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

tokenator/__init__.py

@@ -5,14 +5,9 @@ from .openai.client_openai import tokenator_openai
 from .anthropic.client_anthropic import tokenator_anthropic
 from . import usage
 from .utils import get_default_db_path
-from .migrations import check_and_run_migrations
+from .usage import TokenUsageService
 
-__version__ = "0.1.0"
+usage = TokenUsageService()  # noqa: F811
 __all__ = ["tokenator_openai", "tokenator_anthropic", "usage", "get_default_db_path"]
 
 logger = logging.getLogger(__name__)
-
-try:
-    check_and_run_migrations()
-except Exception as e:
-    logger.warning(f"Failed to run migrations, but continuing anyway: {e}")

tokenator/anthropic/client_anthropic.py

@@ -6,9 +6,13 @@ 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, AnthropicSyncStreamInterceptor
+from .stream_interceptors import (
+    AnthropicAsyncStreamInterceptor,
+    AnthropicSyncStreamInterceptor,
+)
+from ..state import is_tokenator_enabled
 
 logger = logging.getLogger(__name__)
 
@@ -24,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
 
@@ -56,15 +78,23 @@ class BaseAnthropicWrapper(BaseWrapper):
 
 def _create_usage_callback(execution_id, log_usage_fn):
     """Creates a callback function for processing usage statistics from stream chunks."""
+
     def usage_callback(chunks):
         if not chunks:
             return
-        
+
+        # Skip if tokenator is disabled
+        if not is_tokenator_enabled:
+            logger.debug("Tokenator is disabled - skipping stream usage logging")
+            return
+
         usage_data = TokenUsageStats(
-            model=chunks[0].message.model if isinstance(chunks[0], RawMessageStartEvent) else "",
-            usage=Usage(),
+            model=chunks[0].message.model
+            if isinstance(chunks[0], RawMessageStartEvent)
+            else "",
+            usage=TokenMetrics(),
         )
-        
+
         for chunk in chunks:
             if isinstance(chunk, RawMessageStartEvent):
                 usage_data.model = chunk.message.model
@@ -72,8 +102,10 @@ def _create_usage_callback(execution_id, log_usage_fn):
                 usage_data.usage.completion_tokens += chunk.message.usage.output_tokens
             elif isinstance(chunk, RawMessageDeltaEvent):
                 usage_data.usage.completion_tokens += chunk.usage.output_tokens
-        
-        usage_data.usage.total_tokens = usage_data.usage.prompt_tokens + usage_data.usage.completion_tokens
+
+        usage_data.usage.total_tokens = (
+            usage_data.usage.prompt_tokens + usage_data.usage.completion_tokens
+        )
         log_usage_fn(usage_data, execution_id=execution_id)
 
     return usage_callback
@@ -84,7 +116,9 @@ class AnthropicWrapper(BaseAnthropicWrapper):
         self, *args: Any, execution_id: Optional[str] = None, **kwargs: Any
     ) -> Union[Message, Iterator[Message]]:
         """Create a message completion and log token usage."""
-        logger.debug("Creating message completion with args: %s, kwargs: %s", args, kwargs)
+        logger.debug(
+            "Creating message completion with args: %s, kwargs: %s", args, kwargs
+        )
 
         if kwargs.get("stream", False):
             base_stream = self.client.messages.create(*args, **kwargs)
@@ -105,7 +139,9 @@ class AsyncAnthropicWrapper(BaseAnthropicWrapper):
         self, *args: Any, execution_id: Optional[str] = None, **kwargs: Any
     ) -> Union[Message, AsyncIterator[Message]]:
         """Create a message completion and log token usage."""
-        logger.debug("Creating message completion with args: %s, kwargs: %s", args, kwargs)
+        logger.debug(
+            "Creating message completion with args: %s, kwargs: %s", args, kwargs
+        )
 
         if kwargs.get("stream", False):
             base_stream = await self.client.messages.create(*args, **kwargs)

tokenator/base_wrapper.py

@@ -7,6 +7,9 @@ import uuid
 
 from .models import TokenUsageStats
 from .schemas import get_session, TokenUsage
+from . import state
+
+from .migrations import check_and_run_migrations
 
 logger = logging.getLogger(__name__)
 
@@ -16,17 +19,30 @@ ResponseType = TypeVar("ResponseType")
 class BaseWrapper:
     def __init__(self, client: Any, db_path: Optional[str] = None):
         """Initialize the base wrapper."""
-        self.client = client
+        state.is_tokenator_enabled = True
+        try:
+            self.client = client
 
-        if db_path:
-            Path(db_path).parent.mkdir(parents=True, exist_ok=True)
-            logger.info("Created database directory at: %s", Path(db_path).parent)
+            if db_path:
+                Path(db_path).parent.mkdir(parents=True, exist_ok=True)
+                logger.info("Created database directory at: %s", Path(db_path).parent)
+                state.db_path = db_path  # Store db_path in state
 
-        self.Session = get_session(db_path)
+            else:
+                state.db_path = None  # Use default path
 
-        logger.debug(
-            "Initializing %s with db_path: %s", self.__class__.__name__, db_path
-        )
+            self.Session = get_session()
+
+            logger.debug(
+                "Initializing %s with db_path: %s", self.__class__.__name__, db_path
+            )
+
+            check_and_run_migrations(db_path)
+        except Exception as e:
+            state.is_tokenator_enabled = False
+            logger.warning(
+                f"Tokenator initialization failed. Usage tracking will be disabled. Error: {e}"
+            )
 
     def _log_usage_impl(
         self, token_usage_stats: TokenUsageStats, session, execution_id: str
@@ -42,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(
@@ -59,14 +99,20 @@ class BaseWrapper:
         self, token_usage_stats: TokenUsageStats, execution_id: Optional[str] = None
     ):
         """Log token usage to database."""
+        if not state.is_tokenator_enabled:
+            logger.debug("Tokenator is disabled - skipping usage logging")
+            return
+
         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,17 +1,44 @@
 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):
-    total_cost: float = Field(..., description="Total cost in USD")
-    total_tokens: int = Field(..., description="Total tokens used")
-    prompt_tokens: int = Field(..., description="Number of prompt tokens")
-    completion_tokens: int = Field(..., description="Number of completion tokens")
+    total_cost: float = Field(default=0, description="Total cost in USD")
+    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,9 +6,18 @@ 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, OpenAISyncStreamInterceptor
+from .stream_interceptors import (
+    OpenAIAsyncStreamInterceptor,
+    OpenAISyncStreamInterceptor,
+)
+from ..state import is_tokenator_enabled
 
 logger = logging.getLogger(__name__)
 
@@ -26,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),
@@ -58,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."""
@@ -65,10 +109,18 @@ def _create_usage_callback(execution_id, log_usage_fn):
     def usage_callback(chunks):
         if not chunks:
             return
+
+        # Skip if tokenator is disabled
+        if not is_tokenator_enabled:
+            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
@@ -106,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(
@@ -131,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

@@ -1,25 +1,27 @@
 """SQLAlchemy models for tokenator."""
 
 from datetime import datetime
+from typing import Optional
 
 from sqlalchemy import create_engine, Column, Integer, String, DateTime, Index
 from sqlalchemy.orm import sessionmaker, scoped_session, declarative_base
 
 from .utils import get_default_db_path
+from . import state  # Import state to access db_path
 
 Base = declarative_base()
 
 
-def get_engine(db_path: str = None):
+def get_engine(db_path: Optional[str] = None):
     """Create SQLAlchemy engine with the given database path."""
     if db_path is None:
-        db_path = get_default_db_path()
+        db_path = state.db_path or get_default_db_path()  # Use state.db_path if set
     return create_engine(f"sqlite:///{db_path}", echo=False)
 
 
-def get_session(db_path: str = None):
+def get_session():
     """Create a thread-safe session factory."""
-    engine = get_engine(db_path)
+    engine = get_engine()
     # Base.metadata.create_all(engine)
     session_factory = sessionmaker(bind=engine)
     return scoped_session(session_factory)
@@ -38,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/state.py

@@ -0,0 +1,12 @@
+"""Global state for tokenator."""
+
+import logging
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+# Global flag to track if tokenator is properly initialized
+is_tokenator_enabled = True
+
+# Store the database path
+db_path: Optional[str] = None