feat: support wrapping downstream agents as tools (#819)

Author: akotyla

packages/ragbits-agents/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## Unreleased
 
+- Support wrapping downstream agents as tools (#818)
+
 ## 1.3.0 (2025-09-11)
 ### Changed
 

packages/ragbits-agents/src/ragbits/agents/_main.py

@@ -219,6 +219,8 @@ class Agent(
     def __init__(
         self,
         llm: LLM[LLMClientOptionsT],
+        name: str | None = None,
+        description: str | None = None,
         prompt: str | type[Prompt[PromptInputT, PromptOutputT]] | Prompt[PromptInputT, PromptOutputT] | None = None,
         *,
         history: ChatFormat | None = None,
@@ -232,6 +234,8 @@ def __init__(
 
         Args:
             llm: The LLM to run the agent.
+            name: Optional name of the agent. Used to identify the agent instance.
+            description: Optional description of the agent.
             prompt: The prompt for the agent. Can be:
                 - str: A string prompt that will be used as system message when combined with string input,
                     or as the user message when no input is provided during run().
@@ -248,7 +252,17 @@ def __init__(
         self.id = uuid.uuid4().hex[:8]
         self.llm = llm
         self.prompt = prompt
-        self.tools = [Tool.from_callable(tool) for tool in tools or []]
+        self.name = name
+        self.description = description
+        self.tools = []
+        for tool in tools or []:
+            if isinstance(tool, tuple):
+                agent, kwargs = tool
+                self.tools.append(Tool.from_agent(agent, **kwargs))
+            elif isinstance(tool, Agent):
+                self.tools.append(Tool.from_agent(tool))
+            else:
+                self.tools.append(Tool.from_callable(tool))
         self.mcp_servers = mcp_servers or []
         self.history = history or []
         self.keep_history = keep_history
@@ -555,7 +569,10 @@ async def _get_all_tools(self) -> dict[str, Tool]:
         return tools_mapping
 
     async def _execute_tool(
-        self, tool_call: ToolCall, tools_mapping: dict[str, Tool], context: AgentRunContext | None = None
+        self,
+        tool_call: ToolCall,
+        tools_mapping: dict[str, Tool],
+        context: AgentRunContext | None = None,
     ) -> ToolCallResult:
         if tool_call.type != "function":
             raise AgentToolNotSupportedError(tool_call.type)
@@ -577,10 +594,22 @@ async def _execute_tool(
                     else tool.on_tool_call(**call_args)
                 )
 
+                if isinstance(tool_output, AgentResultStreaming):
+                    async for _ in tool_output:
+                        pass
+
+                    tool_output = {
+                        "content": tool_output.content,
+                        "metadata": tool_output.metadata,
+                        "tool_calls": tool_output.tool_calls,
+                        "usage": tool_output.usage,
+                    }
+
                 outputs.result = {
                     "tool_output": tool_output,
                     "tool_call_id": tool_call.id,
                 }
+
             except Exception as e:
                 outputs.result = {
                     "error": str(e),
@@ -758,3 +787,16 @@ def from_pydantic_ai(cls, pydantic_ai_agent: "PydanticAIAgent") -> Self:
             tools=[tool.function for _, tool in pydantic_ai_agent._function_tools.items()],
             mcp_servers=cast(list[MCPServer], mcp_servers),
         )
+
+    def to_tool(self, name: str | None = None, description: str | None = None) -> Tool:
+        """
+        Convert the agent into a Tool instance.
+
+        Args:
+            name: Optional override for the tool name.
+            description: Optional override for the tool description.
+
+        Returns:
+            Tool instance representing the agent.
+        """
+        return Tool.from_agent(self, name=name or self.name, description=description or self.description)

packages/ragbits-agents/src/ragbits/agents/tool.py

@@ -1,13 +1,19 @@
 from collections.abc import Callable
 from contextlib import suppress
 from dataclasses import dataclass
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
+from pydantic import BaseModel
 from typing_extensions import Self
 
+from ragbits.core.llms.base import LLMClientOptionsT
+from ragbits.core.prompt.prompt import PromptInputT, PromptOutputT
 from ragbits.core.utils.decorators import requires_dependencies
 from ragbits.core.utils.function_schema import convert_function_to_function_schema, get_context_variable_name
 
+if TYPE_CHECKING:
+    from ragbits.agents import Agent, AgentResultStreaming
+
 with suppress(ImportError):
     from pydantic_ai import Tool as PydanticAITool
 
@@ -96,5 +102,69 @@ def to_pydantic_ai(self) -> "PydanticAITool":
             description=self.description,
         )
 
+    @classmethod
+    def from_agent(
+        cls,
+        agent: "Agent[LLMClientOptionsT, PromptInputT, PromptOutputT]",
+        name: str | None = None,
+        description: str | None = None,
+    ) -> "Tool":
+        """
+        Wraps a downstream agent as a single tool. The tool parameters are inferred from
+        the downstream agent's prompt input.
+
+        Args:
+            agent: The downstream agent to wrap as a tool.
+            name: Optional override for the tool name.
+            description: Optional override for the tool description.
+
+        Returns:
+            Tool instance representing the agent.
+        """
+        display_name = name or agent.name or "agent"
+        variable_name = display_name.replace(" ", "_").lower()
+        description = description or agent.description
+
+        input_model_cls = getattr(agent.prompt, "input_type", None)
+        if input_model_cls and issubclass(input_model_cls, BaseModel):
+            fields = input_model_cls.model_fields
+            properties = {}
+            required = list(fields.keys())
+
+            for field_name in fields:
+                param_desc = None
+                for t in getattr(agent, "tools", []):
+                    t_params = getattr(t, "parameters", {}).get("properties", {})
+                    if field_name in t_params:
+                        param_desc = t_params[field_name].get("description")
+                        break
+
+                properties[field_name] = {
+                    "type": "string",
+                    "title": field_name.capitalize(),
+                    "description": param_desc,
+                }
+        else:
+            properties = {"input": {"type": "string", "description": "Input for the downstream agent"}}
+            required = ["input"]
+
+        parameters = {"type": "object", "properties": properties, "required": required}
+
+        def _on_tool_call(**kwargs: dict) -> "AgentResultStreaming":
+            if input_model_cls and issubclass(input_model_cls, BaseModel):
+                model_input = input_model_cls(**kwargs)
+            else:
+                model_input = kwargs.get("input")
+
+            return agent.run_streaming(model_input)
+
+        return cls(
+            name=variable_name,
+            description=description,
+            parameters=parameters,
+            on_tool_call=_on_tool_call,
+            context_var_name=get_context_variable_name(agent.run),
+        )
+
 
 ToolChoice = Literal["auto", "none", "required"] | Callable