Refactoring

Author: DouweM

examples/pydantic_ai_examples/chat_app.py

@@ -96,7 +96,7 @@ def options_chat():
 @app.post('/api/chat')
 async def get_chat(request: Request, database: Database = Depends(get_db)) -> Response:
     return await VercelAIAdapter[Deps].dispatch_request(
-        chat_agent, request, deps=Deps(database, 123)
+        request, agent=chat_agent, deps=Deps(database, 123)
     )
 
 

pydantic_ai_slim/pydantic_ai/ag_ui.py

@@ -96,8 +96,8 @@ async def handle_ag_ui_request(
         A streaming Starlette response with AG-UI protocol events.
     """
     return await AGUIAdapter[AgentDepsT].dispatch_request(
-        agent,
         request,
+        agent=agent,
         deps=deps,
         output_type=output_type,
         message_history=message_history,
@@ -112,7 +112,7 @@ async def handle_ag_ui_request(
     )
 
 
-async def run_ag_ui(
+def run_ag_ui(
     agent: AbstractAgent[AgentDepsT, Any],
     run_input: RunAgentInput,
     accept: str = SSE_CONTENT_TYPE,
@@ -153,8 +153,8 @@ async def run_ag_ui(
     Yields:
         Streaming event chunks encoded as strings according to the accept header value.
     """
-    adapter = AGUIAdapter(agent=agent, request=run_input)
-    async for event in adapter.encode_stream(
+    adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
+    return adapter.encode_stream(
         adapter.run_stream(
             output_type=output_type,
             message_history=message_history,
@@ -168,6 +168,4 @@ async def run_ag_ui(
             toolsets=toolsets,
             on_complete=on_complete,
         ),
-        accept=accept,
-    ):
-        yield event
+    )

pydantic_ai_slim/pydantic_ai/agent/abstract.py

@@ -999,6 +999,8 @@ def to_ag_ui(
         *,
         # Agent.iter parameters
         output_type: OutputSpec[OutputDataT] | None = None,
+        message_history: Sequence[_messages.ModelMessage] | None = None,
+        deferred_tool_results: DeferredToolResults | None = None,
         model: models.Model | models.KnownModelName | str | None = None,
         deps: AgentDepsT = None,
         model_settings: ModelSettings | None = None,
@@ -1045,6 +1047,8 @@ def to_ag_ui(
             output_type: Custom output type to use for this run, `output_type` may only be used if the agent has
                 no output validators since output validators would expect an argument that matches the agent's
                 output type.
+            message_history: History of the conversation so far.
+            deferred_tool_results: Optional results for deferred tool calls in the message history.
             model: Optional model to use for this run, required if `model` was not set when creating the agent.
             deps: Optional dependencies to use for this run.
             model_settings: Optional settings to use for this model's request.
@@ -1080,6 +1084,8 @@ def to_ag_ui(
             agent=self,
             # Agent.iter parameters
             output_type=output_type,
+            message_history=message_history,
+            deferred_tool_results=deferred_tool_results,
             model=model,
             deps=deps,
             model_settings=model_settings,