posthog 6.7.2__py3-none-any.whl → 6.9.0__py3-none-any.whl

posthog/ai/langchain/callbacks.py

@@ -20,8 +20,14 @@ from typing import (
 )
 from uuid import UUID
 
-from langchain.callbacks.base import BaseCallbackHandler
-from langchain.schema.agent import AgentAction, AgentFinish
+try:
+    # LangChain 1.0+ and modern 0.x with langchain-core
+    from langchain_core.callbacks.base import BaseCallbackHandler
+    from langchain_core.agents import AgentAction, AgentFinish
+except (ImportError, ModuleNotFoundError):
+    # Fallback for older LangChain versions
+    from langchain.callbacks.base import BaseCallbackHandler
+    from langchain.schema.agent import AgentAction, AgentFinish
 from langchain_core.documents import Document
 from langchain_core.messages import (
     AIMessage,
@@ -486,6 +492,7 @@ class CallbackHandler(BaseCallbackHandler):
             "$ai_latency": run.latency,
             "$ai_span_name": run.name,
             "$ai_span_id": run_id,
+            "$ai_framework": "langchain",
         }
         if parent_run_id is not None:
             event_properties["$ai_parent_id"] = parent_run_id
@@ -556,6 +563,7 @@ class CallbackHandler(BaseCallbackHandler):
             "$ai_http_status": 200,
             "$ai_latency": run.latency,
             "$ai_base_url": run.base_url,
+            "$ai_framework": "langchain",
         }
 
         if run.tools:
@@ -750,12 +758,19 @@ def _parse_usage_model(
         "cache_read": "cache_read_tokens",
         "reasoning": "reasoning_tokens",
     }
-    return ModelUsage(
+    normalized_usage = ModelUsage(
         **{
             dataclass_key: parsed_usage.get(mapped_key) or 0
             for mapped_key, dataclass_key in field_mapping.items()
         },
     )
+    # In LangChain, input_tokens is the sum of input and cache read tokens.
+    # Our cost calculation expects them to be separate, for Anthropic.
+    if normalized_usage.input_tokens and normalized_usage.cache_read_tokens:
+        normalized_usage.input_tokens = max(
+            normalized_usage.input_tokens - normalized_usage.cache_read_tokens, 0
+        )
+    return normalized_usage
 
 
 def _parse_usage(response: LLMResult) -> ModelUsage:

posthog/ai/openai/openai.py

@@ -2,6 +2,8 @@ import time
 import uuid
 from typing import Any, Dict, List, Optional
 
+from posthog.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 posthog.ai.types import StreamingEventData
         from posthog.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 posthog.ai.types import StreamingEventData
         from posthog.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,

posthog/ai/openai/openai_async.py

@@ -2,6 +2,8 @@ import time
 import uuid
 from typing import Any, Dict, List, Optional
 
+from posthog.ai.types import TokenUsage
+
 try:
     import openai
 except ImportError:
@@ -124,9 +126,9 @@ class WrappedResponses:
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {}
+        usage_stats: TokenUsage = TokenUsage()
         final_content = []
-        response = self._original.create(**kwargs)
+        response = await self._original.create(**kwargs)
 
         async def async_generator():
             nonlocal usage_stats
@@ -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,
@@ -211,6 +213,15 @@ class WrappedResponses:
             **(posthog_properties or {}),
         }
 
+        # Add web search count if present
+        web_search_count = usage_stats.get("web_search_count")
+        if (
+            web_search_count is not None
+            and isinstance(web_search_count, int)
+            and web_search_count > 0
+        ):
+            event_properties["$ai_web_search_count"] = web_search_count
+
         if available_tool_calls:
             event_properties["$ai_tools"] = available_tool_calls
 
@@ -336,14 +347,14 @@ 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:
             kwargs["stream_options"] = {}
         kwargs["stream_options"]["include_usage"] = True
-        response = self._original.create(**kwargs)
+        response = await self._original.create(**kwargs)
 
         async def async_generator():
             nonlocal usage_stats
@@ -406,7 +417,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 +441,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
             ),
@@ -442,6 +453,16 @@ class WrappedCompletions:
             **(posthog_properties or {}),
         }
 
+        # Add web search count if present
+        web_search_count = usage_stats.get("web_search_count")
+
+        if (
+            web_search_count is not None
+            and isinstance(web_search_count, int)
+            and web_search_count > 0
+        ):
+            event_properties["$ai_web_search_count"] = web_search_count
+
         if available_tool_calls:
             event_properties["$ai_tools"] = available_tool_calls
 
@@ -497,17 +518,17 @@ class WrappedEmbeddings:
             posthog_trace_id = str(uuid.uuid4())
 
         start_time = time.time()
-        response = self._original.create(**kwargs)
+        response = await self._original.create(**kwargs)
         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 +542,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),

posthog/ai/openai/openai_converter.py

@@ -14,7 +14,6 @@ from posthog.ai.types import (
     FormattedImageContent,
     FormattedMessage,
     FormattedTextContent,
-    StreamingUsageStats,
     TokenUsage,
 )
 
@@ -256,9 +255,180 @@ def format_openai_streaming_content(
     return formatted
 
 
+def extract_openai_web_search_count(response: Any) -> int:
+    """
+    Extract web search count from OpenAI response.
+
+    Uses a two-tier detection strategy:
+    1. Priority 1 (exact count): Check for output[].type == "web_search_call" (Responses API)
+    2. Priority 2 (binary detection): Check for various web search indicators:
+       - Root-level citations, search_results, or usage.search_context_size (Perplexity)
+       - Annotations with type "url_citation" in choices/output (including delta for streaming)
+
+    Args:
+        response: The response from OpenAI API
+
+    Returns:
+        Number of web search requests (exact count or binary 1/0)
+    """
+
+    # Priority 1: Check for exact count in Responses API output
+    if hasattr(response, "output"):
+        web_search_count = 0
+
+        for item in response.output:
+            if hasattr(item, "type") and item.type == "web_search_call":
+                web_search_count += 1
+
+        web_search_count = max(0, web_search_count)
+
+        if web_search_count > 0:
+            return web_search_count
+
+    # Priority 2: Binary detection (returns 1 or 0)
+
+    # Check root-level indicators (Perplexity)
+    if hasattr(response, "citations"):
+        citations = getattr(response, "citations")
+
+        if citations and len(citations) > 0:
+            return 1
+
+    if hasattr(response, "search_results"):
+        search_results = getattr(response, "search_results")
+
+        if search_results and len(search_results) > 0:
+            return 1
+
+    if hasattr(response, "usage") and hasattr(response.usage, "search_context_size"):
+        if response.usage.search_context_size:
+            return 1
+
+    # Check for url_citation annotations in choices (Chat Completions)
+    if hasattr(response, "choices"):
+        for choice in response.choices:
+            # Check message.annotations (non-streaming or final chunk)
+            if hasattr(choice, "message") and hasattr(choice.message, "annotations"):
+                annotations = choice.message.annotations
+
+                if annotations:
+                    for annotation in annotations:
+                        # Support both dict and object formats
+                        annotation_type = (
+                            annotation.get("type")
+                            if isinstance(annotation, dict)
+                            else getattr(annotation, "type", None)
+                        )
+
+                        if annotation_type == "url_citation":
+                            return 1
+
+            # Check delta.annotations (streaming chunks)
+            if hasattr(choice, "delta") and hasattr(choice.delta, "annotations"):
+                annotations = choice.delta.annotations
+
+                if annotations:
+                    for annotation in annotations:
+                        # Support both dict and object formats
+                        annotation_type = (
+                            annotation.get("type")
+                            if isinstance(annotation, dict)
+                            else getattr(annotation, "type", None)
+                        )
+
+                        if annotation_type == "url_citation":
+                            return 1
+
+    # Check for url_citation annotations in output (Responses API)
+    if hasattr(response, "output"):
+        for item in response.output:
+            if hasattr(item, "content") and isinstance(item.content, list):
+                for content_item in item.content:
+                    if hasattr(content_item, "annotations"):
+                        annotations = content_item.annotations
+
+                        if annotations:
+                            for annotation in annotations:
+                                # Support both dict and object formats
+                                annotation_type = (
+                                    annotation.get("type")
+                                    if isinstance(annotation, dict)
+                                    else getattr(annotation, "type", None)
+                                )
+
+                                if annotation_type == "url_citation":
+                                    return 1
+
+    return 0
+
+
+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
+
+    web_search_count = extract_openai_web_search_count(response)
+    if web_search_count > 0:
+        result["web_search_count"] = web_search_count
+
+    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 +442,23 @@ def extract_openai_usage_from_chunk(
         Dictionary of usage statistics
     """
 
-    usage: StreamingUsageStats = {}
+    usage: TokenUsage = TokenUsage()
 
     if provider_type == "chat":
+        # Extract web search count from the chunk before checking for usage
+        # Web search indicators (citations, annotations) can appear on any chunk,
+        # not just those with usage data
+        web_search_count = extract_openai_web_search_count(chunk)
+        if web_search_count > 0:
+            usage["web_search_count"] = web_search_count
+
         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 +487,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(
@@ -328,6 +504,12 @@ def extract_openai_usage_from_chunk(
                         response_usage.output_tokens_details.reasoning_tokens
                     )
 
+                # Extract web search count from the complete response
+                if hasattr(chunk, "response"):
+                    web_search_count = extract_openai_web_search_count(chunk.response)
+                    if web_search_count > 0:
+                        usage["web_search_count"] = web_search_count
+
     return usage
 
 
@@ -535,37 +717,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:
@@ -579,7 +730,6 @@ def format_openai_streaming_input(
     Returns:
         Formatted input ready for PostHog tracking
     """
-    if api_type == "chat":
-        return kwargs.get("messages")
-    else:  # responses API
-        return kwargs.get("input")
+    from posthog.ai.utils import merge_system_prompt
+
+    return merge_system_prompt(kwargs, "openai")

posthog/ai/types.py

@@ -63,6 +63,7 @@ class TokenUsage(TypedDict, total=False):
     cache_read_input_tokens: Optional[int]
     cache_creation_input_tokens: Optional[int]
     reasoning_tokens: Optional[int]
+    web_search_count: Optional[int]
 
 
 class ProviderResponse(TypedDict, total=False):
@@ -77,24 +78,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 +116,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]