|
| 1 | +"""Agent-as-tool adapter. |
| 2 | +
|
| 3 | +This module provides the _AgentAsTool class that wraps an Agent as a tool |
| 4 | +so it can be passed to another agent's tool list. |
| 5 | +""" |
| 6 | + |
| 7 | +from __future__ import annotations |
| 8 | + |
| 9 | +import copy |
| 10 | +import logging |
| 11 | +import threading |
| 12 | +from typing import TYPE_CHECKING, Any |
| 13 | + |
| 14 | +from typing_extensions import override |
| 15 | + |
| 16 | +from ..agent.state import AgentState |
| 17 | +from ..types._events import AgentAsToolStreamEvent, ToolInterruptEvent, ToolResultEvent |
| 18 | +from ..types.content import Messages |
| 19 | +from ..types.interrupt import InterruptResponseContent |
| 20 | +from ..types.tools import AgentTool, ToolGenerator, ToolSpec, ToolUse |
| 21 | + |
| 22 | +if TYPE_CHECKING: |
| 23 | + from .agent import Agent |
| 24 | + |
| 25 | +logger = logging.getLogger(__name__) |
| 26 | + |
| 27 | + |
| 28 | +class _AgentAsTool(AgentTool): |
| 29 | + """Adapter that exposes an Agent as a tool for use by other agents. |
| 30 | +
|
| 31 | + The tool accepts a single ``input`` string parameter, invokes the wrapped |
| 32 | + agent, and returns the text response. |
| 33 | +
|
| 34 | + Example: |
| 35 | + ```python |
| 36 | + from strands import Agent |
| 37 | +
|
| 38 | + researcher = Agent(name="researcher", description="Finds information") |
| 39 | +
|
| 40 | + # Use via convenience method (default: fresh conversation each call) |
| 41 | + tool = researcher.as_tool() |
| 42 | +
|
| 43 | + # Preserve context across invocations |
| 44 | + tool = researcher.as_tool(preserve_context=True) |
| 45 | +
|
| 46 | + writer = Agent(name="writer", tools=[tool]) |
| 47 | + writer("Write about AI agents") |
| 48 | + ``` |
| 49 | + """ |
| 50 | + |
| 51 | + def __init__( |
| 52 | + self, |
| 53 | + agent: Agent, |
| 54 | + *, |
| 55 | + name: str, |
| 56 | + description: str | None = None, |
| 57 | + preserve_context: bool = False, |
| 58 | + ) -> None: |
| 59 | + r"""Initialize the agent-as-tool adapter. |
| 60 | +
|
| 61 | + Args: |
| 62 | + agent: The agent to wrap as a tool. |
| 63 | + name: Tool name. Must match the pattern ``[a-zA-Z0-9_\\-]{1,64}``. |
| 64 | + description: Tool description. Defaults to the agent's description, or a |
| 65 | + generic description if the agent has no description set. |
| 66 | + preserve_context: Whether to preserve the agent's conversation history across |
| 67 | + invocations. When False, the agent's messages and state are reset to the |
| 68 | + values they had at construction time before each call, ensuring every |
| 69 | + invocation starts from the same baseline regardless of any external |
| 70 | + interactions with the agent. Defaults to False. |
| 71 | + """ |
| 72 | + super().__init__() |
| 73 | + self._agent = agent |
| 74 | + self._tool_name = name |
| 75 | + self._description = ( |
| 76 | + description or agent.description or f"Use the {name} agent as a tool by providing a natural language input" |
| 77 | + ) |
| 78 | + self._preserve_context = preserve_context |
| 79 | + |
| 80 | + # When preserve_context=False, we snapshot the agent's initial state so we can |
| 81 | + # restore it before each invocation. This mirrors GraphNode.reset_executor_state(). |
| 82 | + self._initial_messages: Messages = [] |
| 83 | + self._initial_state: AgentState = AgentState() |
| 84 | + # Serialize access so _reset_agent_state + stream_async are atomic. |
| 85 | + # threading.Lock (not asyncio.Lock) because run_async() may create |
| 86 | + # separate event loops in different threads. |
| 87 | + self._lock = threading.Lock() |
| 88 | + |
| 89 | + if not preserve_context: |
| 90 | + if getattr(agent, "_session_manager", None) is not None: |
| 91 | + raise ValueError( |
| 92 | + "preserve_context=False cannot be used with an agent that has a session manager. " |
| 93 | + "The session manager persists conversation history externally, which conflicts with " |
| 94 | + "resetting the agent's state between invocations." |
| 95 | + ) |
| 96 | + self._initial_messages = copy.deepcopy(agent.messages) |
| 97 | + self._initial_state = AgentState(agent.state.get()) |
| 98 | + |
| 99 | + @property |
| 100 | + def agent(self) -> Agent: |
| 101 | + """The wrapped agent instance.""" |
| 102 | + return self._agent |
| 103 | + |
| 104 | + @property |
| 105 | + def tool_name(self) -> str: |
| 106 | + """Get the tool name.""" |
| 107 | + return self._tool_name |
| 108 | + |
| 109 | + @property |
| 110 | + def tool_spec(self) -> ToolSpec: |
| 111 | + """Get the tool specification.""" |
| 112 | + return { |
| 113 | + "name": self._tool_name, |
| 114 | + "description": self._description, |
| 115 | + "inputSchema": { |
| 116 | + "json": { |
| 117 | + "type": "object", |
| 118 | + "properties": { |
| 119 | + "input": { |
| 120 | + "type": "string", |
| 121 | + "description": "The input to send to the agent tool.", |
| 122 | + }, |
| 123 | + }, |
| 124 | + "required": ["input"], |
| 125 | + } |
| 126 | + }, |
| 127 | + } |
| 128 | + |
| 129 | + @property |
| 130 | + def tool_type(self) -> str: |
| 131 | + """Get the tool type.""" |
| 132 | + return "agent" |
| 133 | + |
| 134 | + @override |
| 135 | + async def stream(self, tool_use: ToolUse, invocation_state: dict[str, Any], **kwargs: Any) -> ToolGenerator: |
| 136 | + """Invoke the wrapped agent via streaming and yield events. |
| 137 | +
|
| 138 | + Intermediate agent events are wrapped in AgentAsToolStreamEvent so the caller |
| 139 | + can distinguish sub-agent progress from regular tool events. The final |
| 140 | + AgentResult is yielded as a ToolResultEvent. |
| 141 | +
|
| 142 | + When the sub-agent encounters a hook interrupt (e.g. from BeforeToolCallEvent), |
| 143 | + the interrupts are propagated to the parent agent via ToolInterruptEvent. On |
| 144 | + resume, interrupt responses are forwarded to the sub-agent automatically. |
| 145 | +
|
| 146 | + Args: |
| 147 | + tool_use: The tool use request containing the input parameter. |
| 148 | + invocation_state: Context for the tool invocation. |
| 149 | + **kwargs: Additional keyword arguments. |
| 150 | +
|
| 151 | + Yields: |
| 152 | + AgentAsToolStreamEvent for intermediate events, ToolInterruptEvent if the |
| 153 | + sub-agent is interrupted, or ToolResultEvent with the final response. |
| 154 | + """ |
| 155 | + tool_input = tool_use["input"] |
| 156 | + if isinstance(tool_input, dict): |
| 157 | + prompt = tool_input.get("input", "") |
| 158 | + elif isinstance(tool_input, str): |
| 159 | + prompt = tool_input |
| 160 | + else: |
| 161 | + logger.warning("tool_name=<%s> | unexpected input type: %s", self._tool_name, type(tool_input)) |
| 162 | + prompt = str(tool_input) |
| 163 | + |
| 164 | + tool_use_id = tool_use["toolUseId"] |
| 165 | + |
| 166 | + # Serialize access to the underlying agent. _reset_agent_state() mutates |
| 167 | + # the agent before stream_async acquires its own lock, so a concurrent |
| 168 | + # call would corrupt an in-flight invocation. |
| 169 | + if not self._lock.acquire(blocking=False): |
| 170 | + logger.warning( |
| 171 | + "tool_name=<%s>, tool_use_id=<%s> | agent is already processing a request", |
| 172 | + self._tool_name, |
| 173 | + tool_use_id, |
| 174 | + ) |
| 175 | + yield ToolResultEvent( |
| 176 | + { |
| 177 | + "toolUseId": tool_use_id, |
| 178 | + "status": "error", |
| 179 | + "content": [{"text": f"Agent '{self._tool_name}' is already processing a request"}], |
| 180 | + } |
| 181 | + ) |
| 182 | + return |
| 183 | + |
| 184 | + try: |
| 185 | + # Determine if we are resuming the sub-agent from an interrupt. |
| 186 | + if self._is_sub_agent_interrupted(): |
| 187 | + prompt = self._build_interrupt_responses() |
| 188 | + logger.debug( |
| 189 | + "tool_name=<%s>, tool_use_id=<%s> | resuming sub-agent from interrupt", |
| 190 | + self._tool_name, |
| 191 | + tool_use_id, |
| 192 | + ) |
| 193 | + elif not self._preserve_context: |
| 194 | + self._reset_agent_state(tool_use_id) |
| 195 | + |
| 196 | + logger.debug("tool_name=<%s>, tool_use_id=<%s> | invoking agent", self._tool_name, tool_use_id) |
| 197 | + |
| 198 | + result = None |
| 199 | + async for event in self._agent.stream_async(prompt): |
| 200 | + if "result" in event: |
| 201 | + result = event["result"] |
| 202 | + else: |
| 203 | + yield AgentAsToolStreamEvent(tool_use, event, self) |
| 204 | + |
| 205 | + if result is None: |
| 206 | + yield ToolResultEvent( |
| 207 | + { |
| 208 | + "toolUseId": tool_use_id, |
| 209 | + "status": "error", |
| 210 | + "content": [{"text": "Agent did not produce a result"}], |
| 211 | + } |
| 212 | + ) |
| 213 | + return |
| 214 | + |
| 215 | + # Propagate sub-agent interrupts to the parent agent. |
| 216 | + if result.stop_reason == "interrupt" and result.interrupts: |
| 217 | + yield ToolInterruptEvent(tool_use, list(result.interrupts)) |
| 218 | + return |
| 219 | + |
| 220 | + if result.structured_output: |
| 221 | + yield ToolResultEvent( |
| 222 | + { |
| 223 | + "toolUseId": tool_use_id, |
| 224 | + "status": "success", |
| 225 | + "content": [{"json": result.structured_output.model_dump()}], |
| 226 | + } |
| 227 | + ) |
| 228 | + else: |
| 229 | + yield ToolResultEvent( |
| 230 | + { |
| 231 | + "toolUseId": tool_use_id, |
| 232 | + "status": "success", |
| 233 | + "content": [{"text": str(result)}], |
| 234 | + } |
| 235 | + ) |
| 236 | + |
| 237 | + except Exception as e: |
| 238 | + logger.warning( |
| 239 | + "tool_name=<%s>, tool_use_id=<%s> | agent invocation failed: %s", |
| 240 | + self._tool_name, |
| 241 | + tool_use_id, |
| 242 | + e, |
| 243 | + ) |
| 244 | + yield ToolResultEvent( |
| 245 | + { |
| 246 | + "toolUseId": tool_use_id, |
| 247 | + "status": "error", |
| 248 | + "content": [{"text": f"Agent error: {e}"}], |
| 249 | + } |
| 250 | + ) |
| 251 | + finally: |
| 252 | + self._lock.release() |
| 253 | + |
| 254 | + def _reset_agent_state(self, tool_use_id: str) -> None: |
| 255 | + """Reset the wrapped agent to its initial state. |
| 256 | +
|
| 257 | + Restores messages and state to the values captured at construction time. |
| 258 | + This mirrors the pattern used by ``GraphNode.reset_executor_state()``. |
| 259 | +
|
| 260 | + Args: |
| 261 | + tool_use_id: Tool use ID for logging context. |
| 262 | + """ |
| 263 | + logger.debug( |
| 264 | + "tool_name=<%s>, tool_use_id=<%s> | resetting agent to initial state", |
| 265 | + self._tool_name, |
| 266 | + tool_use_id, |
| 267 | + ) |
| 268 | + self._agent.messages = copy.deepcopy(self._initial_messages) |
| 269 | + self._agent.state = AgentState(self._initial_state.get()) |
| 270 | + |
| 271 | + def _is_sub_agent_interrupted(self) -> bool: |
| 272 | + """Check whether the wrapped agent is in an activated interrupt state.""" |
| 273 | + return self._agent._interrupt_state.activated |
| 274 | + |
| 275 | + def _build_interrupt_responses(self) -> list[InterruptResponseContent]: |
| 276 | + """Build interrupt response payloads from the sub-agent's interrupt state. |
| 277 | +
|
| 278 | + The parent agent's ``_interrupt_state.resume()`` sets ``.response`` on the shared |
| 279 | + ``Interrupt`` objects (registered by the executor), so we re-package them in the |
| 280 | + format expected by ``Agent.stream_async``. |
| 281 | +
|
| 282 | + Returns: |
| 283 | + List of interrupt response content blocks for resuming the sub-agent. |
| 284 | + """ |
| 285 | + return [ |
| 286 | + {"interruptResponse": {"interruptId": interrupt.id, "response": interrupt.response}} |
| 287 | + for interrupt in self._agent._interrupt_state.interrupts.values() |
| 288 | + if interrupt.response is not None |
| 289 | + ] |
| 290 | + |
| 291 | + @override |
| 292 | + def get_display_properties(self) -> dict[str, str]: |
| 293 | + """Get properties for UI display.""" |
| 294 | + properties = super().get_display_properties() |
| 295 | + properties["Agent"] = getattr(self._agent, "name", "unknown") |
| 296 | + return properties |
0 commit comments