tokenator 0.1.14__py3-none-any.whl → 0.1.15__py3-none-any.whl

tokenator/anthropic/client_anthropic.py

@@ -6,7 +6,7 @@ import logging
 from anthropic import Anthropic, AsyncAnthropic
 from anthropic.types import Message, RawMessageStartEvent, RawMessageDeltaEvent
 
-from ..models import Usage, TokenUsageStats
+from ..models import PromptTokenDetails, TokenMetrics, TokenUsageStats
 from ..base_wrapper import BaseWrapper, ResponseType
 from .stream_interceptors import (
     AnthropicAsyncStreamInterceptor,
@@ -28,28 +28,46 @@ class BaseAnthropicWrapper(BaseWrapper):
             if isinstance(response, Message):
                 if not hasattr(response, "usage"):
                     return None
-                usage = Usage(
-                    prompt_tokens=response.usage.input_tokens,
+                usage = TokenMetrics(
+                    prompt_tokens=response.usage.input_tokens
+                    + (getattr(response.usage, "cache_creation_input_tokens", 0) or 0),
                     completion_tokens=response.usage.output_tokens,
                     total_tokens=response.usage.input_tokens
                     + response.usage.output_tokens,
+                    prompt_tokens_details=PromptTokenDetails(
+                        cached_input_tokens=getattr(
+                            response.usage, "cache_read_input_tokens", None
+                        ),
+                        cached_creation_tokens=getattr(
+                            response.usage, "cache_creation_input_tokens", None
+                        ),
+                    ),
                 )
                 return TokenUsageStats(model=response.model, usage=usage)
             elif isinstance(response, dict):
                 usage_dict = response.get("usage")
                 if not usage_dict:
                     return None
-                usage = Usage(
-                    prompt_tokens=usage_dict.get("input_tokens", 0),
+                usage = TokenMetrics(
+                    prompt_tokens=usage_dict.get("input_tokens", 0)
+                    + (getattr(usage_dict, "cache_creation_input_tokens", 0) or 0),
                     completion_tokens=usage_dict.get("output_tokens", 0),
                     total_tokens=usage_dict.get("input_tokens", 0)
                     + usage_dict.get("output_tokens", 0),
+                    prompt_tokens_details=PromptTokenDetails(
+                        cached_input_tokens=getattr(
+                            usage_dict, "cache_read_input_tokens", None
+                        ),
+                        cached_creation_tokens=getattr(
+                            usage_dict, "cache_creation_input_tokens", None
+                        ),
+                    ),
                 )
                 return TokenUsageStats(
                     model=response.get("model", "unknown"), usage=usage
                 )
         except Exception as e:
-            logger.warning("Failed to process usage stats: %s", str(e))
+            logger.warning("Failed to process usage stats: %s", str(e), exc_info=True)
             return None
         return None
 
@@ -74,7 +92,7 @@ def _create_usage_callback(execution_id, log_usage_fn):
             model=chunks[0].message.model
             if isinstance(chunks[0], RawMessageStartEvent)
             else "",
-            usage=Usage(),
+            usage=TokenMetrics(),
         )
 
         for chunk in chunks:

tokenator/base_wrapper.py

@@ -58,9 +58,33 @@ class BaseWrapper:
                 execution_id=execution_id,
                 provider=self.provider,
                 model=token_usage_stats.model,
+                total_cost=0,  # This needs to be calculated based on your rates
                 prompt_tokens=token_usage_stats.usage.prompt_tokens,
                 completion_tokens=token_usage_stats.usage.completion_tokens,
                 total_tokens=token_usage_stats.usage.total_tokens,
+                # Prompt details
+                prompt_cached_input_tokens=token_usage_stats.usage.prompt_tokens_details.cached_input_tokens
+                if token_usage_stats.usage.prompt_tokens_details
+                else None,
+                prompt_cached_creation_tokens=token_usage_stats.usage.prompt_tokens_details.cached_creation_tokens
+                if token_usage_stats.usage.prompt_tokens_details
+                else None,
+                prompt_audio_tokens=token_usage_stats.usage.prompt_tokens_details.audio_tokens
+                if token_usage_stats.usage.prompt_tokens_details
+                else None,
+                # Completion details
+                completion_audio_tokens=token_usage_stats.usage.completion_tokens_details.audio_tokens
+                if token_usage_stats.usage.completion_tokens_details
+                else None,
+                completion_reasoning_tokens=token_usage_stats.usage.completion_tokens_details.reasoning_tokens
+                if token_usage_stats.usage.completion_tokens_details
+                else None,
+                completion_accepted_prediction_tokens=token_usage_stats.usage.completion_tokens_details.accepted_prediction_tokens
+                if token_usage_stats.usage.completion_tokens_details
+                else None,
+                completion_rejected_prediction_tokens=token_usage_stats.usage.completion_tokens_details.rejected_prediction_tokens
+                if token_usage_stats.usage.completion_tokens_details
+                else None,
             )
             session.add(token_usage)
             logger.debug(
@@ -82,11 +106,13 @@ class BaseWrapper:
         if not execution_id:
             execution_id = str(uuid.uuid4())
 
+        logger.debug("Starting token usage logging for execution_id: %s", execution_id)
         session = self.Session()
         try:
             try:
                 self._log_usage_impl(token_usage_stats, session, execution_id)
                 session.commit()
+                logger.debug("Successfully committed token usage for execution_id: %s", execution_id)
             except Exception as e:
                 logger.error("Failed to log token usage: %s", str(e))
                 session.rollback()

tokenator/migrations/versions/f028b8155fed_adding_detailed_input_and_output_token_.py

@@ -0,0 +1,64 @@
+"""Adding detailed input and output token schema
+
+Revision ID: f028b8155fed
+Revises: f6f1f2437513
+Create Date: 2025-01-19 15:41:12.715623
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = "f028b8155fed"
+down_revision: Union[str, None] = "f6f1f2437513"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("token_usage", sa.Column("total_cost", sa.Integer(), nullable=False))
+    op.add_column(
+        "token_usage",
+        sa.Column("prompt_cached_input_tokens", sa.Integer(), nullable=True),
+    )
+    op.add_column(
+        "token_usage",
+        sa.Column("prompt_cached_creation_tokens", sa.Integer(), nullable=True),
+    )
+    op.add_column(
+        "token_usage", sa.Column("prompt_audio_tokens", sa.Integer(), nullable=True)
+    )
+    op.add_column(
+        "token_usage", sa.Column("completion_audio_tokens", sa.Integer(), nullable=True)
+    )
+    op.add_column(
+        "token_usage",
+        sa.Column("completion_reasoning_tokens", sa.Integer(), nullable=True),
+    )
+    op.add_column(
+        "token_usage",
+        sa.Column("completion_accepted_prediction_tokens", sa.Integer(), nullable=True),
+    )
+    op.add_column(
+        "token_usage",
+        sa.Column("completion_rejected_prediction_tokens", sa.Integer(), nullable=True),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("token_usage", "completion_rejected_prediction_tokens")
+    op.drop_column("token_usage", "completion_accepted_prediction_tokens")
+    op.drop_column("token_usage", "completion_reasoning_tokens")
+    op.drop_column("token_usage", "completion_audio_tokens")
+    op.drop_column("token_usage", "prompt_audio_tokens")
+    op.drop_column("token_usage", "prompt_cached_creation_tokens")
+    op.drop_column("token_usage", "prompt_cached_input_tokens")
+    op.drop_column("token_usage", "total_cost")
+    # ### end Alembic commands ###

tokenator/models.py

@@ -1,10 +1,35 @@
 from pydantic import BaseModel, Field
-from typing import List
+from typing import List, Optional
 
 
 class TokenRate(BaseModel):
     prompt: float = Field(..., description="Cost per prompt token")
     completion: float = Field(..., description="Cost per completion token")
+    prompt_audio: Optional[float] = Field(
+        None, description="Cost per audio prompt token"
+    )
+    completion_audio: Optional[float] = Field(
+        None, description="Cost per audio completion token"
+    )
+    prompt_cached_input: Optional[float] = Field(
+        None, description="Cost per cached prompt input token"
+    )
+    prompt_cached_creation: Optional[float] = Field(
+        None, description="Cost per cached prompt creation token"
+    )
+
+
+class PromptTokenDetails(BaseModel):
+    cached_input_tokens: Optional[int] = None
+    cached_creation_tokens: Optional[int] = None
+    audio_tokens: Optional[int] = None
+
+
+class CompletionTokenDetails(BaseModel):
+    reasoning_tokens: Optional[int] = None
+    audio_tokens: Optional[int] = None
+    accepted_prediction_tokens: Optional[int] = None
+    rejected_prediction_tokens: Optional[int] = None
 
 
 class TokenMetrics(BaseModel):
@@ -12,6 +37,8 @@ class TokenMetrics(BaseModel):
     total_tokens: int = Field(default=0, description="Total tokens used")
     prompt_tokens: int = Field(default=0, description="Number of prompt tokens")
     completion_tokens: int = Field(default=0, description="Number of completion tokens")
+    prompt_tokens_details: Optional[PromptTokenDetails] = None
+    completion_tokens_details: Optional[CompletionTokenDetails] = None
 
 
 class ModelUsage(TokenMetrics):
@@ -31,12 +58,6 @@ class TokenUsageReport(TokenMetrics):
     )
 
 
-class Usage(BaseModel):
-    prompt_tokens: int = 0
-    completion_tokens: int = 0
-    total_tokens: int = 0
-
-
 class TokenUsageStats(BaseModel):
     model: str
-    usage: Usage
+    usage: TokenMetrics

tokenator/openai/client_openai.py

@@ -6,7 +6,12 @@ import logging
 from openai import AsyncOpenAI, OpenAI
 from openai.types.chat import ChatCompletion, ChatCompletionChunk
 
-from ..models import Usage, TokenUsageStats
+from ..models import (
+    TokenMetrics,
+    TokenUsageStats,
+    PromptTokenDetails,
+    CompletionTokenDetails,
+)
 from ..base_wrapper import BaseWrapper, ResponseType
 from .stream_interceptors import (
     OpenAIAsyncStreamInterceptor,
@@ -30,18 +35,49 @@ class BaseOpenAIWrapper(BaseWrapper):
             if isinstance(response, ChatCompletion):
                 if response.usage is None:
                     return None
-                usage = Usage(
+                usage = TokenMetrics(
                     prompt_tokens=response.usage.prompt_tokens,
                     completion_tokens=response.usage.completion_tokens,
                     total_tokens=response.usage.total_tokens,
+                    prompt_tokens_details=PromptTokenDetails(
+                        cached_input_tokens=getattr(
+                            response.usage.prompt_tokens_details, "cached_tokens", None
+                        ),
+                        audio_tokens=getattr(
+                            response.usage.prompt_tokens_details, "audio_tokens", None
+                        ),
+                    ),
+                    completion_tokens_details=CompletionTokenDetails(
+                        reasoning_tokens=getattr(
+                            response.usage.completion_tokens_details,
+                            "reasoning_tokens",
+                            None,
+                        ),
+                        audio_tokens=getattr(
+                            response.usage.completion_tokens_details,
+                            "audio_tokens",
+                            None,
+                        ),
+                        accepted_prediction_tokens=getattr(
+                            response.usage.completion_tokens_details,
+                            "accepted_prediction_tokens",
+                            None,
+                        ),
+                        rejected_prediction_tokens=getattr(
+                            response.usage.completion_tokens_details,
+                            "rejected_prediction_tokens",
+                            None,
+                        ),
+                    ),
                 )
+
                 return TokenUsageStats(model=response.model, usage=usage)
 
             elif isinstance(response, dict):
                 usage_dict = response.get("usage")
                 if not usage_dict:
                     return None
-                usage = Usage(
+                usage = TokenMetrics(
                     prompt_tokens=usage_dict.get("prompt_tokens", 0),
                     completion_tokens=usage_dict.get("completion_tokens", 0),
                     total_tokens=usage_dict.get("total_tokens", 0),
@@ -62,6 +98,10 @@ class BaseOpenAIWrapper(BaseWrapper):
     def completions(self):
         return self
 
+    @property
+    def beta(self):
+        return self
+
 
 def _create_usage_callback(execution_id, log_usage_fn):
     """Creates a callback function for processing usage statistics from stream chunks."""
@@ -75,10 +115,12 @@ def _create_usage_callback(execution_id, log_usage_fn):
             logger.debug("Tokenator is disabled - skipping stream usage logging")
             return
 
+        logger.debug("Processing stream usage for execution_id: %s", execution_id)
+
         # Build usage_data from the first chunk's model
         usage_data = TokenUsageStats(
             model=chunks[0].model,
-            usage=Usage(),
+            usage=TokenMetrics(),
         )
         # Sum up usage from all chunks
         has_usage = False
@@ -116,6 +158,26 @@ class OpenAIWrapper(BaseOpenAIWrapper):
 
         return response
 
+    def parse(
+        self, *args: Any, execution_id: Optional[str] = None, **kwargs: Any
+    ) -> Union[ChatCompletion, Iterator[ChatCompletion]]:
+        """Create a chat completion parse and log token usage."""
+        logger.debug("Creating chat completion with args: %s, kwargs: %s", args, kwargs)
+
+        if kwargs.get("stream", False):
+            base_stream = self.client.beta.chat.completions.parse(*args, **kwargs)
+            return OpenAISyncStreamInterceptor(
+                base_stream=base_stream,
+                usage_callback=_create_usage_callback(execution_id, self._log_usage),
+            )
+
+        response = self.client.beta.chat.completions.parse(*args, **kwargs)
+        usage_data = self._process_response_usage(response)
+        if usage_data:
+            self._log_usage(usage_data, execution_id=execution_id)
+
+        return response
+
 
 class AsyncOpenAIWrapper(BaseOpenAIWrapper):
     async def create(
@@ -141,6 +203,26 @@ class AsyncOpenAIWrapper(BaseOpenAIWrapper):
             self._log_usage(usage_data, execution_id=execution_id)
         return response
 
+    async def parse(
+        self, *args: Any, execution_id: Optional[str] = None, **kwargs: Any
+    ) -> Union[ChatCompletion, AsyncIterator[ChatCompletionChunk]]:
+        """Create a chat completion parse and log token usage."""
+        logger.debug("Creating chat completion with args: %s, kwargs: %s", args, kwargs)
+
+        if kwargs.get("stream", False):
+            base_stream = await self.client.beta.chat.completions.parse(*args, **kwargs)
+            return OpenAIAsyncStreamInterceptor(
+                base_stream=base_stream,
+                usage_callback=_create_usage_callback(execution_id, self._log_usage),
+            )
+
+        response = await self.client.beta.chat.completions.parse(*args, **kwargs)
+        usage_data = self._process_response_usage(response)
+        if usage_data:
+            self._log_usage(usage_data, execution_id=execution_id)
+
+        return response
+
 
 @overload
 def tokenator_openai(

tokenator/schemas.py

@@ -40,28 +40,28 @@ class TokenUsage(Base):
     updated_at = Column(
         DateTime, nullable=False, default=datetime.now, onupdate=datetime.now
     )
+
+    # Core metrics (mandatory)
+    total_cost = Column(Integer, nullable=False)
     prompt_tokens = Column(Integer, nullable=False)
     completion_tokens = Column(Integer, nullable=False)
     total_tokens = Column(Integer, nullable=False)
 
-    # Create indexes
+    # Prompt token details (optional)
+    prompt_cached_input_tokens = Column(Integer, nullable=True)
+    prompt_cached_creation_tokens = Column(Integer, nullable=True)
+    prompt_audio_tokens = Column(Integer, nullable=True)
+
+    # Completion token details (optional)
+    completion_audio_tokens = Column(Integer, nullable=True)
+    completion_reasoning_tokens = Column(Integer, nullable=True)
+    completion_accepted_prediction_tokens = Column(Integer, nullable=True)
+    completion_rejected_prediction_tokens = Column(Integer, nullable=True)
+
+    # Keep existing indexes
     __table_args__ = (
         Index("idx_created_at", "created_at"),
         Index("idx_execution_id", "execution_id"),
         Index("idx_provider", "provider"),
         Index("idx_model", "model"),
     )
-
-    def to_dict(self):
-        """Convert model instance to dictionary."""
-        return {
-            "id": self.id,
-            "execution_id": self.execution_id,
-            "provider": self.provider,
-            "model": self.model,
-            "created_at": self.created_at,
-            "updated_at": self.updated_at,
-            "prompt_tokens": self.prompt_tokens,
-            "completion_tokens": self.completion_tokens,
-            "total_tokens": self.total_tokens,
-        }

tokenator/usage.py

@@ -4,7 +4,14 @@ from datetime import datetime, timedelta
 from typing import Dict, Optional, Union
 
 from .schemas import get_session, TokenUsage
-from .models import TokenRate, TokenUsageReport, ModelUsage, ProviderUsage
+from .models import (
+    CompletionTokenDetails,
+    PromptTokenDetails,
+    TokenRate,
+    TokenUsageReport,
+    ModelUsage,
+    ProviderUsage,
+)
 from . import state
 
 import requests
@@ -27,14 +34,25 @@ class TokenUsageService:
         response = requests.get(url)
         data = response.json()
 
-        return {
-            model: TokenRate(
+        model_costs = {}
+        for model, info in data.items():
+            if (
+                "input_cost_per_token" not in info
+                or "output_cost_per_token" not in info
+            ):
+                continue
+
+            rate = TokenRate(
                 prompt=info["input_cost_per_token"],
                 completion=info["output_cost_per_token"],
+                prompt_audio=info.get("input_cost_per_audio_token"),
+                completion_audio=info.get("output_cost_per_audio_token"),
+                prompt_cached_input=info.get("cache_read_input_token_cost") or 0,
+                prompt_cached_creation=info.get("cache_read_creation_token_cost") or 0,
             )
-            for model, info in data.items()
-            if "input_cost_per_token" in info and "output_cost_per_token" in info
-        }
+            model_costs[model] = rate
+
+        return model_costs
 
     def _calculate_cost(
         self, usages: list[TokenUsage], provider: Optional[str] = None
@@ -47,23 +65,26 @@ class TokenUsageService:
             logger.warning("No model costs available.")
             return TokenUsageReport()
 
-        GPT4O_PRICING = self.MODEL_COSTS.get(
-            "gpt-4o", TokenRate(prompt=0.0000025, completion=0.000010)
+        # Default GPT4O pricing updated with provided values
+        GPT4O_PRICING = TokenRate(
+            prompt=0.0000025,
+            completion=0.000010,
+            prompt_audio=0.0001,
+            completion_audio=0.0002,
+            prompt_cached_input=0.00000125,
+            prompt_cached_creation=0.00000125,
         )
 
-        # Existing calculation logic...
         provider_model_usages: Dict[str, Dict[str, list[TokenUsage]]] = {}
         logger.debug(f"usages: {len(usages)}")
 
         for usage in usages:
-            # 1st priority - direct match
+            # Model key resolution logic (unchanged)
             model_key = usage.model
             if model_key in self.MODEL_COSTS:
                 pass
-            # 2nd priority - provider/model format
             elif f"{usage.provider}/{usage.model}" in self.MODEL_COSTS:
                 model_key = f"{usage.provider}/{usage.model}"
-            # 3rd priority - contains search
             else:
                 matched_keys = [k for k in self.MODEL_COSTS.keys() if usage.model in k]
                 if matched_keys:
@@ -72,10 +93,8 @@ class TokenUsageService:
                         f"Model {usage.model} matched with {model_key} in pricing data via contains search"
                     )
                 else:
-                    # Fallback to GPT4O pricing
                     logger.warning(
-                        f"Model {model_key} not found in pricing data. Using gpt-4o pricing as fallback "
-                        f"(prompt: ${GPT4O_PRICING.prompt}/token, completion: ${GPT4O_PRICING.completion}/token)"
+                        f"Model {model_key} not found in pricing data. Using gpt-4o pricing as fallback"
                     )
                     self.MODEL_COSTS[model_key] = GPT4O_PRICING
 
@@ -99,18 +118,93 @@ class TokenUsageService:
                 "total_tokens": 0,
                 "prompt_tokens": 0,
                 "completion_tokens": 0,
+                "prompt_cached_input_tokens": 0,
+                "prompt_cached_creation_tokens": 0,
+                "prompt_audio_tokens": 0,
+                "completion_audio_tokens": 0,
+                "completion_reasoning_tokens": 0,
+                "completion_accepted_prediction_tokens": 0,
+                "completion_rejected_prediction_tokens": 0,
             }
             models_list = []
 
             for model_key, usages in model_usages.items():
-                model_cost = sum(
-                    usage.prompt_tokens * self.MODEL_COSTS[model_key].prompt
-                    + usage.completion_tokens * self.MODEL_COSTS[model_key].completion
-                    for usage in usages
-                )
-                model_total = sum(usage.total_tokens for usage in usages)
-                model_prompt = sum(usage.prompt_tokens for usage in usages)
-                model_completion = sum(usage.completion_tokens for usage in usages)
+                model_rates = self.MODEL_COSTS[model_key]
+                model_cost = 0.0
+                model_total = 0
+                model_prompt = 0
+                model_completion = 0
+
+                for usage in usages:
+                    # Base token costs
+                    prompt_text_tokens = usage.prompt_tokens
+                    if usage.prompt_cached_input_tokens:
+                        prompt_text_tokens = (
+                            usage.prompt_tokens - usage.prompt_cached_input_tokens
+                        )
+                    if usage.prompt_audio_tokens:
+                        prompt_text_tokens = (
+                            usage.prompt_tokens - usage.prompt_audio_tokens
+                        )
+
+                    completion_text_tokens = usage.completion_tokens
+                    if usage.completion_audio_tokens:
+                        completion_text_tokens = (
+                            usage.completion_tokens - usage.completion_audio_tokens
+                        )
+
+                    prompt_cost = prompt_text_tokens * model_rates.prompt
+                    completion_cost = completion_text_tokens * model_rates.completion
+                    model_cost += prompt_cost + completion_cost
+
+                    # Audio token costs
+                    if usage.prompt_audio_tokens:
+                        if model_rates.prompt_audio:
+                            model_cost += (
+                                usage.prompt_audio_tokens * model_rates.prompt_audio
+                            )
+                        else:
+                            logger.warning(
+                                f"Audio prompt tokens present for {model_key} but no audio rate defined"
+                            )
+
+                    if usage.completion_audio_tokens:
+                        if model_rates.completion_audio:
+                            model_cost += (
+                                usage.completion_audio_tokens
+                                * model_rates.completion_audio
+                            )
+                        else:
+                            logger.warning(
+                                f"Audio completion tokens present for {model_key} but no audio rate defined"
+                            )
+
+                    # Cached token costs
+                    if usage.prompt_cached_input_tokens:
+                        if model_rates.prompt_cached_input:
+                            model_cost += (
+                                usage.prompt_cached_input_tokens
+                                * model_rates.prompt_cached_input
+                            )
+                        else:
+                            logger.warning(
+                                f"Cached input tokens present for {model_key} but no cache input rate defined"
+                            )
+
+                    if usage.prompt_cached_creation_tokens:
+                        if model_rates.prompt_cached_creation:
+                            model_cost += (
+                                usage.prompt_cached_creation_tokens
+                                * model_rates.prompt_cached_creation
+                            )
+                        else:
+                            logger.warning(
+                                f"Cached creation tokens present for {model_key} but no cache creation rate defined"
+                            )
+
+                    model_total += usage.total_tokens
+                    model_prompt += usage.prompt_tokens
+                    model_completion += usage.completion_tokens
 
                 models_list.append(
                     ModelUsage(
@@ -119,22 +213,124 @@ class TokenUsageService:
                         total_tokens=model_total,
                         prompt_tokens=model_prompt,
                         completion_tokens=model_completion,
+                        prompt_tokens_details=PromptTokenDetails(
+                            cached_input_tokens=sum(
+                                u.prompt_cached_input_tokens or 0 for u in usages
+                            ),
+                            cached_creation_tokens=sum(
+                                u.prompt_cached_creation_tokens or 0 for u in usages
+                            ),
+                            audio_tokens=sum(
+                                u.prompt_audio_tokens or 0 for u in usages
+                            ),
+                        )
+                        if any(
+                            u.prompt_cached_input_tokens
+                            or u.prompt_cached_creation_tokens
+                            or u.prompt_audio_tokens
+                            for u in usages
+                        )
+                        else None,
+                        completion_tokens_details=CompletionTokenDetails(
+                            audio_tokens=sum(
+                                u.completion_audio_tokens or 0 for u in usages
+                            ),
+                            reasoning_tokens=sum(
+                                u.completion_reasoning_tokens or 0 for u in usages
+                            ),
+                            accepted_prediction_tokens=sum(
+                                u.completion_accepted_prediction_tokens or 0
+                                for u in usages
+                            ),
+                            rejected_prediction_tokens=sum(
+                                u.completion_rejected_prediction_tokens or 0
+                                for u in usages
+                            ),
+                        )
+                        if any(
+                            getattr(u, attr, None)
+                            for u in usages
+                            for attr in [
+                                "completion_audio_tokens",
+                                "completion_reasoning_tokens",
+                                "completion_accepted_prediction_tokens",
+                                "completion_rejected_prediction_tokens",
+                            ]
+                        )
+                        else None,
                     )
                 )
 
+                # Update provider metrics with all token types
                 provider_metrics["total_cost"] += model_cost
                 provider_metrics["total_tokens"] += model_total
                 provider_metrics["prompt_tokens"] += model_prompt
                 provider_metrics["completion_tokens"] += model_completion
+                provider_metrics["prompt_cached_input_tokens"] += sum(
+                    u.prompt_cached_input_tokens or 0 for u in usages
+                )
+                provider_metrics["prompt_cached_creation_tokens"] += sum(
+                    u.prompt_cached_creation_tokens or 0 for u in usages
+                )
+                provider_metrics["prompt_audio_tokens"] += sum(
+                    u.prompt_audio_tokens or 0 for u in usages
+                )
+                provider_metrics["completion_audio_tokens"] += sum(
+                    u.completion_audio_tokens or 0 for u in usages
+                )
+                provider_metrics["completion_reasoning_tokens"] += sum(
+                    u.completion_reasoning_tokens or 0 for u in usages
+                )
+                provider_metrics["completion_accepted_prediction_tokens"] += sum(
+                    u.completion_accepted_prediction_tokens or 0 for u in usages
+                )
+                provider_metrics["completion_rejected_prediction_tokens"] += sum(
+                    u.completion_rejected_prediction_tokens or 0 for u in usages
+                )
 
             providers_list.append(
                 ProviderUsage(
                     provider=provider,
                     models=models_list,
-                    **{
-                        k: (round(v, 6) if k == "total_cost" else v)
-                        for k, v in provider_metrics.items()
-                    },
+                    total_cost=round(provider_metrics["total_cost"], 6),
+                    total_tokens=provider_metrics["total_tokens"],
+                    prompt_tokens=provider_metrics["prompt_tokens"],
+                    completion_tokens=provider_metrics["completion_tokens"],
+                    prompt_tokens_details=PromptTokenDetails(
+                        cached_input_tokens=provider_metrics[
+                            "prompt_cached_input_tokens"
+                        ],
+                        cached_creation_tokens=provider_metrics[
+                            "prompt_cached_creation_tokens"
+                        ],
+                        audio_tokens=provider_metrics["prompt_audio_tokens"],
+                    )
+                    if provider_metrics["prompt_cached_input_tokens"]
+                    or provider_metrics["prompt_cached_creation_tokens"]
+                    or provider_metrics["prompt_audio_tokens"]
+                    else None,
+                    completion_tokens_details=CompletionTokenDetails(
+                        audio_tokens=provider_metrics["completion_audio_tokens"],
+                        reasoning_tokens=provider_metrics[
+                            "completion_reasoning_tokens"
+                        ],
+                        accepted_prediction_tokens=provider_metrics[
+                            "completion_accepted_prediction_tokens"
+                        ],
+                        rejected_prediction_tokens=provider_metrics[
+                            "completion_rejected_prediction_tokens"
+                        ],
+                    )
+                    if any(
+                        provider_metrics[k]
+                        for k in [
+                            "completion_audio_tokens",
+                            "completion_reasoning_tokens",
+                            "completion_accepted_prediction_tokens",
+                            "completion_rejected_prediction_tokens",
+                        ]
+                    )
+                    else None,
                 )
             )
 

tokenator/utils.py

@@ -5,8 +5,21 @@ import platform
 import logging
 from pathlib import Path
 
+
 logger = logging.getLogger(__name__)
 
+def is_notebook() -> bool:
+    try:
+        from IPython import get_ipython # type: ignore
+        shell = get_ipython().__class__.__name__
+        if shell == 'ZMQInteractiveShell':
+            return True   # Jupyter notebook or qtconsole
+        elif shell == 'TerminalInteractiveShell':
+            return False  # Terminal running IPython
+        else:
+            return False  # Other type (?)
+    except NameError:
+        return False
 
 def is_colab() -> bool:
     """Check if running in Google Colab."""
@@ -21,7 +34,7 @@ def is_colab() -> bool:
 def get_default_db_path() -> str:
     """Get the platform-specific default database path."""
     try:
-        if is_colab():
+        if is_colab() or is_notebook():
             # Use in-memory database for Colab
             return "usage.db"
 

{tokenator-0.1.14.dist-info → tokenator-0.1.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: tokenator
-Version: 0.1.14
+Version: 0.1.15
 Summary: Token usage tracking wrapper for LLMs
 License: MIT
 Author: Ujjwal Maheshwari
@@ -15,6 +15,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: alembic (>=1.13.0,<2.0.0)
 Requires-Dist: anthropic (>=0.43.0,<0.44.0)
+Requires-Dist: ipython
 Requires-Dist: openai (>=1.59.0,<2.0.0)
 Requires-Dist: requests (>=2.32.3,<3.0.0)
 Requires-Dist: sqlalchemy (>=2.0.0,<3.0.0)
@@ -22,7 +23,7 @@ Description-Content-Type: text/markdown
 
 # Tokenator : Track and analyze LLM token usage and cost
 
-Have you ever wondered about :
+Have you ever wondered :
 - How many tokens does your AI agent consume? 
 - How much does it cost to do run a complex AI workflow with multiple LLM providers?
 - How much money/tokens did you spend today on developing with LLMs?

tokenator-0.1.15.dist-info/RECORD

@@ -0,0 +1,21 @@
+tokenator/__init__.py,sha256=AEPE73UGB_TeNLhro3eY0hU8yy6T-_6AyDls8vWApnE,465
+tokenator/anthropic/client_anthropic.py,sha256=2oxTLb5-sPK_KL-OumCjE4wPVI8U_eFyRonn9XjGXJw,7196
+tokenator/anthropic/stream_interceptors.py,sha256=4VHC_-WkG3Pa10YizmFLrHcbz0Tm2MR_YB5-uohKp5A,5221
+tokenator/base_wrapper.py,sha256=EQ49xGduEp05-gj1xyZDasrck4RpComaoKslHxQTwuw,4956
+tokenator/create_migrations.py,sha256=k9IHiGK21dLTA8MYNsuhO0-kUVIcMSViMFYtY4WU2Rw,730
+tokenator/migrations/env.py,sha256=JoF5MJ4ae0wJW5kdBHuFlG3ZqeCCDvbMcU8fNA_a6hM,1396
+tokenator/migrations/script.py.mako,sha256=nJL-tbLQE0Qy4P9S4r4ntNAcikPtoFUlvXe6xvm9ot8,635
+tokenator/migrations/versions/f028b8155fed_adding_detailed_input_and_output_token_.py,sha256=WIZN5HdNRXlRdfpUJpJFaPD4G1s-SgRdTMQl4WDB-hA,2189
+tokenator/migrations/versions/f6f1f2437513_initial_migration.py,sha256=4cveHkwSxs-hxOPCm81YfvGZTkJJ2ClAFmyL98-1VCo,1910
+tokenator/migrations.py,sha256=YAf9gZmDzAq36PWWXPtdUQoJFYPXtIDzflC79H6gcJg,1114
+tokenator/models.py,sha256=p4uoFqJYGMlygotxip_HZcfM16Jm4LoyFLFTsM1Z8a4,2132
+tokenator/openai/client_openai.py,sha256=pbdJ-aZPuJs-7OT1VEv0DW36cCYbRAVKhSQEprxVIdY,9686
+tokenator/openai/stream_interceptors.py,sha256=ez1MnjRZW_rEalv2SIPAvrU9oMD6OJoD9vht-057fDM,5243
+tokenator/schemas.py,sha256=kBmShqgpQ3W-ILAP1NuCaFgqFplQM4OH0MmJteLqrwI,2371
+tokenator/state.py,sha256=xdqDC-rlEA88-VgqQqHnAOXQ5pNTpnHcgOtohDIImPY,262
+tokenator/usage.py,sha256=QaudrO6uwnMNRn9aCYVPj9yiQHmbdoAVZ9-G4Q1B0fw,20511
+tokenator/utils.py,sha256=djoWmAhqH-O2Su3qIcuY-_3Vj1-qPwMcdzwq9IlwiDc,2435
+tokenator-0.1.15.dist-info/LICENSE,sha256=wdG-B6-ODk8RQ4jq5uXSn0w1UWTzCH_MMyvh7AwtGns,1074
+tokenator-0.1.15.dist-info/METADATA,sha256=dtws3Qwm2iZLCYZv0meqQP80Q49821HdyZgUmDeqDcg,6035
+tokenator-0.1.15.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
+tokenator-0.1.15.dist-info/RECORD,,

tokenator-0.1.14.dist-info/RECORD

@@ -1,20 +0,0 @@
-tokenator/__init__.py,sha256=AEPE73UGB_TeNLhro3eY0hU8yy6T-_6AyDls8vWApnE,465
-tokenator/anthropic/client_anthropic.py,sha256=uWUrRId7vJlMG6hVKLUzaA3PoOT6mJwTqSRIhAidRFY,6163
-tokenator/anthropic/stream_interceptors.py,sha256=4VHC_-WkG3Pa10YizmFLrHcbz0Tm2MR_YB5-uohKp5A,5221
-tokenator/base_wrapper.py,sha256=UoS3cOuPa3HpuXPTawybvAtwufgZwzzKBj0BhyB-z6w,3160
-tokenator/create_migrations.py,sha256=k9IHiGK21dLTA8MYNsuhO0-kUVIcMSViMFYtY4WU2Rw,730
-tokenator/migrations/env.py,sha256=JoF5MJ4ae0wJW5kdBHuFlG3ZqeCCDvbMcU8fNA_a6hM,1396
-tokenator/migrations/script.py.mako,sha256=nJL-tbLQE0Qy4P9S4r4ntNAcikPtoFUlvXe6xvm9ot8,635
-tokenator/migrations/versions/f6f1f2437513_initial_migration.py,sha256=4cveHkwSxs-hxOPCm81YfvGZTkJJ2ClAFmyL98-1VCo,1910
-tokenator/migrations.py,sha256=YAf9gZmDzAq36PWWXPtdUQoJFYPXtIDzflC79H6gcJg,1114
-tokenator/models.py,sha256=AlNC5NVrycLg0LhDJIww9HXQ3lwM8CoKvRSqXU6iw-k,1225
-tokenator/openai/client_openai.py,sha256=LhD1IbpzPXRK9eSqtcfUfoM9vBsyw6OHA0_a7N_tS9U,6230
-tokenator/openai/stream_interceptors.py,sha256=ez1MnjRZW_rEalv2SIPAvrU9oMD6OJoD9vht-057fDM,5243
-tokenator/schemas.py,sha256=zIgfmSsFJV9ziJdKrpV8p2P1f-BVWUVIpWoqCLpzhEU,2225
-tokenator/state.py,sha256=xdqDC-rlEA88-VgqQqHnAOXQ5pNTpnHcgOtohDIImPY,262
-tokenator/usage.py,sha256=ghnZ7pQuIxeI38O63xDAbEm6jOSmkYE7MChHBGPxbyM,11229
-tokenator/utils.py,sha256=xg9l2GV1yJL1BlxKL1r8CboABWDslf3G5rGQEJSjFrE,1973
-tokenator-0.1.14.dist-info/LICENSE,sha256=wdG-B6-ODk8RQ4jq5uXSn0w1UWTzCH_MMyvh7AwtGns,1074
-tokenator-0.1.14.dist-info/METADATA,sha256=L93LfqCfqvhES92COaQZpX5w9_c2aDaX8pj2wT74Sxw,6018
-tokenator-0.1.14.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
-tokenator-0.1.14.dist-info/RECORD,,