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

tokenator/usage.py

@@ -1,9 +1,8 @@
 """Cost analysis functions for token usage."""
 
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timedelta
 from typing import Dict, Optional, Union
 
-from sqlalchemy import and_
 
 from .schemas import get_session, TokenUsage
 from .models import TokenRate, TokenUsageReport, ModelUsage, ProviderUsage
@@ -13,48 +12,63 @@ 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"]
+            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()
 
-def _calculate_cost(usages: list[TokenUsage], provider: Optional[str] = None) -> TokenUsageReport:
+
+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]]] = {}
 
     print(f"usages: {len(usages)}")
-    
+
     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] = {}
-        
+
         if usage.model not in provider_model_usages[provider]:
             provider_model_usages[provider][usage.model] = []
-            
+
         provider_model_usages[provider][usage.model].append(usage)
 
     # Calculate totals for each level
     providers_list = []
-    total_metrics = {"total_cost": 0.0, "total_tokens": 0, "prompt_tokens": 0, "completion_tokens": 0}
+    total_metrics = {
+        "total_cost": 0.0,
+        "total_tokens": 0,
+        "prompt_tokens": 0,
+        "completion_tokens": 0,
+    }
 
     for provider, model_usages in provider_model_usages.items():
-        provider_metrics = {"total_cost": 0.0, "total_tokens": 0, "prompt_tokens": 0, "completion_tokens": 0}
+        provider_metrics = {
+            "total_cost": 0.0,
+            "total_tokens": 0,
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+        }
         models_list = []
 
         for model, usages in model_usages.items():
@@ -67,17 +81,21 @@ def _calculate_cost(usages: list[TokenUsage], provider: Optional[str] = None) ->
                 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)
-
-            models_list.append(ModelUsage(
-                model=model,
-                total_cost=round(model_cost, 6),
-                total_tokens=model_total,
-                prompt_tokens=model_prompt,
-                completion_tokens=model_completion
-            ))
+
+                model_cost += usage.prompt_tokens * MODEL_COSTS[usage.model].prompt
+                model_cost += (
+                    usage.completion_tokens * MODEL_COSTS[usage.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,
+                )
+            )
 
             # Add to provider totals
             provider_metrics["total_cost"] += model_cost
@@ -85,11 +103,16 @@ def _calculate_cost(usages: list[TokenUsage], provider: Optional[str] = None) ->
             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()}
-        ))
+        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 grand totals
         for key in total_metrics:
@@ -97,76 +120,110 @@ def _calculate_cost(usages: list[TokenUsage], provider: Optional[str] = None) ->
 
     return TokenUsageReport(
         providers=providers_list,
-        **{k: (round(v, 6) if k == "total_cost" else v) for k, v in total_metrics.items()}
+        **{
+            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:
+
+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)
         )
-        
+
         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:
+
+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})")
+    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:
+
+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})")
+    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:
+
+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})")
+    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:
+
+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})")
+    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
+    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})")
-    
+    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
 
@@ -174,12 +231,20 @@ def between(
         try:
             end = datetime.strptime(end_date, "%Y-%m-%d %H:%M:%S")
         except ValueError:
-            end = datetime.strptime(end_date, "%Y-%m-%d") + timedelta(days=1)  # Include the end date
+            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 _query_usage(start, end, provider, model)
 
+
 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}")
@@ -187,6 +252,7 @@ def for_execution(execution_id: str) -> TokenUsageReport:
     query = session.query(TokenUsage).filter(TokenUsage.execution_id == execution_id)
     return _calculate_cost(query.all())
 
+
 def last_execution() -> TokenUsageReport:
     """Get cost analysis for the last execution_id."""
     logger.debug("Getting cost analysis for last execution")
@@ -194,9 +260,10 @@ def last_execution() -> TokenUsageReport:
     query = session.query(TokenUsage).order_by(TokenUsage.created_at.desc()).first()
     return for_execution(query.execution_id)
 
+
 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)
+    return for_execution(query.execution_id)

tokenator/utils.py

@@ -4,27 +4,29 @@ import os
 import platform
 import logging
 from pathlib import Path
-from typing import Optional
 
 logger = logging.getLogger(__name__)
 
+
 def is_colab() -> bool:
     """Check if running in Google Colab."""
     try:
-        import google.colab  # type: ignore
-        return True
+        from importlib.util import find_spec
+
+        return find_spec("google.colab") is not None
     except ImportError:
         return False
 
+
 def get_default_db_path() -> str:
     """Get the platform-specific default database path."""
     try:
         if is_colab():
             # Use in-memory database for Colab
             return "usage.db"
-            
+
         system = platform.system().lower()
-        
+
         if system == "linux" or system == "darwin":
             # Follow XDG Base Directory Specification
             xdg_data_home = os.environ.get("XDG_DATA_HOME", "")
@@ -39,18 +41,21 @@ def get_default_db_path() -> str:
             db_path = os.path.join(local_app_data, "tokenator", "usage.db")
         else:
             db_path = os.path.join(str(Path.home()), ".tokenator", "usage.db")
-        
+
         # Create directory if it doesn't exist
         os.makedirs(os.path.dirname(db_path), exist_ok=True)
         return db_path
     except (OSError, IOError) as e:
         # Fallback to current directory if we can't create the default path
         fallback_path = os.path.join(os.getcwd(), "tokenator_usage.db")
-        logger.warning(f"Could not create default db path, falling back to {fallback_path}. Error: {e}")
-        return fallback_path 
+        logger.warning(
+            f"Could not create default db path, falling back to {fallback_path}. Error: {e}"
+        )
+        return fallback_path
+
 
 __all__ = [
     "get_default_db_path",
     "is_colab",
     # ... other exports ...
-]
+]

{tokenator-0.1.9.dist-info → tokenator-0.1.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tokenator
-Version: 0.1.9
+Version: 0.1.10
 Summary: Token usage tracking wrapper for LLMs
 License: MIT
 Author: Ujjwal Maheshwari
@@ -60,29 +60,31 @@ response = client.chat.completions.create(
 ### Cost Analysis
 
 ```python
-from tokenator import cost
+from tokenator import usage
 
 # Get usage for different time periods
-cost.last_hour()
-cost.last_day()
-cost.last_week()
-cost.last_month()
+usage.last_hour()
+usage.last_day()
+usage.last_week()
+usage.last_month()
 
 # Custom date range
-cost.between("2024-03-01", "2024-03-15")
+usage.between("2024-03-01", "2024-03-15")
 
 # Get usage for different LLM providers
-cost.last_day("openai")
-cost.last_day("anthropic")
-cost.last_day("google")
+usage.last_day("openai")
+usage.last_day("anthropic")
+usage.last_day("google")
 ```
 
-### Example `cost` object
+### Example `usage` object
 
-```json
-# print(cost.last_hour().model_dump_json(indent=4))
+```python
+print(cost.last_hour().model_dump_json(indent=4))
+```
 
-usage : {
+```json
+{
     "total_cost": 0.0004,
     "total_tokens": 79,
     "prompt_tokens": 52,

tokenator-0.1.10.dist-info/RECORD

@@ -0,0 +1,19 @@
+tokenator/__init__.py,sha256=bIAPyGAvWreS2i_5tzxJEyX9JlZgAUNxzVk1iHNUhvU,593
+tokenator/anthropic/client_anthropic.py,sha256=fcKxGsLex99II-WD9SVNI5QVzH0IEWRmVLjyvZd9wKs,5936
+tokenator/anthropic/stream_interceptors.py,sha256=4VHC_-WkG3Pa10YizmFLrHcbz0Tm2MR_YB5-uohKp5A,5221
+tokenator/base_wrapper.py,sha256=VYSkQB1MEudgzBX60T-VAMsNg4fFx7IRzpadzjm4klE,2466
+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=Umfxha3BhBFU_JebPjyuaUZEZuPqJWQo1xTCuAy3R24,5691
+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.10.dist-info/LICENSE,sha256=wdG-B6-ODk8RQ4jq5uXSn0w1UWTzCH_MMyvh7AwtGns,1074
+tokenator-0.1.10.dist-info/METADATA,sha256=ryILkOYlq8V8219sVmK0xUeEEw51msw_FCoF_3VJ_k8,3108
+tokenator-0.1.10.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+tokenator-0.1.10.dist-info/RECORD,,

tokenator/client_anthropic.py

@@ -1,148 +0,0 @@
-"""Anthropic client wrapper with token usage tracking."""
-
-from typing import Any, Dict, Optional, TypeVar, Union, overload, Iterator, AsyncIterator
-import logging
-
-from anthropic import Anthropic, AsyncAnthropic
-from anthropic.types import Message, RawMessageStartEvent, RawMessageDeltaEvent
-
-from .models import Usage, TokenUsageStats
-from .base_wrapper import BaseWrapper, ResponseType
-
-logger = logging.getLogger(__name__)
-
-class BaseAnthropicWrapper(BaseWrapper):
-    provider = "anthropic"
-
-    def _process_response_usage(self, response: ResponseType) -> Optional[TokenUsageStats]:
-        """Process and log usage statistics from a response."""
-        try:
-            if isinstance(response, Message):
-                if not hasattr(response, 'usage'):
-                    return None
-                usage = Usage(
-                    prompt_tokens=response.usage.input_tokens,
-                    completion_tokens=response.usage.output_tokens,
-                    total_tokens=response.usage.input_tokens + response.usage.output_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('input_tokens', 0),
-                    completion_tokens=usage_dict.get('output_tokens', 0),
-                    total_tokens=usage_dict.get('input_tokens', 0) + usage_dict.get('output_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 messages(self):
-        return self
-
-class AnthropicWrapper(BaseAnthropicWrapper):
-    def create(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)
-        
-        response = self.client.messages.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: Iterator[Message], execution_id: Optional[str]) -> Iterator[Message]:
-        """Wrap streaming response to capture final usage stats"""
-        usage_data: TokenUsageStats = TokenUsageStats(model="", usage=Usage())
-        for chunk in response_iter:
-            if isinstance(chunk, RawMessageStartEvent):
-                usage_data.model = chunk.message.model
-                usage_data.usage.prompt_tokens = chunk.message.usage.input_tokens
-                usage_data.usage.completion_tokens = chunk.message.usage.output_tokens
-                usage_data.usage.total_tokens = chunk.message.usage.input_tokens + chunk.message.usage.output_tokens
-            
-            elif isinstance(chunk, RawMessageDeltaEvent):
-                usage_data.usage.prompt_tokens += chunk.usage.input_tokens
-                usage_data.usage.completion_tokens += chunk.usage.output_tokens
-                usage_data.usage.total_tokens += chunk.usage.input_tokens + chunk.usage.output_tokens
-            
-            yield chunk
-
-        self._log_usage(usage_data, execution_id=execution_id)
-
-class AsyncAnthropicWrapper(BaseAnthropicWrapper):
-    async def create(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)
-        
-        if kwargs.get('stream', False):
-            response = await self.client.messages.create(*args, **kwargs)
-            return self._wrap_streaming_response(response, execution_id)
-        
-        response = await self.client.messages.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: AsyncIterator[Message], execution_id: Optional[str]) -> AsyncIterator[Message]:
-        """Wrap streaming response to capture final usage stats"""
-        usage_data: TokenUsageStats = TokenUsageStats(model="", usage=Usage())
-        async for chunk in response_iter:
-            if isinstance(chunk, RawMessageStartEvent):
-                usage_data.model = chunk.message.model
-                usage_data.usage.prompt_tokens = chunk.message.usage.input_tokens
-                usage_data.usage.completion_tokens = chunk.message.usage.output_tokens
-                usage_data.usage.total_tokens = chunk.message.usage.input_tokens + chunk.message.usage.output_tokens
-            
-            elif isinstance(chunk, RawMessageDeltaEvent):
-                usage_data.usage.prompt_tokens += chunk.usage.input_tokens
-                usage_data.usage.completion_tokens += chunk.usage.output_tokens
-                usage_data.usage.total_tokens += chunk.usage.input_tokens + chunk.usage.output_tokens
-            
-            yield chunk
-        
-            
-        self._log_usage(usage_data, execution_id=execution_id)
-
-@overload
-def tokenator_anthropic(
-    client: Anthropic,
-    db_path: Optional[str] = None,
-) -> AnthropicWrapper: ...
-
-@overload
-def tokenator_anthropic(
-    client: AsyncAnthropic,
-    db_path: Optional[str] = None,
-) -> AsyncAnthropicWrapper: ...
-
-def tokenator_anthropic(
-    client: Union[Anthropic, AsyncAnthropic],
-    db_path: Optional[str] = None,
-) -> Union[AnthropicWrapper, AsyncAnthropicWrapper]:
-    """Create a token-tracking wrapper for an Anthropic client.
-    
-    Args:
-        client: Anthropic or AsyncAnthropic client instance
-        db_path: Optional path to SQLite database for token tracking
-    """
-    if isinstance(client, Anthropic):
-        return AnthropicWrapper(client=client, db_path=db_path)
-    
-    if isinstance(client, AsyncAnthropic):
-        return AsyncAnthropicWrapper(client=client, db_path=db_path)
-        
-    raise ValueError("Client must be an instance of Anthropic or AsyncAnthropic")

tokenator/openai/AsyncStreamInterceptor.py

@@ -1,78 +0,0 @@
-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()