posthoganalytics 6.7.2__py3-none-any.whl → 6.7.3__py3-none-any.whl

posthoganalytics/ai/anthropic/anthropic.py

@@ -10,7 +10,7 @@ import time
 import uuid
 from typing import Any, Dict, List, Optional
 
-from posthoganalytics.ai.types import StreamingContentBlock, ToolInProgress
+from posthoganalytics.ai.types import StreamingContentBlock, TokenUsage, ToolInProgress
 from posthoganalytics.ai.utils import (
     call_llm_and_track_usage,
     merge_usage_stats,
@@ -126,7 +126,7 @@ class WrappedMessages(Messages):
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
+        usage_stats: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0)
         accumulated_content = ""
         content_blocks: List[StreamingContentBlock] = []
         tools_in_progress: Dict[str, ToolInProgress] = {}
@@ -210,14 +210,13 @@ class WrappedMessages(Messages):
         posthog_privacy_mode: bool,
         posthog_groups: Optional[Dict[str, Any]],
         kwargs: Dict[str, Any],
-        usage_stats: Dict[str, int],
+        usage_stats: TokenUsage,
         latency: float,
         content_blocks: List[StreamingContentBlock],
         accumulated_content: str,
     ):
         from posthoganalytics.ai.types import StreamingEventData
         from posthoganalytics.ai.anthropic.anthropic_converter import (
-            standardize_anthropic_usage,
             format_anthropic_streaming_input,
             format_anthropic_streaming_output_complete,
         )
@@ -236,7 +235,7 @@ class WrappedMessages(Messages):
             formatted_output=format_anthropic_streaming_output_complete(
                 content_blocks, accumulated_content
             ),
-            usage_stats=standardize_anthropic_usage(usage_stats),
+            usage_stats=usage_stats,
             latency=latency,
             distinct_id=posthog_distinct_id,
             trace_id=posthog_trace_id,

posthoganalytics/ai/anthropic/anthropic_async.py

@@ -11,7 +11,7 @@ import uuid
 from typing import Any, Dict, List, Optional
 
 from posthoganalytics import setup
-from posthoganalytics.ai.types import StreamingContentBlock, ToolInProgress
+from posthoganalytics.ai.types import StreamingContentBlock, TokenUsage, ToolInProgress
 from posthoganalytics.ai.utils import (
     call_llm_and_track_usage_async,
     extract_available_tool_calls,
@@ -131,7 +131,7 @@ class AsyncWrappedMessages(AsyncMessages):
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
+        usage_stats: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0)
         accumulated_content = ""
         content_blocks: List[StreamingContentBlock] = []
         tools_in_progress: Dict[str, ToolInProgress] = {}
@@ -215,7 +215,7 @@ class AsyncWrappedMessages(AsyncMessages):
         posthog_privacy_mode: bool,
         posthog_groups: Optional[Dict[str, Any]],
         kwargs: Dict[str, Any],
-        usage_stats: Dict[str, int],
+        usage_stats: TokenUsage,
         latency: float,
         content_blocks: List[StreamingContentBlock],
         accumulated_content: str,

posthoganalytics/ai/anthropic/anthropic_converter.py

@@ -14,7 +14,6 @@ from posthoganalytics.ai.types import (
     FormattedMessage,
     FormattedTextContent,
     StreamingContentBlock,
-    StreamingUsageStats,
     TokenUsage,
     ToolInProgress,
 )
@@ -164,7 +163,38 @@ def format_anthropic_streaming_content(
     return formatted
 
 
-def extract_anthropic_usage_from_event(event: Any) -> StreamingUsageStats:
+def extract_anthropic_usage_from_response(response: Any) -> TokenUsage:
+    """
+    Extract usage from a full Anthropic response (non-streaming).
+
+    Args:
+        response: The complete response from Anthropic API
+
+    Returns:
+        TokenUsage with standardized usage
+    """
+    if not hasattr(response, "usage"):
+        return TokenUsage(input_tokens=0, output_tokens=0)
+
+    result = TokenUsage(
+        input_tokens=getattr(response.usage, "input_tokens", 0),
+        output_tokens=getattr(response.usage, "output_tokens", 0),
+    )
+
+    if hasattr(response.usage, "cache_read_input_tokens"):
+        cache_read = response.usage.cache_read_input_tokens
+        if cache_read and cache_read > 0:
+            result["cache_read_input_tokens"] = cache_read
+
+    if hasattr(response.usage, "cache_creation_input_tokens"):
+        cache_creation = response.usage.cache_creation_input_tokens
+        if cache_creation and cache_creation > 0:
+            result["cache_creation_input_tokens"] = cache_creation
+
+    return result
+
+
+def extract_anthropic_usage_from_event(event: Any) -> TokenUsage:
     """
     Extract usage statistics from an Anthropic streaming event.
 
@@ -175,7 +205,7 @@ def extract_anthropic_usage_from_event(event: Any) -> StreamingUsageStats:
         Dictionary of usage statistics
     """
 
-    usage: StreamingUsageStats = {}
+    usage: TokenUsage = TokenUsage()
 
     # Handle usage stats from message_start event
     if hasattr(event, "type") and event.type == "message_start":
@@ -329,26 +359,6 @@ def finalize_anthropic_tool_input(
             del tools_in_progress[block["id"]]
 
 
-def standardize_anthropic_usage(usage: Dict[str, Any]) -> TokenUsage:
-    """
-    Standardize Anthropic usage statistics to common TokenUsage format.
-
-    Anthropic already uses standard field names, so this mainly structures the data.
-
-    Args:
-        usage: Raw usage statistics from Anthropic
-
-    Returns:
-        Standardized TokenUsage dict
-    """
-    return TokenUsage(
-        input_tokens=usage.get("input_tokens", 0),
-        output_tokens=usage.get("output_tokens", 0),
-        cache_read_input_tokens=usage.get("cache_read_input_tokens"),
-        cache_creation_input_tokens=usage.get("cache_creation_input_tokens"),
-    )
-
-
 def format_anthropic_streaming_input(kwargs: Dict[str, Any]) -> Any:
     """
     Format Anthropic streaming input using system prompt merging.

posthoganalytics/ai/gemini/gemini.py

@@ -3,6 +3,8 @@ import time
 import uuid
 from typing import Any, Dict, Optional
 
+from posthoganalytics.ai.types import TokenUsage
+
 try:
     from google import genai
 except ImportError:
@@ -294,7 +296,7 @@ class Models:
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
+        usage_stats: TokenUsage = TokenUsage(input_tokens=0, output_tokens=0)
         accumulated_content = []
 
         kwargs_without_stream = {"model": model, "contents": contents, **kwargs}
@@ -350,12 +352,11 @@ class Models:
         privacy_mode: bool,
         groups: Optional[Dict[str, Any]],
         kwargs: Dict[str, Any],
-        usage_stats: Dict[str, int],
+        usage_stats: TokenUsage,
         latency: float,
         output: Any,
     ):
         from posthoganalytics.ai.types import StreamingEventData
-        from posthoganalytics.ai.gemini.gemini_converter import standardize_gemini_usage
 
         # Prepare standardized event data
         formatted_input = self._format_input(contents)
@@ -368,7 +369,7 @@ class Models:
             kwargs=kwargs,
             formatted_input=sanitized_input,
             formatted_output=format_gemini_streaming_output(output),
-            usage_stats=standardize_gemini_usage(usage_stats),
+            usage_stats=usage_stats,
             latency=latency,
             distinct_id=distinct_id,
             trace_id=trace_id,

posthoganalytics/ai/gemini/gemini_converter.py

@@ -10,7 +10,6 @@ from typing import Any, Dict, List, Optional, TypedDict, Union
 from posthoganalytics.ai.types import (
     FormattedContentItem,
     FormattedMessage,
-    StreamingUsageStats,
     TokenUsage,
 )
 
@@ -283,7 +282,54 @@ def format_gemini_input(contents: Any) -> List[FormattedMessage]:
     return [_format_object_message(contents)]
 
 
-def extract_gemini_usage_from_chunk(chunk: Any) -> StreamingUsageStats:
+def _extract_usage_from_metadata(metadata: Any) -> TokenUsage:
+    """
+    Common logic to extract usage from Gemini metadata.
+    Used by both streaming and non-streaming paths.
+
+    Args:
+        metadata: usage_metadata from Gemini response or chunk
+
+    Returns:
+        TokenUsage with standardized usage
+    """
+    usage = TokenUsage(
+        input_tokens=getattr(metadata, "prompt_token_count", 0),
+        output_tokens=getattr(metadata, "candidates_token_count", 0),
+    )
+
+    # Add cache tokens if present (don't add if 0)
+    if hasattr(metadata, "cached_content_token_count"):
+        cache_tokens = metadata.cached_content_token_count
+        if cache_tokens and cache_tokens > 0:
+            usage["cache_read_input_tokens"] = cache_tokens
+
+    # Add reasoning tokens if present (don't add if 0)
+    if hasattr(metadata, "thoughts_token_count"):
+        reasoning_tokens = metadata.thoughts_token_count
+        if reasoning_tokens and reasoning_tokens > 0:
+            usage["reasoning_tokens"] = reasoning_tokens
+
+    return usage
+
+
+def extract_gemini_usage_from_response(response: Any) -> TokenUsage:
+    """
+    Extract usage statistics from a full Gemini response (non-streaming).
+
+    Args:
+        response: The complete response from Gemini API
+
+    Returns:
+        TokenUsage with standardized usage statistics
+    """
+    if not hasattr(response, "usage_metadata") or not response.usage_metadata:
+        return TokenUsage(input_tokens=0, output_tokens=0)
+
+    return _extract_usage_from_metadata(response.usage_metadata)
+
+
+def extract_gemini_usage_from_chunk(chunk: Any) -> TokenUsage:
     """
     Extract usage statistics from a Gemini streaming chunk.
 
@@ -291,21 +337,16 @@ def extract_gemini_usage_from_chunk(chunk: Any) -> StreamingUsageStats:
         chunk: Streaming chunk from Gemini API
 
     Returns:
-        Dictionary of usage statistics
+        TokenUsage with standardized usage statistics
     """
 
-    usage: StreamingUsageStats = {}
+    usage: TokenUsage = TokenUsage()
 
     if not hasattr(chunk, "usage_metadata") or not chunk.usage_metadata:
         return usage
 
-    # Gemini uses prompt_token_count and candidates_token_count
-    usage["input_tokens"] = getattr(chunk.usage_metadata, "prompt_token_count", 0)
-    usage["output_tokens"] = getattr(chunk.usage_metadata, "candidates_token_count", 0)
-
-    # Calculate total if both values are defined (including 0)
-    if "input_tokens" in usage and "output_tokens" in usage:
-        usage["total_tokens"] = usage["input_tokens"] + usage["output_tokens"]
+    # Use the shared helper to extract usage
+    usage = _extract_usage_from_metadata(chunk.usage_metadata)
 
     return usage
 
@@ -417,22 +458,3 @@ def format_gemini_streaming_output(
 
     # Fallback for empty or unexpected input
     return [{"role": "assistant", "content": [{"type": "text", "text": ""}]}]
-
-
-def standardize_gemini_usage(usage: Dict[str, Any]) -> TokenUsage:
-    """
-    Standardize Gemini usage statistics to common TokenUsage format.
-
-    Gemini already uses standard field names (input_tokens/output_tokens).
-
-    Args:
-        usage: Raw usage statistics from Gemini
-
-    Returns:
-        Standardized TokenUsage dict
-    """
-    return TokenUsage(
-        input_tokens=usage.get("input_tokens", 0),
-        output_tokens=usage.get("output_tokens", 0),
-        # Gemini doesn't currently support cache or reasoning tokens
-    )

posthoganalytics/ai/openai/openai.py

@@ -2,6 +2,8 @@ import time
 import uuid
 from typing import Any, Dict, List, Optional
 
+from posthoganalytics.ai.types import TokenUsage
+
 try:
     import openai
 except ImportError:
@@ -120,7 +122,7 @@ class WrappedResponses:
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {}
+        usage_stats: TokenUsage = TokenUsage()
         final_content = []
         response = self._original.create(**kwargs)
 
@@ -171,14 +173,13 @@ class WrappedResponses:
         posthog_privacy_mode: bool,
         posthog_groups: Optional[Dict[str, Any]],
         kwargs: Dict[str, Any],
-        usage_stats: Dict[str, int],
+        usage_stats: TokenUsage,
         latency: float,
         output: Any,
         available_tool_calls: Optional[List[Dict[str, Any]]] = None,
     ):
         from posthoganalytics.ai.types import StreamingEventData
         from posthoganalytics.ai.openai.openai_converter import (
-            standardize_openai_usage,
             format_openai_streaming_input,
             format_openai_streaming_output,
         )
@@ -195,7 +196,7 @@ class WrappedResponses:
             kwargs=kwargs,
             formatted_input=sanitized_input,
             formatted_output=format_openai_streaming_output(output, "responses"),
-            usage_stats=standardize_openai_usage(usage_stats, "responses"),
+            usage_stats=usage_stats,
             latency=latency,
             distinct_id=posthog_distinct_id,
             trace_id=posthog_trace_id,
@@ -316,7 +317,7 @@ class WrappedCompletions:
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {}
+        usage_stats: TokenUsage = TokenUsage()
         accumulated_content = []
         accumulated_tool_calls: Dict[int, Dict[str, Any]] = {}
         if "stream_options" not in kwargs:
@@ -387,7 +388,7 @@ class WrappedCompletions:
         posthog_privacy_mode: bool,
         posthog_groups: Optional[Dict[str, Any]],
         kwargs: Dict[str, Any],
-        usage_stats: Dict[str, int],
+        usage_stats: TokenUsage,
         latency: float,
         output: Any,
         tool_calls: Optional[List[Dict[str, Any]]] = None,
@@ -395,7 +396,6 @@ class WrappedCompletions:
     ):
         from posthoganalytics.ai.types import StreamingEventData
         from posthoganalytics.ai.openai.openai_converter import (
-            standardize_openai_usage,
             format_openai_streaming_input,
             format_openai_streaming_output,
         )
@@ -412,7 +412,7 @@ class WrappedCompletions:
             kwargs=kwargs,
             formatted_input=sanitized_input,
             formatted_output=format_openai_streaming_output(output, "chat", tool_calls),
-            usage_stats=standardize_openai_usage(usage_stats, "chat"),
+            usage_stats=usage_stats,
             latency=latency,
             distinct_id=posthog_distinct_id,
             trace_id=posthog_trace_id,

posthoganalytics/ai/openai/openai_async.py

@@ -2,6 +2,8 @@ import time
 import uuid
 from typing import Any, Dict, List, Optional
 
+from posthoganalytics.ai.types import TokenUsage
+
 try:
     import openai
 except ImportError:
@@ -124,7 +126,7 @@ class WrappedResponses:
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {}
+        usage_stats: TokenUsage = TokenUsage()
         final_content = []
         response = self._original.create(**kwargs)
 
@@ -176,7 +178,7 @@ class WrappedResponses:
         posthog_privacy_mode: bool,
         posthog_groups: Optional[Dict[str, Any]],
         kwargs: Dict[str, Any],
-        usage_stats: Dict[str, int],
+        usage_stats: TokenUsage,
         latency: float,
         output: Any,
         available_tool_calls: Optional[List[Dict[str, Any]]] = None,
@@ -336,7 +338,7 @@ class WrappedCompletions:
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {}
+        usage_stats: TokenUsage = TokenUsage()
         accumulated_content = []
         accumulated_tool_calls: Dict[int, Dict[str, Any]] = {}
 
@@ -406,7 +408,7 @@ class WrappedCompletions:
         posthog_privacy_mode: bool,
         posthog_groups: Optional[Dict[str, Any]],
         kwargs: Dict[str, Any],
-        usage_stats: Dict[str, int],
+        usage_stats: TokenUsage,
         latency: float,
         output: Any,
         tool_calls: Optional[List[Dict[str, Any]]] = None,
@@ -430,8 +432,8 @@ class WrappedCompletions:
                 format_openai_streaming_output(output, "chat", tool_calls),
             ),
             "$ai_http_status": 200,
-            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
-            "$ai_output_tokens": usage_stats.get("completion_tokens", 0),
+            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
+            "$ai_output_tokens": usage_stats.get("output_tokens", 0),
             "$ai_cache_read_input_tokens": usage_stats.get(
                 "cache_read_input_tokens", 0
             ),
@@ -501,13 +503,13 @@ class WrappedEmbeddings:
         end_time = time.time()
 
         # Extract usage statistics if available
-        usage_stats = {}
+        usage_stats: TokenUsage = TokenUsage()
 
         if hasattr(response, "usage") and response.usage:
-            usage_stats = {
-                "prompt_tokens": getattr(response.usage, "prompt_tokens", 0),
-                "total_tokens": getattr(response.usage, "total_tokens", 0),
-            }
+            usage_stats = TokenUsage(
+                input_tokens=getattr(response.usage, "prompt_tokens", 0),
+                output_tokens=getattr(response.usage, "completion_tokens", 0),
+            )
 
         latency = end_time - start_time
 
@@ -521,7 +523,7 @@ class WrappedEmbeddings:
                 sanitize_openai_response(kwargs.get("input")),
             ),
             "$ai_http_status": 200,
-            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
+            "$ai_input_tokens": usage_stats.get("input_tokens", 0),
             "$ai_latency": latency,
             "$ai_trace_id": posthog_trace_id,
             "$ai_base_url": str(self._client.base_url),

posthoganalytics/ai/openai/openai_converter.py

@@ -14,7 +14,6 @@ from posthoganalytics.ai.types import (
     FormattedImageContent,
     FormattedMessage,
     FormattedTextContent,
-    StreamingUsageStats,
     TokenUsage,
 )
 
@@ -256,9 +255,69 @@ def format_openai_streaming_content(
     return formatted
 
 
+def extract_openai_usage_from_response(response: Any) -> TokenUsage:
+    """
+    Extract usage statistics from a full OpenAI response (non-streaming).
+    Handles both Chat Completions and Responses API.
+
+    Args:
+        response: The complete response from OpenAI API
+
+    Returns:
+        TokenUsage with standardized usage statistics
+    """
+    if not hasattr(response, "usage"):
+        return TokenUsage(input_tokens=0, output_tokens=0)
+
+    cached_tokens = 0
+    input_tokens = 0
+    output_tokens = 0
+    reasoning_tokens = 0
+
+    # Responses API format
+    if hasattr(response.usage, "input_tokens"):
+        input_tokens = response.usage.input_tokens
+    if hasattr(response.usage, "output_tokens"):
+        output_tokens = response.usage.output_tokens
+    if hasattr(response.usage, "input_tokens_details") and hasattr(
+        response.usage.input_tokens_details, "cached_tokens"
+    ):
+        cached_tokens = response.usage.input_tokens_details.cached_tokens
+    if hasattr(response.usage, "output_tokens_details") and hasattr(
+        response.usage.output_tokens_details, "reasoning_tokens"
+    ):
+        reasoning_tokens = response.usage.output_tokens_details.reasoning_tokens
+
+    # Chat Completions format
+    if hasattr(response.usage, "prompt_tokens"):
+        input_tokens = response.usage.prompt_tokens
+    if hasattr(response.usage, "completion_tokens"):
+        output_tokens = response.usage.completion_tokens
+    if hasattr(response.usage, "prompt_tokens_details") and hasattr(
+        response.usage.prompt_tokens_details, "cached_tokens"
+    ):
+        cached_tokens = response.usage.prompt_tokens_details.cached_tokens
+    if hasattr(response.usage, "completion_tokens_details") and hasattr(
+        response.usage.completion_tokens_details, "reasoning_tokens"
+    ):
+        reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens
+
+    result = TokenUsage(
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+    )
+
+    if cached_tokens > 0:
+        result["cache_read_input_tokens"] = cached_tokens
+    if reasoning_tokens > 0:
+        result["reasoning_tokens"] = reasoning_tokens
+
+    return result
+
+
 def extract_openai_usage_from_chunk(
     chunk: Any, provider_type: str = "chat"
-) -> StreamingUsageStats:
+) -> TokenUsage:
     """
     Extract usage statistics from an OpenAI streaming chunk.
 
@@ -272,16 +331,16 @@ def extract_openai_usage_from_chunk(
         Dictionary of usage statistics
     """
 
-    usage: StreamingUsageStats = {}
+    usage: TokenUsage = TokenUsage()
 
     if provider_type == "chat":
         if not hasattr(chunk, "usage") or not chunk.usage:
             return usage
 
         # Chat Completions API uses prompt_tokens and completion_tokens
-        usage["prompt_tokens"] = getattr(chunk.usage, "prompt_tokens", 0)
-        usage["completion_tokens"] = getattr(chunk.usage, "completion_tokens", 0)
-        usage["total_tokens"] = getattr(chunk.usage, "total_tokens", 0)
+        # Standardize to input_tokens and output_tokens
+        usage["input_tokens"] = getattr(chunk.usage, "prompt_tokens", 0)
+        usage["output_tokens"] = getattr(chunk.usage, "completion_tokens", 0)
 
         # Handle cached tokens
         if hasattr(chunk.usage, "prompt_tokens_details") and hasattr(
@@ -310,7 +369,6 @@ def extract_openai_usage_from_chunk(
                 response_usage = chunk.response.usage
                 usage["input_tokens"] = getattr(response_usage, "input_tokens", 0)
                 usage["output_tokens"] = getattr(response_usage, "output_tokens", 0)
-                usage["total_tokens"] = getattr(response_usage, "total_tokens", 0)
 
                 # Handle cached tokens
                 if hasattr(response_usage, "input_tokens_details") and hasattr(
@@ -535,37 +593,6 @@ def format_openai_streaming_output(
     ]
 
 
-def standardize_openai_usage(
-    usage: Dict[str, Any], api_type: str = "chat"
-) -> TokenUsage:
-    """
-    Standardize OpenAI usage statistics to common TokenUsage format.
-
-    Args:
-        usage: Raw usage statistics from OpenAI
-        api_type: Either "chat" or "responses" to handle different field names
-
-    Returns:
-        Standardized TokenUsage dict
-    """
-    if api_type == "chat":
-        # Chat API uses prompt_tokens/completion_tokens
-        return TokenUsage(
-            input_tokens=usage.get("prompt_tokens", 0),
-            output_tokens=usage.get("completion_tokens", 0),
-            cache_read_input_tokens=usage.get("cache_read_input_tokens"),
-            reasoning_tokens=usage.get("reasoning_tokens"),
-        )
-    else:  # responses API
-        # Responses API uses input_tokens/output_tokens
-        return TokenUsage(
-            input_tokens=usage.get("input_tokens", 0),
-            output_tokens=usage.get("output_tokens", 0),
-            cache_read_input_tokens=usage.get("cache_read_input_tokens"),
-            reasoning_tokens=usage.get("reasoning_tokens"),
-        )
-
-
 def format_openai_streaming_input(
     kwargs: Dict[str, Any], api_type: str = "chat"
 ) -> Any:

posthoganalytics/ai/types.py

@@ -77,24 +77,6 @@ class ProviderResponse(TypedDict, total=False):
     error: Optional[str]
 
 
-class StreamingUsageStats(TypedDict, total=False):
-    """
-    Usage statistics collected during streaming.
-
-    Different providers populate different fields during streaming.
-    """
-
-    input_tokens: int
-    output_tokens: int
-    cache_read_input_tokens: Optional[int]
-    cache_creation_input_tokens: Optional[int]
-    reasoning_tokens: Optional[int]
-    # OpenAI-specific names
-    prompt_tokens: Optional[int]
-    completion_tokens: Optional[int]
-    total_tokens: Optional[int]
-
-
 class StreamingContentBlock(TypedDict, total=False):
     """
     Content block used during streaming to accumulate content.
@@ -133,7 +115,7 @@ class StreamingEventData(TypedDict):
     kwargs: Dict[str, Any]  # Original kwargs for tool extraction and special handling
     formatted_input: Any  # Provider-formatted input ready for tracking
     formatted_output: Any  # Provider-formatted output ready for tracking
-    usage_stats: TokenUsage  # Standardized token counts
+    usage_stats: TokenUsage
     latency: float
     distinct_id: Optional[str]
     trace_id: Optional[str]

posthoganalytics/ai/utils.py

@@ -2,9 +2,8 @@ import time
 import uuid
 from typing import Any, Callable, Dict, Optional
 
-
 from posthoganalytics.client import Client as PostHogClient
-from posthoganalytics.ai.types import StreamingEventData, StreamingUsageStats
+from posthoganalytics.ai.types import StreamingEventData, TokenUsage
 from posthoganalytics.ai.sanitization import (
     sanitize_openai,
     sanitize_anthropic,
@@ -14,7 +13,7 @@ from posthoganalytics.ai.sanitization import (
 
 
 def merge_usage_stats(
-    target: Dict[str, int], source: StreamingUsageStats, mode: str = "incremental"
+    target: TokenUsage, source: TokenUsage, mode: str = "incremental"
 ) -> None:
     """
     Merge streaming usage statistics into target dict, handling None values.
@@ -25,19 +24,49 @@ def merge_usage_stats(
 
     Args:
         target: Dictionary to update with usage stats
-        source: StreamingUsageStats that may contain None values
+        source: TokenUsage that may contain None values
         mode: Either "incremental" or "cumulative"
     """
     if mode == "incremental":
         # Add new values to existing totals
-        for key, value in source.items():
-            if value is not None and isinstance(value, int):
-                target[key] = target.get(key, 0) + value
+        source_input = source.get("input_tokens")
+        if source_input is not None:
+            current = target.get("input_tokens") or 0
+            target["input_tokens"] = current + source_input
+
+        source_output = source.get("output_tokens")
+        if source_output is not None:
+            current = target.get("output_tokens") or 0
+            target["output_tokens"] = current + source_output
+
+        source_cache_read = source.get("cache_read_input_tokens")
+        if source_cache_read is not None:
+            current = target.get("cache_read_input_tokens") or 0
+            target["cache_read_input_tokens"] = current + source_cache_read
+
+        source_cache_creation = source.get("cache_creation_input_tokens")
+        if source_cache_creation is not None:
+            current = target.get("cache_creation_input_tokens") or 0
+            target["cache_creation_input_tokens"] = current + source_cache_creation
+
+        source_reasoning = source.get("reasoning_tokens")
+        if source_reasoning is not None:
+            current = target.get("reasoning_tokens") or 0
+            target["reasoning_tokens"] = current + source_reasoning
     elif mode == "cumulative":
         # Replace with latest values (already cumulative)
-        for key, value in source.items():
-            if value is not None and isinstance(value, int):
-                target[key] = value
+        if source.get("input_tokens") is not None:
+            target["input_tokens"] = source["input_tokens"]
+        if source.get("output_tokens") is not None:
+            target["output_tokens"] = source["output_tokens"]
+        if source.get("cache_read_input_tokens") is not None:
+            target["cache_read_input_tokens"] = source["cache_read_input_tokens"]
+        if source.get("cache_creation_input_tokens") is not None:
+            target["cache_creation_input_tokens"] = source[
+                "cache_creation_input_tokens"
+            ]
+        if source.get("reasoning_tokens") is not None:
+            target["reasoning_tokens"] = source["reasoning_tokens"]
     else:
         raise ValueError(f"Invalid mode: {mode}. Must be 'incremental' or 'cumulative'")
 
@@ -64,74 +93,31 @@ def get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]:
     return model_params
 
 
-def get_usage(response, provider: str) -> Dict[str, Any]:
+def get_usage(response, provider: str) -> TokenUsage:
+    """
+    Extract usage statistics from response based on provider.
+    Delegates to provider-specific converter functions.
+    """
     if provider == "anthropic":
-        return {
-            "input_tokens": response.usage.input_tokens,
-            "output_tokens": response.usage.output_tokens,
-            "cache_read_input_tokens": response.usage.cache_read_input_tokens,
-            "cache_creation_input_tokens": response.usage.cache_creation_input_tokens,
-        }
+        from posthoganalytics.ai.anthropic.anthropic_converter import (
+            extract_anthropic_usage_from_response,
+        )
+
+        return extract_anthropic_usage_from_response(response)
     elif provider == "openai":
-        cached_tokens = 0
-        input_tokens = 0
-        output_tokens = 0
-        reasoning_tokens = 0
-
-        # responses api
-        if hasattr(response.usage, "input_tokens"):
-            input_tokens = response.usage.input_tokens
-        if hasattr(response.usage, "output_tokens"):
-            output_tokens = response.usage.output_tokens
-        if hasattr(response.usage, "input_tokens_details") and hasattr(
-            response.usage.input_tokens_details, "cached_tokens"
-        ):
-            cached_tokens = response.usage.input_tokens_details.cached_tokens
-        if hasattr(response.usage, "output_tokens_details") and hasattr(
-            response.usage.output_tokens_details, "reasoning_tokens"
-        ):
-            reasoning_tokens = response.usage.output_tokens_details.reasoning_tokens
-
-        # chat completions
-        if hasattr(response.usage, "prompt_tokens"):
-            input_tokens = response.usage.prompt_tokens
-        if hasattr(response.usage, "completion_tokens"):
-            output_tokens = response.usage.completion_tokens
-        if hasattr(response.usage, "prompt_tokens_details") and hasattr(
-            response.usage.prompt_tokens_details, "cached_tokens"
-        ):
-            cached_tokens = response.usage.prompt_tokens_details.cached_tokens
+        from posthoganalytics.ai.openai.openai_converter import (
+            extract_openai_usage_from_response,
+        )
 
-        return {
-            "input_tokens": input_tokens,
-            "output_tokens": output_tokens,
-            "cache_read_input_tokens": cached_tokens,
-            "reasoning_tokens": reasoning_tokens,
-        }
+        return extract_openai_usage_from_response(response)
     elif provider == "gemini":
-        input_tokens = 0
-        output_tokens = 0
+        from posthoganalytics.ai.gemini.gemini_converter import (
+            extract_gemini_usage_from_response,
+        )
 
-        if hasattr(response, "usage_metadata") and response.usage_metadata:
-            input_tokens = getattr(response.usage_metadata, "prompt_token_count", 0)
-            output_tokens = getattr(
-                response.usage_metadata, "candidates_token_count", 0
-            )
+        return extract_gemini_usage_from_response(response)
 
-        return {
-            "input_tokens": input_tokens,
-            "output_tokens": output_tokens,
-            "cache_read_input_tokens": 0,
-            "cache_creation_input_tokens": 0,
-            "reasoning_tokens": 0,
-        }
-    return {
-        "input_tokens": 0,
-        "output_tokens": 0,
-        "cache_read_input_tokens": 0,
-        "cache_creation_input_tokens": 0,
-        "reasoning_tokens": 0,
-    }
+    return TokenUsage(input_tokens=0, output_tokens=0)
 
 
 def format_response(response, provider: str):
@@ -169,6 +155,7 @@ def extract_available_tool_calls(provider: str, kwargs: Dict[str, Any]):
         from posthoganalytics.ai.openai.openai_converter import extract_openai_tools
 
         return extract_openai_tools(kwargs)
+    return None
 
 
 def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
@@ -187,9 +174,9 @@ def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
         contents = kwargs.get("contents", [])
         return format_gemini_input(contents)
     elif provider == "openai":
-        # For OpenAI, handle both Chat Completions and Responses API
         from posthoganalytics.ai.openai.openai_converter import format_openai_input
 
+        # For OpenAI, handle both Chat Completions and Responses API
         messages_param = kwargs.get("messages")
         input_param = kwargs.get("input")
 
@@ -250,7 +237,7 @@ def call_llm_and_track_usage(
     response = None
     error = None
     http_status = 200
-    usage: Dict[str, Any] = {}
+    usage: TokenUsage = TokenUsage()
     error_params: Dict[str, Any] = {}
 
     try:
@@ -305,27 +292,17 @@ def call_llm_and_track_usage(
         if available_tool_calls:
             event_properties["$ai_tools"] = available_tool_calls
 
-        if (
-            usage.get("cache_read_input_tokens") is not None
-            and usage.get("cache_read_input_tokens", 0) > 0
-        ):
-            event_properties["$ai_cache_read_input_tokens"] = usage.get(
-                "cache_read_input_tokens", 0
-            )
+        cache_read = usage.get("cache_read_input_tokens")
+        if cache_read is not None and cache_read > 0:
+            event_properties["$ai_cache_read_input_tokens"] = cache_read
 
-        if (
-            usage.get("cache_creation_input_tokens") is not None
-            and usage.get("cache_creation_input_tokens", 0) > 0
-        ):
-            event_properties["$ai_cache_creation_input_tokens"] = usage.get(
-                "cache_creation_input_tokens", 0
-            )
+        cache_creation = usage.get("cache_creation_input_tokens")
+        if cache_creation is not None and cache_creation > 0:
+            event_properties["$ai_cache_creation_input_tokens"] = cache_creation
 
-        if (
-            usage.get("reasoning_tokens") is not None
-            and usage.get("reasoning_tokens", 0) > 0
-        ):
-            event_properties["$ai_reasoning_tokens"] = usage.get("reasoning_tokens", 0)
+        reasoning = usage.get("reasoning_tokens")
+        if reasoning is not None and reasoning > 0:
+            event_properties["$ai_reasoning_tokens"] = reasoning
 
         if posthog_distinct_id is None:
             event_properties["$process_person_profile"] = False
@@ -367,7 +344,7 @@ async def call_llm_and_track_usage_async(
     response = None
     error = None
     http_status = 200
-    usage: Dict[str, Any] = {}
+    usage: TokenUsage = TokenUsage()
     error_params: Dict[str, Any] = {}
 
     try:
@@ -422,21 +399,13 @@ async def call_llm_and_track_usage_async(
         if available_tool_calls:
             event_properties["$ai_tools"] = available_tool_calls
 
-        if (
-            usage.get("cache_read_input_tokens") is not None
-            and usage.get("cache_read_input_tokens", 0) > 0
-        ):
-            event_properties["$ai_cache_read_input_tokens"] = usage.get(
-                "cache_read_input_tokens", 0
-            )
+        cache_read = usage.get("cache_read_input_tokens")
+        if cache_read is not None and cache_read > 0:
+            event_properties["$ai_cache_read_input_tokens"] = cache_read
 
-        if (
-            usage.get("cache_creation_input_tokens") is not None
-            and usage.get("cache_creation_input_tokens", 0) > 0
-        ):
-            event_properties["$ai_cache_creation_input_tokens"] = usage.get(
-                "cache_creation_input_tokens", 0
-            )
+        cache_creation = usage.get("cache_creation_input_tokens")
+        if cache_creation is not None and cache_creation > 0:
+            event_properties["$ai_cache_creation_input_tokens"] = cache_creation
 
         if posthog_distinct_id is None:
             event_properties["$process_person_profile"] = False

posthoganalytics/version.py

@@ -1,4 +1,4 @@
-VERSION = "6.7.2"
+VERSION = "6.7.3"
 
 if __name__ == "__main__":
     print(VERSION, end="")  # noqa: T201

{posthoganalytics-6.7.2.dist-info → posthoganalytics-6.7.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: posthoganalytics
-Version: 6.7.2
+Version: 6.7.3
 Summary: Integrate PostHog into any python application.
 Home-page: https://github.com/posthog/posthog-python
 Author: Posthog

{posthoganalytics-6.7.2.dist-info → posthoganalytics-6.7.3.dist-info}/RECORD

@@ -11,25 +11,25 @@ posthoganalytics/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 posthoganalytics/request.py,sha256=Bsl2c5WwONKPQzwWMmKPX5VgOlwSiIcSNfhXgoz62Y8,6186
 posthoganalytics/types.py,sha256=Dl3aFGX9XUR0wMmK12r2s5Hjan9jL4HpQ9GHpVcEq5U,10207
 posthoganalytics/utils.py,sha256=-0w-OLcCaoldkbBebPzQyBzLJSo9G9yBOg8NDVz7La8,16088
-posthoganalytics/version.py,sha256=5_SqKJ01JbRPG9x4t8JTHffIB0KktqfXyvL6EK2L4Vg,87
+posthoganalytics/version.py,sha256=pd4U7G30qoIwQl2HZ9MpxMew8kahEVgR-yT4Kp42zcI,87
 posthoganalytics/ai/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 posthoganalytics/ai/sanitization.py,sha256=owipZ4eJYtd4JTI-CM_klatclXaeaIec3XJBOUfsOnQ,5770
-posthoganalytics/ai/types.py,sha256=OsB6u855BdZNl5TyVx6Bxm80fT0dfbfnL9Yr5GsIbOQ,3755
-posthoganalytics/ai/utils.py,sha256=nkqVR8Gpum9cllx8Z-ylhYM7tMy1BENfB58lOlrjIkw,20747
+posthoganalytics/ai/types.py,sha256=ceubs4K9xf8vQx7wokq1NL9hPtxyS7D7sUOuT7Lx1lM,3237
+posthoganalytics/ai/utils.py,sha256=8qhlPZjUzUn3XnFE51TraMAEHLQ3bZ7eG1iuwr9Dre0,20052
 posthoganalytics/ai/anthropic/__init__.py,sha256=8nTvETZzkfW-P3zBMmp06GOHs0N-xyOGu7Oa4di_lno,669
-posthoganalytics/ai/anthropic/anthropic.py,sha256=YA-oZyqnWX1cxZvzMRBq1qwO0_r46Q1NYj2RBNuNAYI,8812
-posthoganalytics/ai/anthropic/anthropic_async.py,sha256=M-LIWqkMRF2LSjIwYU_l9D0S1ZjWh4HOoPPBRuKrY4o,10161
-posthoganalytics/ai/anthropic/anthropic_converter.py,sha256=U310RnrXN_JCjJIabr_zJKgNq14cs_atEI4ENsN-LK0,11415
+posthoganalytics/ai/anthropic/anthropic.py,sha256=njOoVb9vkCdnPWAQuVF0XB0BnT2y1ScIryrCGyt5ur8,8750
+posthoganalytics/ai/anthropic/anthropic_async.py,sha256=nM3oFcNLw6meEtV6RfrvhFcuxD4aS-CXDuepRHycUjM,10169
+posthoganalytics/ai/anthropic/anthropic_converter.py,sha256=LWIQ1kyK3vV3rLBmQIcd-98fet7isK3uhTRmBqBN0lk,11776
 posthoganalytics/ai/anthropic/anthropic_providers.py,sha256=y1_qc8Lbip-YDmpimPGg3DfTm5g-WZk5FrRCXzwF_Ow,2139
 posthoganalytics/ai/gemini/__init__.py,sha256=JV_9-gBR87leHgZW4XAYZP7LSl4YaXeuhqDUpA8HygA,383
-posthoganalytics/ai/gemini/gemini.py,sha256=V_ZHKYIJuRzUIQ-BKX8DzFp4eyXdZbzWUg8WHyPZfOw,14960
-posthoganalytics/ai/gemini/gemini_converter.py,sha256=nKiwtHIKFw3g2KVR47GqFzNMxHUnSF4kqi8BYJIJPPQ,13495
+posthoganalytics/ai/gemini/gemini.py,sha256=mHW20mhFVmEwJfxz2HnjxCrEvvB09Jf7goWEmMNwPgw,14891
+posthoganalytics/ai/gemini/gemini_converter.py,sha256=ylV-6JjzKB47liz2UyP5udRB9vtSS4eEe4rr-LpCtyo,14154
 posthoganalytics/ai/langchain/__init__.py,sha256=9CqAwLynTGj3ASAR80C3PmdTdrYGmu99tz0JL-HPFgI,70
 posthoganalytics/ai/langchain/callbacks.py,sha256=Otha0a6YLBwETfKjDDbdLzNi-RHRgKFJB69GwWCv9lg,29527
 posthoganalytics/ai/openai/__init__.py,sha256=u4OuUT7k1NgFj0TrxjuyegOg7a_UA8nAU6a-Hszr0OM,490
-posthoganalytics/ai/openai/openai.py,sha256=OQxey512e0PfPEDpy5nFQ4fkK2N-aw9kn2YGJaJFHqU,20330
-posthoganalytics/ai/openai/openai_async.py,sha256=jT_C-DgI5eUMhJDtA5NtTJy07vQmIkNSYk9bAO12ZPY,21732
-posthoganalytics/ai/openai/openai_converter.py,sha256=G-VC3JO8yudOduXVoykoL44KI9o4h1CxH13w2uDlRSw,19552
+posthoganalytics/ai/openai/openai.py,sha256=I05NruE9grWezM_EgOZBiG5Ej_gABsDcYKN0pRQWvzU,20235
+posthoganalytics/ai/openai/openai_async.py,sha256=k6bo3LfJ_CAPBZCxAzyM2uLz4BpW2YWEFhNuzVcpJlM,21811
+posthoganalytics/ai/openai/openai_converter.py,sha256=0H_ouclo97y19v9QFn-QC2qnj8HdE0lS4wTJljHcmkM,20484
 posthoganalytics/ai/openai/openai_providers.py,sha256=RPVmj2V0_lAdno_ax5Ul2kwhBA9_rRgAdl_sCqrQc6M,4004
 posthoganalytics/integrations/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 posthoganalytics/integrations/django.py,sha256=KYtBr7CkiZQynRc2TCWWYHe-J3ie8iSUa42WPshYZdc,6795
@@ -47,8 +47,8 @@ posthoganalytics/test/test_request.py,sha256=Zc0VbkjpVmj8mKokQm9rzdgTr0b1U44vvMY
 posthoganalytics/test/test_size_limited_dict.py,sha256=-5IQjIEr_-Dql24M0HusdR_XroOMrtgiT0v6ZQCRvzo,774
 posthoganalytics/test/test_types.py,sha256=bRPHdwVpP7hu7emsplU8UVyzSQptv6PaG5lAoOD_BtM,7595
 posthoganalytics/test/test_utils.py,sha256=sqUTbfweVcxxFRd3WDMFXqPMyU6DvzOBeAOc68Py9aw,9620
-posthoganalytics-6.7.2.dist-info/licenses/LICENSE,sha256=wGf9JBotDkSygFj43m49oiKlFnpMnn97keiZKF-40vE,2450
-posthoganalytics-6.7.2.dist-info/METADATA,sha256=nGtssI8-JFCLJJgoGrwrX7KbGjyaRLigawFm2l7_VSY,6024
-posthoganalytics-6.7.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-posthoganalytics-6.7.2.dist-info/top_level.txt,sha256=8QsNIqIkBh1p2TXvKp0Em9ZLZKwe3uIqCETyW4s1GOE,17
-posthoganalytics-6.7.2.dist-info/RECORD,,
+posthoganalytics-6.7.3.dist-info/licenses/LICENSE,sha256=wGf9JBotDkSygFj43m49oiKlFnpMnn97keiZKF-40vE,2450
+posthoganalytics-6.7.3.dist-info/METADATA,sha256=WJzq5FqyDiFj8cFU0tkAGZeLKF5-OOLVlH1nuBPGmLs,6024
+posthoganalytics-6.7.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+posthoganalytics-6.7.3.dist-info/top_level.txt,sha256=8QsNIqIkBh1p2TXvKp0Em9ZLZKwe3uIqCETyW4s1GOE,17
+posthoganalytics-6.7.3.dist-info/RECORD,,