tokenator 0.1.8__tar.gz → 0.1.10__tar.gz

{tokenator-0.1.8 → tokenator-0.1.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tokenator
-Version: 0.1.8
+Version: 0.1.10
 Summary: Token usage tracking wrapper for LLMs
 License: MIT
 Author: Ujjwal Maheshwari
@@ -27,7 +27,7 @@ Have you ever wondered about :
 - How much does it cost to do run a complex AI workflow with multiple LLM providers?
 - How much money did I spent today on development?
 
-Afraid not, tokenator is here! With tokenator's easy to use API, you can start tracking LLM usage in a matter of minutes and track your LLM usage.
+Afraid not, tokenator is here! With tokenator's easy to use API, you can start tracking LLM usage in a matter of minutes.
 
 Get started with just 3 lines of code!
 
@@ -60,27 +60,54 @@ response = client.chat.completions.create(
 ### Cost Analysis
 
 ```python
-from tokenator import cost
+from tokenator import usage
 
 # Get usage for different time periods
-cost.last_hour()
-cost.last_day()
-cost.last_week()
-cost.last_month()
+usage.last_hour()
+usage.last_day()
+usage.last_week()
+usage.last_month()
 
 # Custom date range
-cost.between("2024-03-01", "2024-03-15")
+usage.between("2024-03-01", "2024-03-15")
 
 # Get usage for different LLM providers
-cost.last_day("openai")
-cost.last_day("anthropic")
-cost.last_day("google")
+usage.last_day("openai")
+usage.last_day("anthropic")
+usage.last_day("google")
 ```
 
-### Example `cost` object
+### Example `usage` object
 
-```json
+```python
+print(cost.last_hour().model_dump_json(indent=4))
+```
 
+```json
+{
+    "total_cost": 0.0004,
+    "total_tokens": 79,
+    "prompt_tokens": 52,
+    "completion_tokens": 27,
+    "providers": [
+        {
+            "total_cost": 0.0004,
+            "total_tokens": 79,
+            "prompt_tokens": 52,
+            "completion_tokens": 27,
+            "provider": "openai",
+            "models": [
+                {
+                    "total_cost": 0.0004,
+                    "total_tokens": 79,
+                    "prompt_tokens": 52,
+                    "completion_tokens": 27,
+                    "model": "gpt-4o-2024-08-06"
+                }
+            ]
+        }
+    ]
+}
 ```
 
 ## Features

{tokenator-0.1.8 → tokenator-0.1.10}/README.md

@@ -5,7 +5,7 @@ Have you ever wondered about :
 - How much does it cost to do run a complex AI workflow with multiple LLM providers?
 - How much money did I spent today on development?
 
-Afraid not, tokenator is here! With tokenator's easy to use API, you can start tracking LLM usage in a matter of minutes and track your LLM usage.
+Afraid not, tokenator is here! With tokenator's easy to use API, you can start tracking LLM usage in a matter of minutes.
 
 Get started with just 3 lines of code!
 
@@ -38,27 +38,54 @@ response = client.chat.completions.create(
 ### Cost Analysis
 
 ```python
-from tokenator import cost
+from tokenator import usage
 
 # Get usage for different time periods
-cost.last_hour()
-cost.last_day()
-cost.last_week()
-cost.last_month()
+usage.last_hour()
+usage.last_day()
+usage.last_week()
+usage.last_month()
 
 # Custom date range
-cost.between("2024-03-01", "2024-03-15")
+usage.between("2024-03-01", "2024-03-15")
 
 # Get usage for different LLM providers
-cost.last_day("openai")
-cost.last_day("anthropic")
-cost.last_day("google")
+usage.last_day("openai")
+usage.last_day("anthropic")
+usage.last_day("google")
 ```
 
-### Example `cost` object
+### Example `usage` object
 
-```json
+```python
+print(cost.last_hour().model_dump_json(indent=4))
+```
 
+```json
+{
+    "total_cost": 0.0004,
+    "total_tokens": 79,
+    "prompt_tokens": 52,
+    "completion_tokens": 27,
+    "providers": [
+        {
+            "total_cost": 0.0004,
+            "total_tokens": 79,
+            "prompt_tokens": 52,
+            "completion_tokens": 27,
+            "provider": "openai",
+            "models": [
+                {
+                    "total_cost": 0.0004,
+                    "total_tokens": 79,
+                    "prompt_tokens": 52,
+                    "completion_tokens": 27,
+                    "model": "gpt-4o-2024-08-06"
+                }
+            ]
+        }
+    ]
+}
 ```
 
 ## Features

{tokenator-0.1.8 → tokenator-0.1.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tokenator"
-version = "0.1.8"
+version = "0.1.10"
 description = "Token usage tracking wrapper for LLMs"
 authors = ["Ujjwal Maheshwari <your.email@example.com>"]
 readme = "README.md"
@@ -19,11 +19,12 @@ anthropic = "^0.40.0"
 pytest = "^8.0.0"
 pytest-asyncio = "^0.23.0"
 pytest-cov = "^4.1.0"
+ruff = "^0.8.4"
 
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.pytest.ini_options]
-pythonpath = ["src"]
-asyncio_mode = "auto" 
+testpaths = ["tests"] 
+pythonpath = "src"

{tokenator-0.1.8 → tokenator-0.1.10}/src/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-0.1.10/src/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-0.1.10/src/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-0.1.8 → tokenator-0.1.10}/src/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-0.1.8 → tokenator-0.1.10}/src/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-0.1.8 → tokenator-0.1.10}/src/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()