pydantic-ai-slim 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

pydantic_ai/models/instrumented.py

@@ -77,6 +77,7 @@ class InstrumentationSettings:
     tracer: Tracer = field(repr=False)
     event_logger: EventLogger = field(repr=False)
     event_mode: Literal['attributes', 'logs'] = 'attributes'
+    include_binary_content: bool = True
 
     def __init__(
         self,
@@ -84,6 +85,7 @@ class InstrumentationSettings:
         event_mode: Literal['attributes', 'logs'] = 'attributes',
         tracer_provider: TracerProvider | None = None,
         event_logger_provider: EventLoggerProvider | None = None,
+        include_binary_content: bool = True,
     ):
         """Create instrumentation options.
 
@@ -97,6 +99,7 @@ class InstrumentationSettings:
                 If not provided, the global event logger provider is used.
                 Calling `logfire.configure()` sets the global event logger provider, so most users don't need this.
                 This is only used if `event_mode='logs'`.
+            include_binary_content: Whether to include binary content in the instrumentation events.
         """
         from pydantic_ai import __version__
 
@@ -105,6 +108,40 @@ class InstrumentationSettings:
         self.tracer = tracer_provider.get_tracer('pydantic-ai', __version__)
         self.event_logger = event_logger_provider.get_event_logger('pydantic-ai', __version__)
         self.event_mode = event_mode
+        self.include_binary_content = include_binary_content
+
+    def messages_to_otel_events(self, messages: list[ModelMessage]) -> list[Event]:
+        """Convert a list of model messages to OpenTelemetry events.
+
+        Args:
+            messages: The messages to convert.
+
+        Returns:
+            A list of OpenTelemetry events.
+        """
+        events: list[Event] = []
+        instructions = InstrumentedModel._get_instructions(messages)  # pyright: ignore [reportPrivateUsage]
+        if instructions is not None:
+            events.append(Event('gen_ai.system.message', body={'content': instructions, 'role': 'system'}))
+
+        for message_index, message in enumerate(messages):
+            message_events: list[Event] = []
+            if isinstance(message, ModelRequest):
+                for part in message.parts:
+                    if hasattr(part, 'otel_event'):
+                        message_events.append(part.otel_event(self))
+            elif isinstance(message, ModelResponse):  # pragma: no branch
+                message_events = message.otel_events()
+            for event in message_events:
+                event.attributes = {
+                    'gen_ai.message.index': message_index,
+                    **(event.attributes or {}),
+                }
+            events.extend(message_events)
+
+        for event in events:
+            event.body = InstrumentedModel.serialize_any(event.body)
+        return events
 
 
 GEN_AI_SYSTEM_ATTRIBUTE = 'gen_ai.system'
@@ -155,7 +192,7 @@ class InstrumentedModel(WrapperModel):
                 ) as response_stream:
                     yield response_stream
             finally:
-                if response_stream:
+                if response_stream:  # pragma: no branch
                     finish(response_stream.get())
 
     @contextmanager
@@ -193,8 +230,8 @@ class InstrumentedModel(WrapperModel):
                 if not span.is_recording():
                     return
 
-                events = self.messages_to_otel_events(messages)
-                for event in self.messages_to_otel_events([response]):
+                events = self.settings.messages_to_otel_events(messages)
+                for event in self.settings.messages_to_otel_events([response]):
                     events.append(
                         Event(
                             'gen_ai.choice',
@@ -253,9 +290,9 @@ class InstrumentedModel(WrapperModel):
             except Exception:  # pragma: no cover
                 pass
             else:
-                if parsed.hostname:
+                if parsed.hostname:  # pragma: no branch
                     attributes['server.address'] = parsed.hostname
-                if parsed.port:
+                if parsed.port:  # pragma: no branch
                     attributes['server.port'] = parsed.port
 
         return attributes
@@ -263,40 +300,13 @@ class InstrumentedModel(WrapperModel):
     @staticmethod
     def event_to_dict(event: Event) -> dict[str, Any]:
         if not event.body:
-            body = {}
+            body = {}  # pragma: no cover
         elif isinstance(event.body, Mapping):
             body = event.body  # type: ignore
         else:
             body = {'body': event.body}
         return {**body, **(event.attributes or {})}
 
-    @staticmethod
-    def messages_to_otel_events(messages: list[ModelMessage]) -> list[Event]:
-        events: list[Event] = []
-        last_model_request: ModelRequest | None = None
-        for message_index, message in enumerate(messages):
-            message_events: list[Event] = []
-            if isinstance(message, ModelRequest):
-                last_model_request = message
-                for part in message.parts:
-                    if hasattr(part, 'otel_event'):
-                        message_events.append(part.otel_event())
-            elif isinstance(message, ModelResponse):
-                message_events = message.otel_events()
-            for event in message_events:
-                event.attributes = {
-                    'gen_ai.message.index': message_index,
-                    **(event.attributes or {}),
-                }
-            events.extend(message_events)
-        if last_model_request and last_model_request.instructions:
-            events.insert(
-                0, Event('gen_ai.system.message', body={'content': last_model_request.instructions, 'role': 'system'})
-            )
-        for event in events:
-            event.body = InstrumentedModel.serialize_any(event.body)
-        return events
-
     @staticmethod
     def serialize_any(value: Any) -> str:
         try:

pydantic_ai/models/mistral.py

@@ -71,7 +71,7 @@ try:
     from mistralai.models.usermessage import UserMessage as MistralUserMessage
     from mistralai.types.basemodel import Unset as MistralUnset
     from mistralai.utils.eventstreaming import EventStreamAsync as MistralEventStreamAsync
-except ImportError as e:
+except ImportError as e:  # pragma: lax no cover
     raise ImportError(
         'Please install `mistral` to use the Mistral model, '
         'you can use the `mistral` optional group — `pip install "pydantic-ai-slim[mistral]"`'
@@ -208,7 +208,7 @@ class MistralModel(Model):
         except SDKError as e:
             if (status_code := e.status_code) >= 400:
                 raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
-            raise
+            raise  # pragma: lax no cover
 
         assert response, 'A unexpected empty response from Mistral.'
         return response
@@ -325,7 +325,9 @@ class MistralModel(Model):
                 tool = self._map_mistral_to_pydantic_tool_call(tool_call=tool_call)
                 parts.append(tool)
 
-        return ModelResponse(parts, usage=_map_usage(response), model_name=response.model, timestamp=timestamp)
+        return ModelResponse(
+            parts, usage=_map_usage(response), model_name=response.model, timestamp=timestamp, vendor_id=response.id
+        )
 
     async def _process_streamed_response(
         self,
@@ -336,7 +338,9 @@ class MistralModel(Model):
         peekable_response = _utils.PeekableAsyncStream(response)
         first_chunk = await peekable_response.peek()
         if isinstance(first_chunk, _utils.Unset):
-            raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')
+            raise UnexpectedModelBehavior(  # pragma: no cover
+                'Streamed response ended without content or tool calls'
+            )
 
         if first_chunk.data.created:
             timestamp = datetime.fromtimestamp(first_chunk.data.created, tz=timezone.utc)
@@ -437,7 +441,7 @@ class MistralModel(Model):
         """Convert a timeout to milliseconds."""
         if timeout is None:
             return None
-        if isinstance(timeout, float):
+        if isinstance(timeout, float):  # pragma: no cover
             return int(1000 * timeout)
         raise NotImplementedError('Timeout object is not yet supported for MistralModel.')
 
@@ -454,7 +458,7 @@ class MistralModel(Model):
                 )
             elif isinstance(part, RetryPromptPart):
                 if part.tool_name is None:
-                    yield MistralUserMessage(content=part.model_response())
+                    yield MistralUserMessage(content=part.model_response())  # pragma: no cover
                 else:
                     yield MistralToolMessage(
                         tool_call_id=part.tool_call_id,
@@ -519,7 +523,7 @@ class MistralModel(Model):
                     else:
                         raise RuntimeError('Only image binary content is supported for Mistral.')
                 elif isinstance(item, DocumentUrl):
-                    raise RuntimeError('DocumentUrl is not supported in Mistral.')
+                    raise RuntimeError('DocumentUrl is not supported in Mistral.')  # pragma: no cover
                 elif isinstance(item, VideoUrl):
                     raise RuntimeError('VideoUrl is not supported in Mistral.')
                 else:  # pragma: no cover
@@ -663,7 +667,7 @@ def _map_usage(response: MistralChatCompletionResponse | MistralCompletionChunk)
             details=None,
         )
     else:
-        return Usage()
+        return Usage()  # pragma: no cover
 
 
 def _map_content(content: MistralOptionalNullable[MistralContent]) -> str | None:
@@ -677,7 +681,9 @@ def _map_content(content: MistralOptionalNullable[MistralContent]) -> str | None
             if isinstance(chunk, MistralTextChunk):
                 output = output or '' + chunk.text
             else:
-                assert False, f'Other data types like (Image, Reference) are not yet supported,  got {type(chunk)}'
+                assert False, (  # pragma: no cover
+                    f'Other data types like (Image, Reference) are not yet supported,  got {type(chunk)}'
+                )
     elif isinstance(content, str):
         output = content
 

pydantic_ai/models/openai.py

@@ -104,6 +104,12 @@ class OpenAIModelSettings(ModelSettings, total=False):
     result in faster responses and fewer tokens used on reasoning in a response.
     """
 
+    openai_logprobs: bool
+    """Include log probabilities in the response."""
+
+    openai_top_logprobs: int
+    """Include log probabilities of the top n tokens in the response."""
+
     openai_user: str
     """A unique identifier representing the end-user, which can help OpenAI monitor and detect abuse.
 
@@ -164,7 +170,7 @@ class OpenAIModel(Model):
         self,
         model_name: OpenAIModelName,
         *,
-        provider: Literal['openai', 'deepseek', 'azure'] | Provider[AsyncOpenAI] = 'openai',
+        provider: Literal['openai', 'deepseek', 'azure', 'openrouter'] | Provider[AsyncOpenAI] = 'openai',
         system_prompt_role: OpenAISystemPromptRole | None = None,
     ):
         """Initialize an OpenAI model.
@@ -287,6 +293,8 @@ class OpenAIModel(Model):
                 frequency_penalty=model_settings.get('frequency_penalty', NOT_GIVEN),
                 logit_bias=model_settings.get('logit_bias', NOT_GIVEN),
                 reasoning_effort=model_settings.get('openai_reasoning_effort', NOT_GIVEN),
+                logprobs=model_settings.get('openai_logprobs', NOT_GIVEN),
+                top_logprobs=model_settings.get('openai_top_logprobs', NOT_GIVEN),
                 user=model_settings.get('openai_user', NOT_GIVEN),
                 extra_headers=extra_headers,
                 extra_body=model_settings.get('extra_body'),
@@ -294,26 +302,54 @@ class OpenAIModel(Model):
         except APIStatusError as e:
             if (status_code := e.status_code) >= 400:
                 raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
-            raise
+            raise  # pragma: lax no cover
 
     def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         timestamp = datetime.fromtimestamp(response.created, tz=timezone.utc)
         choice = response.choices[0]
         items: list[ModelResponsePart] = []
+        vendor_details: dict[str, Any] | None = None
+
+        # Add logprobs to vendor_details if available
+        if choice.logprobs is not None and choice.logprobs.content:
+            # Convert logprobs to a serializable format
+            vendor_details = {
+                'logprobs': [
+                    {
+                        'token': lp.token,
+                        'bytes': lp.bytes,
+                        'logprob': lp.logprob,
+                        'top_logprobs': [
+                            {'token': tlp.token, 'bytes': tlp.bytes, 'logprob': tlp.logprob} for tlp in lp.top_logprobs
+                        ],
+                    }
+                    for lp in choice.logprobs.content
+                ],
+            }
+
         if choice.message.content is not None:
             items.append(TextPart(choice.message.content))
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
                 items.append(ToolCallPart(c.function.name, c.function.arguments, tool_call_id=c.id))
-        return ModelResponse(items, usage=_map_usage(response), model_name=response.model, timestamp=timestamp)
+        return ModelResponse(
+            items,
+            usage=_map_usage(response),
+            model_name=response.model,
+            timestamp=timestamp,
+            vendor_details=vendor_details,
+            vendor_id=response.id,
+        )
 
     async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> OpenAIStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         peekable_response = _utils.PeekableAsyncStream(response)
         first_chunk = await peekable_response.peek()
         if isinstance(first_chunk, _utils.Unset):
-            raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')
+            raise UnexpectedModelBehavior(  # pragma: no cover
+                'Streamed response ended without content or tool calls'
+            )
 
         return OpenAIStreamedResponse(
             _model_name=self._model_name,
@@ -399,7 +435,9 @@ class OpenAIModel(Model):
                 )
             elif isinstance(part, RetryPromptPart):
                 if part.tool_name is None:
-                    yield chat.ChatCompletionUserMessageParam(role='user', content=part.model_response())
+                    yield chat.ChatCompletionUserMessageParam(  # pragma: no cover
+                        role='user', content=part.model_response()
+                    )
                 else:
                     yield chat.ChatCompletionToolMessageParam(
                         role='tool',
@@ -637,7 +675,7 @@ class OpenAIResponsesModel(Model):
         except APIStatusError as e:
             if (status_code := e.status_code) >= 400:
                 raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
-            raise
+            raise  # pragma: lax no cover
 
     def _get_reasoning(self, model_settings: OpenAIResponsesModelSettings) -> Reasoning | NotGiven:
         reasoning_effort = model_settings.get('openai_reasoning_effort', None)
@@ -867,7 +905,7 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
                     args=chunk.delta,
                     tool_call_id=chunk.item_id,
                 )
-                if maybe_event is not None:
+                if maybe_event is not None:  # pragma: no branch
                     yield maybe_event
 
             elif isinstance(chunk, responses.ResponseFunctionCallArgumentsDoneEvent):
@@ -1031,7 +1069,7 @@ class _OpenAIJsonSchema(WalkJsonSchema):
                     notes.append(f'{key}={value}')
                 notes_string = ', '.join(notes)
                 schema['description'] = notes_string if not description else f'{description} ({notes_string})'
-            elif self.strict is None:
+            elif self.strict is None:  # pragma: no branch
                 self.is_strict_compatible = False
 
         schema_type = schema.get('type')

pydantic_ai/models/test.py

@@ -169,7 +169,7 @@ class TestModel(Model):
                 model_name=self._model_name,
             )
 
-        if messages:
+        if messages:  # pragma: no branch
             last_message = messages[-1]
             assert isinstance(last_message, ModelRequest), 'Expected last message to be a `ModelRequest`.'
 

pydantic_ai/models/wrapper.py

@@ -48,4 +48,4 @@ class WrapperModel(Model):
         return self.wrapped.system
 
     def __getattr__(self, item: str):
-        return getattr(self.wrapped, item)
+        return getattr(self.wrapped, item)  # pragma: no cover

pydantic_ai/providers/__init__.py

@@ -52,6 +52,10 @@ def infer_provider(provider: str) -> Provider[Any]:
         from .deepseek import DeepSeekProvider
 
         return DeepSeekProvider()
+    elif provider == 'openrouter':
+        from .openrouter import OpenRouterProvider
+
+        return OpenRouterProvider()
     elif provider == 'azure':
         from .azure import AzureProvider
 

pydantic_ai/providers/google.py

@@ -0,0 +1,143 @@
+from __future__ import annotations as _annotations
+
+import os
+from typing import Literal, overload
+
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.models import get_user_agent
+from pydantic_ai.providers import Provider
+
+try:
+    from google import genai
+    from google.auth.credentials import Credentials
+except ImportError as _import_error:
+    raise ImportError(
+        'Please install the `google-genai` package to use the Google provider, '
+        'you can use the `google` optional group — `pip install "pydantic-ai-slim[google]"`'
+    ) from _import_error
+
+
+class GoogleProvider(Provider[genai.Client]):
+    """Provider for Google."""
+
+    @property
+    def name(self) -> str:
+        return 'google-vertex' if self._client._api_client.vertexai else 'google-gla'  # type: ignore[reportPrivateUsage]
+
+    @property
+    def base_url(self) -> str:
+        return str(self._client._api_client._http_options.base_url)  # type: ignore[reportPrivateUsage]
+
+    @property
+    def client(self) -> genai.Client:
+        return self._client
+
+    @overload
+    def __init__(self, *, api_key: str) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        *,
+        credentials: Credentials | None = None,
+        project: str | None = None,
+        location: VertexAILocation | Literal['global'] | None = None,
+    ) -> None: ...
+
+    @overload
+    def __init__(self, *, client: genai.Client) -> None: ...
+
+    @overload
+    def __init__(self, *, vertexai: bool = False) -> None: ...
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        credentials: Credentials | None = None,
+        project: str | None = None,
+        location: VertexAILocation | Literal['global'] | None = None,
+        client: genai.Client | None = None,
+        vertexai: bool | None = None,
+    ) -> None:
+        """Create a new Google provider.
+
+        Args:
+            api_key: The `API key <https://ai.google.dev/gemini-api/docs/api-key>`_ to
+                use for authentication. It can also be set via the `GOOGLE_API_KEY` environment variable.
+                Applies to the Gemini Developer API only.
+            credentials: The credentials to use for authentication when calling the Vertex AI APIs. Credentials can be
+                obtained from environment variables and default credentials. For more information, see Set up
+                Application Default Credentials. Applies to the Vertex AI API only.
+            project: The Google Cloud project ID to use for quota. Can be obtained from environment variables
+                (for example, GOOGLE_CLOUD_PROJECT). Applies to the Vertex AI API only.
+            location: The location to send API requests to (for example, us-central1). Can be obtained from environment variables.
+                Applies to the Vertex AI API only.
+            client: A pre-initialized client to use.
+            vertexai: Force the use of the Vertex AI API. If `False`, the Google Generative Language API will be used.
+                Defaults to `False`.
+        """
+        if client is None:
+            # NOTE: We are keeping GEMINI_API_KEY for backwards compatibility.
+            api_key = api_key or os.environ.get('GOOGLE_API_KEY')
+
+            if vertexai is None:  # pragma: lax no cover
+                vertexai = bool(location or project or credentials)
+
+            if not vertexai:
+                if api_key is None:
+                    raise UserError(  # pragma: no cover
+                        'Set the `GOOGLE_API_KEY` environment variable or pass it via `GoogleProvider(api_key=...)`'
+                        'to use the Google Generative Language API.'
+                    )
+                self._client = genai.Client(
+                    vertexai=vertexai,
+                    api_key=api_key,
+                    http_options={'headers': {'User-Agent': get_user_agent()}},
+                )
+            else:
+                self._client = genai.Client(
+                    vertexai=vertexai,
+                    project=project or os.environ.get('GOOGLE_CLOUD_PROJECT'),
+                    location=location or os.environ.get('GOOGLE_CLOUD_LOCATION'),
+                    credentials=credentials,
+                    http_options={'headers': {'User-Agent': get_user_agent()}},
+                )
+        else:
+            self._client = client  # pragma: lax no cover
+
+
+VertexAILocation = Literal[
+    'asia-east1',
+    'asia-east2',
+    'asia-northeast1',
+    'asia-northeast3',
+    'asia-south1',
+    'asia-southeast1',
+    'australia-southeast1',
+    'europe-central2',
+    'europe-north1',
+    'europe-southwest1',
+    'europe-west1',
+    'europe-west2',
+    'europe-west3',
+    'europe-west4',
+    'europe-west6',
+    'europe-west8',
+    'europe-west9',
+    'me-central1',
+    'me-central2',
+    'me-west1',
+    'northamerica-northeast1',
+    'southamerica-east1',
+    'us-central1',
+    'us-east1',
+    'us-east4',
+    'us-east5',
+    'us-south1',
+    'us-west1',
+    'us-west4',
+]
+"""Regions available for Vertex AI.
+More details [here](https://cloud.google.com/vertex-ai/generative-ai/docs/learn/locations#genai-locations).
+"""

pydantic_ai/providers/google_vertex.py

@@ -128,7 +128,7 @@ class _VertexAIAuth(httpx.Auth):
         self.credentials = None
 
     async def async_auth_flow(self, request: httpx.Request) -> AsyncGenerator[httpx.Request, httpx.Response]:
-        if self.credentials is None:
+        if self.credentials is None:  # pragma: no branch
             self.credentials = await self._get_credentials()
         if self.credentials.token is None:  # type: ignore[reportUnknownMemberType]
             await self._refresh_token()
@@ -157,9 +157,9 @@ class _VertexAIAuth(httpx.Auth):
             creds, creds_project_id = await _async_google_auth()
             creds_source = '`google.auth.default()`'
 
-        if self.project_id is None:
+        if self.project_id is None:  # pragma: no branch
             if creds_project_id is None:
-                raise UserError(f'No project_id provided and none found in {creds_source}')
+                raise UserError(f'No project_id provided and none found in {creds_source}')  # pragma: no cover
             self.project_id = creds_project_id
         return creds
 

pydantic_ai/providers/openrouter.py

@@ -0,0 +1,69 @@
+from __future__ import annotations as _annotations
+
+import os
+from typing import overload
+
+from httpx import AsyncClient as AsyncHTTPClient
+from openai import AsyncOpenAI
+
+from pydantic_ai.exceptions import UserError
+from pydantic_ai.models import cached_async_http_client
+from pydantic_ai.providers import Provider
+
+try:
+    from openai import AsyncOpenAI
+except ImportError as _import_error:  # pragma: no cover
+    raise ImportError(
+        'Please install the `openai` package to use the OpenRouter provider, '
+        'you can use the `openai` optional group — `pip install "pydantic-ai-slim[openai]"`'
+    ) from _import_error
+
+
+class OpenRouterProvider(Provider[AsyncOpenAI]):
+    """Provider for OpenRouter API."""
+
+    @property
+    def name(self) -> str:
+        return 'openrouter'
+
+    @property
+    def base_url(self) -> str:
+        return 'https://openrouter.ai/api/v1'
+
+    @property
+    def client(self) -> AsyncOpenAI:
+        return self._client
+
+    @overload
+    def __init__(self) -> None: ...
+
+    @overload
+    def __init__(self, *, api_key: str) -> None: ...
+
+    @overload
+    def __init__(self, *, api_key: str, http_client: AsyncHTTPClient) -> None: ...
+
+    @overload
+    def __init__(self, *, openai_client: AsyncOpenAI | None = None) -> None: ...
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        openai_client: AsyncOpenAI | None = None,
+        http_client: AsyncHTTPClient | None = None,
+    ) -> None:
+        api_key = api_key or os.getenv('OPENROUTER_API_KEY')
+        if not api_key and openai_client is None:
+            raise UserError(
+                'Set the `OPENROUTER_API_KEY` environment variable or pass it via `OpenRouterProvider(api_key=...)`'
+                'to use the OpenRouter provider.'
+            )
+
+        if openai_client is not None:
+            self._client = openai_client
+        elif http_client is not None:
+            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)
+        else:
+            http_client = cached_async_http_client(provider='openrouter')
+            self._client = AsyncOpenAI(base_url=self.base_url, api_key=api_key, http_client=http_client)