diff --git a/pydantic_ai_slim/pydantic_ai/_cli.py b/pydantic_ai_slim/pydantic_ai/_cli.py index 95fcf8b520..c265cfab37 100644 --- a/pydantic_ai_slim/pydantic_ai/_cli.py +++ b/pydantic_ai_slim/pydantic_ai/_cli.py @@ -18,7 +18,7 @@ from ._run_context import AgentDepsT from .agent import AbstractAgent, Agent from .exceptions import UserError -from .messages import ModelMessage, ModelResponse +from .messages import FunctionToolCallEvent, FunctionToolResultEvent, ModelMessage, ModelResponse from .models import KnownModelName, infer_model from .output import OutputDataT @@ -274,7 +274,11 @@ async def ask_agent( deps: AgentDepsT = None, messages: list[ModelMessage] | None = None, ) -> list[ModelMessage]: - status = Status('[dim]Working on it…[/dim]', console=console) + MODEL_CALL_STATUS_MSG = '[dim]Calling model…[/dim]' + TOOL_EXECUTION_STATUS_MSG = '[dim]Executing tools…[/dim]' + MAX_TOOL_CALL_RESULT_LEN = 100 + MAX_TOOL_CALL_ID_LEN = 5 + status = Status(MODEL_CALL_STATUS_MSG, console=console) if not stream: with status: @@ -285,15 +289,40 @@ async def ask_agent( with status, ExitStack() as stack: async with agent.iter(prompt, message_history=messages, deps=deps) as agent_run: - live = Live('', refresh_per_second=15, console=console, vertical_overflow='ellipsis') + final_output_live = None async for node in agent_run: if Agent.is_model_request_node(node): + status.update(MODEL_CALL_STATUS_MSG) async with node.stream(agent_run.ctx) as handle_stream: status.stop() # stopping multiple times is idempotent - stack.enter_context(live) # entering multiple times is idempotent async for content in handle_stream.stream_output(debounce_by=None): - live.update(Markdown(str(content), code_theme=code_theme)) + if final_output_live is None: + final_output_live = Live( + '', refresh_per_second=15, console=console, vertical_overflow='ellipsis' + ) + stack.enter_context(final_output_live) # entering multiple times is idempotent + final_output_live.update(Markdown(str(content), code_theme=code_theme)) + elif Agent.is_call_tools_node(node): + status.update(TOOL_EXECUTION_STATUS_MSG) + async with node.stream(agent_run.ctx) as handle_stream: + async for event in handle_stream: + if isinstance(event, FunctionToolCallEvent): + status.stop() # stopping multiple times is idempotent + console.print( + Markdown( + f'[Tool] {event.part.tool_name!r}[{event.part.tool_call_id[-5:]}] called with args={event.part.args}' + ) + ) + status.start() + elif isinstance(event, FunctionToolResultEvent): + status.stop() # stopping multiple times is idempotent + console.print( + Markdown( + f'[Tool] {event.result.tool_name!r}[{event.result.tool_call_id[-MAX_TOOL_CALL_ID_LEN:]}] returned => {event.result.content if len(event.result.content) < MAX_TOOL_CALL_RESULT_LEN else str(event.result.content[:MAX_TOOL_CALL_RESULT_LEN]) + "..."}' + ) + ) + status.start() assert agent_run.result is not None return agent_run.result.all_messages()