pydantic-ai-slim 1.0.1__py3-none-any.whl → 1.0.2__py3-none-any.whl

pydantic_ai/durable_exec/dbos/_mcp_server.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+from abc import ABC
+from collections.abc import Callable
+from typing import Any
+
+from dbos import DBOS
+from typing_extensions import Self
+
+from pydantic_ai.mcp import MCPServer, ToolResult
+from pydantic_ai.tools import AgentDepsT, RunContext
+from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
+from pydantic_ai.toolsets.wrapper import WrapperToolset
+
+from ._utils import StepConfig
+
+
+class DBOSMCPServer(WrapperToolset[AgentDepsT], ABC):
+    """A wrapper for MCPServer that integrates with DBOS, turning call_tool and get_tools to DBOS steps."""
+
+    def __init__(
+        self,
+        wrapped: MCPServer,
+        *,
+        step_name_prefix: str,
+        step_config: StepConfig,
+    ):
+        super().__init__(wrapped)
+        self._step_config = step_config or {}
+        self._step_name_prefix = step_name_prefix
+        id_suffix = f'__{wrapped.id}' if wrapped.id else ''
+        self._name = f'{step_name_prefix}__mcp_server{id_suffix}'
+
+        # Wrap get_tools in a DBOS step.
+        @DBOS.step(
+            name=f'{self._name}.get_tools',
+            **self._step_config,
+        )
+        async def wrapped_get_tools_step(
+            ctx: RunContext[AgentDepsT],
+        ) -> dict[str, ToolsetTool[AgentDepsT]]:
+            return await super(DBOSMCPServer, self).get_tools(ctx)
+
+        self._dbos_wrapped_get_tools_step = wrapped_get_tools_step
+
+        # Wrap call_tool in a DBOS step.
+        @DBOS.step(
+            name=f'{self._name}.call_tool',
+            **self._step_config,
+        )
+        async def wrapped_call_tool_step(
+            name: str,
+            tool_args: dict[str, Any],
+            ctx: RunContext[AgentDepsT],
+            tool: ToolsetTool[AgentDepsT],
+        ) -> ToolResult:
+            return await super(DBOSMCPServer, self).call_tool(name, tool_args, ctx, tool)
+
+        self._dbos_wrapped_call_tool_step = wrapped_call_tool_step
+
+    @property
+    def id(self) -> str | None:
+        return self.wrapped.id
+
+    async def __aenter__(self) -> Self:
+        # The wrapped MCPServer enters itself around listing and calling tools
+        # so we don't need to enter it here (nor could we because we're not inside a DBOS step).
+        return self
+
+    async def __aexit__(self, *args: Any) -> bool | None:
+        return None
+
+    def visit_and_replace(
+        self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]]
+    ) -> AbstractToolset[AgentDepsT]:
+        # DBOS-ified toolsets cannot be swapped out after the fact.
+        return self
+
+    async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
+        return await self._dbos_wrapped_get_tools_step(ctx)
+
+    async def call_tool(
+        self,
+        name: str,
+        tool_args: dict[str, Any],
+        ctx: RunContext[AgentDepsT],
+        tool: ToolsetTool[AgentDepsT],
+    ) -> ToolResult:
+        return await self._dbos_wrapped_call_tool_step(name, tool_args, ctx, tool)

pydantic_ai/durable_exec/dbos/_model.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from datetime import datetime
+from typing import Any
+
+from dbos import DBOS
+
+from pydantic_ai.agent import EventStreamHandler
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelResponse,
+    ModelResponseStreamEvent,
+)
+from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse
+from pydantic_ai.models.wrapper import WrapperModel
+from pydantic_ai.settings import ModelSettings
+from pydantic_ai.tools import RunContext
+from pydantic_ai.usage import RequestUsage
+
+from ._utils import StepConfig
+
+
+class DBOSStreamedResponse(StreamedResponse):
+    def __init__(self, model_request_parameters: ModelRequestParameters, response: ModelResponse):
+        super().__init__(model_request_parameters)
+        self.response = response
+
+    async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
+        return
+        # noinspection PyUnreachableCode
+        yield
+
+    def get(self) -> ModelResponse:
+        return self.response
+
+    def usage(self) -> RequestUsage:
+        return self.response.usage  # pragma: no cover
+
+    @property
+    def model_name(self) -> str:
+        return self.response.model_name or ''  # pragma: no cover
+
+    @property
+    def provider_name(self) -> str:
+        return self.response.provider_name or ''  # pragma: no cover
+
+    @property
+    def timestamp(self) -> datetime:
+        return self.response.timestamp  # pragma: no cover
+
+
+class DBOSModel(WrapperModel):
+    """A wrapper for Model that integrates with DBOS, turning request and request_stream to DBOS steps."""
+
+    def __init__(
+        self,
+        model: Model,
+        *,
+        step_name_prefix: str,
+        step_config: StepConfig,
+        event_stream_handler: EventStreamHandler[Any] | None = None,
+    ):
+        super().__init__(model)
+        self.step_config = step_config
+        self.event_stream_handler = event_stream_handler
+        self._step_name_prefix = step_name_prefix
+
+        # Wrap the request in a DBOS step.
+        @DBOS.step(
+            name=f'{self._step_name_prefix}__model.request',
+            **self.step_config,
+        )
+        async def wrapped_request_step(
+            messages: list[ModelMessage],
+            model_settings: ModelSettings | None,
+            model_request_parameters: ModelRequestParameters,
+        ) -> ModelResponse:
+            return await super(DBOSModel, self).request(messages, model_settings, model_request_parameters)
+
+        self._dbos_wrapped_request_step = wrapped_request_step
+
+        # Wrap the request_stream in a DBOS step.
+        @DBOS.step(
+            name=f'{self._step_name_prefix}__model.request_stream',
+            **self.step_config,
+        )
+        async def wrapped_request_stream_step(
+            messages: list[ModelMessage],
+            model_settings: ModelSettings | None,
+            model_request_parameters: ModelRequestParameters,
+            run_context: RunContext[Any] | None = None,
+        ) -> ModelResponse:
+            async with super(DBOSModel, self).request_stream(
+                messages, model_settings, model_request_parameters, run_context
+            ) as streamed_response:
+                if self.event_stream_handler is not None:
+                    assert run_context is not None, (
+                        'A DBOS model cannot be used with `pydantic_ai.direct.model_request_stream()` as it requires a `run_context`. Set an `event_stream_handler` on the agent and use `agent.run()` instead.'
+                    )
+                    await self.event_stream_handler(run_context, streamed_response)
+
+                async for _ in streamed_response:
+                    pass
+            return streamed_response.get()
+
+        self._dbos_wrapped_request_stream_step = wrapped_request_stream_step
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> ModelResponse:
+        return await self._dbos_wrapped_request_step(messages, model_settings, model_request_parameters)
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+        run_context: RunContext[Any] | None = None,
+    ) -> AsyncIterator[StreamedResponse]:
+        # If not in a workflow (could be in a step), just call the wrapped request_stream method.
+        if DBOS.workflow_id is None or DBOS.step_id is not None:
+            async with super().request_stream(
+                messages, model_settings, model_request_parameters, run_context
+            ) as streamed_response:
+                yield streamed_response
+                return
+
+        response = await self._dbos_wrapped_request_stream_step(
+            messages, model_settings, model_request_parameters, run_context
+        )
+        yield DBOSStreamedResponse(model_request_parameters, response)

pydantic_ai/durable_exec/dbos/_utils.py

@@ -0,0 +1,10 @@
+from typing_extensions import TypedDict
+
+
+class StepConfig(TypedDict, total=False):
+    """Configuration for a step in the DBOS workflow."""
+
+    retries_allowed: bool
+    interval_seconds: float
+    max_attempts: int
+    backoff_rate: float

pydantic_ai/mcp.py

@@ -517,7 +517,7 @@ class MCPServerStdio(MCPServer):
             f'args={self.args!r}',
         ]
         if self.id:
-            repr_args.append(f'id={self.id!r}')  # pragma: no cover
+            repr_args.append(f'id={self.id!r}')
         return f'{self.__class__.__name__}({", ".join(repr_args)})'
 
 

pydantic_ai/messages.py

@@ -52,6 +52,15 @@ ImageFormat: TypeAlias = Literal['jpeg', 'png', 'gif', 'webp']
 DocumentFormat: TypeAlias = Literal['csv', 'doc', 'docx', 'html', 'md', 'pdf', 'txt', 'xls', 'xlsx']
 VideoFormat: TypeAlias = Literal['mkv', 'mov', 'mp4', 'webm', 'flv', 'mpeg', 'mpg', 'wmv', 'three_gp']
 
+FinishReason: TypeAlias = Literal[
+    'stop',
+    'length',
+    'content_filter',
+    'tool_call',
+    'error',
+]
+"""Reason the model finished generating the response, normalized to OpenTelemetry values."""
+
 
 @dataclass(repr=False)
 class SystemPromptPart:
@@ -1032,6 +1041,9 @@ class ModelResponse:
     ] = None
     """request ID as specified by the model provider. This can be used to track the specific request to the model."""
 
+    finish_reason: FinishReason | None = None
+    """Reason the model finished generating the response, normalized to OpenTelemetry values."""
+
     @deprecated('`price` is deprecated, use `cost` instead')
     def price(self) -> genai_types.PriceCalculation:  # pragma: no cover
         return self.cost()

pydantic_ai/models/__init__.py

@@ -28,6 +28,7 @@ from ..exceptions import UserError
 from ..messages import (
     FileUrl,
     FinalResultEvent,
+    FinishReason,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -555,6 +556,10 @@ class StreamedResponse(ABC):
 
     final_result_event: FinalResultEvent | None = field(default=None, init=False)
 
+    provider_response_id: str | None = field(default=None, init=False)
+    provider_details: dict[str, Any] | None = field(default=None, init=False)
+    finish_reason: FinishReason | None = field(default=None, init=False)
+
     _parts_manager: ModelResponsePartsManager = field(default_factory=ModelResponsePartsManager, init=False)
     _event_iterator: AsyncIterator[ModelResponseStreamEvent] | None = field(default=None, init=False)
     _usage: RequestUsage = field(default_factory=RequestUsage, init=False)
@@ -609,6 +614,9 @@ class StreamedResponse(ABC):
             timestamp=self.timestamp,
             usage=self.usage(),
             provider_name=self.provider_name,
+            provider_response_id=self.provider_response_id,
+            provider_details=self.provider_details,
+            finish_reason=self.finish_reason,
         )
 
     def usage(self) -> RequestUsage:

pydantic_ai/models/anthropic.py

@@ -21,6 +21,7 @@ from ..messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     DocumentUrl,
+    FinishReason,
     ImageUrl,
     ModelMessage,
     ModelRequest,
@@ -42,6 +43,16 @@ from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import Model, ModelRequestParameters, StreamedResponse, check_allow_model_requests, download_item, get_user_agent
 
+_FINISH_REASON_MAP: dict[BetaStopReason, FinishReason] = {
+    'end_turn': 'stop',
+    'max_tokens': 'length',
+    'stop_sequence': 'stop',
+    'tool_use': 'tool_call',
+    'pause_turn': 'stop',
+    'refusal': 'content_filter',
+}
+
+
 try:
     from anthropic import NOT_GIVEN, APIStatusError, AsyncStream
     from anthropic.types.beta import (
@@ -70,6 +81,7 @@ try:
         BetaServerToolUseBlock,
         BetaServerToolUseBlockParam,
         BetaSignatureDelta,
+        BetaStopReason,
         BetaTextBlock,
         BetaTextBlockParam,
         BetaTextDelta,
@@ -326,12 +338,20 @@ class AnthropicModel(Model):
                     )
                 )
 
+        finish_reason: FinishReason | None = None
+        provider_details: dict[str, Any] | None = None
+        if raw_finish_reason := response.stop_reason:  # pragma: no branch
+            provider_details = {'finish_reason': raw_finish_reason}
+            finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)
+
         return ModelResponse(
             parts=items,
             usage=_map_usage(response),
             model_name=response.model,
             provider_response_id=response.id,
             provider_name=self._provider.name,
+            finish_reason=finish_reason,
+            provider_details=provider_details,
         )
 
     async def _process_streamed_response(
@@ -583,6 +603,7 @@ class AnthropicStreamedResponse(StreamedResponse):
         async for event in self._response:
             if isinstance(event, BetaRawMessageStartEvent):
                 self._usage = _map_usage(event)
+                self.provider_response_id = event.message.id
 
             elif isinstance(event, BetaRawContentBlockStartEvent):
                 current_block = event.content_block
@@ -646,6 +667,9 @@ class AnthropicStreamedResponse(StreamedResponse):
 
             elif isinstance(event, BetaRawMessageDeltaEvent):
                 self._usage = _map_usage(event)
+                if raw_finish_reason := event.delta.stop_reason:  # pragma: no branch
+                    self.provider_details = {'finish_reason': raw_finish_reason}
+                    self.finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)
 
             elif isinstance(event, BetaRawContentBlockStopEvent | BetaRawMessageStopEvent):  # pragma: no branch
                 current_block = None

pydantic_ai/models/google.py

@@ -20,6 +20,7 @@ from ..messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     FileUrl,
+    FinishReason,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -54,6 +55,7 @@ try:
         ContentUnionDict,
         CountTokensConfigDict,
         ExecutableCodeDict,
+        FinishReason as GoogleFinishReason,
         FunctionCallDict,
         FunctionCallingConfigDict,
         FunctionCallingConfigMode,
@@ -99,6 +101,22 @@ allow any name in the type hints.
 See [the Gemini API docs](https://ai.google.dev/gemini-api/docs/models/gemini#model-variations) for a full list.
 """
 
+_FINISH_REASON_MAP: dict[GoogleFinishReason, FinishReason | None] = {
+    GoogleFinishReason.FINISH_REASON_UNSPECIFIED: None,
+    GoogleFinishReason.STOP: 'stop',
+    GoogleFinishReason.MAX_TOKENS: 'length',
+    GoogleFinishReason.SAFETY: 'content_filter',
+    GoogleFinishReason.RECITATION: 'content_filter',
+    GoogleFinishReason.LANGUAGE: 'error',
+    GoogleFinishReason.OTHER: None,
+    GoogleFinishReason.BLOCKLIST: 'content_filter',
+    GoogleFinishReason.PROHIBITED_CONTENT: 'content_filter',
+    GoogleFinishReason.SPII: 'content_filter',
+    GoogleFinishReason.MALFORMED_FUNCTION_CALL: 'error',
+    GoogleFinishReason.IMAGE_SAFETY: 'content_filter',
+    GoogleFinishReason.UNEXPECTED_TOOL_CALL: 'error',
+}
+
 
 class GoogleModelSettings(ModelSettings, total=False):
     """Settings used for a Gemini model request."""
@@ -129,6 +147,12 @@ class GoogleModelSettings(ModelSettings, total=False):
     See <https://ai.google.dev/api/generate-content#MediaResolution> for more information.
     """
 
+    google_cached_content: str
+    """The name of the cached content to use for the model.
+
+    See <https://ai.google.dev/gemini-api/docs/caching> for more information.
+    """
+
 
 @dataclass(init=False)
 class GoogleModel(Model):
@@ -377,6 +401,7 @@ class GoogleModel(Model):
             thinking_config=model_settings.get('google_thinking_config'),
             labels=model_settings.get('google_labels'),
             media_resolution=model_settings.get('google_video_resolution'),
+            cached_content=model_settings.get('google_cached_content'),
             tools=cast(ToolListUnionDict, tools),
             tool_config=tool_config,
             response_mime_type=response_mime_type,
@@ -396,11 +421,14 @@ class GoogleModel(Model):
                     'Content field missing from Gemini response', str(response)
                 )  # pragma: no cover
         parts = candidate.content.parts or []
-        vendor_id = response.response_id or None
+
+        vendor_id = response.response_id
         vendor_details: dict[str, Any] | None = None
-        finish_reason = candidate.finish_reason
-        if finish_reason:  # pragma: no branch
-            vendor_details = {'finish_reason': finish_reason.value}
+        finish_reason: FinishReason | None = None
+        if raw_finish_reason := candidate.finish_reason:  # pragma: no branch
+            vendor_details = {'finish_reason': raw_finish_reason.value}
+            finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)
+
         usage = _metadata_as_usage(response)
         return _process_response_from_parts(
             parts,
@@ -409,6 +437,7 @@ class GoogleModel(Model):
             usage,
             vendor_id=vendor_id,
             vendor_details=vendor_details,
+            finish_reason=finish_reason,
         )
 
     async def _process_streamed_response(
@@ -543,6 +572,14 @@ class GeminiStreamedResponse(StreamedResponse):
 
             assert chunk.candidates is not None
             candidate = chunk.candidates[0]
+
+            if chunk.response_id:  # pragma: no branch
+                self.provider_response_id = chunk.response_id
+
+            if raw_finish_reason := candidate.finish_reason:
+                self.provider_details = {'finish_reason': raw_finish_reason.value}
+                self.finish_reason = _FINISH_REASON_MAP.get(raw_finish_reason)
+
             if candidate.content is None or candidate.content.parts is None:
                 if candidate.finish_reason == 'STOP':  # pragma: no cover
                     # Normal completion - skip this chunk
@@ -625,6 +662,7 @@ def _process_response_from_parts(
     usage: usage.RequestUsage,
     vendor_id: str | None,
     vendor_details: dict[str, Any] | None = None,
+    finish_reason: FinishReason | None = None,
 ) -> ModelResponse:
     items: list[ModelResponsePart] = []
     for part in parts:
@@ -665,6 +703,7 @@ def _process_response_from_parts(
         provider_response_id=vendor_id,
         provider_details=vendor_details,
         provider_name=provider_name,
+        finish_reason=finish_reason,
     )
 
 

pydantic_ai/models/instrumented.py

@@ -221,7 +221,10 @@ class InstrumentationSettings:
                         _otel_messages.ChatMessage(role='system' if is_system else 'user', parts=message_parts)
                     )
             elif isinstance(message, ModelResponse):  # pragma: no branch
-                result.append(_otel_messages.ChatMessage(role='assistant', parts=message.otel_message_parts(self)))
+                otel_message = _otel_messages.OutputMessage(role='assistant', parts=message.otel_message_parts(self))
+                if message.finish_reason is not None:
+                    otel_message['finish_reason'] = message.finish_reason
+                result.append(otel_message)
         return result
 
     def handle_messages(self, input_messages: list[ModelMessage], response: ModelResponse, system: str, span: Span):
@@ -246,12 +249,10 @@ class InstrumentationSettings:
         else:
             output_messages = self.messages_to_otel_messages([response])
             assert len(output_messages) == 1
-            output_message = cast(_otel_messages.OutputMessage, output_messages[0])
-            if response.provider_details and 'finish_reason' in response.provider_details:
-                output_message['finish_reason'] = response.provider_details['finish_reason']
+            output_message = output_messages[0]
             instructions = InstrumentedModel._get_instructions(input_messages)  # pyright: ignore [reportPrivateUsage]
             system_instructions_attributes = self.system_instructions_attributes(instructions)
-            attributes = {
+            attributes: dict[str, AttributeValue] = {
                 'gen_ai.input.messages': json.dumps(self.messages_to_otel_messages(input_messages)),
                 'gen_ai.output.messages': json.dumps([output_message]),
                 **system_instructions_attributes,
@@ -420,17 +421,25 @@ class InstrumentedModel(WrapperModel):
                         return
 
                     self.instrumentation_settings.handle_messages(messages, response, system, span)
+
+                    attributes_to_set = {
+                        **response.usage.opentelemetry_attributes(),
+                        'gen_ai.response.model': response_model,
+                    }
                     try:
-                        cost_attributes = {'operation.cost': float(response.cost().total_price)}
+                        attributes_to_set['operation.cost'] = float(response.cost().total_price)
                     except LookupError:
-                        cost_attributes = {}
-                    span.set_attributes(
-                        {
-                            **response.usage.opentelemetry_attributes(),
-                            'gen_ai.response.model': response_model,
-                            **cost_attributes,
-                        }
-                    )
+                        # The cost of this provider/model is unknown, which is common.
+                        pass
+                    except Exception as e:
+                        warnings.warn(
+                            f'Failed to get cost from response: {type(e).__name__}: {e}', CostCalculationFailedWarning
+                        )
+                    if response.provider_response_id is not None:
+                        attributes_to_set['gen_ai.response.id'] = response.provider_response_id
+                    if response.finish_reason is not None:
+                        attributes_to_set['gen_ai.response.finish_reasons'] = [response.finish_reason]
+                    span.set_attributes(attributes_to_set)
                     span.update_name(f'{operation} {request_model}')
 
                 yield finish
@@ -478,3 +487,7 @@ class InstrumentedModel(WrapperModel):
                 return str(value)
             except Exception as e:
                 return f'Unable to serialize: {e}'
+
+
+class CostCalculationFailedWarning(Warning):
+    """Warning raised when cost calculation fails."""