Merge pull request #307 from Mirascope/feature/otel-llm-updates

feat: otel llm updates

mirascope/anthropic/types.py

@@ -9,6 +9,7 @@
     ContentBlockStartEvent,
     Message,
     MessageParam,
+    MessageStartEvent,
     MessageStreamEvent,
     TextBlock,
     TextDelta,
@@ -153,6 +154,21 @@ def content(self) -> str:
         block = self.response.content[0]
         return block.text if block.type == "text" else ""
 
+    @property
+    def model(self) -> str:
+        """Returns the name of the response model."""
+        return self.response.model
+
+    @property
+    def id(self) -> str:
+        """Returns the id of the response."""
+        return self.response.id
+
+    @property
+    def finish_reasons(self) -> Optional[list[str]]:
+        """Returns the finish reason of the response."""
+        return [str(self.response.stop_reason)]
+
     @property
     def usage(self) -> Usage:
         """Returns the usage of the message."""
@@ -241,3 +257,45 @@ def content(self) -> str:
                 self.chunk.delta.text if isinstance(self.chunk.delta, TextDelta) else ""
             )
         return ""
+
+    @property
+    def model(self) -> Optional[str]:
+        """Returns the name of the response model."""
+        if isinstance(self.chunk, MessageStartEvent):
+            return self.chunk.message.model
+        return None
+
+    @property
+    def id(self) -> Optional[str]:
+        """Returns the id of the response."""
+        if isinstance(self.chunk, MessageStartEvent):
+            return self.chunk.message.id
+        return None
+
+    @property
+    def finish_reasons(self) -> Optional[list[str]]:
+        """Returns the finish reason of the response."""
+        if isinstance(self.chunk, MessageStartEvent):
+            return [str(self.chunk.message.stop_reason)]
+        return None
+
+    @property
+    def usage(self) -> Optional[Usage]:
+        """Returns the usage of the message."""
+        if isinstance(self.chunk, MessageStartEvent):
+            return self.chunk.message.usage
+        return None
+
+    @property
+    def input_tokens(self) -> Optional[int]:
+        """Returns the number of input tokens."""
+        if self.usage:
+            return self.usage.input_tokens
+        return None
+
+    @property
+    def output_tokens(self) -> Optional[int]:
+        """Returns the number of output tokens."""
+        if self.usage:
+            return self.usage.output_tokens
+        return None

mirascope/base/types.py

@@ -168,6 +168,27 @@ def content(self) -> str:
         """
         ...  # pragma: no cover
 
+    @property
+    @abstractmethod
+    def finish_reasons(self) -> Union[None, list[str]]:
+        """Should return the finish reasons of the response.
+
+        If there is no finish reason, this method must return None.
+        """
+        ...  # pragma: no cover
+
+    @property
+    @abstractmethod
+    def model(self) -> Optional[str]:
+        """Should return the name of the response model."""
+        ...  # pragma: no cover
+
+    @property
+    @abstractmethod
+    def id(self) -> Optional[str]:
+        """Should return the id of the response."""
+        ...  # pragma: no cover
+
     @property
     @abstractmethod
     def usage(self) -> Any:
@@ -223,3 +244,51 @@ def content(self) -> str:
         the empty string.
         """
         ...  # pragma: no cover
+
+    @property
+    @abstractmethod
+    def model(self) -> Optional[str]:
+        """Should return the name of the response model."""
+        ...  # pragma: no cover
+
+    @property
+    @abstractmethod
+    def id(self) -> Optional[str]:
+        """Should return the id of the response."""
+        ...  # pragma: no cover
+
+    @property
+    @abstractmethod
+    def finish_reasons(self) -> Union[None, list[str]]:
+        """Should return the finish reasons of the response.
+
+        If there is no finish reason, this method must return None.
+        """
+        ...  # pragma: no cover
+
+    @property
+    @abstractmethod
+    def usage(self) -> Any:
+        """Should return the usage of the response.
+
+        If there is no usage, this method must return None.
+        """
+        ...  # pragma: no cover
+
+    @property
+    @abstractmethod
+    def input_tokens(self) -> Optional[Union[int, float]]:
+        """Should return the number of input tokens.
+
+        If there is no input_tokens, this method must return None.
+        """
+        ...  # pragma: no cover
+
+    @property
+    @abstractmethod
+    def output_tokens(self) -> Optional[Union[int, float]]:
+        """Should return the number of output tokens.
+
+        If there is no output_tokens, this method must return None.
+        """
+        ...  # pragma: no cover

mirascope/cohere/types.py

@@ -7,6 +7,7 @@
     StreamedChatResponse_SearchQueriesGeneration,
     StreamedChatResponse_SearchResults,
     StreamedChatResponse_StreamEnd,
+    StreamedChatResponse_StreamStart,
     StreamedChatResponse_ToolCallsGeneration,
 )
 from cohere.types import (
@@ -130,6 +131,24 @@ def content(self) -> str:
         """Returns the content of the chat completion for the 0th choice."""
         return self.response.text
 
+    @property
+    def model(self) -> Optional[str]:
+        """Returns the name of the response model.
+
+        Cohere does not return model, so we return None
+        """
+        return None
+
+    @property
+    def id(self) -> Optional[str]:
+        """Returns the id of the response."""
+        return self.response.generation_id
+
+    @property
+    def finish_reasons(self) -> Optional[list[str]]:
+        """Returns the finish reasons of the response."""
+        return [str(self.response.finish_reason)]
+
     @property
     def search_queries(self) -> Optional[list[ChatSearchQuery]]:
         """Returns the search queries for the 0th choice message."""
@@ -307,6 +326,28 @@ def citations(self) -> Optional[list[ChatCitation]]:
             return self.chunk.citations
         return None
 
+    @property
+    def model(self) -> Optional[str]:
+        """Returns the name of the response model.
+
+        Cohere does not return model, so we return None
+        """
+        return None
+
+    @property
+    def id(self) -> Optional[str]:
+        """Returns the id of the response."""
+        if isinstance(self.chunk, StreamedChatResponse_StreamStart):
+            return self.chunk.generation_id
+        return None
+
+    @property
+    def finish_reasons(self) -> Optional[list[str]]:
+        """Returns the finish reasons of the response."""
+        if isinstance(self.chunk, StreamedChatResponse_StreamEnd):
+            return [str(self.chunk.finish_reason)]
+        return None
+
     @property
     def response(self) -> Optional[NonStreamedChatResponse]:
         """Returns the full response for the stream-end event type else None."""
@@ -321,6 +362,28 @@ def tool_calls(self) -> Optional[list[ToolCall]]:
             return self.chunk.tool_calls
         return None
 
+    @property
+    def usage(self) -> Optional[ApiMetaBilledUnits]:
+        """Returns the usage of the response."""
+        if isinstance(self.chunk, StreamedChatResponse_StreamEnd):
+            if self.chunk.response.meta:
+                return self.chunk.response.meta.billed_units
+        return None
+
+    @property
+    def input_tokens(self) -> Optional[float]:
+        """Returns the number of input tokens."""
+        if self.usage:
+            return self.usage.input_tokens
+        return None
+
+    @property
+    def output_tokens(self) -> Optional[float]:
+        """Returns the number of output tokens."""
+        if self.usage:
+            return self.usage.output_tokens
+        return None
+
 
 class CohereEmbeddingResponse(BaseEmbeddingResponse[SkipValidation[EmbedResponse]]):
     """A convenience wrapper around the Cohere `EmbedResponse` response."""

mirascope/gemini/types.py

@@ -126,6 +126,39 @@ def content(self) -> str:
         """Returns the contained string content for the 0th choice."""
         return self.response.candidates[0].content.parts[0].text
 
+    @property
+    def id(self) -> Optional[str]:
+        """Returns the id of the response.
+
+        google.generativeai does not return an id
+        """
+        return None
+
+    @property
+    def finish_reasons(self) -> list[str]:
+        """Returns the finish reasons of the response."""
+        finish_reasons = [
+            "FINISH_REASON_UNSPECIFIED",
+            "STOP",
+            "MAX_TOKENS",
+            "SAFETY",
+            "RECITATION",
+            "OTHER",
+        ]
+
+        return [
+            finish_reasons[candidate.finish_reason]
+            for candidate in self.response.candidates
+        ]
+
+    @property
+    def model(self) -> None:
+        """Returns the model name.
+
+        google.generativeai does not return model, so we return None
+        """
+        return None
+
     @property
     def usage(self) -> None:
         """Returns the usage of the chat completion.
@@ -192,3 +225,54 @@ class Math(GeminiCall):
     def content(self) -> str:
         """Returns the chunk content for the 0th choice."""
         return self.chunk.candidates[0].content.parts[0].text
+
+    @property
+    def id(self) -> Optional[str]:
+        """Returns the id of the response.
+
+        google.generativeai does not return an id
+        """
+        return None
+
+    @property
+    def finish_reasons(self) -> list[str]:
+        """Returns the finish reasons of the response."""
+        finish_reasons = [
+            "FINISH_REASON_UNSPECIFIED",
+            "STOP",
+            "MAX_TOKENS",
+            "SAFETY",
+            "RECITATION",
+            "OTHER",
+        ]
+
+        return [
+            finish_reasons[candidate.finish_reason]
+            for candidate in self.chunk.candidates
+        ]
+
+    @property
+    def model(self) -> None:
+        """Returns the model name.
+
+        google.generativeai does not return model, so we return None
+        """
+        return None
+
+    @property
+    def usage(self) -> None:
+        """Returns the usage of the chat completion.
+
+        google.generativeai does not have Usage, so we return None
+        """
+        return None
+
+    @property
+    def input_tokens(self) -> None:
+        """Returns the number of input tokens."""
+        return None
+
+    @property
+    def output_tokens(self) -> None:
+        """Returns the number of output tokens."""
+        return None