reference/generated/torchrl.envs.llm.transforms.ToolRegistry

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

ToolRegistry

class torchrl.envs.llm.transforms.ToolRegistry(services: Sequence[ToolService] = ())

Registry for managing available tool services.

This class maintains a collection of tool services that can be looked up by name for execution.

Parameters:

services (Sequence[ToolService], optional) – Initial services to register. Defaults to an empty sequence.

Examples

>>> class AddService:
...     name = "add"
...     schema_in = {"a": int, "b": int}
...     schema_out = {"result": int}
...     def __call__(self, a, b, **kwargs):
...         return {"result": a + b}
>>> registry = ToolRegistry([AddService()])
>>> service = registry.get("add")
>>> result = service(a=1, b=2)
>>> print(result)
{"result": 3}

get(name: str) → ToolService

Retrieve a service by name.

Parameters:

name (str) – The name of the service to retrieve.

Returns:

The requested service.

Return type:

ToolService

Raises:

KeyError – If the service is not found.

register(service: ToolService) → None

Register a new service.

Parameters:

service (ToolService) – The service to register.