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
@@ -233,7 +233,9 @@ def as_tool(
233
233
self ,
234
234
tool_name : str | None ,
235
235
tool_description : str | None ,
236
+ * ,
236
237
custom_output_extractor : Callable [[RunResult ], Awaitable [str ]] | None = None ,
238
+ stream_inner_events : bool = False ,
237
239
) -> Tool :
238
240
"""Transform this agent into a tool, callable by other agents.
239
241
@@ -258,17 +260,36 @@ def as_tool(
258
260
async def run_agent (context : RunContextWrapper , input : str ) -> str :
259
261
from .run import Runner
260
262
261
- output = await Runner .run (
262
- starting_agent = self ,
263
- input = input ,
264
- context = context .context ,
265
- )
263
+ output_run : RunResult | RunResultStreaming
264
+ if stream_inner_events :
265
+ from .stream_events import RunItemStreamEvent
266
+
267
+ sub_run = Runner .run_streamed (
268
+ self ,
269
+ input = input ,
270
+ context = context .context ,
271
+ )
272
+ parent_queue = getattr (context , "_event_queue" , None )
273
+ async for ev in sub_run .stream_events ():
274
+ if parent_queue is not None and isinstance (ev , RunItemStreamEvent ):
275
+ if ev .name in ("tool_called" , "tool_output" ):
276
+ parent_queue .put_nowait (ev )
277
+ output_run = sub_run
278
+ else :
279
+ output_run = await Runner .run (
280
+ starting_agent = self ,
281
+ input = input ,
282
+ context = context .context ,
283
+ )
284
+
266
285
if custom_output_extractor :
267
- return await custom_output_extractor (output )
286
+ return await custom_output_extractor (cast ( Any , output_run ) )
268
287
269
- return ItemHelpers .text_message_outputs (output .new_items )
288
+ return ItemHelpers .text_message_outputs (output_run .new_items )
270
289
271
- return run_agent
290
+ tool = run_agent
291
+ tool .stream_inner_events = stream_inner_events
292
+ return tool
272
293
273
294
async def get_system_prompt (self , run_context : RunContextWrapper [TContext ]) -> str | None :
274
295
"""Get the system prompt for the agent."""
0 commit comments