pydantic-ai-slim 0.0.17__tar.gz → 0.0.18__tar.gz

{pydantic_ai_slim-0.0.17 → pydantic_ai_slim-0.0.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.0.17
+Version: 0.0.18
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>
 License-Expression: MIT

{pydantic_ai_slim-0.0.17 → pydantic_ai_slim-0.0.18}/pydantic_ai/_griffe.py

@@ -1,6 +1,8 @@
 from __future__ import annotations as _annotations
 
+import logging
 import re
+from contextlib import contextmanager
 from inspect import Signature
 from typing import Any, Callable, Literal, cast
 
@@ -25,7 +27,8 @@ def doc_descriptions(
     parent = cast(GriffeObject, sig)
 
     docstring = Docstring(doc, lineno=1, parser=style or _infer_docstring_style(doc), parent=parent)
-    sections = docstring.parse()
+    with _disable_griffe_logging():
+        sections = docstring.parse()
 
     params = {}
     if parameters := next((p for p in sections if p.kind == DocstringSectionKind.parameters), None):
@@ -125,3 +128,12 @@ _docstring_style_patterns: list[tuple[str, list[str], DocstringStyle]] = [
         'numpy',
     ),
 ]
+
+
+@contextmanager
+def _disable_griffe_logging():
+    # Hacky, but suggested here: https://github.com/mkdocstrings/griffe/issues/293#issuecomment-2167668117
+    old_level = logging.root.getEffectiveLevel()
+    logging.root.setLevel(logging.ERROR)
+    yield
+    logging.root.setLevel(old_level)

{pydantic_ai_slim-0.0.17 → pydantic_ai_slim-0.0.18}/pydantic_ai/_system_prompt.py

@@ -12,6 +12,7 @@ from .tools import AgentDeps, RunContext, SystemPromptFunc
 @dataclass
 class SystemPromptRunner(Generic[AgentDeps]):
     function: SystemPromptFunc[AgentDeps]
+    dynamic: bool = False
     _takes_ctx: bool = field(init=False)
     _is_async: bool = field(init=False)
 

{pydantic_ai_slim-0.0.17 → pydantic_ai_slim-0.0.18}/pydantic_ai/agent.py

@@ -4,13 +4,13 @@ import asyncio
 import dataclasses
 import inspect
 from collections.abc import AsyncIterator, Awaitable, Iterator, Sequence
-from contextlib import asynccontextmanager, contextmanager
+from contextlib import AbstractAsyncContextManager, asynccontextmanager, contextmanager
 from contextvars import ContextVar
 from types import FrameType
 from typing import Any, Callable, Generic, Literal, cast, final, overload
 
 import logfire_api
-from typing_extensions import assert_never, deprecated
+from typing_extensions import TypeVar, assert_never, deprecated
 
 from . import (
     _result,
@@ -57,6 +57,8 @@ EndStrategy = Literal['early', 'exhaustive']
 - `'early'`: Stop processing other tool calls once a final result is found
 - `'exhaustive'`: Process all tool calls even after finding a final result
 """
+RunResultData = TypeVar('RunResultData')
+"""Type variable for the result data of a run where `result_type` was customized on the run call."""
 
 
 @final
@@ -99,14 +101,17 @@ class Agent(Generic[AgentDeps, ResultData]):
     Note, if `model_settings` is provided by `run`, `run_sync`, or `run_stream`, those settings will
     be merged with this value, with the runtime argument taking priority.
     """
-
+    _result_tool_name: str = dataclasses.field(repr=False)
+    _result_tool_description: str | None = dataclasses.field(repr=False)
     _result_schema: _result.ResultSchema[ResultData] | None = dataclasses.field(repr=False)
     _result_validators: list[_result.ResultValidator[AgentDeps, ResultData]] = dataclasses.field(repr=False)
-    _allow_text_result: bool = dataclasses.field(repr=False)
     _system_prompts: tuple[str, ...] = dataclasses.field(repr=False)
     _function_tools: dict[str, Tool[AgentDeps]] = dataclasses.field(repr=False)
     _default_retries: int = dataclasses.field(repr=False)
     _system_prompt_functions: list[_system_prompt.SystemPromptRunner[AgentDeps]] = dataclasses.field(repr=False)
+    _system_prompt_dynamic_functions: dict[str, _system_prompt.SystemPromptRunner[AgentDeps]] = dataclasses.field(
+        repr=False
+    )
     _deps_type: type[AgentDeps] = dataclasses.field(repr=False)
     _max_result_retries: int = dataclasses.field(repr=False)
     _override_deps: _utils.Option[AgentDeps] = dataclasses.field(default=None, repr=False)
@@ -166,11 +171,11 @@ class Agent(Generic[AgentDeps, ResultData]):
         self.end_strategy = end_strategy
         self.name = name
         self.model_settings = model_settings
+        self._result_tool_name = result_tool_name
+        self._result_tool_description = result_tool_description
         self._result_schema = _result.ResultSchema[result_type].build(
             result_type, result_tool_name, result_tool_description
         )
-        # if the result tool is None, or its schema allows `str`, we allow plain text results
-        self._allow_text_result = self._result_schema is None or self._result_schema.allow_text_result
 
         self._system_prompts = (system_prompt,) if isinstance(system_prompt, str) else tuple(system_prompt)
         self._function_tools = {}
@@ -182,13 +187,31 @@ class Agent(Generic[AgentDeps, ResultData]):
                 self._register_tool(Tool(tool))
         self._deps_type = deps_type
         self._system_prompt_functions = []
+        self._system_prompt_dynamic_functions = {}
         self._max_result_retries = result_retries if result_retries is not None else retries
         self._result_validators = []
 
+    @overload
+    async def run(
+        self,
+        user_prompt: str,
+        *,
+        result_type: None = None,
+        message_history: list[_messages.ModelMessage] | None = None,
+        model: models.Model | models.KnownModelName | None = None,
+        deps: AgentDeps = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.Usage | None = None,
+        infer_name: bool = True,
+    ) -> result.RunResult[ResultData]: ...
+
+    @overload
     async def run(
         self,
         user_prompt: str,
         *,
+        result_type: type[RunResultData],
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDeps = None,
@@ -196,7 +219,21 @@ class Agent(Generic[AgentDeps, ResultData]):
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> result.RunResult[ResultData]:
+    ) -> result.RunResult[RunResultData]: ...
+
+    async def run(
+        self,
+        user_prompt: str,
+        *,
+        message_history: list[_messages.ModelMessage] | None = None,
+        model: models.Model | models.KnownModelName | None = None,
+        deps: AgentDeps = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.Usage | None = None,
+        result_type: type[RunResultData] | None = None,
+        infer_name: bool = True,
+    ) -> result.RunResult[Any]:
         """Run the agent with a user prompt in async mode.
 
         Example:
@@ -211,6 +248,8 @@ class Agent(Generic[AgentDeps, ResultData]):
         ```
 
         Args:
+            result_type: Custom result type to use for this run, `result_type` may only be used if the agent has no
+                result validators since result validators would expect an argument that matches the agent's result type.
             user_prompt: User input to start/continue the conversation.
             message_history: History of the conversation so far.
             model: Optional model to use for this run, required if `model` was not set when creating the agent.
@@ -229,6 +268,7 @@ class Agent(Generic[AgentDeps, ResultData]):
 
         deps = self._get_deps(deps)
         new_message_index = len(message_history) if message_history else 0
+        result_schema = self._prepare_result_schema(result_type)
 
         with _logfire.span(
             '{agent_name} run {prompt=}',
@@ -252,7 +292,7 @@ class Agent(Generic[AgentDeps, ResultData]):
 
                 run_context.run_step += 1
                 with _logfire.span('preparing model and tools {run_step=}', run_step=run_context.run_step):
-                    agent_model = await self._prepare_model(run_context)
+                    agent_model = await self._prepare_model(run_context, result_schema)
 
                 with _logfire.span('model request', run_step=run_context.run_step) as model_req_span:
                     model_response, request_usage = await agent_model.request(messages, model_settings)
@@ -264,7 +304,9 @@ class Agent(Generic[AgentDeps, ResultData]):
                 usage_limits.check_tokens(run_context.usage)
 
                 with _logfire.span('handle model response', run_step=run_context.run_step) as handle_span:
-                    final_result, tool_responses = await self._handle_model_response(model_response, run_context)
+                    final_result, tool_responses = await self._handle_model_response(
+                        model_response, run_context, result_schema
+                    )
 
                     if tool_responses:
                         # Add parts to the conversation as a new message
@@ -287,10 +329,40 @@ class Agent(Generic[AgentDeps, ResultData]):
                         tool_responses_str = ' '.join(r.part_kind for r in tool_responses)
                         handle_span.message = f'handle model response -> {tool_responses_str}'
 
+    @overload
+    def run_sync(
+        self,
+        user_prompt: str,
+        *,
+        message_history: list[_messages.ModelMessage] | None = None,
+        model: models.Model | models.KnownModelName | None = None,
+        deps: AgentDeps = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.Usage | None = None,
+        infer_name: bool = True,
+    ) -> result.RunResult[ResultData]: ...
+
+    @overload
+    def run_sync(
+        self,
+        user_prompt: str,
+        *,
+        result_type: type[RunResultData] | None,
+        message_history: list[_messages.ModelMessage] | None = None,
+        model: models.Model | models.KnownModelName | None = None,
+        deps: AgentDeps = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.Usage | None = None,
+        infer_name: bool = True,
+    ) -> result.RunResult[RunResultData]: ...
+
     def run_sync(
         self,
         user_prompt: str,
         *,
+        result_type: type[RunResultData] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDeps = None,
@@ -298,7 +370,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> result.RunResult[ResultData]:
+    ) -> result.RunResult[Any]:
         """Run the agent with a user prompt synchronously.
 
         This is a convenience method that wraps [`self.run`][pydantic_ai.Agent.run] with `loop.run_until_complete(...)`.
@@ -317,6 +389,8 @@ class Agent(Generic[AgentDeps, ResultData]):
         ```
 
         Args:
+            result_type: Custom result type to use for this run, `result_type` may only be used if the agent has no
+                result validators since result validators would expect an argument that matches the agent's result type.
             user_prompt: User input to start/continue the conversation.
             message_history: History of the conversation so far.
             model: Optional model to use for this run, required if `model` was not set when creating the agent.
@@ -334,6 +408,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         return asyncio.get_event_loop().run_until_complete(
             self.run(
                 user_prompt,
+                result_type=result_type,
                 message_history=message_history,
                 model=model,
                 deps=deps,
@@ -344,11 +419,42 @@ class Agent(Generic[AgentDeps, ResultData]):
             )
         )
 
+    @overload
+    def run_stream(
+        self,
+        user_prompt: str,
+        *,
+        result_type: None = None,
+        message_history: list[_messages.ModelMessage] | None = None,
+        model: models.Model | models.KnownModelName | None = None,
+        deps: AgentDeps = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.Usage | None = None,
+        infer_name: bool = True,
+    ) -> AbstractAsyncContextManager[result.StreamedRunResult[AgentDeps, ResultData]]: ...
+
+    @overload
+    def run_stream(
+        self,
+        user_prompt: str,
+        *,
+        result_type: type[RunResultData],
+        message_history: list[_messages.ModelMessage] | None = None,
+        model: models.Model | models.KnownModelName | None = None,
+        deps: AgentDeps = None,
+        model_settings: ModelSettings | None = None,
+        usage_limits: _usage.UsageLimits | None = None,
+        usage: _usage.Usage | None = None,
+        infer_name: bool = True,
+    ) -> AbstractAsyncContextManager[result.StreamedRunResult[AgentDeps, RunResultData]]: ...
+
     @asynccontextmanager
     async def run_stream(
         self,
         user_prompt: str,
         *,
+        result_type: type[RunResultData] | None = None,
         message_history: list[_messages.ModelMessage] | None = None,
         model: models.Model | models.KnownModelName | None = None,
         deps: AgentDeps = None,
@@ -356,7 +462,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> AsyncIterator[result.StreamedRunResult[AgentDeps, ResultData]]:
+    ) -> AsyncIterator[result.StreamedRunResult[AgentDeps, Any]]:
         """Run the agent with a user prompt in async mode, returning a streamed response.
 
         Example:
@@ -372,6 +478,8 @@ class Agent(Generic[AgentDeps, ResultData]):
         ```
 
         Args:
+            result_type: Custom result type to use for this run, `result_type` may only be used if the agent has no
+                result validators since result validators would expect an argument that matches the agent's result type.
             user_prompt: User input to start/continue the conversation.
             message_history: History of the conversation so far.
             model: Optional model to use for this run, required if `model` was not set when creating the agent.
@@ -392,6 +500,7 @@ class Agent(Generic[AgentDeps, ResultData]):
 
         deps = self._get_deps(deps)
         new_message_index = len(message_history) if message_history else 0
+        result_schema = self._prepare_result_schema(result_type)
 
         with _logfire.span(
             '{agent_name} run stream {prompt=}',
@@ -415,7 +524,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                 usage_limits.check_before_request(run_context.usage)
 
                 with _logfire.span('preparing model and tools {run_step=}', run_step=run_context.run_step):
-                    agent_model = await self._prepare_model(run_context)
+                    agent_model = await self._prepare_model(run_context, result_schema)
 
                 with _logfire.span('model request {run_step=}', run_step=run_context.run_step) as model_req_span:
                     async with agent_model.request_stream(messages, model_settings) as model_response:
@@ -426,7 +535,9 @@ class Agent(Generic[AgentDeps, ResultData]):
                         model_req_span.__exit__(None, None, None)
 
                         with _logfire.span('handle model response') as handle_span:
-                            maybe_final_result = await self._handle_streamed_model_response(model_response, run_context)
+                            maybe_final_result = await self._handle_streamed_model_response(
+                                model_response, run_context, result_schema
+                            )
 
                             # Check if we got a final result
                             if isinstance(maybe_final_result, _MarkFinalResult):
@@ -446,7 +557,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                                         part for part in last_message.parts if isinstance(part, _messages.ToolCallPart)
                                     ]
                                     parts = await self._process_function_tools(
-                                        tool_calls, result_tool_name, run_context
+                                        tool_calls, result_tool_name, run_context, result_schema
                                     )
                                     if parts:
                                         messages.append(_messages.ModelRequest(parts))
@@ -457,7 +568,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                                     new_message_index,
                                     usage_limits,
                                     result_stream,
-                                    self._result_schema,
+                                    result_schema,
                                     run_context,
                                     self._result_validators,
                                     result_tool_name,
@@ -535,17 +646,37 @@ class Agent(Generic[AgentDeps, ResultData]):
     @overload
     def system_prompt(self, func: Callable[[], Awaitable[str]], /) -> Callable[[], Awaitable[str]]: ...
 
+    @overload
+    def system_prompt(
+        self, /, *, dynamic: bool = False
+    ) -> Callable[[_system_prompt.SystemPromptFunc[AgentDeps]], _system_prompt.SystemPromptFunc[AgentDeps]]: ...
+
     def system_prompt(
-        self, func: _system_prompt.SystemPromptFunc[AgentDeps], /
-    ) -> _system_prompt.SystemPromptFunc[AgentDeps]:
+        self,
+        func: _system_prompt.SystemPromptFunc[AgentDeps] | None = None,
+        /,
+        *,
+        dynamic: bool = False,
+    ) -> (
+        Callable[[_system_prompt.SystemPromptFunc[AgentDeps]], _system_prompt.SystemPromptFunc[AgentDeps]]
+        | _system_prompt.SystemPromptFunc[AgentDeps]
+    ):
         """Decorator to register a system prompt function.
 
         Optionally takes [`RunContext`][pydantic_ai.tools.RunContext] as its only argument.
         Can decorate a sync or async functions.
 
+        The decorator can be used either bare (`agent.system_prompt`) or as a function call
+        (`agent.system_prompt(...)`), see the examples below.
+
         Overloads for every possible signature of `system_prompt` are included so the decorator doesn't obscure
         the type of the function, see `tests/typed_agent.py` for tests.
 
+        Args:
+            func: The function to decorate
+            dynamic: If True, the system prompt will be reevaluated even when `messages_history` is provided,
+                see [`SystemPromptPart.dynamic_ref`][pydantic_ai.messages.SystemPromptPart.dynamic_ref]
+
         Example:
         ```python
         from pydantic_ai import Agent, RunContext
@@ -556,17 +687,27 @@ class Agent(Generic[AgentDeps, ResultData]):
         def simple_system_prompt() -> str:
             return 'foobar'
 
-        @agent.system_prompt
+        @agent.system_prompt(dynamic=True)
         async def async_system_prompt(ctx: RunContext[str]) -> str:
             return f'{ctx.deps} is the best'
-
-        result = agent.run_sync('foobar', deps='spam')
-        print(result.data)
-        #> success (no tool calls)
         ```
         """
-        self._system_prompt_functions.append(_system_prompt.SystemPromptRunner(func))
-        return func
+        if func is None:
+
+            def decorator(
+                func_: _system_prompt.SystemPromptFunc[AgentDeps],
+            ) -> _system_prompt.SystemPromptFunc[AgentDeps]:
+                runner = _system_prompt.SystemPromptRunner(func_, dynamic=dynamic)
+                self._system_prompt_functions.append(runner)
+                if dynamic:
+                    self._system_prompt_dynamic_functions[func_.__qualname__] = runner
+                return func_
+
+            return decorator
+        else:
+            assert not dynamic, "dynamic can't be True in this case"
+            self._system_prompt_functions.append(_system_prompt.SystemPromptRunner(func, dynamic=dynamic))
+            return func
 
     @overload
     def result_validator(
@@ -818,7 +959,9 @@ class Agent(Generic[AgentDeps, ResultData]):
 
         return model_
 
-    async def _prepare_model(self, run_context: RunContext[AgentDeps]) -> models.AgentModel:
+    async def _prepare_model(
+        self, run_context: RunContext[AgentDeps], result_schema: _result.ResultSchema[RunResultData] | None
+    ) -> models.AgentModel:
         """Build tools and create an agent model."""
         function_tools: list[ToolDefinition] = []
 
@@ -831,10 +974,39 @@ class Agent(Generic[AgentDeps, ResultData]):
 
         return await run_context.model.agent_model(
             function_tools=function_tools,
-            allow_text_result=self._allow_text_result,
-            result_tools=self._result_schema.tool_defs() if self._result_schema is not None else [],
+            allow_text_result=self._allow_text_result(result_schema),
+            result_tools=result_schema.tool_defs() if result_schema is not None else [],
         )
 
+    async def _reevaluate_dynamic_prompts(
+        self, messages: list[_messages.ModelMessage], run_context: RunContext[AgentDeps]
+    ) -> None:
+        """Reevaluate any `SystemPromptPart` with dynamic_ref in the provided messages by running the associated runner function."""
+        # Only proceed if there's at least one dynamic runner.
+        if self._system_prompt_dynamic_functions:
+            for msg in messages:
+                if isinstance(msg, _messages.ModelRequest):
+                    for i, part in enumerate(msg.parts):
+                        if isinstance(part, _messages.SystemPromptPart) and part.dynamic_ref:
+                            # Look up the runner by its ref
+                            if runner := self._system_prompt_dynamic_functions.get(part.dynamic_ref):
+                                updated_part_content = await runner.run(run_context)
+                                msg.parts[i] = _messages.SystemPromptPart(
+                                    updated_part_content, dynamic_ref=part.dynamic_ref
+                                )
+
+    def _prepare_result_schema(
+        self, result_type: type[RunResultData] | None
+    ) -> _result.ResultSchema[RunResultData] | None:
+        if result_type is not None:
+            if self._result_validators:
+                raise exceptions.UserError('Cannot set a custom run `result_type` when the agent has result validators')
+            return _result.ResultSchema[result_type].build(
+                result_type, self._result_tool_name, self._result_tool_description
+            )
+        else:
+            return self._result_schema  # pyright: ignore[reportReturnType]
+
     async def _prepare_messages(
         self, user_prompt: str, message_history: list[_messages.ModelMessage] | None, run_context: RunContext[AgentDeps]
     ) -> list[_messages.ModelMessage]:
@@ -850,8 +1022,10 @@ class Agent(Generic[AgentDeps, ResultData]):
                 ctx_messages.used = True
 
         if message_history:
-            # shallow copy messages
+            # Shallow copy messages
             messages.extend(message_history)
+            # Reevaluate any dynamic system prompt parts
+            await self._reevaluate_dynamic_prompts(messages, run_context)
             messages.append(_messages.ModelRequest([_messages.UserPromptPart(user_prompt)]))
         else:
             parts = await self._sys_parts(run_context)
@@ -861,8 +1035,11 @@ class Agent(Generic[AgentDeps, ResultData]):
         return messages
 
     async def _handle_model_response(
-        self, model_response: _messages.ModelResponse, run_context: RunContext[AgentDeps]
-    ) -> tuple[_MarkFinalResult[ResultData] | None, list[_messages.ModelRequestPart]]:
+        self,
+        model_response: _messages.ModelResponse,
+        run_context: RunContext[AgentDeps],
+        result_schema: _result.ResultSchema[RunResultData] | None,
+    ) -> tuple[_MarkFinalResult[RunResultData] | None, list[_messages.ModelRequestPart]]:
         """Process a non-streamed response from the model.
 
         Returns:
@@ -883,19 +1060,19 @@ class Agent(Generic[AgentDeps, ResultData]):
         # This accounts for cases like anthropic returns that might contain a text response
         # and a tool call response, where the text response just indicates the tool call will happen.
         if tool_calls:
-            return await self._handle_structured_response(tool_calls, run_context)
+            return await self._handle_structured_response(tool_calls, run_context, result_schema)
         elif texts:
             text = '\n\n'.join(texts)
-            return await self._handle_text_response(text, run_context)
+            return await self._handle_text_response(text, run_context, result_schema)
         else:
             raise exceptions.UnexpectedModelBehavior('Received empty model response')
 
     async def _handle_text_response(
-        self, text: str, run_context: RunContext[AgentDeps]
-    ) -> tuple[_MarkFinalResult[ResultData] | None, list[_messages.ModelRequestPart]]:
+        self, text: str, run_context: RunContext[AgentDeps], result_schema: _result.ResultSchema[RunResultData] | None
+    ) -> tuple[_MarkFinalResult[RunResultData] | None, list[_messages.ModelRequestPart]]:
         """Handle a plain text response from the model for non-streaming responses."""
-        if self._allow_text_result:
-            result_data_input = cast(ResultData, text)
+        if self._allow_text_result(result_schema):
+            result_data_input = cast(RunResultData, text)
             try:
                 result_data = await self._validate_result(result_data_input, run_context, None)
             except _result.ToolRetryError as e:
@@ -911,16 +1088,19 @@ class Agent(Generic[AgentDeps, ResultData]):
             return None, [response]
 
     async def _handle_structured_response(
-        self, tool_calls: list[_messages.ToolCallPart], run_context: RunContext[AgentDeps]
-    ) -> tuple[_MarkFinalResult[ResultData] | None, list[_messages.ModelRequestPart]]:
+        self,
+        tool_calls: list[_messages.ToolCallPart],
+        run_context: RunContext[AgentDeps],
+        result_schema: _result.ResultSchema[RunResultData] | None,
+    ) -> tuple[_MarkFinalResult[RunResultData] | None, list[_messages.ModelRequestPart]]:
         """Handle a structured response containing tool calls from the model for non-streaming responses."""
         assert tool_calls, 'Expected at least one tool call'
 
         # first look for the result tool call
-        final_result: _MarkFinalResult[ResultData] | None = None
+        final_result: _MarkFinalResult[RunResultData] | None = None
 
         parts: list[_messages.ModelRequestPart] = []
-        if result_schema := self._result_schema:
+        if result_schema := result_schema:
             if match := result_schema.find_tool(tool_calls):
                 call, result_tool = match
                 try:
@@ -933,7 +1113,9 @@ class Agent(Generic[AgentDeps, ResultData]):
                     final_result = _MarkFinalResult(result_data, call.tool_name)
 
         # Then build the other request parts based on end strategy
-        parts += await self._process_function_tools(tool_calls, final_result and final_result.tool_name, run_context)
+        parts += await self._process_function_tools(
+            tool_calls, final_result and final_result.tool_name, run_context, result_schema
+        )
 
         return final_result, parts
 
@@ -942,6 +1124,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         tool_calls: list[_messages.ToolCallPart],
         result_tool_name: str | None,
         run_context: RunContext[AgentDeps],
+        result_schema: _result.ResultSchema[RunResultData] | None,
     ) -> list[_messages.ModelRequestPart]:
         """Process function (non-result) tool calls in parallel.
 
@@ -975,7 +1158,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                     )
                 else:
                     tasks.append(asyncio.create_task(tool.run(call, run_context), name=call.tool_name))
-            elif self._result_schema is not None and call.tool_name in self._result_schema.tools:
+            elif result_schema is not None and call.tool_name in result_schema.tools:
                 # if tool_name is in _result_schema, it means we found a result tool but an error occurred in
                 # validation, we don't add another part here
                 if result_tool_name is not None:
@@ -987,7 +1170,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                         )
                     )
             else:
-                parts.append(self._unknown_tool(call.tool_name, run_context))
+                parts.append(self._unknown_tool(call.tool_name, run_context, result_schema))
 
         # Run all tool tasks in parallel
         if tasks:
@@ -1000,6 +1183,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         self,
         model_response: models.EitherStreamedResponse,
         run_context: RunContext[AgentDeps],
+        result_schema: _result.ResultSchema[RunResultData] | None,
     ) -> (
         _MarkFinalResult[models.EitherStreamedResponse]
         | tuple[_messages.ModelResponse, list[_messages.ModelRequestPart]]
@@ -1012,7 +1196,7 @@ class Agent(Generic[AgentDeps, ResultData]):
         """
         if isinstance(model_response, models.StreamTextResponse):
             # plain string response
-            if self._allow_text_result:
+            if self._allow_text_result(result_schema):
                 return _MarkFinalResult(model_response, None)
             else:
                 self._incr_result_retry(run_context)
@@ -1026,7 +1210,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                 text = ''.join(model_response.get(final=True))
                 return _messages.ModelResponse([_messages.TextPart(text)]), [response]
         elif isinstance(model_response, models.StreamStructuredResponse):
-            if self._result_schema is not None:
+            if result_schema is not None:
                 # if there's a result schema, iterate over the stream until we find at least one tool
                 # NOTE: this means we ignore any other tools called here
                 structured_msg = model_response.get()
@@ -1037,7 +1221,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                         break
                     structured_msg = model_response.get()
 
-                if match := self._result_schema.find_tool(structured_msg.parts):
+                if match := result_schema.find_tool(structured_msg.parts):
                     call, _ = match
                     return _MarkFinalResult(model_response, call.tool_name)
 
@@ -1057,7 +1241,7 @@ class Agent(Generic[AgentDeps, ResultData]):
                     if tool := self._function_tools.get(call.tool_name):
                         tasks.append(asyncio.create_task(tool.run(call, run_context), name=call.tool_name))
                     else:
-                        parts.append(self._unknown_tool(call.tool_name, run_context))
+                        parts.append(self._unknown_tool(call.tool_name, run_context, result_schema))
 
             with _logfire.span('running {tools=}', tools=[t.get_name() for t in tasks]):
                 task_results: Sequence[_messages.ModelRequestPart] = await asyncio.gather(*tasks)
@@ -1068,13 +1252,17 @@ class Agent(Generic[AgentDeps, ResultData]):
 
     async def _validate_result(
         self,
-        result_data: ResultData,
+        result_data: RunResultData,
         run_context: RunContext[AgentDeps],
         tool_call: _messages.ToolCallPart | None,
-    ) -> ResultData:
-        for validator in self._result_validators:
-            result_data = await validator.validate(result_data, tool_call, run_context)
-        return result_data
+    ) -> RunResultData:
+        if self._result_validators:
+            agent_result_data = cast(ResultData, result_data)
+            for validator in self._result_validators:
+                agent_result_data = await validator.validate(agent_result_data, tool_call, run_context)
+            return cast(RunResultData, agent_result_data)
+        else:
+            return result_data
 
     def _incr_result_retry(self, run_context: RunContext[AgentDeps]) -> None:
         run_context.retry += 1
@@ -1088,14 +1276,22 @@ class Agent(Generic[AgentDeps, ResultData]):
         messages: list[_messages.ModelRequestPart] = [_messages.SystemPromptPart(p) for p in self._system_prompts]
         for sys_prompt_runner in self._system_prompt_functions:
             prompt = await sys_prompt_runner.run(run_context)
-            messages.append(_messages.SystemPromptPart(prompt))
+            if sys_prompt_runner.dynamic:
+                messages.append(_messages.SystemPromptPart(prompt, dynamic_ref=sys_prompt_runner.function.__qualname__))
+            else:
+                messages.append(_messages.SystemPromptPart(prompt))
         return messages
 
-    def _unknown_tool(self, tool_name: str, run_context: RunContext[AgentDeps]) -> _messages.RetryPromptPart:
+    def _unknown_tool(
+        self,
+        tool_name: str,
+        run_context: RunContext[AgentDeps],
+        result_schema: _result.ResultSchema[RunResultData] | None,
+    ) -> _messages.RetryPromptPart:
         self._incr_result_retry(run_context)
         names = list(self._function_tools.keys())
-        if self._result_schema:
-            names.extend(self._result_schema.tool_names())
+        if result_schema:
+            names.extend(result_schema.tool_names())
         if names:
             msg = f'Available tools: {", ".join(names)}'
         else:
@@ -1133,6 +1329,10 @@ class Agent(Generic[AgentDeps, ResultData]):
                             self.name = name
                             return
 
+    @staticmethod
+    def _allow_text_result(result_schema: _result.ResultSchema[RunResultData] | None) -> bool:
+        return result_schema is None or result_schema.allow_text_result
+
     @property
     @deprecated(
         'The `last_run_messages` attribute has been removed, use `capture_run_messages` instead.', category=None

{pydantic_ai_slim-0.0.17 → pydantic_ai_slim-0.0.18}/pydantic_ai/messages.py

@@ -21,6 +21,12 @@ class SystemPromptPart:
     content: str
     """The content of the prompt."""
 
+    dynamic_ref: str | None = None
+    """The ref of the dynamic system prompt function that generated this part.
+
+    Only set if system prompt is dynamic, see [`system_prompt`][pydantic_ai.Agent.system_prompt] for more information.
+    """
+
     part_kind: Literal['system-prompt'] = 'system-prompt'
     """Part type identifier, this is available on all parts as a discriminator."""
 

{pydantic_ai_slim-0.0.17 → pydantic_ai_slim-0.0.18}/pydantic_ai/models/__init__.py

@@ -48,13 +48,12 @@ KnownModelName = Literal[
     'groq:mixtral-8x7b-32768',
     'groq:gemma2-9b-it',
     'groq:gemma-7b-it',
-    'gemini-1.5-flash',
-    'gemini-1.5-pro',
-    'gemini-2.0-flash-exp',
-    'vertexai:gemini-1.5-flash',
-    'vertexai:gemini-1.5-pro',
-    # since mistral models are supported by other providers (e.g. ollama), and some of their models (e.g. "codestral")
-    # don't start with "mistral", we add the "mistral:" prefix to all to be explicit
+    'google-gla:gemini-1.5-flash',
+    'google-gla:gemini-1.5-pro',
+    'google-gla:gemini-2.0-flash-exp',
+    'google-vertex:gemini-1.5-flash',
+    'google-vertex:gemini-1.5-pro',
+    'google-vertex:gemini-2.0-flash-exp',
     'mistral:mistral-small-latest',
     'mistral:mistral-large-latest',
     'mistral:codestral-latest',
@@ -76,9 +75,9 @@ KnownModelName = Literal[
     'ollama:qwen2',
     'ollama:qwen2.5',
     'ollama:starcoder2',
-    'claude-3-5-haiku-latest',
-    'claude-3-5-sonnet-latest',
-    'claude-3-opus-latest',
+    'anthropic:claude-3-5-haiku-latest',
+    'anthropic:claude-3-5-sonnet-latest',
+    'anthropic:claude-3-opus-latest',
     'test',
 ]
 """Known model names that can be used with the `model` parameter of [`Agent`][pydantic_ai.Agent].
@@ -274,6 +273,15 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .openai import OpenAIModel
 
         return OpenAIModel(model[7:])
+    elif model.startswith(('gpt', 'o1')):
+        from .openai import OpenAIModel
+
+        return OpenAIModel(model)
+    elif model.startswith('google-gla'):
+        from .gemini import GeminiModel
+
+        return GeminiModel(model[11:])  # pyright: ignore[reportArgumentType]
+    # backwards compatibility with old model names (ex, gemini-1.5-flash -> google-gla:gemini-1.5-flash)
     elif model.startswith('gemini'):
         from .gemini import GeminiModel
 
@@ -283,6 +291,11 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .groq import GroqModel
 
         return GroqModel(model[5:])  # pyright: ignore[reportArgumentType]
+    elif model.startswith('google-vertex'):
+        from .vertexai import VertexAIModel
+
+        return VertexAIModel(model[14:])  # pyright: ignore[reportArgumentType]
+    # backwards compatibility with old model names (ex, vertexai:gemini-1.5-flash -> google-vertex:gemini-1.5-flash)
     elif model.startswith('vertexai:'):
         from .vertexai import VertexAIModel
 
@@ -295,6 +308,11 @@ def infer_model(model: Model | KnownModelName) -> Model:
         from .ollama import OllamaModel
 
         return OllamaModel(model[7:])
+    elif model.startswith('anthropic'):
+        from .anthropic import AnthropicModel
+
+        return AnthropicModel(model[10:])
+    # backwards compatibility with old model names (ex, claude-3-5-sonnet-latest -> anthropic:claude-3-5-sonnet-latest)
     elif model.startswith('claude'):
         from .anthropic import AnthropicModel
 

{pydantic_ai_slim-0.0.17 → pydantic_ai_slim-0.0.18}/pydantic_ai/models/anthropic.py

@@ -136,7 +136,7 @@ class AnthropicModel(Model):
         )
 
     def name(self) -> str:
-        return self.model_name
+        return f'anthropic:{self.model_name}'
 
     @staticmethod
     def _map_tool_definition(f: ToolDefinition) -> ToolParam:

{pydantic_ai_slim-0.0.17 → pydantic_ai_slim-0.0.18}/pydantic_ai/models/gemini.py

@@ -111,7 +111,7 @@ class GeminiModel(Model):
         )
 
     def name(self) -> str:
-        return self.model_name
+        return f'google-gla:{self.model_name}'
 
 
 class AuthProtocol(Protocol):
@@ -693,7 +693,7 @@ class _GeminiJsonSchema:
 
     def _simplify(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None:
         schema.pop('title', None)
-        default = schema.pop('default', _utils.UNSET)
+        schema.pop('default', None)
         if ref := schema.pop('$ref', None):
             # noinspection PyTypeChecker
             key = re.sub(r'^#/\$defs/', '', ref)
@@ -708,11 +708,12 @@ class _GeminiJsonSchema:
         if any_of := schema.get('anyOf'):
             for item_schema in any_of:
                 self._simplify(item_schema, refs_stack)
-            if len(any_of) == 2 and {'type': 'null'} in any_of and default is None:
+            if len(any_of) == 2 and {'type': 'null'} in any_of:
                 for item_schema in any_of:
                     if item_schema != {'type': 'null'}:
                         schema.clear()
                         schema.update(item_schema)
+                        schema['nullable'] = True
                         return
 
         type_ = schema.get('type')
@@ -721,6 +722,12 @@ class _GeminiJsonSchema:
             self._object(schema, refs_stack)
         elif type_ == 'array':
             return self._array(schema, refs_stack)
+        elif type_ == 'string' and (fmt := schema.pop('format', None)):
+            description = schema.get('description')
+            if description:
+                schema['description'] = f'{description} (format: {fmt})'
+            else:
+                schema['description'] = f'Format: {fmt}'
 
     def _object(self, schema: dict[str, Any], refs_stack: tuple[str, ...]) -> None:
         ad_props = schema.pop('additionalProperties', None)

{pydantic_ai_slim-0.0.17 → pydantic_ai_slim-0.0.18}/pydantic_ai/models/vertexai.py

@@ -164,7 +164,7 @@ class VertexAIModel(Model):
         return url, auth
 
     def name(self) -> str:
-        return f'vertexai:{self.model_name}'
+        return f'google-vertex:{self.model_name}'
 
 
 # pyright: reportUnknownMemberType=false

{pydantic_ai_slim-0.0.17 → pydantic_ai_slim-0.0.18}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pydantic-ai-slim"
-version = "0.0.17"
+version = "0.0.18"
 description = "Agent Framework / shim to use Pydantic with LLMs, slim package"
 authors = [
     { name = "Samuel Colvin", email = "samuel@pydantic.dev" },
@@ -60,6 +60,7 @@ dev = [
     "pytest-examples>=0.0.14",
     "pytest-mock>=3.14.0",
     "pytest-pretty>=1.2.0",
+    "diff-cover>=9.2.0",
 ]
 
 [tool.hatch.build.targets.wheel]