|
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 |
@@ -356,9 +356,11 @@ def as_tool( |
356 | 356 | self, |
357 | 357 | tool_name: str | None, |
358 | 358 | tool_description: str | None, |
| 359 | + *, |
359 | 360 | custom_output_extractor: Callable[[RunResult], Awaitable[str]] | None = None, |
360 | 361 | is_enabled: bool |
361 | 362 | | Callable[[RunContextWrapper[Any], AgentBase[Any]], MaybeAwaitable[bool]] = True, |
| 363 | + stream_inner_events: bool = False, |
362 | 364 | ) -> Tool: |
363 | 365 | """Transform this agent into a tool, callable by other agents. |
364 | 366 |
|
@@ -387,17 +389,36 @@ def as_tool( |
387 | 389 | async def run_agent(context: RunContextWrapper, input: str) -> str: |
388 | 390 | from .run import Runner |
389 | 391 |
|
390 | | - output = await Runner.run( |
391 | | - starting_agent=self, |
392 | | - input=input, |
393 | | - context=context.context, |
394 | | - ) |
| 392 | + output_run: RunResult | RunResultStreaming |
| 393 | + if stream_inner_events: |
| 394 | + from .stream_events import RunItemStreamEvent |
| 395 | + |
| 396 | + sub_run = Runner.run_streamed( |
| 397 | + self, |
| 398 | + input=input, |
| 399 | + context=context.context, |
| 400 | + ) |
| 401 | + parent_queue = getattr(context, "_event_queue", None) |
| 402 | + async for ev in sub_run.stream_events(): |
| 403 | + if parent_queue is not None and isinstance(ev, RunItemStreamEvent): |
| 404 | + if ev.name in ("tool_called", "tool_output"): |
| 405 | + parent_queue.put_nowait(ev) |
| 406 | + output_run = sub_run |
| 407 | + else: |
| 408 | + output_run = await Runner.run( |
| 409 | + starting_agent=self, |
| 410 | + input=input, |
| 411 | + context=context.context, |
| 412 | + ) |
| 413 | + |
395 | 414 | if custom_output_extractor: |
396 | | - return await custom_output_extractor(output) |
| 415 | + return await custom_output_extractor(cast(Any, output_run)) |
397 | 416 |
|
398 | | - return ItemHelpers.text_message_outputs(output.new_items) |
| 417 | + return ItemHelpers.text_message_outputs(output_run.new_items) |
399 | 418 |
|
400 | | - return run_agent |
| 419 | + tool = run_agent |
| 420 | + tool.stream_inner_events = stream_inner_events |
| 421 | + return tool |
401 | 422 |
|
402 | 423 | async def get_system_prompt(self, run_context: RunContextWrapper[TContext]) -> str | None: |
403 | 424 | if isinstance(self.instructions, str): |
|
0 commit comments