27
27
if TYPE_CHECKING :
28
28
from .lifecycle import AgentHooks
29
29
from .mcp import MCPServer
30
- from .result import RunResult
30
+ from .result import RunResult , RunResultStreaming
31
31
32
32
33
33
@dataclass
@@ -199,7 +199,9 @@ def as_tool(
199
199
self ,
200
200
tool_name : str | None ,
201
201
tool_description : str | None ,
202
+ * ,
202
203
custom_output_extractor : Callable [[RunResult ], Awaitable [str ]] | None = None ,
204
+ stream_inner_events : bool = False ,
203
205
) -> Tool :
204
206
"""Transform this agent into a tool, callable by other agents.
205
207
@@ -224,17 +226,36 @@ def as_tool(
224
226
async def run_agent (context : RunContextWrapper , input : str ) -> str :
225
227
from .run import Runner
226
228
227
- output = await Runner .run (
228
- starting_agent = self ,
229
- input = input ,
230
- context = context .context ,
231
- )
229
+ output_run : RunResult | RunResultStreaming
230
+ if stream_inner_events :
231
+ from .stream_events import RunItemStreamEvent
232
+
233
+ sub_run = Runner .run_streamed (
234
+ self ,
235
+ input = input ,
236
+ context = context .context ,
237
+ )
238
+ parent_queue = getattr (context , "_event_queue" , None )
239
+ async for ev in sub_run .stream_events ():
240
+ if parent_queue is not None and isinstance (ev , RunItemStreamEvent ):
241
+ if ev .name in ("tool_called" , "tool_output" ):
242
+ parent_queue .put_nowait (ev )
243
+ output_run = sub_run
244
+ else :
245
+ output_run = await Runner .run (
246
+ starting_agent = self ,
247
+ input = input ,
248
+ context = context .context ,
249
+ )
250
+
232
251
if custom_output_extractor :
233
- return await custom_output_extractor (output )
252
+ return await custom_output_extractor (cast ( Any , output_run ) )
234
253
235
- return ItemHelpers .text_message_outputs (output .new_items )
254
+ return ItemHelpers .text_message_outputs (output_run .new_items )
236
255
237
- return run_agent
256
+ tool = run_agent
257
+ tool .stream_inner_events = stream_inner_events
258
+ return tool
238
259
239
260
async def get_system_prompt (self , run_context : RunContextWrapper [TContext ]) -> str | None :
240
261
"""Get the system prompt for the agent."""
@@ -256,9 +277,7 @@ async def get_prompt(
256
277
"""Get the prompt for the agent."""
257
278
return await PromptUtil .to_model_input (self .prompt , run_context , self )
258
279
259
- async def get_mcp_tools (
260
- self , run_context : RunContextWrapper [TContext ]
261
- ) -> list [Tool ]:
280
+ async def get_mcp_tools (self , run_context : RunContextWrapper [TContext ]) -> list [Tool ]:
262
281
"""Fetches the available tools from the MCP servers."""
263
282
convert_schemas_to_strict = self .mcp_config .get ("convert_schemas_to_strict" , False )
264
283
return await MCPUtil .get_all_function_tools (
0 commit comments