pydantic-ai-slim 0.0.23__py3-none-any.whl → 0.0.25__py3-none-any.whl

pydantic_ai/models/instrumented.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Iterator
+from contextlib import asynccontextmanager, contextmanager
+from dataclasses import dataclass, field
+from functools import partial
+from typing import Any, Callable, Literal
+
+import logfire_api
+from opentelemetry._events import Event, EventLogger, EventLoggerProvider, get_event_logger_provider
+from opentelemetry.trace import Tracer, TracerProvider, get_tracer_provider
+
+from ..messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelRequestPart,
+    ModelResponse,
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+from ..settings import ModelSettings
+from ..usage import Usage
+from . import KnownModelName, Model, ModelRequestParameters, StreamedResponse
+from .wrapper import WrapperModel
+
+MODEL_SETTING_ATTRIBUTES: tuple[
+    Literal[
+        'max_tokens',
+        'top_p',
+        'seed',
+        'temperature',
+        'presence_penalty',
+        'frequency_penalty',
+    ],
+    ...,
+] = (
+    'max_tokens',
+    'top_p',
+    'seed',
+    'temperature',
+    'presence_penalty',
+    'frequency_penalty',
+)
+
+NOT_GIVEN = object()
+
+
+@dataclass
+class InstrumentedModel(WrapperModel):
+    """Model which is instrumented with logfire."""
+
+    tracer: Tracer = field(repr=False)
+    event_logger: EventLogger = field(repr=False)
+
+    def __init__(
+        self,
+        wrapped: Model | KnownModelName,
+        tracer_provider: TracerProvider | None = None,
+        event_logger_provider: EventLoggerProvider | None = None,
+    ):
+        super().__init__(wrapped)
+        tracer_provider = tracer_provider or get_tracer_provider()
+        event_logger_provider = event_logger_provider or get_event_logger_provider()
+        self.tracer = tracer_provider.get_tracer('pydantic-ai')
+        self.event_logger = event_logger_provider.get_event_logger('pydantic-ai')
+
+    @classmethod
+    def from_logfire(
+        cls,
+        wrapped: Model | KnownModelName,
+        logfire_instance: logfire_api.Logfire = logfire_api.DEFAULT_LOGFIRE_INSTANCE,
+    ) -> InstrumentedModel:
+        if hasattr(logfire_instance.config, 'get_event_logger_provider'):
+            event_provider = logfire_instance.config.get_event_logger_provider()
+        else:
+            event_provider = None
+        tracer_provider = logfire_instance.config.get_tracer_provider()
+        return cls(wrapped, tracer_provider, event_provider)
+
+    async def request(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> tuple[ModelResponse, Usage]:
+        with self._instrument(messages, model_settings) as finish:
+            response, usage = await super().request(messages, model_settings, model_request_parameters)
+            finish(response, usage)
+            return response, usage
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> AsyncIterator[StreamedResponse]:
+        with self._instrument(messages, model_settings) as finish:
+            response_stream: StreamedResponse | None = None
+            try:
+                async with super().request_stream(
+                    messages, model_settings, model_request_parameters
+                ) as response_stream:
+                    yield response_stream
+            finally:
+                if response_stream:
+                    finish(response_stream.get(), response_stream.usage())
+
+    @contextmanager
+    def _instrument(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+    ) -> Iterator[Callable[[ModelResponse, Usage], None]]:
+        operation = 'chat'
+        model_name = self.model_name
+        span_name = f'{operation} {model_name}'
+        system = getattr(self.wrapped, 'system', '') or self.wrapped.__class__.__name__.removesuffix('Model').lower()
+        system = {'google-gla': 'gemini', 'google-vertex': 'vertex_ai', 'mistral': 'mistral_ai'}.get(system, system)
+        # TODO Missing attributes:
+        #  - server.address: requires a Model.base_url abstract method or similar
+        #  - server.port: to parse from the base_url
+        #  - error.type: unclear if we should do something here or just always rely on span exceptions
+        #  - gen_ai.request.stop_sequences/top_k: model_settings doesn't include these
+        attributes: dict[str, Any] = {
+            'gen_ai.operation.name': operation,
+            'gen_ai.system': system,
+            'gen_ai.request.model': model_name,
+        }
+
+        if model_settings:
+            for key in MODEL_SETTING_ATTRIBUTES:
+                if (value := model_settings.get(key, NOT_GIVEN)) is not NOT_GIVEN:
+                    attributes[f'gen_ai.request.{key}'] = value
+
+        emit_event = partial(self._emit_event, system)
+
+        with self.tracer.start_as_current_span(span_name, attributes=attributes) as span:
+            if span.is_recording():
+                for message in messages:
+                    if isinstance(message, ModelRequest):
+                        for part in message.parts:
+                            event_name, body = _request_part_body(part)
+                            if event_name:
+                                emit_event(event_name, body)
+                    elif isinstance(message, ModelResponse):
+                        for body in _response_bodies(message):
+                            emit_event('gen_ai.assistant.message', body)
+
+            def finish(response: ModelResponse, usage: Usage):
+                if not span.is_recording():
+                    return
+
+                for response_body in _response_bodies(response):
+                    if response_body:
+                        emit_event(
+                            'gen_ai.choice',
+                            {
+                                # TODO finish_reason
+                                'index': 0,
+                                'message': response_body,
+                            },
+                        )
+                span.set_attributes(
+                    {
+                        k: v
+                        for k, v in {
+                            # TODO finish_reason (https://github.com/open-telemetry/semantic-conventions/issues/1277), id
+                            #  https://github.com/pydantic/pydantic-ai/issues/886
+                            'gen_ai.response.model': response.model_name or model_name,
+                            'gen_ai.usage.input_tokens': usage.request_tokens,
+                            'gen_ai.usage.output_tokens': usage.response_tokens,
+                        }.items()
+                        if v is not None
+                    }
+                )
+
+            yield finish
+
+    def _emit_event(self, system: str, event_name: str, body: dict[str, Any]) -> None:
+        self.event_logger.emit(Event(event_name, body=body, attributes={'gen_ai.system': system}))
+
+
+def _request_part_body(part: ModelRequestPart) -> tuple[str, dict[str, Any]]:
+    if isinstance(part, SystemPromptPart):
+        return 'gen_ai.system.message', {'content': part.content, 'role': 'system'}
+    elif isinstance(part, UserPromptPart):
+        return 'gen_ai.user.message', {'content': part.content, 'role': 'user'}
+    elif isinstance(part, ToolReturnPart):
+        return 'gen_ai.tool.message', {'content': part.content, 'role': 'tool', 'id': part.tool_call_id}
+    elif isinstance(part, RetryPromptPart):
+        if part.tool_name is None:
+            return 'gen_ai.user.message', {'content': part.model_response(), 'role': 'user'}
+        else:
+            return 'gen_ai.tool.message', {'content': part.model_response(), 'role': 'tool', 'id': part.tool_call_id}
+    else:
+        return '', {}
+
+
+def _response_bodies(message: ModelResponse) -> list[dict[str, Any]]:
+    body: dict[str, Any] = {'role': 'assistant'}
+    result = [body]
+    for part in message.parts:
+        if isinstance(part, ToolCallPart):
+            body.setdefault('tool_calls', []).append(
+                {
+                    'id': part.tool_call_id,
+                    'type': 'function',  # TODO https://github.com/pydantic/pydantic-ai/issues/888
+                    'function': {
+                        'name': part.tool_name,
+                        'arguments': part.args,
+                    },
+                }
+            )
+        elif isinstance(part, TextPart):
+            if body.get('content'):
+                body = {'role': 'assistant'}
+                result.append(body)
+            body['content'] = part.content
+
+    return result

pydantic_ai/models/mistral.py

@@ -134,9 +134,6 @@ class MistralModel(Model):
             api_key = os.getenv('MISTRAL_API_KEY') if api_key is None else api_key
             self.client = Mistral(api_key=api_key, async_client=http_client or cached_async_http_client())
 
-    def name(self) -> str:
-        return f'mistral:{self._model_name}'
-
     async def request(
         self,
         messages: list[ModelMessage],
@@ -165,6 +162,16 @@ class MistralModel(Model):
         async with response:
             yield await self._process_streamed_response(model_request_parameters.result_tools, response)
 
+    @property
+    def model_name(self) -> MistralModelName:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str | None:
+        """The system / model provider."""
+        return self._system
+
     async def _completions_create(
         self,
         messages: list[ModelMessage],
@@ -296,7 +303,7 @@ class MistralModel(Model):
                 tool = self._map_mistral_to_pydantic_tool_call(tool_call=tool_call)
                 parts.append(tool)
 
-        return ModelResponse(parts, model_name=self._model_name, timestamp=timestamp)
+        return ModelResponse(parts, model_name=response.model, timestamp=timestamp)
 
     async def _process_streamed_response(
         self,
@@ -461,6 +468,7 @@ MistralToolCallId = Union[str, None]
 class MistralStreamedResponse(StreamedResponse):
     """Implementation of `StreamedResponse` for Mistral models."""
 
+    _model_name: MistralModelName
     _response: AsyncIterable[MistralCompletionEvent]
     _timestamp: datetime
     _result_tools: dict[str, ToolDefinition]
@@ -502,7 +510,14 @@ class MistralStreamedResponse(StreamedResponse):
                     vendor_part_id=index, tool_name=dtc.function.name, args=dtc.function.arguments, tool_call_id=dtc.id
                 )
 
+    @property
+    def model_name(self) -> MistralModelName:
+        """Get the model name of the response."""
+        return self._model_name
+
+    @property
     def timestamp(self) -> datetime:
+        """Get the timestamp of the response."""
         return self._timestamp
 
     @staticmethod

pydantic_ai/models/openai.py

@@ -119,10 +119,11 @@ class OpenAIModel(Model):
         """
         self._model_name = model_name
         # This is a workaround for the OpenAI client requiring an API key, whilst locally served,
-        # openai compatible models do not always need an API key.
+        # openai compatible models do not always need an API key, but a placeholder (non-empty) key is required.
         if api_key is None and 'OPENAI_API_KEY' not in os.environ and base_url is not None and openai_client is None:
-            api_key = ''
-        elif openai_client is not None:
+            api_key = 'api-key-not-set'
+
+        if openai_client is not None:
             assert http_client is None, 'Cannot provide both `openai_client` and `http_client`'
             assert base_url is None, 'Cannot provide both `openai_client` and `base_url`'
             assert api_key is None, 'Cannot provide both `openai_client` and `api_key`'
@@ -134,9 +135,6 @@ class OpenAIModel(Model):
         self.system_prompt_role = system_prompt_role
         self._system = system
 
-    def name(self) -> str:
-        return f'openai:{self._model_name}'
-
     async def request(
         self,
         messages: list[ModelMessage],
@@ -163,6 +161,16 @@ class OpenAIModel(Model):
         async with response:
             yield await self._process_streamed_response(response)
 
+    @property
+    def model_name(self) -> OpenAIModelName:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str | None:
+        """The system / model provider."""
+        return self._system
+
     @overload
     async def _completions_create(
         self,
@@ -232,7 +240,7 @@ class OpenAIModel(Model):
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
                 items.append(ToolCallPart(c.function.name, c.function.arguments, c.id))
-        return ModelResponse(items, model_name=self._model_name, timestamp=timestamp)
+        return ModelResponse(items, model_name=response.model, timestamp=timestamp)
 
     async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> OpenAIStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
@@ -331,6 +339,7 @@ class OpenAIModel(Model):
 class OpenAIStreamedResponse(StreamedResponse):
     """Implementation of `StreamedResponse` for OpenAI models."""
 
+    _model_name: OpenAIModelName
     _response: AsyncIterable[ChatCompletionChunk]
     _timestamp: datetime
 
@@ -358,7 +367,14 @@ class OpenAIStreamedResponse(StreamedResponse):
                 if maybe_event is not None:
                     yield maybe_event
 
+    @property
+    def model_name(self) -> OpenAIModelName:
+        """Get the model name of the response."""
+        return self._model_name
+
+    @property
     def timestamp(self) -> datetime:
+        """Get the timestamp of the response."""
         return self._timestamp
 
 

pydantic_ai/models/test.py

@@ -107,6 +107,16 @@ class TestModel(Model):
             _model_name=self._model_name, _structured_response=model_response, _messages=messages
         )
 
+    @property
+    def model_name(self) -> str:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str | None:
+        """The system / model provider."""
+        return self._system
+
     def gen_tool_args(self, tool_def: ToolDefinition) -> Any:
         return _JsonSchemaTestData(tool_def.parameters_json_schema, self.seed).generate()
 
@@ -120,15 +130,15 @@ class TestModel(Model):
 
     def _get_result(self, model_request_parameters: ModelRequestParameters) -> _TextResult | _FunctionToolResult:
         if self.custom_result_text is not None:
-            assert (
-                model_request_parameters.allow_text_result
-            ), 'Plain response not allowed, but `custom_result_text` is set.'
+            assert model_request_parameters.allow_text_result, (
+                'Plain response not allowed, but `custom_result_text` is set.'
+            )
             assert self.custom_result_args is None, 'Cannot set both `custom_result_text` and `custom_result_args`.'
             return _TextResult(self.custom_result_text)
         elif self.custom_result_args is not None:
-            assert (
-                model_request_parameters.result_tools is not None
-            ), 'No result tools provided, but `custom_result_args` is set.'
+            assert model_request_parameters.result_tools is not None, (
+                'No result tools provided, but `custom_result_args` is set.'
+            )
             result_tool = model_request_parameters.result_tools[0]
 
             if k := result_tool.outer_typed_dict_key:
@@ -221,9 +231,9 @@ class TestModel(Model):
 class TestStreamedResponse(StreamedResponse):
     """A structured response that streams test data."""
 
+    _model_name: str
     _structured_response: ModelResponse
     _messages: InitVar[Iterable[ModelMessage]]
-
     _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
 
     def __post_init__(self, _messages: Iterable[ModelMessage]):
@@ -249,7 +259,14 @@ class TestStreamedResponse(StreamedResponse):
                     vendor_part_id=i, tool_name=part.tool_name, args=part.args, tool_call_id=part.tool_call_id
                 )
 
+    @property
+    def model_name(self) -> str:
+        """Get the model name of the response."""
+        return self._model_name
+
+    @property
     def timestamp(self) -> datetime:
+        """Get the timestamp of the response."""
         return self._timestamp
 
 

pydantic_ai/models/vertexai.py

@@ -161,6 +161,16 @@ class VertexAIModel(GeminiModel):
         async with super().request_stream(messages, model_settings, model_request_parameters) as value:
             yield value
 
+    @property
+    def model_name(self) -> GeminiModelName:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str | None:
+        """The system / model provider."""
+        return self._system
+
 
 # pyright: reportUnknownMemberType=false
 def _creds_from_file(service_account_file: str | Path) -> ServiceAccountCredentials:

pydantic_ai/models/wrapper.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import Any
+
+from ..messages import ModelMessage, ModelResponse
+from ..settings import ModelSettings
+from ..usage import Usage
+from . import KnownModelName, Model, ModelRequestParameters, StreamedResponse, infer_model
+
+
+@dataclass(init=False)
+class WrapperModel(Model):
+    """Model which wraps another model."""
+
+    wrapped: Model
+
+    def __init__(self, wrapped: Model | KnownModelName):
+        self.wrapped = infer_model(wrapped)
+
+    async def request(self, *args: Any, **kwargs: Any) -> tuple[ModelResponse, Usage]:
+        return await self.wrapped.request(*args, **kwargs)
+
+    @asynccontextmanager
+    async def request_stream(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> AsyncIterator[StreamedResponse]:
+        async with self.wrapped.request_stream(messages, model_settings, model_request_parameters) as response_stream:
+            yield response_stream
+
+    @property
+    def model_name(self) -> str:
+        return self.wrapped.model_name
+
+    @property
+    def system(self) -> str | None:
+        return self.wrapped.system
+
+    def __getattr__(self, item: str):
+        return getattr(self.wrapped, item)