Skip to content

Commit cf5e977

Browse files
mjn96raychen911
authored andcommitted
feat(tools): support GraphAgent in AgentTool
Add GraphAgent support in _run_async_impl so that when the wrapped agent is a GraphAgent, the tool result is retrieved from tool_context.state using STATE_KEY_LAST_RESPONSE.
1 parent 344d3fd commit cf5e977

1 file changed

Lines changed: 3 additions & 0 deletions

File tree

trpc_agent_sdk/tools/_agent_tool.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -136,6 +136,7 @@ async def _run_async_impl(
136136
) -> Any:
137137
try:
138138
from trpc_agent_sdk.agents import LlmAgent
139+
from trpc_agent_sdk.dsl.graph import GraphAgent, STATE_KEY_LAST_RESPONSE
139140
if self.skip_summarization:
140141
tool_context.event_actions.skip_summarization = True
141142
if isinstance(self.agent, LlmAgent) and self.agent.input_schema:
@@ -206,6 +207,8 @@ async def _run_async_impl(
206207
merged_text = '\n'.join([p.text for p in last_event.content.parts if p.text])
207208
repaired = json_repair_string(merged_text)
208209
tool_result = self.agent.output_schema.model_validate_json(repaired).model_dump(exclude_none=True)
210+
elif isinstance(self.agent, GraphAgent):
211+
tool_result = tool_context.state.get(STATE_KEY_LAST_RESPONSE, '')
209212
else:
210213
tool_result = '\n'.join([p.text for p in last_event.content.parts if p.text])
211214
return tool_result

0 commit comments

Comments
 (0)