Use abstract properties for otel semantic attributes on models and stream responses (#904)

Author: sydney-runkle

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -177,9 +177,6 @@ class ModelRequestParameters:
 class Model(ABC):
     """Abstract class for a model."""
 
-    _model_name: str
-    _system: str | None
-
     @abstractmethod
     async def request(
         self,
@@ -205,24 +202,25 @@ async def request_stream(
         yield  # pragma: no cover
 
     @property
+    @abstractmethod
     def model_name(self) -> str:
         """The model name."""
-        return self._model_name
+        raise NotImplementedError()
 
     @property
+    @abstractmethod
     def system(self) -> str | None:
         """The system / model provider, ex: openai."""
-        return self._system
+        raise NotImplementedError()
 
 
 @dataclass
 class StreamedResponse(ABC):
     """Streamed response from an LLM when calling a tool."""
 
-    _model_name: str
-    _usage: Usage = field(default_factory=Usage, init=False)
     _parts_manager: ModelResponsePartsManager = field(default_factory=ModelResponsePartsManager, init=False)
     _event_iterator: AsyncIterator[ModelResponseStreamEvent] | None = field(default=None, init=False)
+    _usage: Usage = field(default_factory=Usage, init=False)
 
     def __aiter__(self) -> AsyncIterator[ModelResponseStreamEvent]:
         """Stream the response as an async iterable of [`ModelResponseStreamEvent`][pydantic_ai.messages.ModelResponseStreamEvent]s."""
@@ -244,17 +242,20 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
     def get(self) -> ModelResponse:
         """Build a [`ModelResponse`][pydantic_ai.messages.ModelResponse] from the data received from the stream so far."""
         return ModelResponse(
-            parts=self._parts_manager.get_parts(), model_name=self._model_name, timestamp=self.timestamp()
+            parts=self._parts_manager.get_parts(), model_name=self.model_name, timestamp=self.timestamp
         )
 
-    def model_name(self) -> str:
-        """Get the model name of the response."""
-        return self._model_name
-
     def usage(self) -> Usage:
         """Get the usage of the response so far. This will not be the final usage until the stream is exhausted."""
         return self._usage
 
+    @property
+    @abstractmethod
+    def model_name(self) -> str:
+        """Get the model name of the response."""
+        raise NotImplementedError()
+
+    @property
     @abstractmethod
     def timestamp(self) -> datetime:
         """Get the timestamp of the response."""

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -162,6 +162,16 @@ async def request_stream(
         async with response:
             yield await self._process_streamed_response(response)
 
+    @property
+    def model_name(self) -> AnthropicModelName:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str | None:
+        """The system / model provider."""
+        return self._system
+
     @overload
     async def _messages_create(
         self,
@@ -362,6 +372,7 @@ def _map_usage(message: AnthropicMessage | RawMessageStreamEvent) -> usage.Usage
 class AnthropicStreamedResponse(StreamedResponse):
     """Implementation of `StreamedResponse` for Anthropic models."""
 
+    _model_name: AnthropicModelName
     _response: AsyncIterable[RawMessageStreamEvent]
     _timestamp: datetime
 
@@ -414,5 +425,12 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
             elif isinstance(event, (RawContentBlockStopEvent, RawMessageStopEvent)):
                 current_block = None
 
+    @property
+    def model_name(self) -> AnthropicModelName:
+        """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_slim/pydantic_ai/models/cohere.py

@@ -136,6 +136,16 @@ async def request(
         response = await self._chat(messages, cast(CohereModelSettings, model_settings or {}), model_request_parameters)
         return self._process_response(response), _map_usage(response)
 
+    @property
+    def model_name(self) -> CohereModelName:
+        """The model name."""
+        return self._model_name
+
+    @property
+    def system(self) -> str | None:
+        """The system / model provider."""
+        return self._system
+
     async def _chat(
         self,
         messages: list[ModelMessage],

pydantic_ai_slim/pydantic_ai/models/function.py

@@ -121,6 +121,16 @@ async def request_stream(
 
         yield FunctionStreamedResponse(_model_name=f'function:{self.stream_function.__name__}', _iter=response_stream)
 
+    @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
+
 
 @dataclass(frozen=True)
 class AgentInfo:
@@ -178,6 +188,7 @@ class DeltaToolCall:
 class FunctionStreamedResponse(StreamedResponse):
     """Implementation of `StreamedResponse` for [FunctionModel][pydantic_ai.models.function.FunctionModel]."""
 
+    _model_name: str
     _iter: AsyncIterator[str | DeltaToolCalls]
     _timestamp: datetime = field(default_factory=_utils.now_utc)
 
@@ -205,7 +216,14 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                     if maybe_event is not None:
                         yield maybe_event
 
+    @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_slim/pydantic_ai/models/gemini.py

@@ -149,6 +149,16 @@ async def request_stream(
         ) as http_response:
             yield await self._process_streamed_response(http_response)
 
+    @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
+
     def _get_tools(self, model_request_parameters: ModelRequestParameters) -> _GeminiTools | None:
         tools = [_function_from_abstract_tool(t) for t in model_request_parameters.function_tools]
         if model_request_parameters.result_tools:
@@ -315,6 +325,7 @@ async def headers(self) -> dict[str, str]:
 class GeminiStreamedResponse(StreamedResponse):
     """Implementation of `StreamedResponse` for the Gemini model."""
 
+    _model_name: GeminiModelName
     _content: bytearray
     _stream: AsyncIterator[bytes]
     _timestamp: datetime = field(default_factory=_utils.now_utc, init=False)
@@ -378,7 +389,14 @@ async def _get_gemini_responses(self) -> AsyncIterator[_GeminiResponse]:
             self._usage += _metadata_as_usage(r)
             yield r
 
+    @property
+    def model_name(self) -> GeminiModelName:
+        """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_slim/pydantic_ai/models/groq.py

@@ -146,6 +146,16 @@ async def request_stream(
         async with response:
             yield await self._process_streamed_response(response)
 
+    @property
+    def model_name(self) -> GroqModelName:
+        """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,
@@ -305,6 +315,7 @@ def _map_user_message(cls, message: ModelRequest) -> Iterable[chat.ChatCompletio
 class GroqStreamedResponse(StreamedResponse):
     """Implementation of `StreamedResponse` for Groq models."""
 
+    _model_name: GroqModelName
     _response: AsyncIterable[ChatCompletionChunk]
     _timestamp: datetime
 
@@ -333,7 +344,14 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                 if maybe_event is not None:
                     yield maybe_event
 
+    @property
+    def model_name(self) -> GroqModelName:
+        """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_slim/pydantic_ai/models/mistral.py

@@ -165,6 +165,16 @@ async def request_stream(
         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],
@@ -461,6 +471,7 @@ def _map_message(cls, message: ModelMessage) -> Iterable[MistralMessages]:
 class MistralStreamedResponse(StreamedResponse):
     """Implementation of `StreamedResponse` for Mistral models."""
 
+    _model_name: MistralModelName
     _response: AsyncIterable[MistralCompletionEvent]
     _timestamp: datetime
     _result_tools: dict[str, ToolDefinition]
@@ -502,7 +513,14 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                     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_slim/pydantic_ai/models/openai.py

@@ -164,6 +164,16 @@ async def request_stream(
         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,
@@ -332,6 +342,7 @@ def _map_user_message(self, message: ModelRequest) -> Iterable[chat.ChatCompleti
 class OpenAIStreamedResponse(StreamedResponse):
     """Implementation of `StreamedResponse` for OpenAI models."""
 
+    _model_name: OpenAIModelName
     _response: AsyncIterable[ChatCompletionChunk]
     _timestamp: datetime
 
@@ -359,7 +370,14 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                 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_slim/pydantic_ai/models/test.py

@@ -107,6 +107,16 @@ async def request_stream(
             _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()
 
@@ -221,9 +231,9 @@ def _request(
 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 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
                     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_slim/pydantic_ai/models/vertexai.py

@@ -161,6 +161,16 @@ async def request_stream(
         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: