wrapper.py•1.67 kB
from __future__ import annotations
from collections.abc import Callable
from dataclasses import dataclass, replace
from typing import Any
from typing_extensions import Self
from .._run_context import AgentDepsT, RunContext
from .abstract import AbstractToolset, ToolsetTool
@dataclass
class WrapperToolset(AbstractToolset[AgentDepsT]):
"""A toolset that wraps another toolset and delegates to it.
See [toolset docs](../toolsets.md#wrapping-a-toolset) for more information.
"""
wrapped: AbstractToolset[AgentDepsT]
@property
def id(self) -> str | None:
return None # pragma: no cover
@property
def label(self) -> str:
return f'{self.__class__.__name__}({self.wrapped.label})'
async def __aenter__(self) -> Self:
await self.wrapped.__aenter__()
return self
async def __aexit__(self, *args: Any) -> bool | None:
return await self.wrapped.__aexit__(*args)
async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
return await self.wrapped.get_tools(ctx)
async def call_tool(
self, name: str, tool_args: dict[str, Any], ctx: RunContext[AgentDepsT], tool: ToolsetTool[AgentDepsT]
) -> Any:
return await self.wrapped.call_tool(name, tool_args, ctx, tool)
def apply(self, visitor: Callable[[AbstractToolset[AgentDepsT]], None]) -> None:
self.wrapped.apply(visitor)
def visit_and_replace(
self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]]
) -> AbstractToolset[AgentDepsT]:
return replace(self, wrapped=self.wrapped.visit_and_replace(visitor))