tokenator 0.1.13__tar.gz → 0.1.15__tar.gz

{tokenator-0.1.13 → tokenator-0.1.15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: tokenator
-Version: 0.1.13
+Version: 0.1.15
 Summary: Token usage tracking wrapper for LLMs
 License: MIT
 Author: Ujjwal Maheshwari
@@ -14,15 +14,16 @@ Classifier: Programming Language :: Python :: 3.11
 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.40.0,<0.41.0)
-Requires-Dist: openai (>=1.57.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)
 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.13 → tokenator-0.1.15}/README.md

@@ -1,6 +1,6 @@
 # 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.13 → tokenator-0.1.15}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tokenator"
-version = "0.1.13"
+version = "0.1.15"
 description = "Token usage tracking wrapper for LLMs"
 authors = ["Ujjwal Maheshwari <your.email@example.com>"]
 readme = "README.md"
@@ -9,11 +9,12 @@ packages = [{include = "tokenator", from = "src"}]
 
 [tool.poetry.dependencies]
 python = "^3.9"
-openai = "^1.57.0"
+openai = "^1.59.0"
 sqlalchemy = "^2.0.0"
 requests = "^2.32.3"
 alembic = "^1.13.0"
-anthropic = "^0.40.0"
+anthropic = "^0.43.0"
+ipython = "*"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.0.0"

{tokenator-0.1.13 → tokenator-0.1.15}/src/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-0.1.13 → tokenator-0.1.15}/src/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-0.1.15/src/tokenator/base_wrapper.py

@@ -0,0 +1,120 @@
+"""Base wrapper class for token usage tracking."""
+
+from pathlib import Path
+from typing import Any, Optional, TypeVar
+import logging
+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__)
+
+ResponseType = TypeVar("ResponseType")
+
+
+class BaseWrapper:
+    def __init__(self, client: Any, db_path: Optional[str] = None):
+        """Initialize the base wrapper."""
+        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)
+                state.db_path = db_path  # Store db_path in state
+
+            else:
+                state.db_path = None  # Use default 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
+    ) -> None:
+        """Implementation of token usage logging."""
+        logger.debug(
+            "Logging usage for model %s: %s",
+            token_usage_stats.model,
+            token_usage_stats.usage.model_dump(),
+        )
+        try:
+            token_usage = TokenUsage(
+                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(
+                "Logged token usage: model=%s, total_tokens=%d",
+                token_usage_stats.model,
+                token_usage_stats.usage.total_tokens,
+            )
+        except Exception as e:
+            logger.error("Failed to log token usage: %s", str(e))
+
+    def _log_usage(
+        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()
+        finally:
+            session.close()

tokenator-0.1.15/src/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-0.1.15/src/tokenator/models.py

@@ -0,0 +1,63 @@
+from pydantic import BaseModel, Field
+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(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):
+    model: str = Field(..., description="Model name")
+
+
+class ProviderUsage(TokenMetrics):
+    provider: str = Field(..., description="Provider name")
+    models: List[ModelUsage] = Field(
+        default_factory=list, description="Usage breakdown by model"
+    )
+
+
+class TokenUsageReport(TokenMetrics):
+    providers: List[ProviderUsage] = Field(
+        default_factory=list, description="Usage breakdown by provider"
+    )
+
+
+class TokenUsageStats(BaseModel):
+    model: str
+    usage: TokenMetrics

{tokenator-0.1.13 → tokenator-0.1.15}/src/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(