Repository: Mirascope/mirascope

Issue:
Add provider-agnostic usage abstraction
### Description

The docs state [that all BaseCallResponse objects have a usage property](https://mirascope.com/learn/calls/#common-response-properties-and-methods) for retrieving information about token usage. While this is true, the usage object returned is provider-specific without any shared properties across providers. Which means users who want to get the token count and are using multiple providers need code like this:

```py
def usage_to_tokens(usage: OpenAIUsage | AnthropicUsage) -> int:
    if isinstance(usage, OpenAIUsage):
        return usage.total_tokens
    elif isinstance(usage, AnthropicUsage):
        return usage.input_tokens + usage.output_tokens
    raise NotImplementedError()
```

Especially as users are guided to provider-agnostic calls by default (#811), it seems desirable to have provider-agnostic usage. I think the cleanest API would be to have usage return an agnostic usage class with a property for retrieving the provider-specific usage; however, that would break the existing API. An alternative approach would be adding a new field, but that seems kind of ugly. Or, we could attach provider-agnostic fields onto the provider-agnostic usages, but that also seems a bit frankenstein.

Discussion:
@willbakst: For this, I think adding a provider-agnostic `Usage` would be best. I would not consider adding that "breaking" since we generally would expect users to use the `input_tokens` / `output_tokens` / `cost` properties since there is no provider-agnostic `Usage` but those fields are provider-agnostic already.

We would implement this like we do for `finish_reason` where each provider would implement a `common_usage` property and `CallResponse` would use that to implement the `usage` property.

We would also need to update `input_tokens` and `output_tokens` to use the `Usage` type since the underlying implementations rely on the provider-specific object and would otherwise break.@willbakst: Honestly I'm going to update this to a "bug" since I think that's more fitting than an enhancement.@willbakst: Fix is released in `v1.18.4`!

Solution patch:
diff --git a/mirascope/core/base/__init__.py b/mirascope/core/base/__init__.py
index 62f0ffbb..704f9150 100644
--- a/mirascope/core/base/__init__.py
+++ b/mirascope/core/base/__init__.py
@@ -30,7 +30,7 @@ from .stream import BaseStream
 from .structured_stream import BaseStructuredStream
 from .tool import BaseTool, GenerateJsonSchemaNoTitles, ToolConfig
 from .toolkit import BaseToolKit, toolkit_tool
-from .types import AudioSegment, JsonableType
+from .types import AudioSegment, JsonableType, Usage
 
 __all__ = [
     "AudioPart",
@@ -69,6 +69,7 @@ __all__ = [
     "ToolResultPart",
     "toolkit_tool",
     "transform_tool_outputs",
+    "Usage",
     "_partial",
     "_utils",
 ]
diff --git a/mirascope/core/base/_utils/_get_common_usage.py b/mirascope/core/base/_utils/_get_common_usage.py
index f531dbdd..e0c1de1b 100644
--- a/mirascope/core/base/_utils/_get_common_usage.py
+++ b/mirascope/core/base/_utils/_get_common_usage.py
@@ -10,7 +10,7 @@ def get_common_usage(
     input_tokens = int(input_tokens or 0)
     output_tokens = int(output_tokens or 0)
     return Usage(
-        prompt_tokens=input_tokens,
-        completion_tokens=output_tokens,
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
         total_tokens=input_tokens + output_tokens,
     )
diff --git a/mirascope/core/base/types.py b/mirascope/core/base/types.py
index 440fca21..1588951c 100644
--- a/mirascope/core/base/types.py
+++ b/mirascope/core/base/types.py
@@ -43,12 +43,12 @@ 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 = 0
     """Number of tokens in the prompt."""
 
+    output_tokens: int = 0
+    """Number of tokens in the generated output."""
+
     total_tokens: int = 0
     """Total number of tokens used in the request (prompt + completion)."""
 
diff --git a/mirascope/llm/call_response.py b/mirascope/llm/call_response.py
index 04efa738..cedf9350 100644
--- a/mirascope/llm/call_response.py
+++ b/mirascope/llm/call_response.py
@@ -13,6 +13,7 @@ from mirascope.core.base import (
     BaseCallResponse,
     BaseMessageParam,
     BaseTool,
+    Usage,
     transform_tool_outputs,
 )
 from mirascope.core.base.message_param import ToolResultPart
@@ -22,8 +23,6 @@ from mirascope.llm.tool import Tool
 
 _ResponseT = TypeVar("_ResponseT")
 
-_ToolMessageParamT = TypeVar("_ToolMessageParamT")
-
 
 class CallResponse(
     BaseCallResponse[
@@ -64,7 +63,13 @@ class CallResponse(
     def __getattribute__(self, name: str) -> Any:  # noqa: ANN401
         special_names = {
             "_response",
+            "finish_reasons",
+            "usage",
+            "message_param",
             "user_message_param",
+            "tools",
+            "tool",
+            "tool_message_params",
             "__dict__",
             "__class__",
             "model_fields",
@@ -96,6 +101,11 @@ class CallResponse(
     def finish_reasons(self) -> list[FinishReason] | None:  # pyright: ignore [reportIncompatibleMethodOverride]
         return self._response.common_finish_reasons
 
+    @property
+    def usage(self) -> Usage | None:
+        """Returns the usage of the chat completion."""
+        return self._response.common_usage
+
     @computed_field
     @cached_property
     def message_param(self) -> BaseMessageParam:


Test patch:
diff --git a/tests/llm/test_call_response.py b/tests/llm/test_call_response.py
index e676b678..9e96108b 100644
--- a/tests/llm/test_call_response.py
+++ b/tests/llm/test_call_response.py
@@ -13,6 +13,7 @@ from mirascope.core.base import (
     BaseTool,
     Metadata,
     ToolResultPart,
+    Usage,
 )
 from mirascope.core.base.types import FinishReason
 from mirascope.llm.call_response import CallResponse
@@ -104,7 +105,8 @@ class DummyProviderCallResponse(
         return BaseMessageParam(role="user", content="common_user_message")
 
     @property
-    def common_usage(self): ...
+    def common_usage(self) -> Usage | None:
+        return Usage(input_tokens=1, output_tokens=1, total_tokens=2)
 
     def common_construct_call_response(self): ...
 
@@ -141,6 +143,9 @@ def test_call_response(dummy_call_response_instance):
     assert dummy_call_response_instance.tool is not None
     assert str(dummy_call_response_instance) == "dummy_content"
     assert dummy_call_response_instance._response.common_finish_reasons == ["finish"]
+    assert dummy_call_response_instance.usage == Usage(
+        input_tokens=1, output_tokens=1, total_tokens=2
+    )
 
 
 def test_call_response_attribute_fallback_on_instance(dummy_call_response_instance):
diff --git a/tests/llm/test_call_response_chunk.py b/tests/llm/test_call_response_chunk.py
index 0c85f5b9..0afbb222 100644
--- a/tests/llm/test_call_response_chunk.py
+++ b/tests/llm/test_call_response_chunk.py
@@ -40,8 +40,8 @@ class DummyCallResponseChunk(BaseCallResponseChunk[Any, str]):
         input_tokens = int(self.input_tokens or 0)
         output_tokens = int(self.output_tokens or 0)
         return Usage(
-            prompt_tokens=input_tokens,
-            completion_tokens=output_tokens,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
             total_tokens=input_tokens + output_tokens,
         )