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

tokenator/openai/client_openai.py

@@ -1,21 +1,24 @@
 """OpenAI client wrapper with token usage tracking."""
 
-from typing import Any, Dict, Optional, TypeVar, Union, overload, Iterator, AsyncIterator
+from typing import Any, Optional, Union, overload, Iterator, AsyncIterator
 import logging
 
-from openai import AsyncOpenAI, AsyncStream, OpenAI, Stream
+from openai import AsyncOpenAI, OpenAI
 from openai.types.chat import ChatCompletion, ChatCompletionChunk
 
 from ..models import Usage, TokenUsageStats
 from ..base_wrapper import BaseWrapper, ResponseType
-from .AsyncStreamInterceptor import AsyncStreamInterceptor
+from .stream_interceptors import OpenAIAsyncStreamInterceptor, OpenAISyncStreamInterceptor
 
 logger = logging.getLogger(__name__)
 
+
 class BaseOpenAIWrapper(BaseWrapper):
     provider = "openai"
 
-    def _process_response_usage(self, response: ResponseType) -> Optional[TokenUsageStats]:
+    def _process_response_usage(
+        self, response: ResponseType
+    ) -> Optional[TokenUsageStats]:
         """Process and log usage statistics from a response."""
         try:
             if isinstance(response, ChatCompletion):
@@ -27,19 +30,18 @@ class BaseOpenAIWrapper(BaseWrapper):
                     total_tokens=response.usage.total_tokens,
                 )
                 return TokenUsageStats(model=response.model, usage=usage)
-                
+
             elif isinstance(response, dict):
-                usage_dict = response.get('usage')
+                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)
+                    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
+                    model=response.get("model", "unknown"), usage=usage
                 )
         except Exception as e:
             logger.warning("Failed to process usage stats: %s", str(e))
@@ -54,45 +56,58 @@ class BaseOpenAIWrapper(BaseWrapper):
     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]]:
+    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)
-        
-        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
-            
+        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
+        self, *args: Any, execution_id: Optional[str] = None, **kwargs: Any
     ) -> Union[ChatCompletion, AsyncIterator[ChatCompletionChunk]]:
         """
         Create a chat completion and log token usage.
@@ -102,32 +117,9 @@ class AsyncOpenAIWrapper(BaseOpenAIWrapper):
         # If user wants a stream, return an interceptor
         if kwargs.get("stream", False):
             base_stream = await self.client.chat.completions.create(*args, **kwargs)
-
-            # Define a callback that will get called once the stream ends
-            def usage_callback(chunks):
-                # Mimic your old logic to gather usage from chunk.usage
-                # e.g. ChatCompletionChunk.usage
-                # Then call self._log_usage(...)
-                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
-                for ch in chunks:
-                    if ch.usage:
-                        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
-
-                self._log_usage(usage_data, execution_id=execution_id)
-
-            # Return the interceptor that wraps the real AsyncStream
-            return AsyncStreamInterceptor(
+            return OpenAIAsyncStreamInterceptor(
                 base_stream=base_stream,
-                usage_callback=usage_callback,
+                usage_callback=_create_usage_callback(execution_id, self._log_usage),
             )
 
         # Non-streaming path remains unchanged
@@ -136,32 +128,36 @@ class AsyncOpenAIWrapper(BaseOpenAIWrapper):
         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,
+        }