|
32 | 32 | if TYPE_CHECKING: |
33 | 33 | from .lifecycle import AgentHooks |
34 | 34 | from .mcp import MCPServer |
35 | | - from .result import RunResult |
| 35 | + from .result import RunResult, RunResultStreaming |
36 | 36 |
|
37 | 37 |
|
38 | 38 | @dataclass |
@@ -381,9 +381,11 @@ def as_tool( |
381 | 381 | self, |
382 | 382 | tool_name: str | None, |
383 | 383 | tool_description: str | None, |
| 384 | + *, |
384 | 385 | custom_output_extractor: Callable[[RunResult], Awaitable[str]] | None = None, |
385 | 386 | is_enabled: bool |
386 | 387 | | Callable[[RunContextWrapper[Any], AgentBase[Any]], MaybeAwaitable[bool]] = True, |
| 388 | + stream_inner_events: bool = False, |
387 | 389 | ) -> Tool: |
388 | 390 | """Transform this agent into a tool, callable by other agents. |
389 | 391 |
|
@@ -412,17 +414,36 @@ def as_tool( |
412 | 414 | async def run_agent(context: RunContextWrapper, input: str) -> str: |
413 | 415 | from .run import Runner |
414 | 416 |
|
415 | | - output = await Runner.run( |
416 | | - starting_agent=self, |
417 | | - input=input, |
418 | | - context=context.context, |
419 | | - ) |
| 417 | + output_run: RunResult | RunResultStreaming |
| 418 | + if stream_inner_events: |
| 419 | + from .stream_events import RunItemStreamEvent |
| 420 | + |
| 421 | + sub_run = Runner.run_streamed( |
| 422 | + self, |
| 423 | + input=input, |
| 424 | + context=context.context, |
| 425 | + ) |
| 426 | + parent_queue = getattr(context, "_event_queue", None) |
| 427 | + async for ev in sub_run.stream_events(): |
| 428 | + if parent_queue is not None and isinstance(ev, RunItemStreamEvent): |
| 429 | + if ev.name in ("tool_called", "tool_output"): |
| 430 | + parent_queue.put_nowait(ev) |
| 431 | + output_run = sub_run |
| 432 | + else: |
| 433 | + output_run = await Runner.run( |
| 434 | + starting_agent=self, |
| 435 | + input=input, |
| 436 | + context=context.context, |
| 437 | + ) |
| 438 | + |
420 | 439 | if custom_output_extractor: |
421 | | - return await custom_output_extractor(output) |
| 440 | + return await custom_output_extractor(cast(Any, output_run)) |
422 | 441 |
|
423 | | - return ItemHelpers.text_message_outputs(output.new_items) |
| 442 | + return ItemHelpers.text_message_outputs(output_run.new_items) |
424 | 443 |
|
425 | | - return run_agent |
| 444 | + tool = run_agent |
| 445 | + tool.stream_inner_events = stream_inner_events |
| 446 | + return tool |
426 | 447 |
|
427 | 448 | async def get_system_prompt(self, run_context: RunContextWrapper[TContext]) -> str | None: |
428 | 449 | if isinstance(self.instructions, str): |
|
0 commit comments