pydantic-ai-slim 0.2.4__tar.gz → 0.2.6__tar.gz

{pydantic_ai_slim-0.2.4 → pydantic_ai_slim-0.2.6}/.gitignore

@@ -17,3 +17,5 @@ examples/pydantic_ai_examples/.chat_app_messages.sqlite
 /docs-site/.wrangler/
 /CLAUDE.md
 node_modules/
+**.idea/
+.coverage*

pydantic_ai_slim-0.2.6/LICENSE

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) Pydantic Services Inc. 2024 to present
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

{pydantic_ai_slim-0.2.4 → pydantic_ai_slim-0.2.6}/PKG-INFO

@@ -1,9 +1,10 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-slim
-Version: 0.2.4
+Version: 0.2.6
 Summary: Agent Framework / shim to use Pydantic with LLMs, slim package
 Author-email: Samuel Colvin <samuel@pydantic.dev>, Marcelo Trylesinski <marcelotryle@gmail.com>, David Montague <david@pydantic.dev>, Alex Hall <alex@pydantic.dev>
 License-Expression: MIT
+License-File: LICENSE
 Classifier: Development Status :: 4 - Beta
 Classifier: Environment :: Console
 Classifier: Environment :: MacOS X
@@ -29,11 +30,11 @@ Requires-Dist: exceptiongroup; python_version < '3.11'
 Requires-Dist: griffe>=1.3.2
 Requires-Dist: httpx>=0.27
 Requires-Dist: opentelemetry-api>=1.28.0
-Requires-Dist: pydantic-graph==0.2.4
+Requires-Dist: pydantic-graph==0.2.6
 Requires-Dist: pydantic>=2.10
 Requires-Dist: typing-inspection>=0.4.0
 Provides-Extra: a2a
-Requires-Dist: fasta2a==0.2.4; extra == 'a2a'
+Requires-Dist: fasta2a==0.2.6; extra == 'a2a'
 Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.49.0; extra == 'anthropic'
 Provides-Extra: bedrock
@@ -47,7 +48,9 @@ Requires-Dist: cohere>=5.13.11; (platform_system != 'Emscripten') and extra == '
 Provides-Extra: duckduckgo
 Requires-Dist: duckduckgo-search>=7.0.0; extra == 'duckduckgo'
 Provides-Extra: evals
-Requires-Dist: pydantic-evals==0.2.4; extra == 'evals'
+Requires-Dist: pydantic-evals==0.2.6; extra == 'evals'
+Provides-Extra: google
+Requires-Dist: google-genai>=1.15.0; extra == 'google'
 Provides-Extra: groq
 Requires-Dist: groq>=0.15.0; extra == 'groq'
 Provides-Extra: logfire

{pydantic_ai_slim-0.2.4 → pydantic_ai_slim-0.2.6}/pydantic_ai/_agent_graph.py

@@ -26,7 +26,7 @@ from . import (
 )
 from .result import OutputDataT, ToolOutput
 from .settings import ModelSettings, merge_model_settings
-from .tools import RunContext, Tool, ToolDefinition
+from .tools import RunContext, Tool, ToolDefinition, ToolsPrepareFunc
 
 if TYPE_CHECKING:
     from .mcp import MCPServer
@@ -97,6 +97,8 @@ class GraphAgentDeps(Generic[DepsT, OutputDataT]):
 
     tracer: Tracer
 
+    prepare_tools: ToolsPrepareFunc[DepsT] | None = None
+
 
 class AgentNode(BaseNode[GraphAgentState, GraphAgentDeps[DepsT, Any], result.FinalResult[NodeRunEndT]]):
     """The base class for all agent nodes.
@@ -196,7 +198,9 @@ class UserPromptNode(AgentNode[DepsT, NodeRunEndT]):
                     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):
+                            if runner := self.system_prompt_dynamic_functions.get(  # pragma: lax no cover
+                                part.dynamic_ref
+                            ):
                                 updated_part_content = await runner.run(run_context)
                                 msg.parts[i] = _messages.SystemPromptPart(
                                     updated_part_content, dynamic_ref=part.dynamic_ref
@@ -239,6 +243,11 @@ async def _prepare_request_parameters(
         *map(add_mcp_server_tools, ctx.deps.mcp_servers),
     )
 
+    if ctx.deps.prepare_tools:
+        # Prepare the tools using the provided function
+        # This also acts over tool definitions pulled from MCP servers
+        function_tool_defs = await ctx.deps.prepare_tools(run_context, function_tool_defs) or []
+
     output_schema = ctx.deps.output_schema
     return models.ModelRequestParameters(
         function_tools=function_tool_defs,
@@ -265,7 +274,7 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
         if self._did_stream:
             # `self._result` gets set when exiting the `stream` contextmanager, so hitting this
             # means that the stream was started but not finished before `run()` was called
-            raise exceptions.AgentRunError('You must finish streaming before calling run()')
+            raise exceptions.AgentRunError('You must finish streaming before calling run()')  # pragma: no cover
 
         return await self._make_request(ctx)
 
@@ -316,7 +325,7 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
         self, ctx: GraphRunContext[GraphAgentState, GraphAgentDeps[DepsT, NodeRunEndT]]
     ) -> CallToolsNode[DepsT, NodeRunEndT]:
         if self._result is not None:
-            return self._result
+            return self._result  # pragma: no cover
 
         model_settings, model_request_parameters = await self._prepare_request(ctx)
         model_request_parameters = ctx.deps.model.customize_request_parameters(model_request_parameters)
@@ -333,7 +342,7 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
         ctx.state.message_history.append(self.request)
 
         # Check usage
-        if ctx.deps.usage_limits:
+        if ctx.deps.usage_limits:  # pragma: no branch
             ctx.deps.usage_limits.check_before_request(ctx.state.usage)
 
         # Increment run_step
@@ -350,7 +359,7 @@ class ModelRequestNode(AgentNode[DepsT, NodeRunEndT]):
     ) -> CallToolsNode[DepsT, NodeRunEndT]:
         # Update usage
         ctx.state.usage.incr(response.usage)
-        if ctx.deps.usage_limits:
+        if ctx.deps.usage_limits:  # pragma: no branch
             ctx.deps.usage_limits.check_tokens(ctx.state.usage)
 
         # Append the model response to state.message_history
@@ -735,7 +744,7 @@ async def _tool_from_mcp_server(
 
     for server in ctx.deps.mcp_servers:
         tools = await server.list_tools()
-        if tool_name in {tool.name for tool in tools}:
+        if tool_name in {tool.name for tool in tools}:  # pragma: no branch
             return Tool(name=tool_name, function=run_tool, takes_ctx=True, max_retries=ctx.deps.default_retries)
     return None
 

{pydantic_ai_slim-0.2.4 → pydantic_ai_slim-0.2.6}/pydantic_ai/_cli.py

@@ -53,11 +53,11 @@ PYDANTIC_AI_HOME = Path.home() / '.pydantic-ai'
 This folder is used to store the prompt history and configuration.
 """
 
-PROMPT_HISTORY_PATH = PYDANTIC_AI_HOME / 'prompt-history.txt'
+PROMPT_HISTORY_FILENAME = 'prompt-history.txt'
 
 
 class SimpleCodeBlock(CodeBlock):
-    """Customised code blocks in markdown.
+    """Customized code blocks in markdown.
 
     This avoids a background color which messes up copy-pasting and sets the language name as dim prefix and suffix.
     """
@@ -70,7 +70,7 @@ class SimpleCodeBlock(CodeBlock):
 
 
 class LeftHeading(Heading):
-    """Customised headings in markdown to stop centering and prepend markdown style hashes."""
+    """Customized headings in markdown to stop centering and prepend markdown style hashes."""
 
     def __rich_console__(self, console: Console, options: ConsoleOptions) -> RenderResult:
         # note we use `Style(bold=True)` not `self.style_name` here to disable underlining which is ugly IMHO
@@ -202,7 +202,7 @@ Special prompts:
     elif args.code_theme == 'dark':
         code_theme = 'monokai'
     else:
-        code_theme = args.code_theme
+        code_theme = args.code_theme  # pragma: no cover
 
     if prompt := cast(str, args.prompt):
         try:
@@ -211,27 +211,26 @@ Special prompts:
             pass
         return 0
 
-    # Ensure the history directory and file exist
-    PROMPT_HISTORY_PATH.parent.mkdir(parents=True, exist_ok=True)
-    PROMPT_HISTORY_PATH.touch(exist_ok=True)
-
-    # doing this instead of `PromptSession[Any](history=` allows mocking of PromptSession in tests
-    session: PromptSession[Any] = PromptSession(history=FileHistory(str(PROMPT_HISTORY_PATH)))
     try:
-        return asyncio.run(run_chat(session, stream, agent, console, code_theme, prog_name))
+        return asyncio.run(run_chat(stream, agent, console, code_theme, prog_name))
     except KeyboardInterrupt:  # pragma: no cover
         return 0
 
 
 async def run_chat(
-    session: PromptSession[Any],
     stream: bool,
     agent: Agent[AgentDepsT, OutputDataT],
     console: Console,
     code_theme: str,
     prog_name: str,
+    config_dir: Path | None = None,
     deps: AgentDepsT = None,
 ) -> int:
+    prompt_history_path = (config_dir or PYDANTIC_AI_HOME) / PROMPT_HISTORY_FILENAME
+    prompt_history_path.parent.mkdir(parents=True, exist_ok=True)
+    prompt_history_path.touch(exist_ok=True)
+    session: PromptSession[Any] = PromptSession(history=FileHistory(str(prompt_history_path)))
+
     multiline = False
     messages: list[ModelMessage] = []
 

{pydantic_ai_slim-0.2.4 → pydantic_ai_slim-0.2.6}/pydantic_ai/_output.py

@@ -140,8 +140,8 @@ class OutputSchema(Generic[OutputDataT]):
         self, parts: Iterable[_messages.ModelResponsePart], tool_name: str
     ) -> tuple[_messages.ToolCallPart, OutputSchemaTool[OutputDataT]] | None:
         """Find a tool that matches one of the calls, with a specific name."""
-        for part in parts:
-            if isinstance(part, _messages.ToolCallPart):
+        for part in parts:  # pragma: no branch
+            if isinstance(part, _messages.ToolCallPart):  # pragma: no branch
                 if part.tool_name == tool_name:
                     return part, self.tools[tool_name]
 
@@ -151,7 +151,7 @@ class OutputSchema(Generic[OutputDataT]):
     ) -> Iterator[tuple[_messages.ToolCallPart, OutputSchemaTool[OutputDataT]]]:
         """Find a tool that matches one of the calls."""
         for part in parts:
-            if isinstance(part, _messages.ToolCallPart):
+            if isinstance(part, _messages.ToolCallPart):  # pragma: no branch
                 if result := self.tools.get(part.tool_name):
                     yield part, result
 
@@ -201,7 +201,7 @@ class OutputSchemaTool(Generic[OutputDataT]):
             if description is None:
                 tool_description = json_schema_description
             else:
-                tool_description = f'{description}. {json_schema_description}'
+                tool_description = f'{description}. {json_schema_description}'  # pragma: no cover
         else:
             tool_description = description or DEFAULT_DESCRIPTION
             if multiple:
@@ -243,7 +243,7 @@ class OutputSchemaTool(Generic[OutputDataT]):
                 )
                 raise ToolRetryError(m) from e
             else:
-                raise
+                raise  # pragma: lax no cover
         else:
             if k := self.tool_def.outer_typed_dict_key:
                 output = output[k]
@@ -269,11 +269,11 @@ def extract_str_from_union(output_type: Any) -> _utils.Option[Any]:
                 includes_str = True
             else:
                 remain_args.append(arg)
-        if includes_str:
+        if includes_str:  # pragma: no branch
             if len(remain_args) == 1:
                 return _utils.Some(remain_args[0])
             else:
-                return _utils.Some(Union[tuple(remain_args)])
+                return _utils.Some(Union[tuple(remain_args)])  # pragma: no cover
 
 
 def get_union_args(tp: Any) -> tuple[Any, ...]:

{pydantic_ai_slim-0.2.4 → pydantic_ai_slim-0.2.6}/pydantic_ai/_parts_manager.py

@@ -164,7 +164,7 @@ class ModelResponsePartsManager:
             if tool_name is None and self._parts:
                 part_index = len(self._parts) - 1
                 latest_part = self._parts[part_index]
-                if isinstance(latest_part, (ToolCallPart, ToolCallPartDelta)):
+                if isinstance(latest_part, (ToolCallPart, ToolCallPartDelta)):  # pragma: no branch
                     existing_matching_part_and_index = latest_part, part_index
         else:
             # vendor_part_id is provided, so look up the corresponding part or delta

{pydantic_ai_slim-0.2.4 → pydantic_ai_slim-0.2.6}/pydantic_ai/agent.py

@@ -42,6 +42,7 @@ from .tools import (
     ToolFuncPlain,
     ToolParams,
     ToolPrepareFunc,
+    ToolsPrepareFunc,
 )
 
 # Re-exporting like this improves auto-import behavior in PyCharm
@@ -148,6 +149,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
     _system_prompt_dynamic_functions: dict[str, _system_prompt.SystemPromptRunner[AgentDepsT]] = dataclasses.field(
         repr=False
     )
+    _prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = dataclasses.field(repr=False)
     _function_tools: dict[str, Tool[AgentDepsT]] = dataclasses.field(repr=False)
     _mcp_servers: Sequence[MCPServer] = dataclasses.field(repr=False)
     _default_retries: int = dataclasses.field(repr=False)
@@ -172,6 +174,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         retries: int = 1,
         output_retries: int | None = None,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
+        prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = None,
         mcp_servers: Sequence[MCPServer] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
@@ -200,6 +203,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         result_tool_description: str | None = None,
         result_retries: int | None = None,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
+        prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = None,
         mcp_servers: Sequence[MCPServer] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
@@ -223,6 +227,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         retries: int = 1,
         output_retries: int | None = None,
         tools: Sequence[Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...]] = (),
+        prepare_tools: ToolsPrepareFunc[AgentDepsT] | None = None,
         mcp_servers: Sequence[MCPServer] = (),
         defer_model_check: bool = False,
         end_strategy: EndStrategy = 'early',
@@ -251,6 +256,9 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
             output_retries: The maximum number of retries to allow for result validation, defaults to `retries`.
             tools: Tools to register with the agent, you can also register tools via the decorators
                 [`@agent.tool`][pydantic_ai.Agent.tool] and [`@agent.tool_plain`][pydantic_ai.Agent.tool_plain].
+            prepare_tools: custom method to prepare the tool definition of all tools for each step.
+                This is useful if you want to customize the definition of multiple tools or you want to register
+                a subset of tools for a given step. See [`ToolsPrepareFunc`][pydantic_ai.tools.ToolsPrepareFunc]
             mcp_servers: MCP servers to register with the agent. You should register a [`MCPServer`][pydantic_ai.mcp.MCPServer]
                 for each server you want the agent to connect to.
             defer_model_check: by default, if you provide a [named][pydantic_ai.models.KnownModelName] model,
@@ -334,6 +342,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         self._default_retries = retries
         self._max_result_retries = output_retries if output_retries is not None else retries
         self._mcp_servers = mcp_servers
+        self._prepare_tools = prepare_tools
         for tool in tools:
             if isinstance(tool, Tool):
                 self._register_tool(tool)
@@ -585,6 +594,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
                         model_name='gpt-4o',
                         timestamp=datetime.datetime(...),
                         kind='response',
+                        vendor_id=None,
                     )
                 ),
                 End(data=FinalResult(output='Paris', tool_name=None, tool_call_id=None)),
@@ -654,8 +664,10 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
         usage_limits = usage_limits or _usage.UsageLimits()
 
         if isinstance(model_used, InstrumentedModel):
+            instrumentation_settings = model_used.settings
             tracer = model_used.settings.tracer
         else:
+            instrumentation_settings = None
             tracer = NoOpTracer()
         agent_name = self.name or 'agent'
         run_span = tracer.start_span(
@@ -691,6 +703,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
             mcp_servers=self._mcp_servers,
             default_retries=self._default_retries,
             tracer=tracer,
+            prepare_tools=self._prepare_tools,
             get_instructions=get_instructions,
         )
         start_node = _agent_graph.UserPromptNode[AgentDepsT](
@@ -723,19 +736,18 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
                     )
         finally:
             try:
-                if run_span.is_recording():
-                    run_span.set_attributes(self._run_span_end_attributes(state, usage))
+                if instrumentation_settings and run_span.is_recording():
+                    run_span.set_attributes(self._run_span_end_attributes(state, usage, instrumentation_settings))
             finally:
                 run_span.end()
 
-    def _run_span_end_attributes(self, state: _agent_graph.GraphAgentState, usage: _usage.Usage):
+    def _run_span_end_attributes(
+        self, state: _agent_graph.GraphAgentState, usage: _usage.Usage, settings: InstrumentationSettings
+    ):
         return {
             **usage.opentelemetry_attributes(),
             'all_messages_events': json.dumps(
-                [
-                    InstrumentedModel.event_to_dict(e)
-                    for e in InstrumentedModel.messages_to_otel_events(state.message_history)
-                ]
+                [InstrumentedModel.event_to_dict(e) for e in settings.messages_to_otel_events(state.message_history)]
             ),
             'logfire.json_schema': json.dumps(
                 {
@@ -1001,7 +1013,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
                         final_result_details = await stream_to_final(streamed_response)
                         if final_result_details is not None:
                             if yielded:
-                                raise exceptions.AgentRunError('Agent run produced final results')
+                                raise exceptions.AgentRunError('Agent run produced final results')  # pragma: no cover
                             yielded = True
 
                             messages = graph_ctx.state.message_history.copy()
@@ -1048,11 +1060,13 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
                             break
                 next_node = await agent_run.next(node)
                 if not isinstance(next_node, _agent_graph.AgentNode):
-                    raise exceptions.AgentRunError('Should have produced a StreamedRunResult before getting here')
+                    raise exceptions.AgentRunError(  # pragma: no cover
+                        'Should have produced a StreamedRunResult before getting here'
+                    )
                 node = cast(_agent_graph.AgentNode[Any, Any], next_node)
 
         if not yielded:
-            raise exceptions.AgentRunError('Agent run finished without producing a final result')
+            raise exceptions.AgentRunError('Agent run finished without producing a final result')  # pragma: no cover
 
     @contextmanager
     def override(
@@ -1226,7 +1240,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
             ) -> _system_prompt.SystemPromptFunc[AgentDepsT]:
                 runner = _system_prompt.SystemPromptRunner[AgentDepsT](func_, dynamic=dynamic)
                 self._system_prompt_functions.append(runner)
-                if dynamic:
+                if dynamic:  # pragma: lax no cover
                     self._system_prompt_dynamic_functions[func_.__qualname__] = runner
                 return func_
 
@@ -1608,7 +1622,7 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
                     if item is self:
                         self.name = name
                         return
-                if parent_frame.f_locals != parent_frame.f_globals:
+                if parent_frame.f_locals != parent_frame.f_globals:  # pragma: no branch
                     # if we couldn't find the agent in locals and globals are a different dict, try globals
                     for name, item in parent_frame.f_globals.items():
                         if item is self:
@@ -1759,18 +1773,14 @@ class Agent(Generic[AgentDepsT, OutputDataT]):
             await agent.to_cli()
         ```
         """
-        from prompt_toolkit import PromptSession
-        from prompt_toolkit.history import FileHistory
         from rich.console import Console
 
-        from pydantic_ai._cli import PROMPT_HISTORY_PATH, run_chat
+        from pydantic_ai._cli import run_chat
 
         # TODO(Marcelo): We need to refactor the CLI code to be able to be able to just pass `agent`, `deps` and
         # `prog_name` from here.
 
-        session: PromptSession[Any] = PromptSession(history=FileHistory(str(PROMPT_HISTORY_PATH)))
         await run_chat(
-            session=session,
             stream=True,
             agent=self,
             deps=deps,
@@ -1851,6 +1861,7 @@ class AgentRun(Generic[AgentDepsT, OutputDataT]):
                     model_name='gpt-4o',
                     timestamp=datetime.datetime(...),
                     kind='response',
+                    vendor_id=None,
                 )
             ),
             End(data=FinalResult(output='Paris', tool_name=None, tool_call_id=None)),
@@ -1996,6 +2007,7 @@ class AgentRun(Generic[AgentDepsT, OutputDataT]):
                             model_name='gpt-4o',
                             timestamp=datetime.datetime(...),
                             kind='response',
+                            vendor_id=None,
                         )
                     ),
                     End(data=FinalResult(output='Paris', tool_name=None, tool_call_id=None)),
@@ -2024,7 +2036,7 @@ class AgentRun(Generic[AgentDepsT, OutputDataT]):
         """Get usage statistics for the run so far, including token usage, model requests, and so on."""
         return self._graph_run.state.usage
 
-    def __repr__(self) -> str:
+    def __repr__(self) -> str:  # pragma: no cover
         result = self._graph_run.result
         result_repr = '<run not finished>' if result is None else repr(result.output)
         return f'<{type(self).__name__} result={result_repr} usage={self.usage()}>'

{pydantic_ai_slim-0.2.4 → pydantic_ai_slim-0.2.6}/pydantic_ai/direct.py

@@ -52,6 +52,7 @@ async def model_request(
             model_name='claude-3-5-haiku-latest',
             timestamp=datetime.datetime(...),
             kind='response',
+            vendor_id=None,
         )
         '''
     ```
@@ -108,6 +109,7 @@ def model_request_sync(
         model_name='claude-3-5-haiku-latest',
         timestamp=datetime.datetime(...),
         kind='response',
+        vendor_id=None,
     )
     '''
     ```

{pydantic_ai_slim-0.2.4 → pydantic_ai_slim-0.2.6}/pydantic_ai/exceptions.py

@@ -4,9 +4,9 @@ import json
 import sys
 
 if sys.version_info < (3, 11):
-    from exceptiongroup import ExceptionGroup
+    from exceptiongroup import ExceptionGroup  # pragma: lax no cover
 else:
-    ExceptionGroup = ExceptionGroup
+    ExceptionGroup = ExceptionGroup  # pragma: lax no cover
 
 __all__ = (
     'ModelRetry',

{pydantic_ai_slim-0.2.4 → pydantic_ai_slim-0.2.6}/pydantic_ai/messages.py

@@ -6,7 +6,7 @@ from collections.abc import Sequence
 from dataclasses import dataclass, field, replace
 from datetime import datetime
 from mimetypes import guess_type
-from typing import Annotated, Any, Literal, Union, cast, overload
+from typing import TYPE_CHECKING, Annotated, Any, Literal, Union, cast, overload
 
 import pydantic
 import pydantic_core
@@ -17,6 +17,10 @@ from ._utils import generate_tool_call_id as _generate_tool_call_id, now_utc as
 from .exceptions import UnexpectedModelBehavior
 from .usage import Usage
 
+if TYPE_CHECKING:
+    from .models.instrumented import InstrumentationSettings
+
+
 AudioMediaType: TypeAlias = Literal['audio/wav', 'audio/mpeg']
 ImageMediaType: TypeAlias = Literal['image/jpeg', 'image/png', 'image/gif', 'image/webp']
 DocumentMediaType: TypeAlias = Literal[
@@ -68,7 +72,7 @@ class SystemPromptPart:
     part_kind: Literal['system-prompt'] = 'system-prompt'
     """Part type identifier, this is available on all parts as a discriminator."""
 
-    def otel_event(self) -> Event:
+    def otel_event(self, _settings: InstrumentationSettings) -> Event:
         return Event('gen_ai.system.message', body={'content': self.content, 'role': 'system'})
 
 
@@ -305,7 +309,7 @@ class UserPromptPart:
     part_kind: Literal['user-prompt'] = 'user-prompt'
     """Part type identifier, this is available on all parts as a discriminator."""
 
-    def otel_event(self) -> Event:
+    def otel_event(self, settings: InstrumentationSettings) -> Event:
         content: str | list[dict[str, Any] | str]
         if isinstance(self.content, str):
             content = self.content
@@ -317,10 +321,12 @@ class UserPromptPart:
                 elif isinstance(part, (ImageUrl, AudioUrl, DocumentUrl, VideoUrl)):
                     content.append({'kind': part.kind, 'url': part.url})
                 elif isinstance(part, BinaryContent):
-                    base64_data = base64.b64encode(part.data).decode()
-                    content.append({'kind': part.kind, 'content': base64_data, 'media_type': part.media_type})
+                    converted_part = {'kind': part.kind, 'media_type': part.media_type}
+                    if settings.include_binary_content:
+                        converted_part['binary_content'] = base64.b64encode(part.data).decode()
+                    content.append(converted_part)
                 else:
-                    content.append({'kind': part.kind})
+                    content.append({'kind': part.kind})  # pragma: no cover
         return Event('gen_ai.user.message', body={'content': content, 'role': 'user'})
 
 
@@ -357,11 +363,11 @@ class ToolReturnPart:
         """Return a dictionary representation of the content, wrapping non-dict types appropriately."""
         # gemini supports JSON dict return values, but no other JSON types, hence we wrap anything else in a dict
         if isinstance(self.content, dict):
-            return tool_return_ta.dump_python(self.content, mode='json')  # pyright: ignore[reportUnknownMemberType]
+            return tool_return_ta.dump_python(self.content, mode='json')  # pyright: ignore[reportUnknownMemberType]  # pragma: no cover
         else:
             return {'return_value': tool_return_ta.dump_python(self.content, mode='json')}
 
-    def otel_event(self) -> Event:
+    def otel_event(self, _settings: InstrumentationSettings) -> Event:
         return Event(
             'gen_ai.tool.message',
             body={'content': self.content, 'role': 'tool', 'id': self.tool_call_id, 'name': self.tool_name},
@@ -418,7 +424,7 @@ class RetryPromptPart:
             description = f'{len(self.content)} validation errors: {json_errors.decode()}'
         return f'{description}\n\nFix the errors and try again.'
 
-    def otel_event(self) -> Event:
+    def otel_event(self, _settings: InstrumentationSettings) -> Event:
         if self.tool_name is None:
             return Event('gen_ai.user.message', body={'content': self.model_response(), 'role': 'user'})
         else:
@@ -556,6 +562,16 @@ class ModelResponse:
     kind: Literal['response'] = 'response'
     """Message type identifier, this is available on all parts as a discriminator."""
 
+    vendor_details: dict[str, Any] | None = field(default=None, repr=False)
+    """Additional vendor-specific details in a serializable format.
+
+    This allows storing selected vendor-specific data that isn't mapped to standard ModelResponse fields.
+    For OpenAI models, this may include 'logprobs', 'finish_reason', etc.
+    """
+
+    vendor_id: str | None = None
+    """Vendor ID as specified by the model provider. This can be used to track the specific request to the model."""
+
     def otel_events(self) -> list[Event]:
         """Return OpenTelemetry events for the response."""
         result: list[Event] = []
@@ -619,7 +635,7 @@ class TextPartDelta:
             ValueError: If `part` is not a `TextPart`.
         """
         if not isinstance(part, TextPart):
-            raise ValueError('Cannot apply TextPartDeltas to non-TextParts')
+            raise ValueError('Cannot apply TextPartDeltas to non-TextParts')  # pragma: no cover
         return replace(part, content=part.content + self.content_delta)
 
 
@@ -682,7 +698,9 @@ class ToolCallPartDelta:
         if isinstance(part, ToolCallPartDelta):
             return self._apply_to_delta(part)
 
-        raise ValueError(f'Can only apply ToolCallPartDeltas to ToolCallParts or ToolCallPartDeltas, not {part}')
+        raise ValueError(  # pragma: no cover
+            f'Can only apply ToolCallPartDeltas to ToolCallParts or ToolCallPartDeltas, not {part}'
+        )
 
     def _apply_to_delta(self, delta: ToolCallPartDelta) -> ToolCallPart | ToolCallPartDelta:
         """Internal helper to apply this delta to another delta."""