tokenator 0.1.8__py3-none-any.whl → 0.1.10__py3-none-any.whl

tokenator/__init__.py

@@ -1,8 +1,8 @@
 """Tokenator - Track and analyze your OpenAI API token usage and costs."""
 
 import logging
-from .client_openai import tokenator_openai
-from .client_anthropic import tokenator_anthropic
+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
@@ -15,4 +15,4 @@ logger = logging.getLogger(__name__)
 try:
     check_and_run_migrations()
 except Exception as e:
-    logger.warning(f"Failed to run migrations, but continuing anyway: {e}") 
+    logger.warning(f"Failed to run migrations, but continuing anyway: {e}")

tokenator/anthropic/client_anthropic.py

@@ -0,0 +1,155 @@
+"""Anthropic client wrapper with token usage tracking."""
+
+from typing import Any, Optional, Union, overload, Iterator, AsyncIterator
+import logging
+
+from anthropic import Anthropic, AsyncAnthropic
+from anthropic.types import Message, RawMessageStartEvent, RawMessageDeltaEvent
+
+from ..models import Usage, TokenUsageStats
+from ..base_wrapper import BaseWrapper, ResponseType
+from .stream_interceptors import AnthropicAsyncStreamInterceptor, AnthropicSyncStreamInterceptor
+
+logger = logging.getLogger(__name__)
+
+
+class BaseAnthropicWrapper(BaseWrapper):
+    provider = "anthropic"
+
+    def _process_response_usage(
+        self, response: ResponseType
+    ) -> Optional[TokenUsageStats]:
+        """Process and log usage statistics from a response."""
+        try:
+            if isinstance(response, Message):
+                if not hasattr(response, "usage"):
+                    return None
+                usage = Usage(
+                    prompt_tokens=response.usage.input_tokens,
+                    completion_tokens=response.usage.output_tokens,
+                    total_tokens=response.usage.input_tokens
+                    + response.usage.output_tokens,
+                )
+                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),
+                    completion_tokens=usage_dict.get("output_tokens", 0),
+                    total_tokens=usage_dict.get("input_tokens", 0)
+                    + usage_dict.get("output_tokens", 0),
+                )
+                return TokenUsageStats(
+                    model=response.get("model", "unknown"), usage=usage
+                )
+        except Exception as e:
+            logger.warning("Failed to process usage stats: %s", str(e))
+            return None
+        return None
+
+    @property
+    def messages(self):
+        return self
+
+
+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
+        
+        usage_data = TokenUsageStats(
+            model=chunks[0].message.model if isinstance(chunks[0], RawMessageStartEvent) else "",
+            usage=Usage(),
+        )
+        
+        for chunk in chunks:
+            if isinstance(chunk, RawMessageStartEvent):
+                usage_data.model = chunk.message.model
+                usage_data.usage.prompt_tokens += chunk.message.usage.input_tokens
+                usage_data.usage.completion_tokens += chunk.message.usage.output_tokens
+            elif isinstance(chunk, RawMessageDeltaEvent):
+                usage_data.usage.prompt_tokens += chunk.usage.input_tokens
+                usage_data.usage.completion_tokens += chunk.usage.output_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
+
+
+class AnthropicWrapper(BaseAnthropicWrapper):
+    def create(
+        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)
+
+        if kwargs.get("stream", False):
+            base_stream = self.client.messages.create(*args, **kwargs)
+            return AnthropicSyncStreamInterceptor(
+                base_stream=base_stream,
+                usage_callback=_create_usage_callback(execution_id, self._log_usage),
+            )
+
+        response = self.client.messages.create(*args, **kwargs)
+        usage_data = self._process_response_usage(response)
+        if usage_data:
+            self._log_usage(usage_data, execution_id=execution_id)
+        return response
+
+
+class AsyncAnthropicWrapper(BaseAnthropicWrapper):
+    async def create(
+        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)
+
+        if kwargs.get("stream", False):
+            base_stream = await self.client.messages.create(*args, **kwargs)
+            return AnthropicAsyncStreamInterceptor(
+                base_stream=base_stream,
+                usage_callback=_create_usage_callback(execution_id, self._log_usage),
+            )
+
+        response = await self.client.messages.create(*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_anthropic(
+    client: Anthropic,
+    db_path: Optional[str] = None,
+) -> AnthropicWrapper: ...
+
+
+@overload
+def tokenator_anthropic(
+    client: AsyncAnthropic,
+    db_path: Optional[str] = None,
+) -> AsyncAnthropicWrapper: ...
+
+
+def tokenator_anthropic(
+    client: Union[Anthropic, AsyncAnthropic],
+    db_path: Optional[str] = None,
+) -> Union[AnthropicWrapper, AsyncAnthropicWrapper]:
+    """Create a token-tracking wrapper for an Anthropic client.
+
+    Args:
+        client: Anthropic or AsyncAnthropic client instance
+        db_path: Optional path to SQLite database for token tracking
+    """
+    if isinstance(client, Anthropic):
+        return AnthropicWrapper(client=client, db_path=db_path)
+
+    if isinstance(client, AsyncAnthropic):
+        return AsyncAnthropicWrapper(client=client, db_path=db_path)
+
+    raise ValueError("Client must be an instance of Anthropic or AsyncAnthropic")

tokenator/anthropic/stream_interceptors.py

@@ -0,0 +1,146 @@
+import logging
+from typing import AsyncIterator, Callable, List, Optional, TypeVar, Iterator
+
+from anthropic import AsyncStream, Stream
+
+logger = logging.getLogger(__name__)
+
+_T = TypeVar("_T")
+
+
+class AnthropicAsyncStreamInterceptor(AsyncStream[_T]):
+    """
+    A wrapper around anthropic.AsyncStream that delegates all functionality
+    to the 'base_stream' but intercepts each chunk to handle usage or
+    logging logic. This preserves .response and other methods.
+
+    You can store aggregated usage in a local list and process it when
+    the stream ends (StopAsyncIteration).
+    """
+
+    def __init__(
+        self,
+        base_stream: AsyncStream[_T],
+        usage_callback: Optional[Callable[[List[_T]], None]] = None,
+    ):
+        # We do NOT call super().__init__() because anthropic.AsyncStream
+        # expects constructor parameters we don't want to re-initialize.
+        # Instead, we just store the base_stream and delegate everything to it.
+        self._base_stream = base_stream
+        self._usage_callback = usage_callback
+        self._chunks: List[_T] = []
+
+    @property
+    def response(self):
+        """Expose the original stream's 'response' so user code can do stream.response, etc."""
+        return self._base_stream.response
+
+    def __aiter__(self) -> AsyncIterator[_T]:
+        """
+        Called when we do 'async for chunk in wrapped_stream:'
+        We simply return 'self'. Then __anext__ does the rest.
+        """
+        return self
+
+    async def __anext__(self) -> _T:
+        """
+        Intercept iteration. We pull the next chunk from the base_stream.
+        If it's the end, do any final usage logging, then raise StopAsyncIteration.
+        Otherwise, we can accumulate usage info or do whatever we need with the chunk.
+        """
+        try:
+            chunk = await self._base_stream.__anext__()
+        except StopAsyncIteration:
+            # Once the base stream is fully consumed, we can do final usage/logging.
+            if self._usage_callback and self._chunks:
+                self._usage_callback(self._chunks)
+            raise
+
+        # Intercept each chunk
+        self._chunks.append(chunk)
+        return chunk
+
+    async def __aenter__(self) -> "AnthropicAsyncStreamInterceptor[_T]":
+        """Support async with ... : usage."""
+        await self._base_stream.__aenter__()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """
+        Ensure we propagate __aexit__ to the base stream,
+        so connections are properly closed.
+        """
+        return await self._base_stream.__aexit__(exc_type, exc_val, exc_tb)
+
+    async def close(self) -> None:
+        """Delegate close to the base_stream."""
+        await self._base_stream.close()
+
+
+class AnthropicSyncStreamInterceptor(Stream[_T]):
+    """
+    A wrapper around anthropic.Stream that delegates all functionality
+    to the 'base_stream' but intercepts each chunk to handle usage or
+    logging logic. This preserves .response and other methods.
+
+    You can store aggregated usage in a local list and process it when
+    the stream ends (StopIteration).
+    """
+
+    def __init__(
+        self,
+        base_stream: Stream[_T],
+        usage_callback: Optional[Callable[[List[_T]], None]] = None,
+    ):
+        # We do NOT call super().__init__() because openai.SyncStream
+        # expects constructor parameters we don't want to re-initialize.
+        # Instead, we just store the base_stream and delegate everything to it.
+        self._base_stream = base_stream
+        self._usage_callback = usage_callback
+        self._chunks: List[_T] = []
+
+    @property
+    def response(self):
+        """Expose the original stream's 'response' so user code can do stream.response, etc."""
+        return self._base_stream.response
+
+    def __iter__(self) -> Iterator[_T]:
+        """
+        Called when we do 'for chunk in wrapped_stream:'
+        We simply return 'self'. Then __next__ does the rest.
+        """
+        return self
+
+    def __next__(self) -> _T:
+        """
+        Intercept iteration. We pull the next chunk from the base_stream.
+        If it's the end, do any final usage logging, then raise StopIteration.
+        Otherwise, we can accumulate usage info or do whatever we need with the chunk.
+        """
+        try:
+            chunk = self._base_stream.__next__()
+        except StopIteration:
+            # Once the base stream is fully consumed, we can do final usage/logging.
+            if self._usage_callback and self._chunks:
+                self._usage_callback(self._chunks)
+            raise
+
+        # Intercept each chunk
+        self._chunks.append(chunk)
+        return chunk
+
+    def __enter__(self) -> "AnthropicSyncStreamInterceptor[_T]":
+        """Support with ... : usage."""
+        self._base_stream.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        Ensure we propagate __aexit__ to the base stream,
+        so connections are properly closed.
+        """
+        return self._base_stream.__exit__(exc_type, exc_val, exc_tb)
+
+    async def close(self) -> None:
+        """Delegate close to the base_stream."""
+        self._base_stream.close()

tokenator/base_wrapper.py

@@ -1,16 +1,17 @@
 """Base wrapper class for token usage tracking."""
 
 from pathlib import Path
-from typing import Any, Dict, Optional, TypeVar, Union
+from typing import Any, Optional, TypeVar
 import logging
 import uuid
 
-from .models import Usage, TokenUsageStats
+from .models import TokenUsageStats
 from .schemas import get_session, TokenUsage
 
 logger = logging.getLogger(__name__)
 
-ResponseType = TypeVar('ResponseType')
+ResponseType = TypeVar("ResponseType")
+
 
 class BaseWrapper:
     def __init__(self, client: Any, db_path: Optional[str] = None):
@@ -22,13 +23,20 @@ class BaseWrapper:
             logger.info("Created database directory at: %s", Path(db_path).parent)
 
         self.Session = get_session(db_path)
-        
-        logger.debug("Initializing %s with db_path: %s", 
-                    self.__class__.__name__, db_path)
 
-    def _log_usage_impl(self, token_usage_stats: TokenUsageStats, session, execution_id: str) -> None:
+        logger.debug(
+            "Initializing %s with db_path: %s", self.__class__.__name__, db_path
+        )
+
+    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())
+        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,
@@ -36,15 +44,20 @@ class BaseWrapper:
                 model=token_usage_stats.model,
                 prompt_tokens=token_usage_stats.usage.prompt_tokens,
                 completion_tokens=token_usage_stats.usage.completion_tokens,
-                total_tokens=token_usage_stats.usage.total_tokens
+                total_tokens=token_usage_stats.usage.total_tokens,
             )
             session.add(token_usage)
-            logger.info("Logged token usage: model=%s, total_tokens=%d", 
-                       token_usage_stats.model, token_usage_stats.usage.total_tokens)
+            logger.info(
+                "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):
+    def _log_usage(
+        self, token_usage_stats: TokenUsageStats, execution_id: Optional[str] = None
+    ):
         """Log token usage to database."""
         if not execution_id:
             execution_id = str(uuid.uuid4())
@@ -58,4 +71,4 @@ class BaseWrapper:
                 logger.error("Failed to log token usage: %s", str(e))
                 session.rollback()
         finally:
-            session.close() 
+            session.close()

tokenator/create_migrations.py

@@ -1,25 +1,26 @@
 """Development utilities for tokenator."""
 
-import os
 import sys
 from pathlib import Path
 from alembic import command
 from tokenator.migrations import get_alembic_config
 
+
 def create_migration(message: str):
     """Create a new migration based on model changes."""
     config = get_alembic_config()
-    
+
     # Get the migrations directory
     migrations_dir = Path(__file__).parent / "migrations" / "versions"
     migrations_dir.mkdir(parents=True, exist_ok=True)
-    
+
     # Generate migration with custom message
-    command.revision(config, autogenerate=True, message=message) 
+    command.revision(config, autogenerate=True, message=message)
+
 
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         msg = " ".join(sys.argv[1:])
     else:
         msg = "auto generated migration"
-    create_migration(msg)
+    create_migration(msg)

tokenator/migrations/env.py

@@ -18,6 +18,7 @@ if config.config_file_name is not None:
 # add your model's MetaData object here
 target_metadata = Base.metadata
 
+
 def run_migrations_offline() -> None:
     """Run migrations in 'offline' mode."""
     url = config.get_main_option("sqlalchemy.url")
@@ -31,6 +32,7 @@ def run_migrations_offline() -> None:
     with context.begin_transaction():
         context.run_migrations()
 
+
 def run_migrations_online() -> None:
     """Run migrations in 'online' mode."""
     connectable = engine_from_config(
@@ -40,14 +42,13 @@ def run_migrations_online() -> None:
     )
 
     with connectable.connect() as connection:
-        context.configure(
-            connection=connection, target_metadata=target_metadata
-        )
+        context.configure(connection=connection, target_metadata=target_metadata)
 
         with context.begin_transaction():
             context.run_migrations()
 
+
 if context.is_offline_mode():
     run_migrations_offline()
 else:
-    run_migrations_online() 
+    run_migrations_online()

tokenator/migrations/versions/f6f1f2437513_initial_migration.py

@@ -1,10 +1,11 @@
 """Initial migration
 
 Revision ID: f6f1f2437513
-Revises: 
+Revises:
 Create Date: 2024-12-21 17:33:27.187221
 
 """
+
 from typing import Sequence, Union
 
 from alembic import op
@@ -12,7 +13,7 @@ import sqlalchemy as sa
 
 
 # revision identifiers, used by Alembic.
-revision: str = 'f6f1f2437513'
+revision: str = "f6f1f2437513"
 down_revision: Union[str, None] = None
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
@@ -20,30 +21,31 @@ depends_on: Union[str, Sequence[str], None] = None
 
 def upgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    op.create_table('token_usage',
-    sa.Column('id', sa.Integer(), nullable=False),
-    sa.Column('execution_id', sa.String(), nullable=False),
-    sa.Column('provider', sa.String(), nullable=False),
-    sa.Column('model', sa.String(), nullable=False),
-    sa.Column('created_at', sa.DateTime(), nullable=False),
-    sa.Column('updated_at', sa.DateTime(), nullable=False),
-    sa.Column('prompt_tokens', sa.Integer(), nullable=False),
-    sa.Column('completion_tokens', sa.Integer(), nullable=False),
-    sa.Column('total_tokens', sa.Integer(), nullable=False),
-    sa.PrimaryKeyConstraint('id')
+    op.create_table(
+        "token_usage",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("execution_id", sa.String(), nullable=False),
+        sa.Column("provider", sa.String(), nullable=False),
+        sa.Column("model", sa.String(), nullable=False),
+        sa.Column("created_at", sa.DateTime(), nullable=False),
+        sa.Column("updated_at", sa.DateTime(), nullable=False),
+        sa.Column("prompt_tokens", sa.Integer(), nullable=False),
+        sa.Column("completion_tokens", sa.Integer(), nullable=False),
+        sa.Column("total_tokens", sa.Integer(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
     )
-    op.create_index('idx_created_at', 'token_usage', ['created_at'], unique=False)
-    op.create_index('idx_execution_id', 'token_usage', ['execution_id'], unique=False)
-    op.create_index('idx_model', 'token_usage', ['model'], unique=False)
-    op.create_index('idx_provider', 'token_usage', ['provider'], unique=False)
+    op.create_index("idx_created_at", "token_usage", ["created_at"], unique=False)
+    op.create_index("idx_execution_id", "token_usage", ["execution_id"], unique=False)
+    op.create_index("idx_model", "token_usage", ["model"], unique=False)
+    op.create_index("idx_provider", "token_usage", ["provider"], unique=False)
     # ### end Alembic commands ###
 
 
 def downgrade() -> None:
     # ### commands auto generated by Alembic - please adjust! ###
-    op.drop_index('idx_provider', table_name='token_usage')
-    op.drop_index('idx_model', table_name='token_usage')
-    op.drop_index('idx_execution_id', table_name='token_usage')
-    op.drop_index('idx_created_at', table_name='token_usage')
-    op.drop_table('token_usage')
-    # ### end Alembic commands ### 
+    op.drop_index("idx_provider", table_name="token_usage")
+    op.drop_index("idx_model", table_name="token_usage")
+    op.drop_index("idx_execution_id", table_name="token_usage")
+    op.drop_index("idx_created_at", table_name="token_usage")
+    op.drop_table("token_usage")
+    # ### end Alembic commands ###

tokenator/migrations.py

@@ -6,6 +6,7 @@ from alembic.config import Config
 from alembic import command
 from .utils import get_default_db_path
 
+
 def get_alembic_config(db_path: str = None) -> Config:
     """Get Alembic config for migrations."""
     if db_path is None:
@@ -13,27 +14,29 @@ def get_alembic_config(db_path: str = None) -> Config:
 
     # Get the directory containing this file
     migrations_dir = Path(__file__).parent / "migrations"
-    
+
     # Create Config object
     config = Config()
     config.set_main_option("script_location", str(migrations_dir))
     config.set_main_option("sqlalchemy.url", f"sqlite:///{db_path}")
-    
+
     return config
 
+
 def check_and_run_migrations(db_path: str = None):
     """Check and run any pending database migrations."""
     if db_path is None:
         db_path = get_default_db_path()
-        
+
     dirname = os.path.dirname(db_path)
     if dirname:
         os.makedirs(dirname, exist_ok=True)
-    
+
     # Initialize database
     import sqlite3
+
     conn = sqlite3.connect(db_path)
     conn.close()
-    
+
     config = get_alembic_config(db_path)
-    command.upgrade(config, "head")
+    command.upgrade(config, "head")

tokenator/models.py

@@ -1,31 +1,42 @@
 from pydantic import BaseModel, Field
-from typing import Dict, List
+from typing import List
+
 
 class TokenRate(BaseModel):
     prompt: float = Field(..., description="Cost per prompt token")
     completion: float = Field(..., description="Cost per completion token")
 
+
 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")
 
+
 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")
+    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")
+    providers: List[ProviderUsage] = Field(
+        default_factory=list, description="Usage breakdown by provider"
+    )
+
 
 class Usage(BaseModel):
     prompt_tokens: int = 0
     completion_tokens: int = 0
     total_tokens: int = 0
 
+
 class TokenUsageStats(BaseModel):
     model: str
-    usage: Usage
+    usage: Usage