tokenator 0.1.0__tar.gz

tokenator-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Ujjwal Maheshwari
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

tokenator-0.1.0/PKG-INFO

@@ -0,0 +1,100 @@
+Metadata-Version: 2.1
+Name: tokenator
+Version: 0.1.0
+Summary: Token usage tracking wrapper for LLMs
+License: MIT
+Author: Ujjwal Maheshwari
+Author-email: your.email@example.com
+Requires-Python: >=3.9,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+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: 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
+
+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?
+
+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.
+
+Get started with just 3 lines of code!
+
+## Installation
+
+```bash
+pip install tokenator
+```
+
+## Usage
+
+### OpenAI
+
+```python
+from openai import OpenAI
+from tokenator import tokenator_openai
+
+openai_client = OpenAI(api_key="your-api-key")
+
+# Wrap it with Tokenator
+client = tokenator_openai(openai_client)
+
+# Use it exactly like the OpenAI client
+response = client.chat.completions.create(
+    model="gpt-4o",
+    messages=[{"role": "user", "content": "Hello!"}]
+)
+```
+
+### Cost Analysis
+
+```python
+from tokenator import cost
+
+# Get usage for different time periods
+cost.last_hour()
+cost.last_day()
+cost.last_week()
+cost.last_month()
+
+# Custom date range
+cost.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")
+```
+
+### Example `cost` object
+
+```json
+
+```
+
+## Features
+
+- Drop-in replacement for OpenAI, Anthropic client
+- Automatic token usage tracking
+- Cost analysis for different time periods
+- SQLite storage with zero configuration
+- Thread-safe operations
+- Minimal memory footprint
+- Minimal latency footprint
+
+Most importantly, none of your data is ever sent to any server.
+
+## License
+
+MIT 

tokenator-0.1.0/README.md

@@ -0,0 +1,78 @@
+# Tokenator : Easiest way to 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?
+
+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.
+
+Get started with just 3 lines of code!
+
+## Installation
+
+```bash
+pip install tokenator
+```
+
+## Usage
+
+### OpenAI
+
+```python
+from openai import OpenAI
+from tokenator import tokenator_openai
+
+openai_client = OpenAI(api_key="your-api-key")
+
+# Wrap it with Tokenator
+client = tokenator_openai(openai_client)
+
+# Use it exactly like the OpenAI client
+response = client.chat.completions.create(
+    model="gpt-4o",
+    messages=[{"role": "user", "content": "Hello!"}]
+)
+```
+
+### Cost Analysis
+
+```python
+from tokenator import cost
+
+# Get usage for different time periods
+cost.last_hour()
+cost.last_day()
+cost.last_week()
+cost.last_month()
+
+# Custom date range
+cost.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")
+```
+
+### Example `cost` object
+
+```json
+
+```
+
+## Features
+
+- Drop-in replacement for OpenAI, Anthropic client
+- Automatic token usage tracking
+- Cost analysis for different time periods
+- SQLite storage with zero configuration
+- Thread-safe operations
+- Minimal memory footprint
+- Minimal latency footprint
+
+Most importantly, none of your data is ever sent to any server.
+
+## License
+
+MIT 

tokenator-0.1.0/pyproject.toml

@@ -0,0 +1,29 @@
+[tool.poetry]
+name = "tokenator"
+version = "0.1.0"
+description = "Token usage tracking wrapper for LLMs"
+authors = ["Ujjwal Maheshwari <your.email@example.com>"]
+readme = "README.md"
+license = "MIT"
+packages = [{include = "tokenator", from = "src"}]
+
+[tool.poetry.dependencies]
+python = "^3.9"
+openai = "^1.57.0"
+sqlalchemy = "^2.0.0"
+requests = "^2.32.3"
+alembic = "^1.13.0"
+anthropic = "^0.40.0"
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^8.0.0"
+pytest-asyncio = "^0.23.0"
+pytest-cov = "^4.1.0"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.pytest.ini_options]
+pythonpath = ["src"]
+asyncio_mode = "auto" 

tokenator-0.1.0/src/tokenator/__init__.py

@@ -0,0 +1,12 @@
+"""Tokenator - Track and analyze your OpenAI API token usage and costs."""
+
+from .client_openai import OpenAIWrapper
+from . import usage
+from .utils import get_default_db_path
+from .migrations import check_and_run_migrations
+
+__version__ = "0.1.0"
+__all__ = ["OpenAIWrapper", "usage", "get_default_db_path"]
+
+# Run migrations on import
+check_and_run_migrations() 

tokenator-0.1.0/src/tokenator/base_wrapper.py

@@ -0,0 +1,61 @@
+"""Base wrapper class for token usage tracking."""
+
+from pathlib import Path
+from typing import Any, Dict, Optional, TypeVar, Union
+import logging
+import uuid
+
+from .models import Usage, TokenUsageStats
+from .schemas import get_session, TokenUsage
+
+logger = logging.getLogger(__name__)
+
+ResponseType = TypeVar('ResponseType')
+
+class BaseWrapper:
+    def __init__(self, client: Any, db_path: Optional[str] = None):
+        """Initialize the base wrapper."""
+        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)
+
+        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:
+        """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,
+                prompt_tokens=token_usage_stats.usage.prompt_tokens,
+                completion_tokens=token_usage_stats.usage.completion_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)
+        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 execution_id:
+            execution_id = str(uuid.uuid4())
+
+        session = self.Session()
+        try:
+            try:
+                self._log_usage_impl(token_usage_stats, session, execution_id)
+                session.commit()
+            except Exception as e:
+                logger.error("Failed to log token usage: %s", str(e))
+                session.rollback()
+        finally:
+            session.close() 

tokenator-0.1.0/src/tokenator/client_anthropic.py

@@ -0,0 +1,148 @@
+"""Anthropic client wrapper with token usage tracking."""
+
+from typing import Any, Dict, Optional, TypeVar, 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
+
+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
+
+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)
+        
+        response = self.client.messages.create(*args, **kwargs)
+        
+        if not kwargs.get('stream', False):
+            usage_data = self._process_response_usage(response)
+            if usage_data:
+                self._log_usage(usage_data, execution_id=execution_id)
+            return response
+        
+        return self._wrap_streaming_response(response, execution_id)
+    
+    def _wrap_streaming_response(self, response_iter: Iterator[Message], execution_id: Optional[str]) -> Iterator[Message]:
+        """Wrap streaming response to capture final usage stats"""
+        usage_data: TokenUsageStats = TokenUsageStats(model="", usage=Usage())
+        for chunk in response_iter:
+            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
+                usage_data.usage.total_tokens = chunk.message.usage.input_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 += chunk.usage.input_tokens + chunk.usage.output_tokens
+            
+            yield chunk
+
+        self._log_usage(usage_data, execution_id=execution_id)
+
+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):
+            response = await self.client.messages.create(*args, **kwargs)
+            return self._wrap_streaming_response(response, execution_id)
+        
+        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
+
+    async def _wrap_streaming_response(self, response_iter: AsyncIterator[Message], execution_id: Optional[str]) -> AsyncIterator[Message]:
+        """Wrap streaming response to capture final usage stats"""
+        usage_data: TokenUsageStats = TokenUsageStats(model="", usage=Usage())
+        async for chunk in response_iter:
+            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
+                usage_data.usage.total_tokens = chunk.message.usage.input_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 += chunk.usage.input_tokens + chunk.usage.output_tokens
+            
+            yield chunk
+        
+            
+        self._log_usage(usage_data, execution_id=execution_id)
+
+@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.0/src/tokenator/client_openai.py

@@ -0,0 +1,149 @@
+"""OpenAI client wrapper with token usage tracking."""
+
+from typing import Any, Dict, Optional, TypeVar, Union, overload, Iterator, AsyncIterator
+import logging
+
+from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream
+from openai.types.chat import ChatCompletion, ChatCompletionChunk
+
+from .models import Usage, TokenUsageStats
+from .base_wrapper import BaseWrapper, ResponseType
+
+logger = logging.getLogger(__name__)
+
+class BaseOpenAIWrapper(BaseWrapper):
+    provider = "openai"
+
+    def _process_response_usage(self, response: ResponseType) -> Optional[TokenUsageStats]:
+        """Process and log usage statistics from a response."""
+        try:
+            if isinstance(response, ChatCompletion):
+                if response.usage is None:
+                    return None
+                usage = Usage(
+                    prompt_tokens=response.usage.prompt_tokens,
+                    completion_tokens=response.usage.completion_tokens,
+                    total_tokens=response.usage.total_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('prompt_tokens', 0),
+                    completion_tokens=usage_dict.get('completion_tokens', 0),
+                    total_tokens=usage_dict.get('total_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 chat(self):
+        return self
+
+    @property
+    def completions(self):
+        return self
+
+class OpenAIWrapper(BaseOpenAIWrapper):
+    def create(self, *args: Any, execution_id: Optional[str] = None, **kwargs: Any) -> Union[ChatCompletion, Iterator[ChatCompletion]]:
+        """Create a chat completion and log token usage."""
+        logger.debug("Creating chat completion with args: %s, kwargs: %s", args, kwargs)
+        
+        response = self.client.chat.completions.create(*args, **kwargs)
+        
+        if not kwargs.get('stream', False):
+            usage_data = self._process_response_usage(response)
+            if usage_data:
+                self._log_usage(usage_data, execution_id=execution_id)
+            return response
+        
+        return self._wrap_streaming_response(response, execution_id)
+    
+    def _wrap_streaming_response(self, response_iter: Stream[ChatCompletionChunk], execution_id: Optional[str]) -> Iterator[ChatCompletionChunk]:
+        """Wrap streaming response to capture final usage stats"""
+        chunks_with_usage = []
+        for chunk in response_iter:
+            if isinstance(chunk, ChatCompletionChunk) and chunk.usage is not None:
+                chunks_with_usage.append(chunk)
+            yield chunk
+            
+        if len(chunks_with_usage) > 0:
+            usage_data: TokenUsageStats = TokenUsageStats(model=chunks_with_usage[0].model, usage=Usage())
+            for chunk in chunks_with_usage:
+                usage_data.usage.prompt_tokens += chunk.usage.prompt_tokens
+                usage_data.usage.completion_tokens += chunk.usage.completion_tokens
+                usage_data.usage.total_tokens += chunk.usage.total_tokens
+            
+            self._log_usage(usage_data, execution_id=execution_id)
+            
+
+class AsyncOpenAIWrapper(BaseOpenAIWrapper):
+    async def create(self, *args: Any, execution_id: Optional[str] = None, **kwargs: Any) -> Union[ChatCompletion, AsyncIterator[ChatCompletion]]:
+        """Create a chat completion and log token usage."""
+        logger.debug("Creating chat completion with args: %s, kwargs: %s", args, kwargs)
+        
+        if kwargs.get('stream', False):
+            response = await self.client.chat.completions.create(*args, **kwargs)
+            return self._wrap_streaming_response(response, execution_id)
+        
+        response = await self.client.chat.completions.create(*args, **kwargs)
+        usage_data = self._process_response_usage(response)
+        if usage_data:
+            self._log_usage(usage_data, execution_id=execution_id)
+        return response
+
+    async def _wrap_streaming_response(self, response_iter: AsyncStream[ChatCompletionChunk], execution_id: Optional[str]) -> AsyncIterator[ChatCompletionChunk]:
+        """Wrap streaming response to capture final usage stats"""
+        chunks_with_usage = []
+        async for chunk in response_iter:
+            if isinstance(chunk, ChatCompletionChunk) and chunk.usage is not None:
+                chunks_with_usage.append(chunk)
+            yield chunk
+
+        if len(chunks_with_usage) > 0:
+            usage_data: TokenUsageStats = TokenUsageStats(model=chunks_with_usage[0].model, usage=Usage())
+            for chunk in chunks_with_usage:
+                usage_data.usage.prompt_tokens += chunk.usage.prompt_tokens
+                usage_data.usage.completion_tokens += chunk.usage.completion_tokens
+                usage_data.usage.total_tokens += chunk.usage.total_tokens
+            
+            self._log_usage(usage_data, execution_id=execution_id)
+
+@overload
+def tokenator_openai(
+    client: OpenAI,
+    db_path: Optional[str] = None,
+) -> OpenAIWrapper: ...
+
+@overload
+def tokenator_openai(
+    client: AsyncOpenAI,
+    db_path: Optional[str] = None,
+) -> AsyncOpenAIWrapper: ...
+
+def tokenator_openai(
+    client: Union[OpenAI, AsyncOpenAI],
+    db_path: Optional[str] = None,
+) -> Union[OpenAIWrapper, AsyncOpenAIWrapper]:
+    """Create a token-tracking wrapper for an OpenAI client.
+    
+    Args:
+        client: OpenAI or AsyncOpenAI client instance
+        db_path: Optional path to SQLite database for token tracking
+    """
+    if isinstance(client, OpenAI):
+        return OpenAIWrapper(client=client, db_path=db_path)
+    
+    if isinstance(client, AsyncOpenAI):
+        return AsyncOpenAIWrapper(client=client, db_path=db_path)
+        
+    raise ValueError("Client must be an instance of OpenAI or AsyncOpenAI")

tokenator-0.1.0/src/tokenator/create_migrations.py

@@ -0,0 +1,21 @@
+"""Development utilities for tokenator."""
+
+import os
+from pathlib import Path
+from alembic import command
+from tokenator.migrations import get_alembic_config
+
+def create_migration():
+    """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
+    command.revision(config, autogenerate=True, message="auto generated migration") 
+
+
+if __name__ == "__main__":
+    create_migration()

tokenator-0.1.0/src/tokenator/migrations/env.py

@@ -0,0 +1,53 @@
+from logging.config import fileConfig
+
+from sqlalchemy import engine_from_config
+from sqlalchemy import pool
+
+from alembic import context
+
+from tokenator.models import Base
+
+# this is the Alembic Config object
+config = context.config
+
+# Interpret the config file for Python logging.
+# This line sets up loggers basically.
+if config.config_file_name is not None:
+    fileConfig(config.config_file_name)
+
+# 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")
+    context.configure(
+        url=url,
+        target_metadata=target_metadata,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+def run_migrations_online() -> None:
+    """Run migrations in 'online' mode."""
+    connectable = engine_from_config(
+        config.get_section(config.config_ini_section, {}),
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+    )
+
+    with connectable.connect() as connection:
+        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() 

tokenator-0.1.0/src/tokenator/migrations/script.py.mako

@@ -0,0 +1,26 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision: str = ${repr(up_revision)}
+down_revision: Union[str, None] = ${repr(down_revision)}
+branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
+depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    ${downgrades if downgrades else "pass"} 

tokenator-0.1.0/src/tokenator/migrations.py

@@ -0,0 +1,40 @@
+"""Automatic database migrations manager."""
+
+import os
+from pathlib import Path
+from alembic import command
+from alembic.config import Config
+from alembic.runtime.migration import MigrationContext
+from alembic.script import ScriptDirectory
+from sqlalchemy import create_engine
+
+from .utils import get_default_db_path
+
+def get_alembic_config():
+    """Get Alembic config pointing to the package's migrations."""
+    package_dir = Path(__file__).parent
+    migrations_dir = package_dir / "migrations"
+    
+    alembic_cfg = Config()
+    alembic_cfg.set_main_option("script_location", str(migrations_dir))
+    alembic_cfg.set_main_option("sqlalchemy.url", f"sqlite:///{get_default_db_path()}")
+    
+    return alembic_cfg
+
+def check_and_run_migrations():
+    """Check if migrations are needed and run them automatically."""
+    engine = create_engine(f"sqlite:///{get_default_db_path()}")
+    
+    # Create migrations table if it doesn't exist
+    with engine.connect() as conn:
+        context = MigrationContext.configure(conn)
+        current_rev = context.get_current_revision()
+    
+    # Get latest available revision
+    config = get_alembic_config()
+    script = ScriptDirectory.from_config(config)
+    head_rev = script.get_current_head()
+    
+    # Run migrations if needed
+    if current_rev != head_rev:
+        command.upgrade(config, "head") 

tokenator-0.1.0/src/tokenator/models.py

@@ -0,0 +1,31 @@
+from pydantic import BaseModel, Field
+from typing import Dict, 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")
+
+class TokenUsageReport(TokenMetrics):
+    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

tokenator-0.1.0/src/tokenator/schemas.py

@@ -0,0 +1,60 @@
+"""SQLAlchemy models for tokenator."""
+
+import uuid
+from datetime import datetime
+
+from sqlalchemy import create_engine, Column, Integer, String, DateTime, Float, Index
+from sqlalchemy.orm import sessionmaker, scoped_session, declarative_base
+
+from .utils import get_default_db_path
+
+Base = declarative_base()
+
+def get_engine(db_path: str = None):
+    """Create SQLAlchemy engine with the given database path."""
+    db_path = db_path or get_default_db_path()
+    return create_engine(f"sqlite:///{db_path}", echo=False)
+
+def get_session(db_path: str = None):
+    """Create a thread-safe session factory."""
+    engine = get_engine(db_path)
+    Base.metadata.create_all(engine)
+    session_factory = sessionmaker(bind=engine)
+    return scoped_session(session_factory)
+
+class TokenUsage(Base):
+    """Model for tracking token usage."""
+    
+    __tablename__ = "token_usage"
+    
+    id = Column(Integer, primary_key=True)
+    execution_id = Column(String, nullable=False)
+    provider = Column(String, nullable=False)
+    model = Column(String, nullable=False)
+    created_at = Column(DateTime, nullable=False, default=datetime.now)
+    updated_at = Column(DateTime, nullable=False, default=datetime.now, onupdate=datetime.now)
+    prompt_tokens = Column(Integer, nullable=False)
+    completion_tokens = Column(Integer, nullable=False)
+    total_tokens = Column(Integer, nullable=False)
+    
+    # Create indexes
+    __table_args__ = (
+        Index('idx_created_at', 'created_at'),
+        Index('idx_execution_id', 'execution_id'),
+        Index('idx_provider', 'provider'),
+        Index('idx_model', 'model'),
+    )
+    
+    def to_dict(self):
+        """Convert model instance to dictionary."""
+        return {
+            'id': self.id,
+            'execution_id': self.execution_id,
+            'provider': self.provider,
+            'model': self.model,
+            'created_at': self.created_at,
+            'updated_at': self.updated_at,
+            'prompt_tokens': self.prompt_tokens,
+            'completion_tokens': self.completion_tokens,
+            'total_tokens': self.total_tokens
+        } 

tokenator-0.1.0/src/tokenator/usage.py

@@ -0,0 +1,202 @@
+"""Cost analysis functions for token usage."""
+
+from datetime import datetime, timedelta, timezone
+from typing import Dict, Optional, Union
+
+from sqlalchemy import and_
+
+from .schemas import get_session, TokenUsage
+from .models import TokenRate, TokenUsageReport, ModelUsage, ProviderUsage
+
+import requests
+import logging
+
+logger = logging.getLogger(__name__)
+
+def _get_model_costs() -> Dict[str, TokenRate]:
+    url = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
+    response = requests.get(url)
+    data = response.json()
+    
+    return {
+        model: TokenRate(
+            prompt=info["input_cost_per_token"],
+            completion=info["output_cost_per_token"]
+        )
+        for model, info in data.items()
+        if "input_cost_per_token" in info and "output_cost_per_token" in info
+    }
+
+MODEL_COSTS = _get_model_costs()
+
+def _calculate_cost(usages: list[TokenUsage], provider: Optional[str] = None) -> TokenUsageReport:
+    """Calculate cost from token usage records."""
+    # Group usages by provider and model
+    provider_model_usages: Dict[str, Dict[str, list[TokenUsage]]] = {}
+
+    print(f"usages: {len(usages)}")
+    
+    for usage in usages:
+        if usage.model not in MODEL_COSTS:
+            continue
+            
+        provider = usage.provider
+        if provider not in provider_model_usages:
+            provider_model_usages[provider] = {}
+        
+        if usage.model not in provider_model_usages[provider]:
+            provider_model_usages[provider][usage.model] = []
+            
+        provider_model_usages[provider][usage.model].append(usage)
+
+    # Calculate totals for each level
+    providers_list = []
+    total_metrics = {"total_cost": 0.0, "total_tokens": 0, "prompt_tokens": 0, "completion_tokens": 0}
+
+    for provider, model_usages in provider_model_usages.items():
+        provider_metrics = {"total_cost": 0.0, "total_tokens": 0, "prompt_tokens": 0, "completion_tokens": 0}
+        models_list = []
+
+        for model, usages in model_usages.items():
+            model_cost = 0.0
+            model_total = 0
+            model_prompt = 0
+            model_completion = 0
+
+            for usage in usages:
+                model_prompt += usage.prompt_tokens
+                model_completion += usage.completion_tokens
+                model_total += usage.total_tokens
+                
+                model_cost += (usage.prompt_tokens * MODEL_COSTS[usage.model].prompt)
+                model_cost += (usage.completion_tokens * MODEL_COSTS[usage.model].completion)
+
+            models_list.append(ModelUsage(
+                model=model,
+                total_cost=round(model_cost, 6),
+                total_tokens=model_total,
+                prompt_tokens=model_prompt,
+                completion_tokens=model_completion
+            ))
+
+            # Add to provider totals
+            provider_metrics["total_cost"] += model_cost
+            provider_metrics["total_tokens"] += model_total
+            provider_metrics["prompt_tokens"] += model_prompt
+            provider_metrics["completion_tokens"] += model_completion
+
+        providers_list.append(ProviderUsage(
+            provider=provider,
+            models=models_list,
+            **{k: (round(v, 6) if k == "total_cost" else v) for k, v in provider_metrics.items()}
+        ))
+
+        # Add to grand totals
+        for key in total_metrics:
+            total_metrics[key] += provider_metrics[key]
+
+    return TokenUsageReport(
+        providers=providers_list,
+        **{k: (round(v, 6) if k == "total_cost" else v) for k, v in total_metrics.items()}
+    )
+
+def _query_usage(start_date: datetime, end_date: datetime, 
+                provider: Optional[str] = None, 
+                model: Optional[str] = None) -> TokenUsageReport:
+    """Query token usage for a specific time period."""
+    session = get_session()()
+    try:
+        query = session.query(TokenUsage).filter(
+            TokenUsage.created_at.between(start_date, end_date)
+        )
+        
+        if provider:
+            query = query.filter(TokenUsage.provider == provider)
+        if model:
+            query = query.filter(TokenUsage.model == model)
+            
+        usages = query.all()
+        return _calculate_cost(usages, provider or "all")
+    finally:
+        session.close()
+
+def last_hour(provider: Optional[str] = None, model: Optional[str] = None) -> TokenUsageReport:
+    """Get cost analysis for the last hour."""
+    logger.debug(f"Getting cost analysis for last hour (provider={provider}, model={model})")
+    end = datetime.now()
+    start = end - timedelta(hours=1)
+    return _query_usage(start, end, provider, model)
+
+def last_day(provider: Optional[str] = None, model: Optional[str] = None) -> TokenUsageReport:
+    """Get cost analysis for the last 24 hours."""
+    logger.debug(f"Getting cost analysis for last 24 hours (provider={provider}, model={model})")
+    end = datetime.now()
+    start = end - timedelta(days=1)
+    return _query_usage(start, end, provider, model)
+
+def last_week(provider: Optional[str] = None, model: Optional[str] = None) -> TokenUsageReport:
+    """Get cost analysis for the last 7 days."""
+    logger.debug(f"Getting cost analysis for last 7 days (provider={provider}, model={model})")
+    end = datetime.now()
+    start = end - timedelta(weeks=1)
+    return _query_usage(start, end, provider, model)
+
+def last_month(provider: Optional[str] = None, model: Optional[str] = None) -> TokenUsageReport:
+    """Get cost analysis for the last 30 days."""
+    logger.debug(f"Getting cost analysis for last 30 days (provider={provider}, model={model})")
+    end = datetime.now()
+    start = end - timedelta(days=30)
+    return _query_usage(start, end, provider, model)
+
+def between(
+    start_date: Union[datetime, str],
+    end_date: Union[datetime, str],
+    provider: Optional[str] = None,
+    model: Optional[str] = None
+) -> TokenUsageReport:
+    """Get cost analysis between two dates.
+    
+    Args:
+        start_date: datetime object or string (format: YYYY-MM-DD or YYYY-MM-DD HH:MM:SS)
+        end_date: datetime object or string (format: YYYY-MM-DD or YYYY-MM-DD HH:MM:SS)
+    """
+    logger.debug(f"Getting cost analysis between {start_date} and {end_date} (provider={provider}, model={model})")
+    
+    if isinstance(start_date, str):
+        try:
+            start = datetime.strptime(start_date, "%Y-%m-%d %H:%M:%S")
+        except ValueError:
+            start = datetime.strptime(start_date, "%Y-%m-%d")
+    else:
+        start = start_date
+
+    if isinstance(end_date, str):
+        try:
+            end = datetime.strptime(end_date, "%Y-%m-%d %H:%M:%S")
+        except ValueError:
+            end = datetime.strptime(end_date, "%Y-%m-%d") + timedelta(days=1)  # Include the end date
+    else:
+        end = end_date
+
+    return _query_usage(start, end, provider, model)
+
+def for_execution(execution_id: str) -> TokenUsageReport:
+    """Get cost analysis for a specific execution."""
+    logger.debug(f"Getting cost analysis for execution_id={execution_id}")
+    session = get_session()()
+    query = session.query(TokenUsage).filter(TokenUsage.execution_id == execution_id)
+    return _calculate_cost(query.all())
+
+def last_execution() -> TokenUsageReport:
+    """Get cost analysis for the last execution_id."""
+    logger.debug("Getting cost analysis for last execution")
+    session = get_session()()
+    query = session.query(TokenUsage).order_by(TokenUsage.created_at.desc()).first()
+    return for_execution(query.execution_id)
+
+def all_time() -> TokenUsageReport:
+    """Get cost analysis for all time."""
+    logger.warning("Getting cost analysis for all time. This may take a while...")
+    session = get_session()()
+    query = session.query(TokenUsage).all()
+    return for_execution(query.execution_id)

tokenator-0.1.0/src/tokenator/utils.py

@@ -0,0 +1,35 @@
+"""Shared utility functions for tokenator."""
+
+import os
+import platform
+from pathlib import Path
+from typing import Optional
+
+def get_default_db_path() -> str:
+    """Get the platform-specific default database path.
+    
+    Returns:
+        str: Path to the SQLite database file
+        
+    The path follows platform conventions:
+        - Linux/macOS: ~/.local/share/tokenator/usage.db (XDG spec)
+        - Windows: %LOCALAPPDATA%\\tokenator\\usage.db
+        - Others: ~/.tokenator/usage.db
+    """
+    system = platform.system().lower()
+    
+    if system == "linux" or system == "darwin":
+        # Follow XDG Base Directory Specification
+        xdg_data_home = os.environ.get("XDG_DATA_HOME", "")
+        if not xdg_data_home:
+            xdg_data_home = os.path.join(str(Path.home()), ".local", "share")
+        return os.path.join(xdg_data_home, "tokenator", "usage.db")
+    elif system == "windows":
+        # Use %LOCALAPPDATA% on Windows
+        local_app_data = os.environ.get("LOCALAPPDATA", "")
+        if not local_app_data:
+            local_app_data = os.path.join(str(Path.home()), "AppData", "Local")
+        return os.path.join(local_app_data, "tokenator", "usage.db")
+    else:
+        # Fallback for other systems
+        return os.path.join(str(Path.home()), ".tokenator", "usage.db")