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

tokenator/client_openai.py

@@ -1,151 +0,0 @@
-"""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")
-
-__all__ = ["tokenator_openai"]

tokenator-0.1.8.dist-info/RECORD

@@ -1,17 +0,0 @@
-tokenator/__init__.py,sha256=ZKe0zMGa_AqOeXUVgYqivUavht_byk03XNFEvAnxqsA,576
-tokenator/base_wrapper.py,sha256=vSu_pStKYulho7_5g0jMCNf84KRxC4kTKep0v8YE61M,2377
-tokenator/client_anthropic.py,sha256=1ejWIZBxtk-mWTVaKWeMUvS2hZ_Dn-vNKYa3yopdjAU,6714
-tokenator/client_openai.py,sha256=_4jvchKzpCFhpioMZTYIWV7_ephQp1abMCtswUDJv1M,6339
-tokenator/create_migrations.py,sha256=n1OVbWrdwvBdaN-Aqqt1gLCPQidfoQfeJtGsab_epGk,746
-tokenator/migrations/env.py,sha256=LR_hONDa8Saiq9CyNUpH8kZCi5PtXLaDlfABs_CePkk,1415
-tokenator/migrations/script.py.mako,sha256=nJL-tbLQE0Qy4P9S4r4ntNAcikPtoFUlvXe6xvm9ot8,635
-tokenator/migrations/versions/f6f1f2437513_initial_migration.py,sha256=DvHcjnREmUHZVX9q1e6PS4wNK_d4qGw-8pz0eS4_3mE,1860
-tokenator/migrations.py,sha256=BFgZRsdIx-Qs_WwDaH6cyi2124mLf5hA8VrIlW7f7Mg,1134
-tokenator/models.py,sha256=EprE_MMJxDS-YXlcIQLZzfekH7xTYbeOC3bx3B2osVw,1171
-tokenator/schemas.py,sha256=V7NYfY9eZvH3J6uOwXJz4dSAU6WYzINRnfFi1wWsTcc,2280
-tokenator/usage.py,sha256=aHjGwzDzaiVznahNk5HqVyk3IxDo5FtFVfOUCeE7DZ4,7833
-tokenator/utils.py,sha256=5mDiGHgt4koCY0onHwkRjwZIuAgP6QvrDZCwD20Sdk8,1969
-tokenator-0.1.8.dist-info/LICENSE,sha256=wdG-B6-ODk8RQ4jq5uXSn0w1UWTzCH_MMyvh7AwtGns,1074
-tokenator-0.1.8.dist-info/METADATA,sha256=1xgNdiPKTJlnBCPH6iMfi7-LoOl-t9soFzH_5V_eYIk,2444
-tokenator-0.1.8.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-tokenator-0.1.8.dist-info/RECORD,,