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

tokenator/__init__.py

@@ -1,7 +1,7 @@
 """Tokenator - Track and analyze your OpenAI API token usage and costs."""
 
 import logging
-from .client_openai import tokenator_openai
+from .openai.client_openai import tokenator_openai
 from .client_anthropic import tokenator_anthropic
 from . import usage
 from .utils import get_default_db_path

tokenator/openai/AsyncStreamInterceptor.py

@@ -0,0 +1,78 @@
+import logging
+from typing import AsyncIterator, Callable, Generic, List, Optional, TypeVar
+
+from openai import AsyncStream, AsyncOpenAI
+from openai.types.chat import ChatCompletionChunk
+
+logger = logging.getLogger(__name__)
+
+_T = TypeVar("_T")  # or you might specifically do _T = ChatCompletionChunk
+
+
+class AsyncStreamInterceptor(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) -> "AsyncStreamInterceptor[_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()

tokenator/{client_openai.py → openai/client_openai.py}

@@ -6,8 +6,9 @@ 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
+from ..models import Usage, TokenUsageStats
+from ..base_wrapper import BaseWrapper, ResponseType
+from .AsyncStreamInterceptor import AsyncStreamInterceptor
 
 logger = logging.getLogger(__name__)
 
@@ -87,37 +88,54 @@ class OpenAIWrapper(BaseOpenAIWrapper):
             
 
 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."""
+    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 kwargs.get('stream', False):
-            response = await self.client.chat.completions.create(*args, **kwargs)
-            return self._wrap_streaming_response(response, execution_id)
-        
+
+        # 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(
+                base_stream=base_stream,
+                usage_callback=usage_callback,
+            )
+
+        # 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
-
-    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,
@@ -147,5 +165,3 @@ def tokenator_openai(
         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 → tokenator-0.1.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tokenator
-Version: 0.1.8
+Version: 0.1.9
 Summary: Token usage tracking wrapper for LLMs
 License: MIT
 Author: Ujjwal Maheshwari
@@ -27,7 +27,7 @@ Have you ever wondered about :
 - 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.
+Afraid not, tokenator is here! With tokenator's easy to use API, you can start tracking LLM usage in a matter of minutes.
 
 Get started with just 3 lines of code!
 
@@ -80,7 +80,32 @@ cost.last_day("google")
 ### Example `cost` object
 
 ```json
-
+# print(cost.last_hour().model_dump_json(indent=4))
+
+usage : {
+    "total_cost": 0.0004,
+    "total_tokens": 79,
+    "prompt_tokens": 52,
+    "completion_tokens": 27,
+    "providers": [
+        {
+            "total_cost": 0.0004,
+            "total_tokens": 79,
+            "prompt_tokens": 52,
+            "completion_tokens": 27,
+            "provider": "openai",
+            "models": [
+                {
+                    "total_cost": 0.0004,
+                    "total_tokens": 79,
+                    "prompt_tokens": 52,
+                    "completion_tokens": 27,
+                    "model": "gpt-4o-2024-08-06"
+                }
+            ]
+        }
+    ]
+}
 ```
 
 ## Features

{tokenator-0.1.8.dist-info → tokenator-0.1.9.dist-info}/RECORD

@@ -1,17 +1,18 @@
-tokenator/__init__.py,sha256=ZKe0zMGa_AqOeXUVgYqivUavht_byk03XNFEvAnxqsA,576
+tokenator/__init__.py,sha256=mYwK5EJTlbh_7WvylzxXcL-yzWe_fESSL6FLrlY1qck,583
 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/openai/AsyncStreamInterceptor.py,sha256=estfEFBFyo5BWqTNwHlCZ-wE0dRjtGeyQ0ihBeW3jrU,2842
+tokenator/openai/client_openai.py,sha256=q-0abTq54zRORPLeushdHx1UYq-hOAlp6qY8wAOP2GQ,6682
 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,,
+tokenator-0.1.9.dist-info/LICENSE,sha256=wdG-B6-ODk8RQ4jq5uXSn0w1UWTzCH_MMyvh7AwtGns,1074
+tokenator-0.1.9.dist-info/METADATA,sha256=A7x7gEjbTwOBoR7mxGiHKiZVvKL8ZD6ecL7Wd0y6jfM,3093
+tokenator-0.1.9.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+tokenator-0.1.9.dist-info/RECORD,,