mirascope 1.18.3__py3-none-any.whl → 1.18.4__py3-none-any.whl

mirascope/core/google/call_response_chunk.py

@@ -6,7 +6,10 @@ usage docs: learn/streams.md#handling-streamed-responses
 from typing import cast
 
 from google.genai.types import FinishReason as GoogleFinishReason
-from google.genai.types import GenerateContentResponse
+from google.genai.types import (
+    GenerateContentResponse,
+    GenerateContentResponseUsageMetadata,
+)
 
 from mirascope.core.base.types import FinishReason
 
@@ -57,12 +60,12 @@ class GoogleCallResponseChunk(
         ]
 
     @property
-    def model(self) -> None:
+    def model(self) -> str | None:
         """Returns the model name.
 
         google.generativeai does not return model, so we return None
         """
-        return None
+        return self.chunk.model_version
 
     @property
     def id(self) -> str | None:
@@ -73,22 +76,27 @@ class GoogleCallResponseChunk(
         return None
 
     @property
-    def usage(self) -> None:
+    def usage(self) -> GenerateContentResponseUsageMetadata | None:
         """Returns the usage of the chat completion.
 
         google.generativeai does not have Usage, so we return None
         """
-        return None
+        return self.chunk.usage_metadata
 
     @property
-    def input_tokens(self) -> None:
+    def input_tokens(self) -> int | None:
         """Returns the number of input tokens."""
-        return None
+        return self.usage.prompt_token_count if self.usage else None
 
     @property
-    def output_tokens(self) -> None:
+    def cached_tokens(self) -> int | None:
+        """Returns the number of cached tokens."""
+        return self.usage.cached_content_token_count if self.usage else None
+
+    @property
+    def output_tokens(self) -> int | None:
         """Returns the number of output tokens."""
-        return None
+        return self.usage.candidates_token_count if self.usage else None
 
     @property
     def common_finish_reasons(self) -> list[FinishReason] | None:

mirascope/core/google/stream.py

@@ -14,6 +14,7 @@ from google.genai.types import (
     FinishReason,
     FunctionCall,
     GenerateContentResponse,
+    GenerateContentResponseUsageMetadata,
     PartDict,
     Tool,
 )
@@ -66,7 +67,9 @@ class GoogleStream(
     @property
     def cost(self) -> float | None:
         """Returns the cost of the call."""
-        return calculate_cost(self.input_tokens, self.output_tokens, self.model)
+        return calculate_cost(
+            self.input_tokens, self.cached_tokens, self.output_tokens, self.model
+        )
 
     def _construct_message_param(
         self, tool_calls: list[FunctionCall] | None = None, content: str | None = None
@@ -98,6 +101,15 @@ class GoogleStream(
             raise ValueError(
                 "No stream response, check if the stream has been consumed."
             )
+        candidates_token_count = (
+            int(self.output_tokens) if self.output_tokens is not None else None
+        )
+        prompt_token_count = (
+            int(self.input_tokens) if self.input_tokens is not None else None
+        )
+        total_token_count = int(candidates_token_count or 0) + int(
+            prompt_token_count or 0
+        )
         response = GenerateContentResponse(
             candidates=[
                 Candidate(
@@ -109,7 +121,13 @@ class GoogleStream(
                         parts=self.message_param["parts"],  # pyright: ignore [reportTypedDictNotRequiredAccess, reportArgumentType]
                     ),
                 )
-            ]
+            ],
+            model_version=self.model,
+            usage_metadata=GenerateContentResponseUsageMetadata(
+                candidates_token_count=candidates_token_count,
+                prompt_token_count=prompt_token_count,
+                total_token_count=total_token_count,
+            ),
         )
 
         return GoogleCallResponse(

mirascope/core/groq/__init__.py

@@ -18,13 +18,13 @@ GroqMessageParam: TypeAlias = ChatCompletionMessageParam | BaseMessageParam
 
 __all__ = [
     "AsyncGroqDynamicConfig",
-    "call",
-    "GroqDynamicConfig",
     "GroqCallParams",
     "GroqCallResponse",
     "GroqCallResponseChunk",
+    "GroqDynamicConfig",
     "GroqMessageParam",
     "GroqStream",
     "GroqTool",
+    "call",
     "groq_call",
 ]

mirascope/core/groq/_utils/_calculate_cost.py

@@ -3,6 +3,7 @@
 
 def calculate_cost(
     input_tokens: int | float | None,
+    cached_tokens: int | float | None,
     output_tokens: int | float | None,
     model: str = "mixtral-8x7b-32768",
 ) -> float | None:
@@ -10,17 +11,17 @@ def calculate_cost(
 
     https://wow.groq.com/
 
-    Model                                  Input               Output
-    llama-3.1-405b-reasoning               N/A                 N/A
-    llama-3.1-70b-versatile                N/A                 N/A
-    llama-3.1-8b-instant                   N/A                 N/A
-    llama3-groq-70b-8192-tool-use-preview  $0.89 / 1M tokens   $0.89 / 1M tokens
-    llama3-groq-8b-8192-tool-use-preview   $0.19 / 1M tokens   $0.19 / 1M tokens
-    llama3-70b-8192                        $0.59 / 1M tokens   $0.79 / 1M tokens
-    llama3-8b-8192                         $0.05 / 1M tokens   $0.08 / 1M tokens
-    mixtral-8x7b-32768                     $0.27 / 1M tokens   $0.27 / 1M tokens
-    gemma-7b-it                            $0.07 / 1M tokens   $0.07 / 1M tokens
-    gemma2-9b-it                           $0.20 / 1M tokens   $0.20 / 1M tokens
+    Model                                  Input                Cached     Output
+    llama-3.1-405b-reasoning               N/A                             N/A
+    llama-3.1-70b-versatile                N/A                             N/A
+    llama-3.1-8b-instant                   N/A                             N/A
+    llama3-groq-70b-8192-tool-use-preview  $0.89 / 1M tokens               $0.89 / 1M tokens
+    llama3-groq-8b-8192-tool-use-preview   $0.19 / 1M tokens               $0.19 / 1M tokens
+    llama3-70b-8192                        $0.59 / 1M tokens               $0.79 / 1M tokens
+    llama3-8b-8192                         $0.05 / 1M tokens               $0.08 / 1M tokens
+    mixtral-8x7b-32768                     $0.27 / 1M tokens               $0.27 / 1M tokens
+    gemma-7b-it                            $0.07 / 1M tokens               $0.07 / 1M tokens
+    gemma2-9b-it                           $0.20 / 1M tokens               $0.20 / 1M tokens
     """
     pricing = {
         "llama3-groq-70b-8192-tool-use-preview": {

mirascope/core/groq/call_response.py

@@ -98,6 +98,12 @@ class GroqCallResponse(
         """Returns the number of input tokens."""
         return self.usage.prompt_tokens if self.usage else None
 
+    @computed_field
+    @property
+    def cached_tokens(self) -> int | None:
+        """Returns the number of cached tokens."""
+        return 0
+
     @computed_field
     @property
     def output_tokens(self) -> int | None:
@@ -108,7 +114,9 @@ class GroqCallResponse(
     @property
     def cost(self) -> float | None:
         """Returns the cost of the call."""
-        return calculate_cost(self.input_tokens, self.output_tokens, self.model)
+        return calculate_cost(
+            self.input_tokens, self.cached_tokens, self.output_tokens, self.model
+        )
 
     @computed_field
     @cached_property

mirascope/core/groq/call_response_chunk.py

@@ -81,6 +81,11 @@ class GroqCallResponseChunk(BaseCallResponseChunk[ChatCompletionChunk, FinishRea
             return self.usage.prompt_tokens
         return None
 
+    @property
+    def cached_tokens(self) -> int | None:
+        """Returns the number of cached tokens."""
+        return 0
+
     @property
     def output_tokens(self) -> int | None:
         """Returns the number of output tokens."""

mirascope/core/groq/stream.py

@@ -66,7 +66,9 @@ class GroqStream(
     @property
     def cost(self) -> float | None:
         """Returns the cost of the call."""
-        return calculate_cost(self.input_tokens, self.output_tokens, self.model)
+        return calculate_cost(
+            self.input_tokens, self.cached_tokens, self.output_tokens, self.model
+        )
 
     def _construct_message_param(
         self,

mirascope/core/litellm/__init__.py

@@ -16,7 +16,6 @@ LiteLLMMessageParam: TypeAlias = OpenAIMessageParam
 
 __all__ = [
     "AsyncLiteLLMDynamicConfig",
-    "call",
     "LiteLLMCallParams",
     "LiteLLMCallResponse",
     "LiteLLMCallResponseChunk",
@@ -24,5 +23,6 @@ __all__ = [
     "LiteLLMMessageParam",
     "LiteLLMStream",
     "LiteLLMTool",
+    "call",
     "litellm_call",
 ]

mirascope/core/litellm/_utils/_setup_call.py

@@ -1,7 +1,7 @@
 """This module contains the setup_call function for OpenAI tools."""
 
 from collections.abc import Awaitable, Callable
-from typing import Any, cast, overload
+from typing import Any, TypeAlias, cast, overload
 
 from litellm import acompletion, completion
 from openai import OpenAI
@@ -21,12 +21,16 @@ from ...openai import (
 from ...openai._call_kwargs import OpenAICallKwargs
 from ...openai._utils import setup_call as setup_call_openai
 
+# Note: MyPy doesn't like `client: ...` so we use these aliases instead.
+_AsyncLiteLLMClient: TypeAlias = Any
+_SyncLiteLLMClient: TypeAlias = Any
+
 
 @overload
 def setup_call(
     *,
     model: str,
-    client: ...,
+    client: _AsyncLiteLLMClient | None,
     fn: Callable[..., Awaitable[AsyncOpenAIDynamicConfig]],
     fn_args: dict[str, Any],
     dynamic_config: AsyncOpenAIDynamicConfig,
@@ -48,7 +52,7 @@ def setup_call(
 def setup_call(
     *,
     model: str,
-    client: ...,
+    client: _SyncLiteLLMClient | None,
     fn: Callable[..., OpenAIDynamicConfig],
     fn_args: dict[str, Any],
     dynamic_config: OpenAIDynamicConfig,

mirascope/core/mistral/__init__.py

@@ -24,13 +24,13 @@ MistralMessageParam: TypeAlias = (
 )
 
 __all__ = [
-    "call",
-    "MistralDynamicConfig",
     "MistralCallParams",
     "MistralCallResponse",
     "MistralCallResponseChunk",
+    "MistralDynamicConfig",
     "MistralMessageParam",
     "MistralStream",
     "MistralTool",
+    "call",
     "mistral_call",
 ]

mirascope/core/mistral/_utils/_calculate_cost.py

@@ -3,6 +3,7 @@
 
 def calculate_cost(
     input_tokens: int | float | None,
+    cached_tokens: int | float | None,
     output_tokens: int | float | None,
     model: str = "open-mistral-7b",
 ) -> float | None:
@@ -10,15 +11,15 @@ def calculate_cost(
 
     https://mistral.ai/technology/#pricing
 
-    Model                     Input               Output
-    open-mistral-nemo         $0.3/1M tokens	  $0.3/1M tokens
-    mistral-large-latest      $3/1M tokens	      $9/1M tokens
-    codestral-2405            $1/1M tokens	      $3/1M tokens
-    open-mistral-7b	          $0.25/1M tokens	  $0.25/1M tokens
-    open-mixtral-8x7b	      $0.7/1M tokens	  $0.7/1M tokens
-    open-mixtral-8x22b	      $2/1M tokens	      $6/1M tokens
-    mistral-small-latest	  $2/1M tokens	      $6/1M tokens
-    mistral-medium-latest     $2.75/1M tokens	  $8.1/1M tokens
+    Model                     Input               Cached     Output
+    open-mistral-nemo         $0.3/1M tokens	             $0.3/1M tokens
+    mistral-large-latest      $3/1M tokens	                 $9/1M tokens
+    codestral-2405            $1/1M tokens	                 $3/1M tokens
+    open-mistral-7b	          $0.25/1M tokens	             $0.25/1M tokens
+    open-mixtral-8x7b	      $0.7/1M tokens	             $0.7/1M tokens
+    open-mixtral-8x22b	      $2/1M tokens	                 $6/1M tokens
+    mistral-small-latest	  $2/1M tokens	                 $6/1M tokens
+    mistral-medium-latest     $2.75/1M tokens	             $8.1/1M tokens
     """
     pricing = {
         "open-mistral-nemo": {"prompt": 0.000_000_3, "completion": 0.000_000_3},

mirascope/core/mistral/call_response.py

@@ -107,6 +107,12 @@ class MistralCallResponse(
         """Returns the number of input tokens."""
         return self.usage.prompt_tokens
 
+    @computed_field
+    @property
+    def cached_tokens(self) -> int:
+        """Returns the number of cached tokens."""
+        return 0
+
     @computed_field
     @property
     def output_tokens(self) -> int | None:
@@ -117,7 +123,9 @@ class MistralCallResponse(
     @property
     def cost(self) -> float | None:
         """Returns the cost of the call."""
-        return calculate_cost(self.input_tokens, self.output_tokens, self.model)
+        return calculate_cost(
+            self.input_tokens, self.cached_tokens, self.output_tokens, self.model
+        )
 
     @computed_field
     @cached_property

mirascope/core/mistral/call_response_chunk.py

@@ -80,6 +80,11 @@ class MistralCallResponseChunk(BaseCallResponseChunk[CompletionChunk, FinishReas
             return self.usage.prompt_tokens
         return None
 
+    @property
+    def cached_tokens(self) -> int:
+        """Returns the number of cached tokens."""
+        return 0
+
     @property
     def output_tokens(self) -> int | None:
         """Returns the number of output tokens."""

mirascope/core/mistral/stream.py

@@ -65,7 +65,9 @@ class MistralStream(
     @property
     def cost(self) -> float | None:
         """Returns the cost of the call."""
-        return calculate_cost(self.input_tokens, self.output_tokens, self.model)
+        return calculate_cost(
+            self.input_tokens, self.cached_tokens, self.output_tokens, self.model
+        )
 
     def _construct_message_param(
         self, tool_calls: list | None = None, content: str | None = None

mirascope/core/openai/__init__.py

@@ -18,14 +18,14 @@ OpenAIMessageParam: TypeAlias = ChatCompletionMessageParam | BaseMessageParam
 
 __all__ = [
     "AsyncOpenAIDynamicConfig",
-    "call",
-    "OpenAIDynamicConfig",
     "OpenAICallParams",
     "OpenAICallResponse",
     "OpenAICallResponseChunk",
+    "OpenAIDynamicConfig",
     "OpenAIMessageParam",
     "OpenAIStream",
     "OpenAITool",
     "OpenAIToolConfig",
+    "call",
     "openai_call",
 ]