A node that wraps an ADK Tool.
| 33 | |
| 34 | |
| 35 | class _ToolNode(BaseNode): |
| 36 | """A node that wraps an ADK Tool.""" |
| 37 | |
| 38 | model_config = ConfigDict(arbitrary_types_allowed=True) |
| 39 | tool: BaseTool = Field(...) |
| 40 | |
| 41 | def __init__( |
| 42 | self, |
| 43 | *, |
| 44 | tool: BaseTool, |
| 45 | name: str | None = None, |
| 46 | retry_config: RetryConfig | None = None, |
| 47 | timeout: float | None = None, |
| 48 | ): |
| 49 | super().__init__( |
| 50 | tool=tool, |
| 51 | name=name or tool.name, |
| 52 | rerun_on_resume=False, |
| 53 | retry_config=retry_config, |
| 54 | timeout=timeout, |
| 55 | ) |
| 56 | |
| 57 | @override |
| 58 | async def _run_impl( |
| 59 | self, |
| 60 | *, |
| 61 | ctx: Context, |
| 62 | node_input: Any, |
| 63 | ) -> AsyncGenerator[Any, None]: |
| 64 | tool_context = ToolContext( |
| 65 | invocation_context=ctx.get_invocation_context(), |
| 66 | function_call_id=str(uuid.uuid4()), |
| 67 | ) |
| 68 | |
| 69 | args = node_input |
| 70 | if args is None: |
| 71 | args = {} |
| 72 | elif not isinstance(args, dict): |
| 73 | raise TypeError( |
| 74 | 'The input to ToolNode must be a dictionary of tool arguments or' |
| 75 | f' None, but got {type(args)}.' |
| 76 | ) |
| 77 | |
| 78 | response = await self.tool.run_async(args=args, tool_context=tool_context) |
| 79 | state_delta = ( |
| 80 | dict(tool_context.actions.state_delta) |
| 81 | if tool_context.actions.state_delta |
| 82 | else None |
| 83 | ) |
| 84 | if response is not None: |
| 85 | yield Event( |
| 86 | output=response, |
| 87 | state=state_delta, |
| 88 | ) |
| 89 | elif state_delta: |
| 90 | yield Event(state=state_delta) |