tokenator 0.1.14__py3-none-any.whl → 0.1.16__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,
-        }