revenium-python-sdk 0.1.7__tar.gz → 0.1.8__tar.gz

{revenium_python_sdk-0.1.7/revenium_python_sdk.egg-info → revenium_python_sdk-0.1.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: revenium-python-sdk
-Version: 0.1.7
+Version: 0.1.8
 Summary: The official Revenium Python SDK — unified AI metering middleware for OpenAI, Anthropic, Google, Ollama, LiteLLM, Perplexity, and fal.ai.
 Author-email: Revenium <support@revenium.io>
 License: MIT

{revenium_python_sdk-0.1.7 → revenium_python_sdk-0.1.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "revenium-python-sdk"
-version = "0.1.7"
+version = "0.1.8"
 description = "The official Revenium Python SDK — unified AI metering middleware for OpenAI, Anthropic, Google, Ollama, LiteLLM, Perplexity, and fal.ai."
 readme = "README.md"
 requires-python = ">=3.8"

{revenium_python_sdk-0.1.7 → revenium_python_sdk-0.1.8}/revenium_middleware/anthropic/bedrock_adapter.py

@@ -169,7 +169,74 @@ def _model_id(model_name: str) -> str:
     return _MODEL_MAP.get(model_name, f"anthropic.{model_name}")
 
 
-def bedrock_invoke(model: str, payload: dict, region: Optional[str] = None) -> Tuple[str, int, int]:
+def _as_dict(value: Any) -> dict:
+    """Return value when it is a dict, otherwise an empty dict."""
+    return value if isinstance(value, dict) else {}
+
+
+def _token_count(value: Any) -> Optional[int]:
+    """Return a non-negative integer token count, or None for invalid values."""
+    if isinstance(value, bool):
+        return None
+    if isinstance(value, int):
+        return value if value >= 0 else None
+    if isinstance(value, str) and value.isdigit():
+        return int(value)
+    return None
+
+
+def _first_token_count(source: Any, keys: Tuple[str, ...]) -> Optional[int]:
+    """Return the first valid token count for any of the supplied keys."""
+    source = _as_dict(source)
+    for key in keys:
+        if key in source:
+            value = _token_count(source.get(key))
+            if value is not None:
+                return value
+    return None
+
+
+def _extract_cache_tokens(usage: Any) -> Tuple[Optional[int], Optional[int]]:
+    """Return (cache_creation_tokens, cache_read_tokens) from a Bedrock usage dict."""
+    cache_creation = _first_token_count(
+        usage,
+        (
+            "cacheWriteInputTokensCount",
+            "cacheWriteInputTokens",
+            "cache_write_input_tokens_count",
+            "cache_write_input_tokens",
+            "cache_creation_input_tokens_count",
+            "cache_creation_input_tokens",
+        )
+    )
+    cache_read = _first_token_count(
+        usage,
+        (
+            "cacheReadInputTokensCount",
+            "cacheReadInputTokens",
+            "cache_read_input_tokens_count",
+            "cache_read_input_tokens",
+        )
+    )
+    return cache_creation, cache_read
+
+
+def _extract_token_counts(usage: Any, metrics: Optional[Any] = None) -> Tuple[Optional[int], Optional[int]]:
+    """Return (input_tokens, output_tokens) from Bedrock usage and metrics blocks."""
+    usage = _as_dict(usage)
+    metrics = _as_dict(metrics)
+    input_tokens = _first_token_count(usage, ("inputTokens", "input_tokens", "prompt_tokens"))
+    if input_tokens is None:
+        input_tokens = _first_token_count(metrics, ("inputTokenCount",))
+
+    output_tokens = _first_token_count(usage, ("outputTokens", "output_tokens", "completion_tokens"))
+    if output_tokens is None:
+        output_tokens = _first_token_count(metrics, ("outputTokenCount",))
+
+    return input_tokens, output_tokens
+
+
+def bedrock_invoke(model: str, payload: dict, region: Optional[str] = None) -> Tuple[str, int, int, int, int]:
     """
     Invoke Bedrock model with Anthropic-compatible parameters.
 
@@ -179,7 +246,7 @@ def bedrock_invoke(model: str, payload: dict, region: Optional[str] = None) -> T
         region: AWS region (defaults to AWS_REGION env var or us-east-1)
 
     Returns:
-        Tuple of (text_content, input_tokens, output_tokens)
+        Tuple of (text_content, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens)
 
     Raises:
         BedrockValidationError: For invalid input parameters
@@ -225,22 +292,15 @@ def bedrock_invoke(model: str, payload: dict, region: Optional[str] = None) -> T
             if c.get("type") == "text"
         )
 
-        # Try multiple token field formats to find the right one
-        input_tokens = usage.get("inputTokens", 0)  # camelCase (AWS standard)
-        output_tokens = usage.get("outputTokens", 0)
-
-        # Try snake_case if camelCase returns 0
-        if input_tokens == 0 and output_tokens == 0:
-            input_tokens = usage.get("input_tokens", 0)
-            output_tokens = usage.get("output_tokens", 0)
-
-        # Try other possible field names
-        if input_tokens == 0 and output_tokens == 0:
-            input_tokens = usage.get("prompt_tokens", 0)
-            output_tokens = usage.get("completion_tokens", 0)
+        input_tokens, output_tokens = _extract_token_counts(usage)
+        cache_creation_tokens, cache_read_tokens = _extract_cache_tokens(usage)
+        input_tokens = input_tokens if input_tokens is not None else 0
+        output_tokens = output_tokens if output_tokens is not None else 0
+        cache_creation_tokens = cache_creation_tokens if cache_creation_tokens is not None else 0
+        cache_read_tokens = cache_read_tokens if cache_read_tokens is not None else 0
 
         logger.debug(f"Bedrock invoke successful: {input_tokens} input tokens, {output_tokens} output tokens")
-        return text, input_tokens, output_tokens
+        return text, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens
 
     except BedrockValidationError:
         # Re-raise validation errors as-is
@@ -267,6 +327,8 @@ class BedrockStreamIterator:
         self.accumulated_text = ""
         self.input_tokens = 0
         self.output_tokens = 0
+        self.cache_creation_tokens = 0
+        self.cache_read_tokens = 0
         self._stream = None
         self._started = False
 
@@ -314,7 +376,17 @@ class BedrockStreamIterator:
     def _process_stream(self, stream):
         """Process the Bedrock stream and yield text chunks."""
         for event in stream:
-            chunk = event.get("chunk")
+            if not hasattr(event, "get"):
+                continue
+
+            metadata = _as_dict(event.get("metadata"))
+            if metadata:
+                self._update_usage_from_blocks(
+                    metadata.get("usage") or {},
+                    metadata.get("metrics") or {}
+                )
+
+            chunk = _as_dict(event.get("chunk"))
             if not chunk:
                 continue
 
@@ -324,13 +396,25 @@ class BedrockStreamIterator:
 
             try:
                 chunk_data = json.loads(chunk_bytes.decode("utf-8"))
+                if not isinstance(chunk_data, dict):
+                    continue
+
                 logger.debug(f"Bedrock stream chunk: {chunk_data}")
 
                 chunk_type = chunk_data.get("type")
+                usage = _as_dict(chunk_data.get("usage"))
+                metrics = _as_dict(chunk_data.get("amazon-bedrock-invocationMetrics"))
+                if usage or metrics:
+                    self._update_usage_from_blocks(usage, metrics)
+
+                if chunk_type == "message_start":
+                    message_usage = _as_dict(_as_dict(chunk_data.get("message")).get("usage"))
+                    if message_usage:
+                        self._update_usage_from_blocks(message_usage)
 
                 # Handle content block delta (text chunk)
                 if chunk_type == "content_block_delta":
-                    delta = chunk_data.get("delta", {})
+                    delta = _as_dict(chunk_data.get("delta"))
                     text = delta.get("text", "")
                     if text:
                         self.accumulated_text += text
@@ -338,27 +422,14 @@ class BedrockStreamIterator:
 
                 # Handle message stop (final message with usage)
                 elif chunk_type == "message_stop":
-                    # Check for usage in the chunk data
-                    usage = chunk_data.get("usage", {})
-                    # Also check for Amazon Bedrock invocation metrics
-                    metrics = chunk_data.get("amazon-bedrock-invocationMetrics", {})
-
-                    # Try multiple token field formats
-                    self.input_tokens = (
-                        usage.get("inputTokens", 0) or
-                        usage.get("input_tokens", 0) or
-                        metrics.get("inputTokenCount", 0)
-                    )
-                    self.output_tokens = (
-                        usage.get("outputTokens", 0) or
-                        usage.get("output_tokens", 0) or
-                        metrics.get("outputTokenCount", 0)
-                    )
-
                     logger.debug(f"Bedrock streaming completed. Input tokens: {self.input_tokens}, Output tokens: {self.output_tokens}")
 
+                # Handle message delta usage-only updates.
+                elif chunk_type == "message_delta":
+                    logger.debug("Bedrock stream: message_delta")
+
                 # Handle other chunk types (content_block_start, etc.)
-                elif chunk_type in ["content_block_start", "message_start"]:
+                elif chunk_type == "content_block_start":
                     logger.debug(f"Bedrock stream: {chunk_type}")
 
             except json.JSONDecodeError as e:
@@ -367,6 +438,22 @@ class BedrockStreamIterator:
 
         logger.debug(f"Bedrock streaming finished. Total text length: {len(self.accumulated_text)}")
 
+    def _update_usage_from_blocks(self, usage: Any, metrics: Optional[Any] = None):
+        """Update token counters from a stream usage block without resetting missing fields."""
+        usage = _as_dict(usage)
+        metrics = _as_dict(metrics)
+        input_tokens, output_tokens = _extract_token_counts(usage, metrics)
+        if input_tokens is not None:
+            self.input_tokens = input_tokens
+        if output_tokens is not None:
+            self.output_tokens = output_tokens
+
+        cache_creation_tokens, cache_read_tokens = _extract_cache_tokens(usage)
+        if cache_creation_tokens is not None:
+            self.cache_creation_tokens = cache_creation_tokens
+        if cache_read_tokens is not None:
+            self.cache_read_tokens = cache_read_tokens
+
 
 def bedrock_invoke_stream(model: str, payload: dict, region: Optional[str] = None) -> BedrockStreamIterator:
     """
@@ -451,7 +538,8 @@ def create_bedrock_payload(messages: Union[list, List[Dict[str, Any]]], **kwargs
 
 
 def create_anthropic_response(text: str, input_tokens: int, output_tokens: int,
-                            model: str, request_id: Optional[str] = None):
+                            model: str, request_id: Optional[str] = None,
+                            cache_creation_tokens: int = 0, cache_read_tokens: int = 0):
     """
     Create an Anthropic-compatible response object.
 
@@ -461,6 +549,8 @@ def create_anthropic_response(text: str, input_tokens: int, output_tokens: int,
         output_tokens: Number of output tokens
         model: Model name
         request_id: Optional request ID
+        cache_creation_tokens: Number of cache creation input tokens (default 0)
+        cache_read_tokens: Number of cache read input tokens (default 0)
 
     Returns:
         Object that mimics Anthropic's Message response structure with both
@@ -510,26 +600,25 @@ def create_anthropic_response(text: str, input_tokens: int, output_tokens: int,
             self.text = text
 
     class Usage(HybridAccessMixin):
-        def __init__(self, input_tokens, output_tokens):
+        def __init__(self, input_tokens, output_tokens, cache_creation_tokens=0, cache_read_tokens=0):
             self.input_tokens = input_tokens
             self.output_tokens = output_tokens
             self.total_tokens = input_tokens + output_tokens
-            # Add cache token attributes for compatibility
-            self.cache_creation_input_tokens = 0
-            self.cache_read_input_tokens = 0
+            self.cache_creation_input_tokens = cache_creation_tokens
+            self.cache_read_input_tokens = cache_read_tokens
 
     class Message(HybridAccessMixin):
-        def __init__(self, text, input_tokens, output_tokens, model, request_id):
+        def __init__(self, text, input_tokens, output_tokens, model, request_id, cache_creation_tokens=0, cache_read_tokens=0):
             self.id = request_id or _generate_safe_id("msg_bedrock", text)
             self.type = "message"
             self.role = "assistant"
             self.model = model
             self.content = [TextBlock(text)]
-            self.usage = Usage(input_tokens, output_tokens)
+            self.usage = Usage(input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens)
             self.stop_reason = "end_turn"
             self.stop_sequence = None
 
-    return Message(text, input_tokens, output_tokens, model, request_id)
+    return Message(text, input_tokens, output_tokens, model, request_id, cache_creation_tokens, cache_read_tokens)
 
 
 class BedrockStreamWrapper:
@@ -558,6 +647,8 @@ class BedrockStreamWrapper:
         self.first_token_time = None
         self.request_start_time = time.time() * 1000  # Convert to milliseconds
         self.accumulated_text = ""
+        self.cache_creation_tokens = 0
+        self.cache_read_tokens = 0
 
     def __enter__(self):
         """Enter the context manager and initialize the stream."""
@@ -635,6 +726,8 @@ class BedrockStreamWrapper:
         # Get token counts from the stream iterator
         input_tokens = getattr(self.stream_iterator, 'input_tokens', 0)
         output_tokens = getattr(self.stream_iterator, 'output_tokens', 0)
+        self.cache_creation_tokens = getattr(self.stream_iterator, 'cache_creation_tokens', 0)
+        self.cache_read_tokens = getattr(self.stream_iterator, 'cache_read_tokens', 0)
 
         # Generate a response ID
         self.response_id = _generate_safe_id("msg_bedrock_stream", self.accumulated_text)
@@ -645,7 +738,9 @@ class BedrockStreamWrapper:
             input_tokens=input_tokens,
             output_tokens=output_tokens,
             model=self.model,
-            request_id=self.response_id
+            request_id=self.response_id,
+            cache_creation_tokens=self.cache_creation_tokens,
+            cache_read_tokens=self.cache_read_tokens
         )
 
     def _send_metering_data(self, request_duration: float):
@@ -707,8 +802,8 @@ class BedrockStreamWrapper:
                     meta = extract_common_metadata(self.usage_metadata)
 
                     result = submit_ai_event("completion", {
-                        "cache_creation_token_count": 0,
-                        "cache_read_token_count": 0,
+                        "cache_creation_token_count": self.cache_creation_tokens,
+                        "cache_read_token_count": self.cache_read_tokens,
                         "input_token_cost": None,
                         "output_token_cost": None,
                         "total_cost": None,

{revenium_python_sdk-0.1.7 → revenium_python_sdk-0.1.8}/revenium_middleware/anthropic/middleware.py

@@ -165,14 +165,16 @@ def _handle_bedrock_request(args, kwargs, usage_metadata, request_time_dt, reque
     payload = create_bedrock_payload(messages, **bedrock_kwargs)
 
     # Invoke Bedrock
-    text, input_tokens, output_tokens = bedrock_invoke(model, payload)
+    text, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens = bedrock_invoke(model, payload)
 
     # Create Anthropic-compatible response
     response = create_anthropic_response(
         text=text,
         input_tokens=input_tokens,
         output_tokens=output_tokens,
-        model=model
+        model=model,
+        cache_creation_tokens=cache_creation_tokens,
+        cache_read_tokens=cache_read_tokens
     )
 
     # Calculate timing
@@ -182,7 +184,8 @@ def _handle_bedrock_request(args, kwargs, usage_metadata, request_time_dt, reque
 
     # Create metering call for Bedrock (pass kwargs for vision detection)
     _create_bedrock_metering_call(
-        response, usage_metadata, request_time, response_time, request_duration, kwargs
+        response, usage_metadata, request_time, response_time, request_duration, kwargs,
+        cache_creation_tokens=cache_creation_tokens, cache_read_tokens=cache_read_tokens
     )
 
     return response
@@ -288,7 +291,7 @@ def _extract_trace_fields(usage_metadata, request_body=None):
     }
 
 
-def _create_bedrock_metering_call(response, usage_metadata, request_time, response_time, request_duration, request_kwargs=None):
+def _create_bedrock_metering_call(response, usage_metadata, request_time, response_time, request_duration, request_kwargs=None, cache_creation_tokens=0, cache_read_tokens=0):
     """Create a metering call for Bedrock usage."""
 
     # Get provider metadata
@@ -336,8 +339,8 @@ def _create_bedrock_metering_call(response, usage_metadata, request_time, respon
             meta = extract_common_metadata(usage_metadata)
 
             result = submit_ai_event("completion", {
-                "cache_creation_token_count": 0,
-                "cache_read_token_count": 0,
+                "cache_creation_token_count": cache_creation_tokens,
+                "cache_read_token_count": cache_read_tokens,
                 "input_token_cost": None,
                 "output_token_cost": None,
                 "total_cost": None,

{revenium_python_sdk-0.1.7 → revenium_python_sdk-0.1.8}/revenium_middleware/openai/langchain/unified_handler.py

@@ -217,18 +217,15 @@ class UnifiedReveniumCallbackHandler(AsyncCallbackHandler):
                 logger.debug(f"Found usage_metadata: {usage_data}")
 
         # Check for response_metadata with token_usage (LangChain v0.1+)
-        elif hasattr(response, 'response_metadata') and response.response_metadata:
-            if 'token_usage' in response.response_metadata:
-                usage_data = response.response_metadata['token_usage']
-                if self.enable_debug_logging:
-                    logger.debug(f"Found token_usage in response_metadata: {usage_data}")
-            elif 'usage' in response.response_metadata:
-                usage_data = response.response_metadata['usage']
-                if self.enable_debug_logging:
-                    logger.debug(f"Found usage in response_metadata: {usage_data}")
+        if not usage_data and hasattr(response, 'response_metadata') and response.response_metadata:
+            usage_data = self._extract_usage_from_response_metadata(
+                response.response_metadata
+            )
+            if usage_data and self.enable_debug_logging:
+                logger.debug(f"Found usage in response_metadata: {usage_data}")
 
         # Check for llm_output (older LangChain versions)
-        elif hasattr(response, 'llm_output') and response.llm_output:
+        if not usage_data and hasattr(response, 'llm_output') and response.llm_output:
             if 'token_usage' in response.llm_output:
                 usage_data = response.llm_output['token_usage']
                 if self.enable_debug_logging:
@@ -239,7 +236,7 @@ class UnifiedReveniumCallbackHandler(AsyncCallbackHandler):
                     logger.debug(f"Found usage in llm_output: {usage_data}")
 
         # For streaming responses, check if there's accumulated usage data
-        elif hasattr(response, 'content') and hasattr(response, 'additional_kwargs'):
+        if not usage_data and hasattr(response, 'content') and hasattr(response, 'additional_kwargs'):
             # This might be a streaming chunk with accumulated data
             if 'usage' in response.additional_kwargs:
                 usage_data = response.additional_kwargs['usage']
@@ -247,11 +244,17 @@ class UnifiedReveniumCallbackHandler(AsyncCallbackHandler):
                     logger.debug(f"Found usage in additional_kwargs: {usage_data}")
 
         # Last resort: check for any 'usage' attribute directly
-        elif hasattr(response, 'usage'):
+        if not usage_data and hasattr(response, 'usage'):
             usage_data = response.usage
             if self.enable_debug_logging:
                 logger.debug(f"Found direct usage attribute: {usage_data}")
 
+        # LLMResult often stores AIMessage usage on generations[0][0].message.
+        if not usage_data:
+            usage_data = self._extract_usage_from_generations(response)
+            if usage_data and self.enable_debug_logging:
+                logger.debug(f"Found usage in generation message: {usage_data}")
+
         if not usage_data and self.enable_debug_logging:
             logger.debug("No usage data found in response")
             # Log response structure for debugging
@@ -260,6 +263,44 @@ class UnifiedReveniumCallbackHandler(AsyncCallbackHandler):
 
         return usage_data
 
+    def _extract_usage_from_response_metadata(
+        self, response_metadata: Dict[str, Any]
+    ) -> Optional[Dict[str, Any]]:
+        if not isinstance(response_metadata, dict):
+            return None
+        if 'token_usage' in response_metadata:
+            return response_metadata['token_usage']
+        if 'usage' in response_metadata:
+            return response_metadata['usage']
+        return None
+
+    def _extract_usage_from_generations(self, response: Any) -> Optional[Dict[str, Any]]:
+        generations = getattr(response, 'generations', None)
+        if not generations:
+            return None
+
+        for generation_group in generations:
+            if isinstance(generation_group, (list, tuple)):
+                candidates = generation_group
+            else:
+                candidates = [generation_group]
+
+            for generation in candidates:
+                message = getattr(generation, 'message', None) or generation
+
+                usage_metadata = getattr(message, 'usage_metadata', None)
+                if usage_metadata:
+                    return usage_metadata
+
+                response_metadata = getattr(message, 'response_metadata', None)
+                usage_data = self._extract_usage_from_response_metadata(
+                    response_metadata
+                )
+                if usage_data:
+                    return usage_data
+
+        return None
+
     def _process_metering_call(self, usage_data: Dict[str, Any], run_info: Dict[str, Any]) -> None:
         """
         Create a Revenium metering call with the usage data.
@@ -316,20 +357,37 @@ class UnifiedReveniumCallbackHandler(AsyncCallbackHandler):
         """
         class MockResponse:
             def __init__(self, usage_data, model_name):
+                def has_usage_key(key):
+                    if isinstance(usage_data, dict):
+                        return key in usage_data
+                    return hasattr(usage_data, key)
+
+                def get_usage_value(key, default=0):
+                    if isinstance(usage_data, dict):
+                        return usage_data.get(key, default)
+                    return getattr(usage_data, key, default)
+
+                def detail_object(details):
+                    if not details:
+                        return None
+                    if isinstance(details, dict):
+                        return type('TokenDetails', (), details)()
+                    return details
+
                 # Set usage information in the format expected by middleware
-                if 'input_tokens' in usage_data:
+                if has_usage_key('input_tokens'):
                     # LangChain v0.2+ format
                     self.usage = type('Usage', (), {
-                        'prompt_tokens': usage_data.get('input_tokens', 0),
-                        'completion_tokens': usage_data.get('output_tokens', 0),
-                        'total_tokens': usage_data.get('total_tokens', 0)
+                        'prompt_tokens': get_usage_value('input_tokens', 0),
+                        'completion_tokens': get_usage_value('output_tokens', 0),
+                        'total_tokens': get_usage_value('total_tokens', 0)
                     })()
-                elif 'prompt_tokens' in usage_data:
+                elif has_usage_key('prompt_tokens'):
                     # OpenAI format
                     self.usage = type('Usage', (), {
-                        'prompt_tokens': usage_data.get('prompt_tokens', 0),
-                        'completion_tokens': usage_data.get('completion_tokens', 0),
-                        'total_tokens': usage_data.get('total_tokens', 0)
+                        'prompt_tokens': get_usage_value('prompt_tokens', 0),
+                        'completion_tokens': get_usage_value('completion_tokens', 0),
+                        'total_tokens': get_usage_value('total_tokens', 0)
                     })()
                 else:
                     # Fallback - create minimal usage
@@ -339,6 +397,27 @@ class UnifiedReveniumCallbackHandler(AsyncCallbackHandler):
                         'total_tokens': 0
                     })()
 
+                # Preserve provider-specific cache detail shapes for middleware.
+                for detail_name in (
+                    'prompt_tokens_details',
+                    'input_tokens_details',
+                    'input_token_details',
+                ):
+                    details = get_usage_value(detail_name, None)
+                    details = detail_object(details)
+                    if details is not None:
+                        setattr(self.usage, detail_name, details)
+
+                # Preserve Anthropic's native flat cache keys. A live
+                # ChatAnthropic LLMResult reports cache tokens here rather than
+                # under an `input_token_details` sub-object.
+                for flat_key in (
+                    'cache_creation_input_tokens',
+                    'cache_read_input_tokens',
+                ):
+                    if has_usage_key(flat_key):
+                        setattr(self.usage, flat_key, get_usage_value(flat_key, 0))
+
                 # Set model name
                 self.model = model_name
 

{revenium_python_sdk-0.1.7 → revenium_python_sdk-0.1.8}/revenium_middleware/openai/middleware.py

@@ -2,6 +2,7 @@ import datetime
 import logging
 import os
 import uuid
+from numbers import Number
 from typing import Dict, Any, Optional, Tuple
 from enum import Enum
 
@@ -43,6 +44,73 @@ logger = logging.getLogger("revenium_middleware.extension")
 SENSITIVE_FIELDS = SecurityConfig.SENSITIVE_FIELDS
 
 
+def _get_value(source: Any, key: str, default: Any = None) -> Any:
+    """Read a value from either a dict-like or object-like SDK payload."""
+    if source is None:
+        return default
+    if isinstance(source, dict):
+        return source.get(key, default)
+    return getattr(source, key, default)
+
+
+def _coerce_int(value: Any, default: int = 0) -> int:
+    if value is None:
+        return default
+    if isinstance(value, bool):
+        return default
+    if isinstance(value, Number):
+        return int(value)
+    if isinstance(value, str):
+        try:
+            return int(value)
+        except ValueError:
+            return default
+    return default
+
+
+def _extract_cache_token_counts(usage: Any) -> Tuple[int, int]:
+    """
+    Return (cache_creation_token_count, cache_read_token_count).
+
+    OpenAI reports cached prompt/input tokens as cache reads. Anthropic-shaped
+    LangChain metadata reports read and creation counts separately.
+    """
+    input_token_details = _get_value(usage, "input_token_details")
+    cache_creation_token_count = _coerce_int(
+        _get_value(input_token_details, "cache_creation", 0)
+    )
+    cache_read_token_count = _coerce_int(
+        _get_value(input_token_details, "cache_read", 0)
+    )
+    if cache_creation_token_count or cache_read_token_count:
+        return cache_creation_token_count, cache_read_token_count
+
+    # Anthropic native usage shape: a live ChatAnthropic LLMResult exposes the
+    # raw `response_metadata.usage` block with these flat keys rather than the
+    # normalized `input_token_details` shape.
+    cache_creation_token_count = _coerce_int(
+        _get_value(usage, "cache_creation_input_tokens", 0)
+    )
+    cache_read_token_count = _coerce_int(
+        _get_value(usage, "cache_read_input_tokens", 0)
+    )
+    if cache_creation_token_count or cache_read_token_count:
+        return cache_creation_token_count, cache_read_token_count
+
+    input_tokens_details = _get_value(usage, "input_tokens_details")
+    cache_read_token_count = _coerce_int(
+        _get_value(input_tokens_details, "cached_tokens", 0)
+    )
+    if cache_read_token_count:
+        return 0, cache_read_token_count
+
+    prompt_details = _get_value(usage, "prompt_tokens_details")
+    cache_read_token_count = _coerce_int(
+        _get_value(prompt_details, "cached_tokens", 0)
+    )
+    return 0, cache_read_token_count
+
+
 def extract_prompt_data_if_enabled(
     request_body: Optional[Dict[str, Any]],
     response: Any = None,
@@ -303,10 +371,10 @@ def extract_usage_data(response, operation_type: OperationType, request_time: st
             # Responses API doesn't have choices, use default
             stop_reason = "END"
 
-    # Extract cached tokens (only available for chat completions)
-    cached_tokens = 0
-    if operation_type == OperationType.CHAT and hasattr(response.usage, 'prompt_tokens_details'):
-        cached_tokens = getattr(response.usage.prompt_tokens_details, 'cached_tokens', 0)
+    # Extract provider-specific cache token details.
+    cache_creation_token_count, cache_read_token_count = _extract_cache_token_counts(
+        response.usage
+    )
 
     # Build unified usage data structure
     usage_data = {
@@ -321,8 +389,8 @@ def extract_usage_data(response, operation_type: OperationType, request_time: st
         "model_source": provider_metadata["model_source"],
         "is_streamed": False,  # Will be overridden for streaming
         "time_to_first_token": 0,  # Will be set by caller if applicable
-        "cache_creation_token_count": cached_tokens,
-        "cache_read_token_count": 0,
+        "cache_creation_token_count": cache_creation_token_count,
+        "cache_read_token_count": cache_read_token_count,
         "reasoning_token_count": 0,
         "request_time": request_time,
         "response_time": response_time,
@@ -379,6 +447,8 @@ async def log_token_usage(
         input_messages: Optional[str] = None,
         output_response: Optional[str] = None,
         prompts_truncated: Optional[bool] = None,
+        cache_creation_token_count: Optional[int] = None,
+        cache_read_token_count: Optional[int] = None,
 ) -> None:
     """Log token usage to Revenium."""
     if client is None:
@@ -403,9 +473,16 @@ async def log_token_usage(
 
     # Prepare arguments for create_completion
     # Build completion args, only including non-None values for optional fields
+    if cache_creation_token_count is None:
+        cache_creation_tokens = 0
+    else:
+        cache_creation_tokens = cache_creation_token_count
+    if cache_read_token_count is None:
+        cache_read_token_count = _coerce_int(cached_tokens)
+
     completion_args = {
-        "cache_creation_token_count": cached_tokens,
-        "cache_read_token_count": 0,
+        "cache_creation_token_count": cache_creation_tokens,
+        "cache_read_token_count": cache_read_token_count,
         "output_token_count": completion_tokens,
         "cost_type": "AI",
         "model": model,
@@ -634,6 +711,8 @@ def create_metering_call(
             completion_tokens=usage_data["output_token_count"],
             total_tokens=usage_data["total_token_count"],
             cached_tokens=usage_data["cache_creation_token_count"],
+            cache_creation_token_count=usage_data["cache_creation_token_count"],
+            cache_read_token_count=usage_data["cache_read_token_count"],
             stop_reason=usage_data["stop_reason"],
             request_time=usage_data["request_time"],
             response_time=usage_data["response_time"],
@@ -1044,17 +1123,17 @@ def handle_streaming_response(
             prompt_tokens = 0
             completion_tokens = 0
             total_tokens = 0
-            cached_tokens = 0
+            cache_creation_token_count = 0
+            cache_read_token_count = 0
 
             # First check if we have the final usage data from the special chunk
             if self.final_usage:
                 prompt_tokens = self.final_usage.prompt_tokens
                 completion_tokens = self.final_usage.completion_tokens
                 total_tokens = self.final_usage.total_tokens
-                # Check if we have cached tokens info
-                if hasattr(self.final_usage, 'prompt_tokens_details') and hasattr(
-                        self.final_usage.prompt_tokens_details, 'cached_tokens'):
-                    cached_tokens = self.final_usage.prompt_tokens_details.cached_tokens
+                cache_creation_token_count, cache_read_token_count = _extract_cache_token_counts(
+                    self.final_usage
+                )
                 logger.debug(
                     f"Using token usage from final chunk: prompt={prompt_tokens}, completion={completion_tokens}, total={total_tokens}")
             else:
@@ -1168,7 +1247,9 @@ def handle_streaming_response(
                         prompt_tokens=prompt_tokens,
                         completion_tokens=completion_tokens,
                         total_tokens=total_tokens,
-                        cached_tokens=cached_tokens,
+                        cached_tokens=0,
+                        cache_creation_token_count=cache_creation_token_count,
+                        cache_read_token_count=cache_read_token_count,
                         stop_reason=stop_reason,
                         request_time=self.request_time_dt.strftime(
                             "%Y-%m-%dT%H:%M:%SZ"
@@ -1359,12 +1440,17 @@ def handle_streaming_responses(stream, request_time_dt, usage_metadata,
             input_tokens = 0
             output_tokens = 0
             total_tokens = 0
+            cache_creation_token_count = 0
+            cache_read_token_count = 0
 
             # Get usage data from the final chunk or last chunk
             if self.final_usage:
                 input_tokens = self.final_usage.input_tokens
                 output_tokens = self.final_usage.output_tokens
                 total_tokens = self.final_usage.total_tokens
+                cache_creation_token_count, cache_read_token_count = _extract_cache_token_counts(
+                    self.final_usage
+                )
                 logger.debug(
                     f"Using token usage from Responses API stream final chunk: input={input_tokens}, "
                     f"output={output_tokens}, total={total_tokens}")
@@ -1373,6 +1459,9 @@ def handle_streaming_responses(stream, request_time_dt, usage_metadata,
                 input_tokens = self.last_chunk.usage.input_tokens
                 output_tokens = self.last_chunk.usage.output_tokens
                 total_tokens = self.last_chunk.usage.total_tokens
+                cache_creation_token_count, cache_read_token_count = _extract_cache_token_counts(
+                    self.last_chunk.usage
+                )
                 logger.debug(
                     f"Using token usage from Responses API last chunk: input={input_tokens}, "
                     f"output={output_tokens}, total={total_tokens}")
@@ -1463,6 +1552,8 @@ def handle_streaming_responses(stream, request_time_dt, usage_metadata,
                         completion_tokens=output_tokens,
                         total_tokens=total_tokens,
                         cached_tokens=0,
+                        cache_creation_token_count=cache_creation_token_count,
+                        cache_read_token_count=cache_read_token_count,
                         stop_reason="END",
                         request_time=self.request_time_dt.strftime("%Y-%m-%dT%H:%M:%SZ"),
                         response_time=response_time,
@@ -1541,24 +1632,11 @@ def _wrap_async_stream(stream, request_time_dt, usage_metadata, client_instance=
             if not self.chunks and not self.final_usage:
                 return
 
-            prompt_tokens = 0
-            completion_tokens = 0
-            cached_tokens = 0
-
-            if self.final_usage:
-                prompt_tokens = getattr(self.final_usage, 'prompt_tokens', 0) or 0
-                completion_tokens = getattr(self.final_usage, 'completion_tokens', 0) or 0
-                if hasattr(self.final_usage, 'prompt_tokens_details') and self.final_usage.prompt_tokens_details:
-                    cached_tokens = getattr(self.final_usage.prompt_tokens_details, 'cached_tokens', 0) or 0
+            if self.final_usage is None:
+                return
 
-            total_tokens = prompt_tokens + completion_tokens
             time_to_first_token = int((self.first_token_time - request_time_dt).total_seconds() * 1000) if self.first_token_time else 0
 
-            response_time_dt = datetime.datetime.now(datetime.timezone.utc)
-            request_time = request_time_dt.strftime("%Y-%m-%dT%H:%M:%SZ")
-            response_time = response_time_dt.strftime("%Y-%m-%dT%H:%M:%SZ")
-            request_duration = (response_time_dt - request_time_dt).total_seconds() * 1000
-
             try:
                 create_metering_call(
                     type('Response', (), {

{revenium_python_sdk-0.1.7 → revenium_python_sdk-0.1.8/revenium_python_sdk.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: revenium-python-sdk
-Version: 0.1.7
+Version: 0.1.8
 Summary: The official Revenium Python SDK — unified AI metering middleware for OpenAI, Anthropic, Google, Ollama, LiteLLM, Perplexity, and fal.ai.
 Author-email: Revenium <support@revenium.io>
 License: MIT