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

tokenator/openai/client_openai.py

@@ -0,0 +1,163 @@
+"""OpenAI client wrapper with token usage tracking."""
+
+from typing import Any, Optional, Union, overload, Iterator, AsyncIterator
+import logging
+
+from openai import AsyncOpenAI, OpenAI
+from openai.types.chat import ChatCompletion, ChatCompletionChunk
+
+from ..models import Usage, TokenUsageStats
+from ..base_wrapper import BaseWrapper, ResponseType
+from .stream_interceptors import OpenAIAsyncStreamInterceptor, OpenAISyncStreamInterceptor
+
+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
+
+
+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
+        # Build usage_data from the first chunk's model
+        usage_data = TokenUsageStats(
+            model=chunks[0].model,
+            usage=Usage(),
+        )
+        # Sum up usage from all chunks
+        has_usage = False
+        for ch in chunks:
+            if ch.usage:
+                has_usage = True
+                usage_data.usage.prompt_tokens += ch.usage.prompt_tokens
+                usage_data.usage.completion_tokens += ch.usage.completion_tokens
+                usage_data.usage.total_tokens += ch.usage.total_tokens
+
+        if has_usage:
+            log_usage_fn(usage_data, execution_id=execution_id)
+
+    return usage_callback
+
+
+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)
+
+        if kwargs.get("stream", False):
+            base_stream = self.client.chat.completions.create(*args, **kwargs)
+            return OpenAISyncStreamInterceptor(
+                base_stream=base_stream,
+                usage_callback=_create_usage_callback(execution_id, self._log_usage),
+            )
+
+        response = 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
+
+
+class AsyncOpenAIWrapper(BaseOpenAIWrapper):
+    async def create(
+        self, *args: Any, execution_id: Optional[str] = None, **kwargs: Any
+    ) -> Union[ChatCompletion, AsyncIterator[ChatCompletionChunk]]:
+        """
+        Create a chat completion and log token usage.
+        """
+        logger.debug("Creating chat completion with args: %s, kwargs: %s", args, kwargs)
+
+        # If user wants a stream, return an interceptor
+        if kwargs.get("stream", False):
+            base_stream = await self.client.chat.completions.create(*args, **kwargs)
+            return OpenAIAsyncStreamInterceptor(
+                base_stream=base_stream,
+                usage_callback=_create_usage_callback(execution_id, self._log_usage),
+            )
+
+        # Non-streaming path remains unchanged
+        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
+
+
+@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/openai/stream_interceptors.py

@@ -0,0 +1,146 @@
+import logging
+from typing import AsyncIterator, Callable, List, Optional, TypeVar, Iterator
+
+from openai import AsyncStream, Stream
+
+logger = logging.getLogger(__name__)
+
+_T = TypeVar("_T")  # or you might specifically do _T = ChatCompletionChunk
+
+
+class OpenAIAsyncStreamInterceptor(AsyncStream[_T]):
+    """
+    A wrapper around openai.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 openai.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) -> "OpenAIAsyncStreamInterceptor[_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 OpenAISyncStreamInterceptor(Stream[_T]):
+    """
+    A wrapper around openai.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.Stream
+        # 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) -> "OpenAISyncStreamInterceptor[_T]":
+        """Support with ... : usage."""
+        self._base_stream.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        Ensure we propagate __exit__ to the base stream,
+        so connections are properly closed.
+        """
+        return self._base_stream.__exit__(exc_type, exc_val, exc_tb)
+
+    def close(self) -> None:
+        """Delegate close to the base_stream."""
+        self._base_stream.close()

tokenator/schemas.py

@@ -1,26 +1,22 @@
 """SQLAlchemy models for tokenator."""
 
-import uuid
 from datetime import datetime
-import os
 
-from sqlalchemy import create_engine, Column, Integer, String, DateTime, Float, Index
+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
 
 Base = declarative_base()
 
+
 def get_engine(db_path: str = None):
     """Create SQLAlchemy engine with the given database path."""
     if db_path is None:
-        try:
-            import google.colab  # type: ignore
-            db_path = '/content/tokenator.db'
-        except ImportError:
-            db_path = get_default_db_path()
+        db_path = 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)
@@ -28,39 +24,42 @@ def get_session(db_path: str = None):
     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)
+    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'),
+        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
-        } 
+            "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,
+        }