tokenator 0.1.12__py3-none-any.whl → 0.1.14__py3-none-any.whl

tokenator/__init__.py

@@ -5,14 +5,9 @@ from .openai.client_openai import tokenator_openai
 from .anthropic.client_anthropic import tokenator_anthropic
 from . import usage
 from .utils import get_default_db_path
-from .migrations import check_and_run_migrations
+from .usage import TokenUsageService
 
-__version__ = "0.1.0"
+usage = TokenUsageService()  # noqa: F811
 __all__ = ["tokenator_openai", "tokenator_anthropic", "usage", "get_default_db_path"]
 
 logger = logging.getLogger(__name__)
-
-try:
-    check_and_run_migrations()
-except Exception as e:
-    logger.warning(f"Failed to run migrations, but continuing anyway: {e}")

tokenator/anthropic/client_anthropic.py

@@ -8,7 +8,11 @@ from anthropic.types import Message, RawMessageStartEvent, RawMessageDeltaEvent
 
 from ..models import Usage, TokenUsageStats
 from ..base_wrapper import BaseWrapper, ResponseType
-from .stream_interceptors import AnthropicAsyncStreamInterceptor, AnthropicSyncStreamInterceptor
+from .stream_interceptors import (
+    AnthropicAsyncStreamInterceptor,
+    AnthropicSyncStreamInterceptor,
+)
+from ..state import is_tokenator_enabled
 
 logger = logging.getLogger(__name__)
 
@@ -56,15 +60,23 @@ class BaseAnthropicWrapper(BaseWrapper):
 
 def _create_usage_callback(execution_id, log_usage_fn):
     """Creates a callback function for processing usage statistics from stream chunks."""
+
     def usage_callback(chunks):
         if not chunks:
             return
-        
+
+        # Skip if tokenator is disabled
+        if not is_tokenator_enabled:
+            logger.debug("Tokenator is disabled - skipping stream usage logging")
+            return
+
         usage_data = TokenUsageStats(
-            model=chunks[0].message.model if isinstance(chunks[0], RawMessageStartEvent) else "",
+            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
@@ -72,8 +84,10 @@ def _create_usage_callback(execution_id, log_usage_fn):
                 usage_data.usage.completion_tokens += chunk.message.usage.output_tokens
             elif isinstance(chunk, RawMessageDeltaEvent):
                 usage_data.usage.completion_tokens += chunk.usage.output_tokens
-        
-        usage_data.usage.total_tokens = usage_data.usage.prompt_tokens + usage_data.usage.completion_tokens
+
+        usage_data.usage.total_tokens = (
+            usage_data.usage.prompt_tokens + usage_data.usage.completion_tokens
+        )
         log_usage_fn(usage_data, execution_id=execution_id)
 
     return usage_callback
@@ -84,7 +98,9 @@ class AnthropicWrapper(BaseAnthropicWrapper):
         self, *args: Any, execution_id: Optional[str] = None, **kwargs: Any
     ) -> Union[Message, Iterator[Message]]:
         """Create a message completion and log token usage."""
-        logger.debug("Creating message completion with args: %s, kwargs: %s", args, kwargs)
+        logger.debug(
+            "Creating message completion with args: %s, kwargs: %s", args, kwargs
+        )
 
         if kwargs.get("stream", False):
             base_stream = self.client.messages.create(*args, **kwargs)
@@ -105,7 +121,9 @@ class AsyncAnthropicWrapper(BaseAnthropicWrapper):
         self, *args: Any, execution_id: Optional[str] = None, **kwargs: Any
     ) -> Union[Message, AsyncIterator[Message]]:
         """Create a message completion and log token usage."""
-        logger.debug("Creating message completion with args: %s, kwargs: %s", args, kwargs)
+        logger.debug(
+            "Creating message completion with args: %s, kwargs: %s", args, kwargs
+        )
 
         if kwargs.get("stream", False):
             base_stream = await self.client.messages.create(*args, **kwargs)

tokenator/base_wrapper.py

@@ -7,6 +7,9 @@ import uuid
 
 from .models import TokenUsageStats
 from .schemas import get_session, TokenUsage
+from . import state
+
+from .migrations import check_and_run_migrations
 
 logger = logging.getLogger(__name__)
 
@@ -16,17 +19,30 @@ ResponseType = TypeVar("ResponseType")
 class BaseWrapper:
     def __init__(self, client: Any, db_path: Optional[str] = None):
         """Initialize the base wrapper."""
-        self.client = client
+        state.is_tokenator_enabled = True
+        try:
+            self.client = client
 
-        if db_path:
-            Path(db_path).parent.mkdir(parents=True, exist_ok=True)
-            logger.info("Created database directory at: %s", Path(db_path).parent)
+            if db_path:
+                Path(db_path).parent.mkdir(parents=True, exist_ok=True)
+                logger.info("Created database directory at: %s", Path(db_path).parent)
+                state.db_path = db_path  # Store db_path in state
 
-        self.Session = get_session(db_path)
+            else:
+                state.db_path = None  # Use default path
 
-        logger.debug(
-            "Initializing %s with db_path: %s", self.__class__.__name__, db_path
-        )
+            self.Session = get_session()
+
+            logger.debug(
+                "Initializing %s with db_path: %s", self.__class__.__name__, db_path
+            )
+
+            check_and_run_migrations(db_path)
+        except Exception as e:
+            state.is_tokenator_enabled = False
+            logger.warning(
+                f"Tokenator initialization failed. Usage tracking will be disabled. Error: {e}"
+            )
 
     def _log_usage_impl(
         self, token_usage_stats: TokenUsageStats, session, execution_id: str
@@ -59,6 +75,10 @@ class BaseWrapper:
         self, token_usage_stats: TokenUsageStats, execution_id: Optional[str] = None
     ):
         """Log token usage to database."""
+        if not state.is_tokenator_enabled:
+            logger.debug("Tokenator is disabled - skipping usage logging")
+            return
+
         if not execution_id:
             execution_id = str(uuid.uuid4())
 

tokenator/models.py

@@ -8,10 +8,10 @@ class TokenRate(BaseModel):
 
 
 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")
+    total_cost: float = Field(default=0, description="Total cost in USD")
+    total_tokens: int = Field(default=0, description="Total tokens used")
+    prompt_tokens: int = Field(default=0, description="Number of prompt tokens")
+    completion_tokens: int = Field(default=0, description="Number of completion tokens")
 
 
 class ModelUsage(TokenMetrics):

tokenator/openai/client_openai.py

@@ -8,7 +8,11 @@ from openai.types.chat import ChatCompletion, ChatCompletionChunk
 
 from ..models import Usage, TokenUsageStats
 from ..base_wrapper import BaseWrapper, ResponseType
-from .stream_interceptors import OpenAIAsyncStreamInterceptor, OpenAISyncStreamInterceptor
+from .stream_interceptors import (
+    OpenAIAsyncStreamInterceptor,
+    OpenAISyncStreamInterceptor,
+)
+from ..state import is_tokenator_enabled
 
 logger = logging.getLogger(__name__)
 
@@ -65,6 +69,12 @@ def _create_usage_callback(execution_id, log_usage_fn):
     def usage_callback(chunks):
         if not chunks:
             return
+
+        # Skip if tokenator is disabled
+        if not is_tokenator_enabled:
+            logger.debug("Tokenator is disabled - skipping stream usage logging")
+            return
+
         # Build usage_data from the first chunk's model
         usage_data = TokenUsageStats(
             model=chunks[0].model,

tokenator/schemas.py

@@ -1,25 +1,27 @@
 """SQLAlchemy models for tokenator."""
 
 from datetime import datetime
+from typing import Optional
 
 from sqlalchemy import create_engine, Column, Integer, String, DateTime, Index
 from sqlalchemy.orm import sessionmaker, scoped_session, declarative_base
 
 from .utils import get_default_db_path
+from . import state  # Import state to access db_path
 
 Base = declarative_base()
 
 
-def get_engine(db_path: str = None):
+def get_engine(db_path: Optional[str] = None):
     """Create SQLAlchemy engine with the given database path."""
     if db_path is None:
-        db_path = get_default_db_path()
+        db_path = state.db_path or get_default_db_path()  # Use state.db_path if set
     return create_engine(f"sqlite:///{db_path}", echo=False)
 
 
-def get_session(db_path: str = None):
+def get_session():
     """Create a thread-safe session factory."""
-    engine = get_engine(db_path)
+    engine = get_engine()
     # Base.metadata.create_all(engine)
     session_factory = sessionmaker(bind=engine)
     return scoped_session(session_factory)

tokenator/state.py

@@ -0,0 +1,12 @@
+"""Global state for tokenator."""
+
+import logging
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+# Global flag to track if tokenator is properly initialized
+is_tokenator_enabled = True
+
+# Store the database path
+db_path: Optional[str] = None

tokenator/usage.py

@@ -3,9 +3,9 @@
 from datetime import datetime, timedelta
 from typing import Dict, Optional, Union
 
-
 from .schemas import get_session, TokenUsage
 from .models import TokenRate, TokenUsageReport, ModelUsage, ProviderUsage
+from . import state
 
 import requests
 import logging
@@ -13,257 +13,295 @@ 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()
-
+class TokenUsageService:
+    def __init__(self):
+        if not state.is_tokenator_enabled:
+            logger.info("Tokenator is disabled. Database access is unavailable.")
 
-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]]] = {}
+        self.MODEL_COSTS = self._get_model_costs()
 
-    print(f"usages: {len(usages)}")
+    def _get_model_costs(self) -> Dict[str, TokenRate]:
+        if not state.is_tokenator_enabled:
+            return {}
+        url = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
+        response = requests.get(url)
+        data = response.json()
 
-    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] = {}
+        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
+        }
 
-        if usage.model not in provider_model_usages[provider]:
-            provider_model_usages[provider][usage.model] = []
+    def _calculate_cost(
+        self, usages: list[TokenUsage], provider: Optional[str] = None
+    ) -> TokenUsageReport:
+        if not state.is_tokenator_enabled:
+            logger.warning("Tokenator is disabled. Skipping cost calculation.")
+            return TokenUsageReport()
 
-        provider_model_usages[provider][usage.model].append(usage)
+        if not self.MODEL_COSTS:
+            logger.warning("No model costs available.")
+            return TokenUsageReport()
 
-    # Calculate totals for each level
-    providers_list = []
-    total_metrics = {
-        "total_cost": 0.0,
-        "total_tokens": 0,
-        "prompt_tokens": 0,
-        "completion_tokens": 0,
-    }
+        GPT4O_PRICING = self.MODEL_COSTS.get(
+            "gpt-4o", TokenRate(prompt=0.0000025, completion=0.000010)
+        )
 
-    for provider, model_usages in provider_model_usages.items():
-        provider_metrics = {
+        # Existing calculation logic...
+        provider_model_usages: Dict[str, Dict[str, list[TokenUsage]]] = {}
+        logger.debug(f"usages: {len(usages)}")
+
+        for usage in usages:
+            # 1st priority - direct match
+            model_key = usage.model
+            if model_key in self.MODEL_COSTS:
+                pass
+            # 2nd priority - provider/model format
+            elif f"{usage.provider}/{usage.model}" in self.MODEL_COSTS:
+                model_key = f"{usage.provider}/{usage.model}"
+            # 3rd priority - contains search
+            else:
+                matched_keys = [k for k in self.MODEL_COSTS.keys() if usage.model in k]
+                if matched_keys:
+                    model_key = matched_keys[0]
+                    logger.warning(
+                        f"Model {usage.model} matched with {model_key} in pricing data via contains search"
+                    )
+                else:
+                    # Fallback to GPT4O pricing
+                    logger.warning(
+                        f"Model {model_key} not found in pricing data. Using gpt-4o pricing as fallback "
+                        f"(prompt: ${GPT4O_PRICING.prompt}/token, completion: ${GPT4O_PRICING.completion}/token)"
+                    )
+                    self.MODEL_COSTS[model_key] = GPT4O_PRICING
+
+            provider_key = usage.provider or "default"
+            provider_model_usages.setdefault(provider_key, {}).setdefault(
+                model_key, []
+            ).append(usage)
+
+        # Calculate totals for each level
+        providers_list = []
+        total_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
+
+        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_key, usages in model_usages.items():
+                model_cost = sum(
+                    usage.prompt_tokens * self.MODEL_COSTS[model_key].prompt
+                    + usage.completion_tokens * self.MODEL_COSTS[model_key].completion
+                    for usage in usages
+                )
+                model_total = sum(usage.total_tokens for usage in usages)
+                model_prompt = sum(usage.prompt_tokens for usage in usages)
+                model_completion = sum(usage.completion_tokens for usage in usages)
+
+                models_list.append(
+                    ModelUsage(
+                        model=model_key,
+                        total_cost=round(model_cost, 6),
+                        total_tokens=model_total,
+                        prompt_tokens=model_prompt,
+                        completion_tokens=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,
+                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 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()
-                },
-            )
-        )
+            for key in total_metrics:
+                total_metrics[key] += provider_metrics[key]
 
-        # 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)
+        return TokenUsageReport(
+            providers=providers_list,
+            **{
+                k: (round(v, 6) if k == "total_cost" else v)
+                for k, v in total_metrics.items()
+            },
         )
 
-        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:
-            logger.warning(
-                f"Date-only string provided for start_date: {start_date}. Setting time to 00:00:00"
-            )
-            start = datetime.strptime(start_date, "%Y-%m-%d")
-
-    else:
-        start = start_date
-
-    if isinstance(end_date, str):
+    def _query_usage(
+        self,
+        start_date: datetime,
+        end_date: datetime,
+        provider: Optional[str] = None,
+        model: Optional[str] = None,
+    ) -> TokenUsageReport:
+        if not state.is_tokenator_enabled:
+            logger.warning("Tokenator is disabled. Skipping usage query.")
+            return TokenUsageReport()
+
+        session = get_session()()
         try:
-            end = datetime.strptime(end_date, "%Y-%m-%d %H:%M:%S")
-        except ValueError:
-            logger.warning(
-                f"Date-only string provided for end_date: {end_date}. Setting time to 23:59:59"
+            query = session.query(TokenUsage).filter(
+                TokenUsage.created_at.between(start_date, end_date)
             )
-            end = (
-                datetime.strptime(end_date, "%Y-%m-%d")
-                + timedelta(days=1)
-                - timedelta(seconds=1)
-            )
-    else:
-        end = end_date
 
-    return _query_usage(start, end, provider, model)
+            if provider:
+                query = query.filter(TokenUsage.provider == provider)
+            if model:
+                query = query.filter(TokenUsage.model == model)
 
+            usages = query.all()
 
-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())
+            return self._calculate_cost(usages, provider or "all")
+        finally:
+            session.close()
 
+    def last_hour(
+        self, provider: Optional[str] = None, model: Optional[str] = None
+    ) -> TokenUsageReport:
+        if not state.is_tokenator_enabled:
+            return TokenUsageReport()
+        logger.debug(
+            f"Getting cost analysis for last hour (provider={provider}, model={model})"
+        )
+        end = datetime.now()
+        start = end - timedelta(hours=1)
+        return self._query_usage(start, end, provider, model)
+
+    def last_day(
+        self, provider: Optional[str] = None, model: Optional[str] = None
+    ) -> TokenUsageReport:
+        if not state.is_tokenator_enabled:
+            return TokenUsageReport()
+        logger.debug(
+            f"Getting cost analysis for last 24 hours (provider={provider}, model={model})"
+        )
+        end = datetime.now()
+        start = end - timedelta(days=1)
+        return self._query_usage(start, end, provider, model)
+
+    def last_week(
+        self, provider: Optional[str] = None, model: Optional[str] = None
+    ) -> TokenUsageReport:
+        if not state.is_tokenator_enabled:
+            return TokenUsageReport()
+        logger.debug(
+            f"Getting cost analysis for last 7 days (provider={provider}, model={model})"
+        )
+        end = datetime.now()
+        start = end - timedelta(weeks=1)
+        return self._query_usage(start, end, provider, model)
+
+    def last_month(
+        self, provider: Optional[str] = None, model: Optional[str] = None
+    ) -> TokenUsageReport:
+        if not state.is_tokenator_enabled:
+            return TokenUsageReport()
+        logger.debug(
+            f"Getting cost analysis for last 30 days (provider={provider}, model={model})"
+        )
+        end = datetime.now()
+        start = end - timedelta(days=30)
+        return self._query_usage(start, end, provider, model)
+
+    def between(
+        self,
+        start_date: Union[datetime, str],
+        end_date: Union[datetime, str],
+        provider: Optional[str] = None,
+        model: Optional[str] = None,
+    ) -> TokenUsageReport:
+        if not state.is_tokenator_enabled:
+            return TokenUsageReport()
+        logger.debug(
+            f"Getting cost analysis between {start_date} and {end_date} (provider={provider}, model={model})"
+        )
 
-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)
+        if isinstance(start_date, str):
+            try:
+                start = datetime.strptime(start_date, "%Y-%m-%d %H:%M:%S")
+            except ValueError:
+                logger.warning(
+                    f"Date-only string provided for start_date: {start_date}. Setting time to 00:00:00"
+                )
+                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:
+                logger.warning(
+                    f"Date-only string provided for end_date: {end_date}. Setting time to 23:59:59"
+                )
+                end = (
+                    datetime.strptime(end_date, "%Y-%m-%d")
+                    + timedelta(days=1)
+                    - timedelta(seconds=1)
+                )
+        else:
+            end = end_date
 
+        return self._query_usage(start, end, provider, model)
 
-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)
+    def for_execution(self, execution_id: str) -> TokenUsageReport:
+        if not state.is_tokenator_enabled:
+            return TokenUsageReport()
+        logger.debug(f"Getting cost analysis for execution_id={execution_id}")
+        session = get_session()()
+        try:
+            query = session.query(TokenUsage).filter(
+                TokenUsage.execution_id == execution_id
+            )
+            return self._calculate_cost(query.all())
+        finally:
+            session.close()
+
+    def last_execution(self) -> TokenUsageReport:
+        if not state.is_tokenator_enabled:
+            return TokenUsageReport()
+        logger.debug("Getting cost analysis for last execution")
+        session = get_session()()
+        try:
+            query = (
+                session.query(TokenUsage).order_by(TokenUsage.created_at.desc()).first()
+            )
+            if query:
+                return self.for_execution(query.execution_id)
+            return TokenUsageReport()
+        finally:
+            session.close()
+
+    def all_time(self) -> TokenUsageReport:
+        if not state.is_tokenator_enabled:
+            return TokenUsageReport()
+
+        logger.warning("Getting cost analysis for all time. This may take a while...")
+        session = get_session()()
+        try:
+            query = session.query(TokenUsage)
+            return self._calculate_cost(query.all())
+        finally:
+            session.close()

{tokenator-0.1.12.dist-info → tokenator-0.1.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: tokenator
-Version: 0.1.12
+Version: 0.1.14
 Summary: Token usage tracking wrapper for LLMs
 License: MIT
 Author: Ujjwal Maheshwari
@@ -14,8 +14,8 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: alembic (>=1.13.0,<2.0.0)
-Requires-Dist: anthropic (>=0.40.0,<0.41.0)
-Requires-Dist: openai (>=1.57.0,<2.0.0)
+Requires-Dist: anthropic (>=0.43.0,<0.44.0)
+Requires-Dist: openai (>=1.59.0,<2.0.0)
 Requires-Dist: requests (>=2.32.3,<3.0.0)
 Requires-Dist: sqlalchemy (>=2.0.0,<3.0.0)
 Description-Content-Type: text/markdown
@@ -186,7 +186,7 @@ xai_client = OpenAI(
         )
 
 # Wrap it with Tokenator
-client = tokenator_openai(client, db_path=temp_db, provider="xai")
+client = tokenator_openai(xai_client, db_path=temp_db, provider="xai")
 
 # Use it exactly like the OpenAI client but with xAI models
 response = client.chat.completions.create(
@@ -210,17 +210,17 @@ For example, let's see how we can track usage of `perplexity` tokens.
 from openai import OpenAI
 from tokenator import tokenator_openai
 
-xai_client = OpenAI(
+perplexity_client = OpenAI(
             api_key=os.getenv("PERPLEXITY_API_KEY"),
             base_url="https://api.perplexity.ai"
         )
 
 # Wrap it with Tokenator
-client = tokenator_openai(client, db_path=temp_db, provider="perplexity")
+client = tokenator_openai(perplexity_client, db_path=temp_db, provider="perplexity")
 
-# Use it exactly like the OpenAI client but with xAI models
+# Use it exactly like the OpenAI client but with perplexity models
 response = client.chat.completions.create(
-    model="grok-2-latest",
+    model="llama-3.1-sonar-small-128k-online",
     messages=[{"role": "user", "content": "Hello!"}]
 )
 

tokenator-0.1.14.dist-info/RECORD

@@ -0,0 +1,20 @@
+tokenator/__init__.py,sha256=AEPE73UGB_TeNLhro3eY0hU8yy6T-_6AyDls8vWApnE,465
+tokenator/anthropic/client_anthropic.py,sha256=uWUrRId7vJlMG6hVKLUzaA3PoOT6mJwTqSRIhAidRFY,6163
+tokenator/anthropic/stream_interceptors.py,sha256=4VHC_-WkG3Pa10YizmFLrHcbz0Tm2MR_YB5-uohKp5A,5221
+tokenator/base_wrapper.py,sha256=UoS3cOuPa3HpuXPTawybvAtwufgZwzzKBj0BhyB-z6w,3160
+tokenator/create_migrations.py,sha256=k9IHiGK21dLTA8MYNsuhO0-kUVIcMSViMFYtY4WU2Rw,730
+tokenator/migrations/env.py,sha256=JoF5MJ4ae0wJW5kdBHuFlG3ZqeCCDvbMcU8fNA_a6hM,1396
+tokenator/migrations/script.py.mako,sha256=nJL-tbLQE0Qy4P9S4r4ntNAcikPtoFUlvXe6xvm9ot8,635
+tokenator/migrations/versions/f6f1f2437513_initial_migration.py,sha256=4cveHkwSxs-hxOPCm81YfvGZTkJJ2ClAFmyL98-1VCo,1910
+tokenator/migrations.py,sha256=YAf9gZmDzAq36PWWXPtdUQoJFYPXtIDzflC79H6gcJg,1114
+tokenator/models.py,sha256=AlNC5NVrycLg0LhDJIww9HXQ3lwM8CoKvRSqXU6iw-k,1225
+tokenator/openai/client_openai.py,sha256=LhD1IbpzPXRK9eSqtcfUfoM9vBsyw6OHA0_a7N_tS9U,6230
+tokenator/openai/stream_interceptors.py,sha256=ez1MnjRZW_rEalv2SIPAvrU9oMD6OJoD9vht-057fDM,5243
+tokenator/schemas.py,sha256=zIgfmSsFJV9ziJdKrpV8p2P1f-BVWUVIpWoqCLpzhEU,2225
+tokenator/state.py,sha256=xdqDC-rlEA88-VgqQqHnAOXQ5pNTpnHcgOtohDIImPY,262
+tokenator/usage.py,sha256=ghnZ7pQuIxeI38O63xDAbEm6jOSmkYE7MChHBGPxbyM,11229
+tokenator/utils.py,sha256=xg9l2GV1yJL1BlxKL1r8CboABWDslf3G5rGQEJSjFrE,1973
+tokenator-0.1.14.dist-info/LICENSE,sha256=wdG-B6-ODk8RQ4jq5uXSn0w1UWTzCH_MMyvh7AwtGns,1074
+tokenator-0.1.14.dist-info/METADATA,sha256=L93LfqCfqvhES92COaQZpX5w9_c2aDaX8pj2wT74Sxw,6018
+tokenator-0.1.14.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
+tokenator-0.1.14.dist-info/RECORD,,

tokenator-0.1.12.dist-info/RECORD

@@ -1,19 +0,0 @@
-tokenator/__init__.py,sha256=bIAPyGAvWreS2i_5tzxJEyX9JlZgAUNxzVk1iHNUhvU,593
-tokenator/anthropic/client_anthropic.py,sha256=fnjWz_Kf8D0GUTudkZNeSmH9ueCGFLDSBDz1U8Jri3Y,5861
-tokenator/anthropic/stream_interceptors.py,sha256=4VHC_-WkG3Pa10YizmFLrHcbz0Tm2MR_YB5-uohKp5A,5221
-tokenator/base_wrapper.py,sha256=IO344KWbRswQy4vG_pBxWPR7Wp7K-4mlgmS3SCYGep8,2467
-tokenator/create_migrations.py,sha256=k9IHiGK21dLTA8MYNsuhO0-kUVIcMSViMFYtY4WU2Rw,730
-tokenator/migrations/env.py,sha256=JoF5MJ4ae0wJW5kdBHuFlG3ZqeCCDvbMcU8fNA_a6hM,1396
-tokenator/migrations/script.py.mako,sha256=nJL-tbLQE0Qy4P9S4r4ntNAcikPtoFUlvXe6xvm9ot8,635
-tokenator/migrations/versions/f6f1f2437513_initial_migration.py,sha256=4cveHkwSxs-hxOPCm81YfvGZTkJJ2ClAFmyL98-1VCo,1910
-tokenator/migrations.py,sha256=YAf9gZmDzAq36PWWXPtdUQoJFYPXtIDzflC79H6gcJg,1114
-tokenator/models.py,sha256=MhYwCvmqposUNDRxFZNAVnzCqBTHxNL3Hp0MNFXM5ck,1201
-tokenator/openai/client_openai.py,sha256=Ffa3ujLh5PuPe1W8KSISGH3NonZ_AC6ZpKhO6kTupTU,5996
-tokenator/openai/stream_interceptors.py,sha256=ez1MnjRZW_rEalv2SIPAvrU9oMD6OJoD9vht-057fDM,5243
-tokenator/schemas.py,sha256=Ye8hqZlrm3Gh2FyvOVX-hWCpKynWxS58QQRQMfDtIAQ,2114
-tokenator/usage.py,sha256=eTWfcRrTLop-30FmwHpi7_GwCJxU6Qfji374hG1Qptw,8476
-tokenator/utils.py,sha256=xg9l2GV1yJL1BlxKL1r8CboABWDslf3G5rGQEJSjFrE,1973
-tokenator-0.1.12.dist-info/LICENSE,sha256=wdG-B6-ODk8RQ4jq5uXSn0w1UWTzCH_MMyvh7AwtGns,1074
-tokenator-0.1.12.dist-info/METADATA,sha256=VdJVlwESY2_QbiterwI1lH9dng4r4WwWYd6MwXlT9V4,5969
-tokenator-0.1.12.dist-info/WHEEL,sha256=IYZQI976HJqqOpQU6PHkJ8fb3tMNBFjg-Cn-pwAbaFM,88
-tokenator-0.1.12.dist-info/RECORD,,