Module livekit.agents.beta.toolsets

class ToolProxyToolset(ToolSearchToolset):
    """Exposes exactly two fixed tools: search_tools and call_tool.

    Unlike ToolSearchToolset which dynamically modifies the tool list,
    ToolProxyToolset keeps a constant tool list. ``search_tools`` returns
    tool schemas as text, and ``call_tool`` executes tools by name.

    This is useful for maximizing prompt cache hit rates with providers
    that cache based on tool definitions (e.g. Anthropic, OpenAI).
    """

    def __init__(
        self,
        *,
        id: str,
        tools: list[Tool | Toolset] | None = None,
        max_results: int = 5,
        search_strategy: NotGivenOr[SearchStrategy] = NOT_GIVEN,
        search_description: NotGivenOr[str] = NOT_GIVEN,
        query_description: NotGivenOr[str] = NOT_GIVEN,
        call_description: NotGivenOr[str] = NOT_GIVEN,
    ) -> None:
        super().__init__(
            id=id,
            tools=tools,
            max_results=max_results,
            search_strategy=search_strategy,
            search_description=search_description or _DEFAULT_SEARCH_DESCRIPTION,
            query_description=query_description,
        )
        self._tool_ctx: ToolContext | None = None

        call_description = call_description or _DEFAULT_CALL_DESCRIPTION
        self._call_tool = function_tool(
            self._handle_call,
            raw_schema={
                "name": "call_tool",
                "description": call_description,
                "parameters": {
                    "type": "object",
                    "properties": {
                        "name": {
                            "type": "string",
                            "description": "The name of the tool to call",
                        },
                        "parameters": {
                            "type": "object",
                            "description": "The parameters to pass to the tool",
                        },
                    },
                    "required": ["name", "parameters"],
                },
            },
        )

    @property
    def tools(self) -> list[Tool | Toolset]:
        # constant tool list — only search_tools and call_tool
        return [self._search_tool, self._call_tool]

    async def setup(self, *, reload: bool = False) -> Self:
        await super().setup(reload=reload)

        # build a ToolContext from all wrapped tools for call_tool execution
        self._tool_ctx = ToolContext(self._tools)
        return self

    async def _handle_search(self, raw_arguments: dict[str, object]) -> str:
        query = str(raw_arguments.get("query", ""))
        if not query:
            raise ToolError("query cannot be empty")

        tools = await self._search_tools(query)
        if not tools:
            raise ToolError(f"No tools found matching '{query}'.")

        tool_ctx = ToolContext(tools)
        schemas = [_build_tool_schema(tool) for tool in tool_ctx.function_tools.values()]
        return "\n".join(json.dumps(schema) for schema in schemas)

    async def _handle_call(self, ctx: RunContext[Any], raw_arguments: dict[str, object]) -> Any:
        name = str(raw_arguments.get("name", ""))
        parameters = raw_arguments.get("parameters")

        if not name:
            raise ToolError("tool name cannot be empty")

        if parameters is None:
            raise ToolError("parameters is required")

        if self._tool_ctx is None:
            raise RuntimeError("toolset not initialized, call setup() first")

        fnc_tool = self._tool_ctx.get_function_tool(name)
        if fnc_tool is None:
            raise ToolError(f"unknown tool '{name}', use search_tools to discover available tools")

        try:
            json_args = json.dumps(parameters) if isinstance(parameters, dict) else str(parameters)
            fnc_args, fnc_kwargs = prepare_function_arguments(
                fnc=fnc_tool,
                json_arguments=json_args,
                call_ctx=ctx,
            )
        except ValidationError as e:
            raise ToolError(
                f"invalid parameters for tool '{name}': {e.json(include_url=False)}"
            ) from e
        except ToolError:
            raise
        except Exception as e:
            logger.exception(
                f"error parsing arguments for tool '{name}'",
                extra={"tool": name, "arguments": parameters},
            )
            raise ToolError(f"error calling '{name}': {e}") from e

        return await fnc_tool(*fnc_args, **fnc_kwargs)

class ToolSearchToolset(Toolset):
    """Wraps tools/toolsets and exposes a tool_search function for dynamic loading.

    Instead of loading all tool definitions into LLM context, this exposes a single
    ``tool_search`` function. When the LLM calls it, matching tools are dynamically
    loaded into the context.

    Each tool (FunctionTool, RawFunctionTool, ProviderTool) is indexed as its own
    SearchItem. If a matched tool belongs to a Toolset, the entire Toolset is loaded
    atomically.
    """

    def __init__(
        self,
        *,
        id: str,
        tools: list[Tool | Toolset] | None = None,
        max_results: int = 5,
        search_strategy: NotGivenOr[SearchStrategy] = NOT_GIVEN,
        search_description: NotGivenOr[str] = NOT_GIVEN,
        query_description: NotGivenOr[str] = NOT_GIVEN,
    ) -> None:
        super().__init__(id=id, tools=tools)
        self._strategy = search_strategy or BM25SearchStrategy()
        self._max_results = max_results
        self._loaded_tools: list[Tool | Toolset] = []

        self._search_items: list[SearchItem] = []
        self._initialized = False
        self._lock = asyncio.Lock()

        search_description = search_description or _DEFAULT_SEARCH_DESCRIPTION
        query_description = query_description or _DEFAULT_QUERY_DESCRIPTION
        self._search_tool = function_tool(
            self._handle_search,
            raw_schema={
                "name": "tool_search",
                "description": search_description,
                "parameters": {
                    "type": "object",
                    "properties": {"query": {"type": "string", "description": query_description}},
                    "required": ["query"],
                },
            },
        )

    @property
    def tools(self) -> list[Tool | Toolset]:
        return [self._search_tool, *self._loaded_tools]

    async def setup(self, *, reload: bool = False) -> Self:
        await super().setup()
        async with self._lock:
            if not reload and self._initialized:
                return self

            # setup wrapped toolsets
            toolsets = [t for t in self._tools if isinstance(t, Toolset)]
            if toolsets:
                await asyncio.gather(*(ts.setup() for ts in toolsets))

            self._search_items = []

            def _index_tool(tool: Tool | Toolset, source: Tool | Toolset) -> None:
                if isinstance(tool, Toolset):
                    tool_ctx = ToolContext([tool])
                    for tool in tool_ctx.flatten():
                        _index_tool(tool, source)
                elif isinstance(tool, (FunctionTool, RawFunctionTool)):
                    self._search_items.append(
                        SearchItem(
                            name=tool.id,
                            description=_get_tool_description(tool),
                            parameters=_get_tool_params(tool),
                            source=source,
                        )
                    )
                elif isinstance(tool, ProviderTool):
                    self._search_items.append(
                        SearchItem(name=tool.id, description="", parameters={}, source=source)
                    )
                else:
                    raise ValueError(f"Unsupported tool type: {type(tool)}")

            for tool in self._tools:
                _index_tool(tool, tool)

            result = self._strategy.build_index(self._search_items)
            if inspect.isawaitable(result):
                await result
            self._initialized = True
            return self

    async def _handle_search(self, raw_arguments: dict[str, object]) -> str:
        query = str(raw_arguments.get("query", ""))
        tools = await self._search_tools(query)
        if not tools:
            raise ToolError(f"No tools found matching '{query}'.")

        self._loaded_tools = tools
        return "Tools loaded successfully."

    async def _search_tools(self, query: str) -> list[Tool | Toolset]:
        if not query:
            raise ToolError("query cannot be empty")

        results = self._strategy.search(query, self._search_items, self._max_results)
        if inspect.isawaitable(results):
            results = await results

        return list(dict.fromkeys(result.source for result in results))

    async def aclose(self) -> None:
        await super().aclose()
        self._initialized = False
        self._search_items.clear()
        self._loaded_tools.clear()

        result = self._strategy.cleanup()
        if inspect.isawaitable(result):
            await result