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

mirascope/core/base/_utils/__init__.py

@@ -43,21 +43,26 @@ from ._setup_call import setup_call
 from ._setup_extract_tool import setup_extract_tool
 
 __all__ = [
+    "DEFAULT_TOOL_DOCSTRING",
     "AsyncCreateFn",
-    "SameSyncAndAsyncClientSetupCall",
     "BaseType",
     "CalculateCost",
+    "CreateFn",
+    "GetJsonOutput",
+    "HandleStream",
+    "HandleStreamAsync",
+    "LLMFunctionDecorator",
+    "MessagesDecorator",
+    "SameSyncAndAsyncClientSetupCall",
+    "SetupCall",
     "convert_base_model_to_base_tool",
     "convert_base_type_to_base_tool",
     "convert_function_to_base_tool",
-    "CreateFn",
-    "DEFAULT_TOOL_DOCSTRING",
     "extract_tool_return",
     "fn_is_async",
     "format_template",
-    "GetJsonOutput",
-    "get_audio_type",
     "get_async_create_fn",
+    "get_audio_type",
     "get_common_usage",
     "get_create_fn",
     "get_document_type",
@@ -70,18 +75,13 @@ __all__ = [
     "get_template_values",
     "get_template_variables",
     "get_unsupported_tool_config_keys",
-    "HandleStream",
-    "HandleStreamAsync",
     "is_base_type",
     "is_prompt_template",
     "json_mode_content",
-    "LLMFunctionDecorator",
-    "MessagesDecorator",
     "messages_decorator",
     "parse_content_template",
     "parse_prompt_messages",
     "pil_image_to_bytes",
-    "SetupCall",
     "setup_call",
     "setup_extract_tool",
 ]

mirascope/core/base/_utils/_get_common_usage.py

@@ -2,15 +2,19 @@ from mirascope.core.base.types import Usage
 
 
 def get_common_usage(
-    input_tokens: int | float | None, output_tokens: int | float | None
+    input_tokens: int | float | None,
+    cached_tokens: int | float | None,
+    output_tokens: int | float | None,
 ) -> Usage | None:
     """Get common usage from input and output tokens."""
-    if input_tokens is None and output_tokens is None:
+    if input_tokens is None and cached_tokens is None and output_tokens is None:
         return None
     input_tokens = int(input_tokens or 0)
+    cached_tokens = int(cached_tokens or 0)
     output_tokens = int(output_tokens or 0)
     return Usage(
-        prompt_tokens=input_tokens,
-        completion_tokens=output_tokens,
+        input_tokens=input_tokens,
+        cached_tokens=cached_tokens,
+        output_tokens=output_tokens,
         total_tokens=input_tokens + output_tokens,
     )

mirascope/core/base/_utils/_get_create_fn_or_async_create_fn.py

@@ -53,7 +53,7 @@ def get_async_create_fn(
     async_generator_func: _AsyncGeneratorFunc[_StreamedResponse] | None = None,
 ) -> AsyncCreateFn[_NonStreamedResponse, _StreamedResponse]:
     @overload
-    def create_or_stream(
+    def create_or_stream(  # pyright: ignore[reportOverlappingOverload]
         *,
         stream: Literal[True] | StreamConfig = True,
         **kwargs: Any,  # noqa: ANN401
@@ -100,7 +100,7 @@ def get_create_fn(
     sync_generator_func: _SyncGeneratorFunc[_StreamedResponse] | None = None,
 ) -> CreateFn[_NonStreamedResponse, _StreamedResponse]:
     @overload
-    def create_or_stream(
+    def create_or_stream(  # pyright: ignore[reportOverlappingOverload]
         *,
         stream: Literal[True] | StreamConfig = True,
         **kwargs: Any,  # noqa: ANN401

mirascope/core/base/_utils/_protocols.py

@@ -129,7 +129,7 @@ class LLMFunctionDecorator(
 
 class AsyncCreateFn(Protocol[_ResponseT, _ResponseChunkT]):
     @overload
-    def __call__(
+    def __call__(  # pyright: ignore[reportOverlappingOverload]
         self,
         *,
         stream: Literal[False] = False,
@@ -153,7 +153,7 @@ class AsyncCreateFn(Protocol[_ResponseT, _ResponseChunkT]):
 
 class CreateFn(Protocol[_ResponseT, _ResponseChunkT]):
     @overload
-    def __call__(
+    def __call__(  # pyright: ignore[reportOverlappingOverload]
         self,
         *,
         stream: Literal[False] = False,
@@ -371,6 +371,7 @@ class CalculateCost(Protocol):
     def __call__(
         self,
         input_tokens: int | float | None,
+        cached_tokens: int | float | None,
         output_tokens: int | float | None,
         model: str,
     ) -> float | None: ...  # pragma: no cover
@@ -390,7 +391,7 @@ class CallDecorator(
     ],
 ):
     @overload
-    def __call__(
+    def __call__(  # pyright: ignore[reportOverlappingOverload]
         self,
         model: str,
         *,
@@ -409,7 +410,7 @@ class CallDecorator(
     ]: ...
 
     @overload
-    def __call__(
+    def __call__(  # pyright: ignore[reportOverlappingOverload]
         self,
         model: str,
         *,
@@ -437,7 +438,7 @@ class CallDecorator(
     ) -> SyncLLMFunctionDecorator[_BaseDynamicConfigT, _BaseCallResponseT]: ...
 
     @overload
-    def __call__(
+    def __call__(  # pyright: ignore[reportOverlappingOverload]
         self,
         model: str,
         *,
@@ -576,7 +577,7 @@ class CallDecorator(
     ) -> NoReturn: ...
 
     @overload
-    def __call__(
+    def __call__(  # pyright: ignore[reportOverlappingOverload]
         self,
         model: str,
         *,
@@ -664,7 +665,7 @@ class CallDecorator(
     ) -> SyncLLMFunctionDecorator[_BaseDynamicConfigT, _ParsedOutputT]: ...
 
     @overload
-    def __call__(
+    def __call__(  # pyright: ignore[reportOverlappingOverload]
         self,
         model: str,
         *,
@@ -774,7 +775,7 @@ class CallDecorator(
     ####
 
     @overload
-    def __call__(
+    def __call__(  # pyright: ignore[reportOverlappingOverload]
         self,
         model: str,
         *,

mirascope/core/base/call_response.py

@@ -7,7 +7,7 @@ import json
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from functools import cached_property, wraps
-from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeAlias, TypeVar
+from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar
 
 from pydantic import (
     BaseModel,
@@ -24,7 +24,7 @@ from .call_params import BaseCallParams
 from .dynamic_config import BaseDynamicConfig
 from .metadata import Metadata
 from .tool import BaseTool
-from .types import FinishReason, Usage
+from .types import FinishReason, JsonableType, Usage
 
 if TYPE_CHECKING:
     from ...llm.tool import Tool
@@ -35,36 +35,26 @@ _BaseToolT = TypeVar("_BaseToolT", bound=BaseTool)
 _ToolSchemaT = TypeVar("_ToolSchemaT")
 _BaseDynamicConfigT = TypeVar("_BaseDynamicConfigT", bound=BaseDynamicConfig)
 _MessageParamT = TypeVar("_MessageParamT", bound=Any)
+_ToolMessageParamT = TypeVar("_ToolMessageParamT", bound=Any)
 _CallParamsT = TypeVar("_CallParamsT", bound=BaseCallParams)
 _UserMessageParamT = TypeVar("_UserMessageParamT", bound=Any)
 _BaseCallResponseT = TypeVar("_BaseCallResponseT", bound="BaseCallResponse")
 
 
-JsonableType: TypeAlias = (
-    str
-    | int
-    | float
-    | bool
-    | bytes
-    | list["JsonableType"]
-    | set["JsonableType"]
-    | tuple["JsonableType", ...]
-    | dict[str, "JsonableType"]
-    | BaseModel
-)
-
-
 def transform_tool_outputs(
-    fn: Callable[[type[_BaseCallResponseT], list[tuple[_BaseToolT, str]]], list[Any]],
+    fn: Callable[
+        [type[_BaseCallResponseT], list[tuple[_BaseToolT, str]]],
+        list[_ToolMessageParamT],
+    ],
 ) -> Callable[
     [type[_BaseCallResponseT], list[tuple[_BaseToolT, JsonableType]]],
-    list[Any],
+    list[_ToolMessageParamT],
 ]:
     @wraps(fn)
     def wrapper(
         cls: type[_BaseCallResponseT],
         tools_and_outputs: list[tuple[_BaseToolT, JsonableType]],
-    ) -> list[Any]:
+    ) -> list[_ToolMessageParamT]:
         def recursive_serializer(value: JsonableType) -> BaseType:
             if isinstance(value, str):
                 return value
@@ -212,6 +202,16 @@ class BaseCallResponse(
         """
         ...
 
+    @computed_field
+    @property
+    @abstractmethod
+    def cached_tokens(self) -> int | float | None:
+        """Should return the number of cached tokens.
+
+        If there is no cached_tokens, this method must return None.
+        """
+        ...
+
     @computed_field
     @property
     @abstractmethod
@@ -239,14 +239,12 @@ class BaseCallResponse(
         """Returns the assistant's response as a message parameter."""
         ...
 
-    @computed_field
     @cached_property
     @abstractmethod
     def tools(self) -> list[_BaseToolT] | None:
         """Returns the tools for the 0th choice message."""
         ...
 
-    @computed_field
     @cached_property
     @abstractmethod
     def tool(self) -> _BaseToolT | None:
@@ -297,4 +295,6 @@ class BaseCallResponse(
     @property
     def common_usage(self) -> Usage | None:
         """Provider-agnostic usage info."""
-        return get_common_usage(self.input_tokens, self.output_tokens)
+        return get_common_usage(
+            self.input_tokens, self.cached_tokens, self.output_tokens
+        )

mirascope/core/base/call_response_chunk.py

@@ -84,6 +84,15 @@ class BaseCallResponseChunk(BaseModel, Generic[_ChunkT, _FinishReasonT], ABC):
         """
         ...
 
+    @property
+    @abstractmethod
+    def cached_tokens(self) -> int | float | None:
+        """Should return the number of cached tokens.
+
+        If there is no cached_tokens, this method must return None.
+        """
+        ...
+
     @property
     @abstractmethod
     def output_tokens(self) -> int | float | None:
@@ -102,4 +111,6 @@ class BaseCallResponseChunk(BaseModel, Generic[_ChunkT, _FinishReasonT], ABC):
     @property
     def common_usage(self) -> Usage | None:
         """Provider-agnostic usage info."""
-        return get_common_usage(self.input_tokens, self.output_tokens)
+        return get_common_usage(
+            self.input_tokens, self.cached_tokens, self.output_tokens
+        )

mirascope/core/base/stream.py

@@ -92,6 +92,7 @@ class BaseStream(
     user_message_param: _UserMessageParamT | None = None
     message_param: _AssistantMessageParamT
     input_tokens: int | float | None = None
+    cached_tokens: int | float | None = None
     output_tokens: int | float | None = None
     id: str | None = None
     finish_reasons: list[_FinishReason] | None = None
@@ -138,9 +139,9 @@ class BaseStream(
         self,
     ) -> Generator[tuple[_BaseCallResponseChunkT, _BaseToolT | None], None, None]:
         """Iterator over the stream and stores useful information."""
-        assert isinstance(
-            self.stream, Generator
-        ), "Stream must be a generator for __iter__"
+        assert isinstance(self.stream, Generator), (
+            "Stream must be a generator for __iter__"
+        )
         self.content, tool_calls = "", []
         self.start_time = datetime.datetime.now().timestamp() * 1000
         for chunk, tool in self.stream:
@@ -161,12 +162,12 @@ class BaseStream(
         """Iterates over the stream and stores useful information."""
         self.content = ""
 
-        async def generator() -> (
-            AsyncGenerator[tuple[_BaseCallResponseChunkT, _BaseToolT | None], None]
-        ):
-            assert isinstance(
-                self.stream, AsyncGenerator
-            ), "Stream must be an async generator for __aiter__"
+        async def generator() -> AsyncGenerator[
+            tuple[_BaseCallResponseChunkT, _BaseToolT | None], None
+        ]:
+            assert isinstance(self.stream, AsyncGenerator), (
+                "Stream must be an async generator for __aiter__"
+            )
             tool_calls = []
             async for chunk, tool in self.stream:
                 self._update_properties(chunk)
@@ -190,6 +191,12 @@ class BaseStream(
                 if not self.input_tokens
                 else self.input_tokens + chunk.input_tokens
             )
+        if chunk.cached_tokens is not None:
+            self.cached_tokens = (
+                chunk.cached_tokens
+                if not self.cached_tokens
+                else self.cached_tokens + chunk.cached_tokens
+            )
         if chunk.output_tokens is not None:
             self.output_tokens = (
                 chunk.output_tokens
@@ -373,11 +380,9 @@ def stream_factory(  # noqa: ANN201
                     stream=True,
                 )
 
-                async def generator() -> (
-                    AsyncGenerator[
-                        tuple[_BaseCallResponseChunkT, _BaseToolT | None], None
-                    ]
-                ):
+                async def generator() -> AsyncGenerator[
+                    tuple[_BaseCallResponseChunkT, _BaseToolT | None], None
+                ]:
                     async for chunk, tool in handle_stream_async(
                         await create(stream=True, **call_kwargs),
                         tool_types,
@@ -422,13 +427,11 @@ def stream_factory(  # noqa: ANN201
                     stream=True,
                 )
 
-                def generator() -> (
-                    Generator[
-                        tuple[_BaseCallResponseChunkT, _BaseToolT | None],
-                        None,
-                        None,
-                    ]
-                ):
+                def generator() -> Generator[
+                    tuple[_BaseCallResponseChunkT, _BaseToolT | None],
+                    None,
+                    None,
+                ]:
                     yield from handle_stream(
                         create(stream=True, **call_kwargs),
                         tool_types,

mirascope/core/base/tool.py

@@ -36,17 +36,19 @@ class ToolConfig(TypedDict, total=False):
 class GenerateJsonSchemaNoTitles(GenerateJsonSchema):
     _openai_strict: ClassVar[bool] = False
 
-    def _remove_title(self, obj: Any) -> Any:  # noqa: ANN401
+    def _remove_title(self, key: str | None, obj: Any) -> Any:  # noqa: ANN401
         if isinstance(obj, dict):
             if self._openai_strict and "type" in obj and obj["type"] == "object":
                 obj["additionalProperties"] = False
             if "type" in obj or "$ref" in obj or "properties" in obj:
-                obj.pop("title", None)
+                title = obj.pop("title", None)
+                if key and title and key.lower() != title.lower():
+                    obj["title"] = title
 
             for key, value in list(obj.items()):
-                obj[key] = self._remove_title(value)
+                obj[key] = self._remove_title(key, value)
         elif isinstance(obj, list):
-            return [self._remove_title(item) for item in obj]
+            return [self._remove_title(None, item) for item in obj]
 
         return obj
 
@@ -56,7 +58,7 @@ class GenerateJsonSchemaNoTitles(GenerateJsonSchema):
         json_schema = super().generate(schema, mode=mode)
         json_schema.pop("title", None)
         json_schema.pop("description", None)
-        json_schema = self._remove_title(json_schema)
+        json_schema = self._remove_title(None, json_schema)
         return json_schema
 
 

mirascope/core/base/types.py

@@ -43,11 +43,28 @@ FinishReason: TypeAlias = Literal["stop", "length", "tool_calls", "content_filte
 
 
 class Usage(BaseModel):
-    completion_tokens: int = 0
-    """Number of tokens in the generated completion."""
-
-    prompt_tokens: int = 0
+    input_tokens: int
     """Number of tokens in the prompt."""
 
-    total_tokens: int = 0
+    cached_tokens: int
+    """Number of tokens used that were previously cached (and thus cheaper)."""
+
+    output_tokens: int
+    """Number of tokens in the generated output."""
+
+    total_tokens: int
     """Total number of tokens used in the request (prompt + completion)."""
+
+
+JsonableType: TypeAlias = (
+    str
+    | int
+    | float
+    | bool
+    | bytes
+    | list["JsonableType"]
+    | set["JsonableType"]
+    | tuple["JsonableType", ...]
+    | dict[str, "JsonableType"]
+    | BaseModel
+)

mirascope/core/bedrock/__init__.py

@@ -18,9 +18,8 @@ from .tool import BedrockTool, BedrockToolConfig
 BedrockMessageParam: TypeAlias = InternalBedrockMessageParam | BaseMessageParam
 
 __all__ = [
-    "AsyncBedrockDynamicConfig",
     "AssistantMessageTypeDef",
-    "bedrock_call",
+    "AsyncBedrockDynamicConfig",
     "BedrockCallParams",
     "BedrockCallResponse",
     "BedrockCallResponseChunk",
@@ -29,6 +28,7 @@ __all__ = [
     "BedrockStream",
     "BedrockTool",
     "BedrockToolConfig",
-    "call",
     "UserMessageTypeDef",
+    "bedrock_call",
+    "call",
 ]

mirascope/core/bedrock/_utils/_calculate_cost.py

@@ -2,7 +2,10 @@
 
 
 def calculate_cost(
-    input_tokens: int | float | None, output_tokens: int | float | None, model: str
+    input_tokens: int | float | None,
+    cached_tokens: int | float | None,
+    output_tokens: int | float | None,
+    model: str,
 ) -> float | None:
     """Calculate the cost of a completion using the Bedrock API."""
     # NOTE: We are currently investigating a dynamic approach to determine costs

mirascope/core/bedrock/call_response.py

@@ -129,6 +129,11 @@ class BedrockCallResponse(
         """Returns the number of input tokens."""
         return self.usage["inputTokens"] if self.usage else None
 
+    @property
+    def cached_tokens(self) -> int | None:
+        """Returns the number of cached tokens."""
+        return None
+
     @computed_field
     @property
     def output_tokens(self) -> int | None:
@@ -139,7 +144,9 @@ class BedrockCallResponse(
     @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
@@ -150,7 +157,6 @@ class BedrockCallResponse(
             return AssistantMessageTypeDef(role="assistant", content=[])
         return AssistantMessageTypeDef(role="assistant", content=message["content"])
 
-    @computed_field
     @cached_property
     def tools(self) -> list[BedrockTool] | None:
         """Returns any available tool calls as their `BedrockTool` definition.
@@ -180,7 +186,6 @@ class BedrockCallResponse(
 
         return extracted_tools
 
-    @computed_field
     @cached_property
     def tool(self) -> BedrockTool | None:
         """Returns the 0th tool for the 0th choice message.

mirascope/core/bedrock/call_response_chunk.py

@@ -86,6 +86,11 @@ class BedrockCallResponseChunk(
             return self.usage["inputTokens"]
         return None
 
+    @property
+    def cached_tokens(self) -> int | None:
+        """Returns the number of cached tokens."""
+        return None
+
     @property
     def output_tokens(self) -> int | None:
         """Returns the number of output tokens."""

mirascope/core/bedrock/stream.py

@@ -94,7 +94,9 @@ class BedrockStream(
     @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/cohere/__init__.py

@@ -18,13 +18,13 @@ CohereMessageParam: TypeAlias = ChatMessage | ToolResult | BaseMessageParam
 
 __all__ = [
     "AsyncCohereDynamicConfig",
-    "call",
-    "CohereDynamicConfig",
     "CohereCallParams",
     "CohereCallResponse",
     "CohereCallResponseChunk",
+    "CohereDynamicConfig",
     "CohereMessageParam",
     "CohereStream",
     "CohereTool",
+    "call",
     "cohere_call",
 ]

mirascope/core/cohere/_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 = "command-r-plus",
 ) -> float | None:
@@ -10,9 +11,9 @@ def calculate_cost(
 
     https://cohere.com/pricing
 
-    Model              Input               Output
-    command-r          $0.5 / 1M tokens	   $1.5 / 1M tokens
-    command-r-plus     $3 / 1M tokens	   $15 / 1M tokens
+    Model              Input               Cached     Output
+    command-r          $0.5 / 1M tokens	              $1.5 / 1M tokens
+    command-r-plus     $3 / 1M tokens	              $15 / 1M tokens
     """
     pricing = {
         "command-r": {

mirascope/core/cohere/call_response.py

@@ -105,6 +105,12 @@ class CohereCallResponse(
             return self.usage.input_tokens
         return None
 
+    @computed_field
+    @property
+    def cached_tokens(self) -> float | None:
+        """Returns the number of cached tokens."""
+        return None
+
     @computed_field
     @property
     def output_tokens(self) -> float | None:
@@ -117,7 +123,9 @@ class CohereCallResponse(
     @property
     def cost(self) -> float | None:
         """Returns the cost of the response."""
-        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
@@ -129,7 +137,6 @@ class CohereCallResponse(
             role="assistant",  # pyright: ignore [reportCallIssue]
         )
 
-    @computed_field
     @cached_property
     def tools(self) -> list[CohereTool] | None:
         """Returns the tools for the 0th choice message.
@@ -147,7 +154,6 @@ class CohereCallResponse(
                     break
         return extracted_tools
 
-    @computed_field
     @cached_property
     def tool(self) -> CohereTool | None:
         """Returns the 0th tool for the 0th choice message.

mirascope/core/cohere/call_response_chunk.py

@@ -100,6 +100,11 @@ class CohereCallResponseChunk(
             return self.usage.input_tokens
         return None
 
+    @property
+    def cached_tokens(self) -> float | None:
+        """Returns the number of cached tokens."""
+        return None
+
     @property
     def output_tokens(self) -> float | None:
         """Returns the number of output tokens."""

mirascope/core/cohere/stream.py

@@ -62,7 +62,9 @@ class CohereStream(
     @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[ToolCall] | None = None, content: str | None = None

mirascope/core/gemini/__init__.py

@@ -28,13 +28,13 @@ warnings.warn(
 )
 
 __all__ = [
-    "call",
-    "GeminiDynamicConfig",
     "GeminiCallParams",
     "GeminiCallResponse",
     "GeminiCallResponseChunk",
+    "GeminiDynamicConfig",
     "GeminiMessageParam",
     "GeminiStream",
     "GeminiTool",
+    "call",
     "gemini_call",
 ]