tokenator 0.1.9__tar.gz → 0.1.11__tar.gz

{tokenator-0.1.9 → tokenator-0.1.11}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: tokenator
-Version: 0.1.9
+Version: 0.1.11
 Summary: Token usage tracking wrapper for LLMs
 License: MIT
 Author: Ujjwal Maheshwari
@@ -20,12 +20,12 @@ Requires-Dist: requests (>=2.32.3,<3.0.0)
 Requires-Dist: sqlalchemy (>=2.0.0,<3.0.0)
 Description-Content-Type: text/markdown
 
-# Tokenator : Easiest way to track and analyze LLM token usage and cost
+# Tokenator : Track and analyze LLM token usage and cost
 
 Have you ever wondered about :
 - 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 did I spent today on development?
+- How much money/tokens did you spend today on developing with LLMs?
 
 Afraid not, tokenator is here! With tokenator's easy to use API, you can start tracking LLM usage in a matter of minutes.
 
@@ -57,32 +57,37 @@ response = client.chat.completions.create(
 )
 ```
 
+Works with AsyncOpenAI and `streaming=True` as well!
+Note : When streaming, don't forget to add `stream_options={"include_usage": True}` to the `create()` call!
+
 ### 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
-# print(cost.last_hour().model_dump_json(indent=4))
+```python
+print(cost.last_hour().model_dump_json(indent=4))
+```
 
-usage : {
+```json
+{
     "total_cost": 0.0004,
     "total_tokens": 79,
     "prompt_tokens": 52,
@@ -118,6 +123,56 @@ usage : {
 - Minimal memory footprint
 - Minimal latency footprint
 
+### Anthropic
+
+```python
+from anthropic import Anthropic, AsyncAnthropic
+from tokenator import tokenator_anthropic
+
+anthropic_client = AsyncAnthropic(api_key="your-api-key")
+
+# Wrap it with Tokenator
+client = tokenator_anthropic(anthropic_client)
+
+# Use it exactly like the Anthropic client
+response = await client.messages.create(
+    model="claude-3-5-haiku-20241022",
+    messages=[{"role": "user", "content": "hello how are you"}],
+    max_tokens=20,
+)
+
+print(response)
+
+print(usage.last_execution().model_dump_json(indent=4))
+"""
+{
+    "total_cost": 0.0001,
+    "total_tokens": 23,
+    "prompt_tokens": 10,
+    "completion_tokens": 13,
+    "providers": [
+        {
+            "total_cost": 0.0001,
+            "total_tokens": 23,
+            "prompt_tokens": 10,
+            "completion_tokens": 13,
+            "provider": "anthropic",
+            "models": [
+                {
+                    "total_cost": 0.0004,
+                    "total_tokens": 79,
+                    "prompt_tokens": 52,
+                    "completion_tokens": 27,
+                    "model": "claude-3-5-haiku-20241022"
+                }
+            ]
+        }
+    ]
+}
+"""
+```
+---
+
 Most importantly, none of your data is ever sent to any server.
 
 ## License

{tokenator-0.1.9 → tokenator-0.1.11}/README.md

@@ -1,9 +1,9 @@
-# Tokenator : Easiest way to track and analyze LLM token usage and cost
+# Tokenator : Track and analyze LLM token usage and cost
 
 Have you ever wondered about :
 - 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 did I spent today on development?
+- How much money/tokens did you spend today on developing with LLMs?
 
 Afraid not, tokenator is here! With tokenator's easy to use API, you can start tracking LLM usage in a matter of minutes.
 
@@ -35,32 +35,37 @@ response = client.chat.completions.create(
 )
 ```
 
+Works with AsyncOpenAI and `streaming=True` as well!
+Note : When streaming, don't forget to add `stream_options={"include_usage": True}` to the `create()` call!
+
 ### 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
-# print(cost.last_hour().model_dump_json(indent=4))
+```python
+print(cost.last_hour().model_dump_json(indent=4))
+```
 
-usage : {
+```json
+{
     "total_cost": 0.0004,
     "total_tokens": 79,
     "prompt_tokens": 52,
@@ -96,6 +101,56 @@ usage : {
 - Minimal memory footprint
 - Minimal latency footprint
 
+### Anthropic
+
+```python
+from anthropic import Anthropic, AsyncAnthropic
+from tokenator import tokenator_anthropic
+
+anthropic_client = AsyncAnthropic(api_key="your-api-key")
+
+# Wrap it with Tokenator
+client = tokenator_anthropic(anthropic_client)
+
+# Use it exactly like the Anthropic client
+response = await client.messages.create(
+    model="claude-3-5-haiku-20241022",
+    messages=[{"role": "user", "content": "hello how are you"}],
+    max_tokens=20,
+)
+
+print(response)
+
+print(usage.last_execution().model_dump_json(indent=4))
+"""
+{
+    "total_cost": 0.0001,
+    "total_tokens": 23,
+    "prompt_tokens": 10,
+    "completion_tokens": 13,
+    "providers": [
+        {
+            "total_cost": 0.0001,
+            "total_tokens": 23,
+            "prompt_tokens": 10,
+            "completion_tokens": 13,
+            "provider": "anthropic",
+            "models": [
+                {
+                    "total_cost": 0.0004,
+                    "total_tokens": 79,
+                    "prompt_tokens": 52,
+                    "completion_tokens": 27,
+                    "model": "claude-3-5-haiku-20241022"
+                }
+            ]
+        }
+    ]
+}
+"""
+```
+---
+
 Most importantly, none of your data is ever sent to any server.
 
 ## License

{tokenator-0.1.9 → tokenator-0.1.11}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "tokenator"
-version = "0.1.9"
+version = "0.1.11"
 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.9 → tokenator-0.1.11}/src/tokenator/__init__.py

@@ -2,7 +2,7 @@
 
 import logging
 from .openai.client_openai import tokenator_openai
-from .client_anthropic import tokenator_anthropic
+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.11/src/tokenator/anthropic/client_anthropic.py

@@ -0,0 +1,154 @@
+"""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.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.11/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.9 → tokenator-0.1.11}/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.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):
+    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()