|
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