|
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