posthoganalytics 6.7.5__py3-none-any.whl → 7.4.3__py3-none-any.whl

posthoganalytics/ai/langchain/callbacks.py

@@ -1,8 +1,8 @@
 try:
-    import langchain  # noqa: F401
+    import langchain_core  # noqa: F401
 except ImportError:
     raise ModuleNotFoundError(
-        "Please install LangChain to use this feature: 'pip install langchain'"
+        "Please install LangChain to use this feature: 'pip install langchain-core'"
     )
 
 import json
@@ -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,
@@ -73,6 +79,8 @@ class GenerationMetadata(SpanMetadata):
     """Base URL of the provider's API used in the run."""
     tools: Optional[List[Dict[str, Any]]] = None
     """Tools provided to the model."""
+    posthog_properties: Optional[Dict[str, Any]] = None
+    """PostHog properties of the run."""
 
 
 RunMetadata = Union[SpanMetadata, GenerationMetadata]
@@ -414,6 +422,8 @@ class CallbackHandler(BaseCallbackHandler):
                 generation.model = model
             if provider := metadata.get("ls_provider"):
                 generation.provider = provider
+
+            generation.posthog_properties = metadata.get("posthog_properties")
         try:
             base_url = serialized["kwargs"]["openai_api_base"]
             if base_url is not None:
@@ -486,6 +496,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,8 +567,12 @@ class CallbackHandler(BaseCallbackHandler):
             "$ai_http_status": 200,
             "$ai_latency": run.latency,
             "$ai_base_url": run.base_url,
+            "$ai_framework": "langchain",
         }
 
+        if isinstance(run.posthog_properties, dict):
+            event_properties.update(run.posthog_properties)
+
         if run.tools:
             event_properties["$ai_tools"] = run.tools
 
@@ -567,7 +582,7 @@ class CallbackHandler(BaseCallbackHandler):
             event_properties["$ai_is_error"] = True
         else:
             # Add usage
-            usage = _parse_usage(output)
+            usage = _parse_usage(output, run.provider, run.model)
             event_properties["$ai_input_tokens"] = usage.input_tokens
             event_properties["$ai_output_tokens"] = usage.output_tokens
             event_properties["$ai_cache_creation_input_tokens"] = (
@@ -688,6 +703,8 @@ class ModelUsage:
 
 def _parse_usage_model(
     usage: Union[BaseModel, dict],
+    provider: Optional[str] = None,
+    model: Optional[str] = None,
 ) -> ModelUsage:
     if isinstance(usage, BaseModel):
         usage = usage.__dict__
@@ -750,15 +767,38 @@ 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()
         },
     )
+    # For Anthropic providers, LangChain reports input_tokens as the sum of all input tokens.
+    # Our cost calculation expects them to be separate for Anthropic, so we subtract cache tokens.
+    # Both cache_read and cache_write tokens should be subtracted since Anthropic's raw API
+    # reports input_tokens as tokens NOT read from or used to create a cache.
+    # For other providers (OpenAI, etc.), input_tokens already excludes cache tokens as expected.
+    # Match logic consistent with plugin-server: exact match on provider OR substring match on model
+    is_anthropic = False
+    if provider and provider.lower() == "anthropic":
+        is_anthropic = True
+    elif model and "anthropic" in model.lower():
+        is_anthropic = True
+
+    if is_anthropic and normalized_usage.input_tokens:
+        cache_tokens = (normalized_usage.cache_read_tokens or 0) + (
+            normalized_usage.cache_write_tokens or 0
+        )
+        if cache_tokens > 0:
+            normalized_usage.input_tokens = max(
+                normalized_usage.input_tokens - cache_tokens, 0
+            )
+    return normalized_usage
 
 
-def _parse_usage(response: LLMResult) -> ModelUsage:
+def _parse_usage(
+    response: LLMResult, provider: Optional[str] = None, model: Optional[str] = None
+) -> ModelUsage:
     # langchain-anthropic uses the usage field
     llm_usage_keys = ["token_usage", "usage"]
     llm_usage: ModelUsage = ModelUsage(
@@ -772,13 +812,15 @@ def _parse_usage(response: LLMResult) -> ModelUsage:
     if response.llm_output is not None:
         for key in llm_usage_keys:
             if response.llm_output.get(key):
-                llm_usage = _parse_usage_model(response.llm_output[key])
+                llm_usage = _parse_usage_model(
+                    response.llm_output[key], provider, model
+                )
                 break
 
     if hasattr(response, "generations"):
         for generation in response.generations:
             if "usage" in generation:
-                llm_usage = _parse_usage_model(generation["usage"])
+                llm_usage = _parse_usage_model(generation["usage"], provider, model)
                 break
 
             for generation_chunk in generation:
@@ -786,7 +828,9 @@ def _parse_usage(response: LLMResult) -> ModelUsage:
                     "usage_metadata" in generation_chunk.generation_info
                 ):
                     llm_usage = _parse_usage_model(
-                        generation_chunk.generation_info["usage_metadata"]
+                        generation_chunk.generation_info["usage_metadata"],
+                        provider,
+                        model,
                     )
                     break
 
@@ -813,7 +857,7 @@ def _parse_usage(response: LLMResult) -> ModelUsage:
                     bedrock_anthropic_usage or bedrock_titan_usage or ollama_usage
                 )
                 if chunk_usage:
-                    llm_usage = _parse_usage_model(chunk_usage)
+                    llm_usage = _parse_usage_model(chunk_usage, provider, model)
                     break
 
     return llm_usage

posthoganalytics/ai/openai/openai.py

@@ -124,14 +124,23 @@ class WrappedResponses:
         start_time = time.time()
         usage_stats: TokenUsage = TokenUsage()
         final_content = []
+        model_from_response: Optional[str] = None
         response = self._original.create(**kwargs)
 
         def generator():
             nonlocal usage_stats
             nonlocal final_content  # noqa: F824
+            nonlocal model_from_response
 
             try:
                 for chunk in response:
+                    # Extract model from response object in chunk (for stored prompts)
+                    if hasattr(chunk, "response") and chunk.response:
+                        if model_from_response is None and hasattr(
+                            chunk.response, "model"
+                        ):
+                            model_from_response = chunk.response.model
+
                     # Extract usage stats from chunk
                     chunk_usage = extract_openai_usage_from_chunk(chunk, "responses")
 
@@ -161,6 +170,7 @@ class WrappedResponses:
                     latency,
                     output,
                     None,  # Responses API doesn't have tools
+                    model_from_response,
                 )
 
         return generator()
@@ -177,6 +187,7 @@ class WrappedResponses:
         latency: float,
         output: Any,
         available_tool_calls: Optional[List[Dict[str, Any]]] = None,
+        model_from_response: Optional[str] = None,
     ):
         from posthoganalytics.ai.types import StreamingEventData
         from posthoganalytics.ai.openai.openai_converter import (
@@ -189,9 +200,12 @@ class WrappedResponses:
         formatted_input = format_openai_streaming_input(kwargs, "responses")
         sanitized_input = sanitize_openai_response(formatted_input)
 
+        # Use model from kwargs, fallback to model from response
+        model = kwargs.get("model") or model_from_response or "unknown"
+
         event_data = StreamingEventData(
             provider="openai",
-            model=kwargs.get("model", "unknown"),
+            model=model,
             base_url=str(self._client.base_url),
             kwargs=kwargs,
             formatted_input=sanitized_input,
@@ -320,6 +334,7 @@ class WrappedCompletions:
         usage_stats: TokenUsage = TokenUsage()
         accumulated_content = []
         accumulated_tool_calls: Dict[int, Dict[str, Any]] = {}
+        model_from_response: Optional[str] = None
         if "stream_options" not in kwargs:
             kwargs["stream_options"] = {}
         kwargs["stream_options"]["include_usage"] = True
@@ -329,9 +344,14 @@ class WrappedCompletions:
             nonlocal usage_stats
             nonlocal accumulated_content  # noqa: F824
             nonlocal accumulated_tool_calls
+            nonlocal model_from_response
 
             try:
                 for chunk in response:
+                    # Extract model from chunk (Chat Completions chunks have model field)
+                    if model_from_response is None and hasattr(chunk, "model"):
+                        model_from_response = chunk.model
+
                     # Extract usage stats from chunk
                     chunk_usage = extract_openai_usage_from_chunk(chunk, "chat")
 
@@ -376,6 +396,7 @@ class WrappedCompletions:
                     accumulated_content,
                     tool_calls_list,
                     extract_available_tool_calls("openai", kwargs),
+                    model_from_response,
                 )
 
         return generator()
@@ -393,6 +414,7 @@ class WrappedCompletions:
         output: Any,
         tool_calls: Optional[List[Dict[str, Any]]] = None,
         available_tool_calls: Optional[List[Dict[str, Any]]] = None,
+        model_from_response: Optional[str] = None,
     ):
         from posthoganalytics.ai.types import StreamingEventData
         from posthoganalytics.ai.openai.openai_converter import (
@@ -405,9 +427,12 @@ class WrappedCompletions:
         formatted_input = format_openai_streaming_input(kwargs, "chat")
         sanitized_input = sanitize_openai(formatted_input)
 
+        # Use model from kwargs, fallback to model from response
+        model = kwargs.get("model") or model_from_response or "unknown"
+
         event_data = StreamingEventData(
             provider="openai",
-            model=kwargs.get("model", "unknown"),
+            model=model,
             base_url=str(self._client.base_url),
             kwargs=kwargs,
             formatted_input=sanitized_input,

posthoganalytics/ai/openai/openai_async.py

@@ -128,14 +128,23 @@ class WrappedResponses:
         start_time = time.time()
         usage_stats: TokenUsage = TokenUsage()
         final_content = []
-        response = self._original.create(**kwargs)
+        model_from_response: Optional[str] = None
+        response = await self._original.create(**kwargs)
 
         async def async_generator():
             nonlocal usage_stats
             nonlocal final_content  # noqa: F824
+            nonlocal model_from_response
 
             try:
                 async for chunk in response:
+                    # Extract model from response object in chunk (for stored prompts)
+                    if hasattr(chunk, "response") and chunk.response:
+                        if model_from_response is None and hasattr(
+                            chunk.response, "model"
+                        ):
+                            model_from_response = chunk.response.model
+
                     # Extract usage stats from chunk
                     chunk_usage = extract_openai_usage_from_chunk(chunk, "responses")
 
@@ -166,6 +175,7 @@ class WrappedResponses:
                     latency,
                     output,
                     extract_available_tool_calls("openai", kwargs),
+                    model_from_response,
                 )
 
         return async_generator()
@@ -182,13 +192,17 @@ class WrappedResponses:
         latency: float,
         output: Any,
         available_tool_calls: Optional[List[Dict[str, Any]]] = None,
+        model_from_response: Optional[str] = None,
     ):
         if posthog_trace_id is None:
             posthog_trace_id = str(uuid.uuid4())
 
+        # Use model from kwargs, fallback to model from response
+        model = kwargs.get("model") or model_from_response or "unknown"
+
         event_properties = {
             "$ai_provider": "openai",
-            "$ai_model": kwargs.get("model"),
+            "$ai_model": model,
             "$ai_model_parameters": get_model_params(kwargs),
             "$ai_input": with_privacy_mode(
                 self._client._ph_client,
@@ -213,6 +227,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
 
@@ -341,19 +364,25 @@ class WrappedCompletions:
         usage_stats: TokenUsage = TokenUsage()
         accumulated_content = []
         accumulated_tool_calls: Dict[int, Dict[str, Any]] = {}
+        model_from_response: Optional[str] = None
 
         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
             nonlocal accumulated_content  # noqa: F824
             nonlocal accumulated_tool_calls
+            nonlocal model_from_response
 
             try:
                 async for chunk in response:
+                    # Extract model from chunk (Chat Completions chunks have model field)
+                    if model_from_response is None and hasattr(chunk, "model"):
+                        model_from_response = chunk.model
+
                     # Extract usage stats from chunk
                     chunk_usage = extract_openai_usage_from_chunk(chunk, "chat")
                     if chunk_usage:
@@ -396,6 +425,7 @@ class WrappedCompletions:
                     accumulated_content,
                     tool_calls_list,
                     extract_available_tool_calls("openai", kwargs),
+                    model_from_response,
                 )
 
         return async_generator()
@@ -413,13 +443,17 @@ class WrappedCompletions:
         output: Any,
         tool_calls: Optional[List[Dict[str, Any]]] = None,
         available_tool_calls: Optional[List[Dict[str, Any]]] = None,
+        model_from_response: Optional[str] = None,
     ):
         if posthog_trace_id is None:
             posthog_trace_id = str(uuid.uuid4())
 
+        # Use model from kwargs, fallback to model from response
+        model = kwargs.get("model") or model_from_response or "unknown"
+
         event_properties = {
             "$ai_provider": "openai",
-            "$ai_model": kwargs.get("model"),
+            "$ai_model": model,
             "$ai_model_parameters": get_model_params(kwargs),
             "$ai_input": with_privacy_mode(
                 self._client._ph_client,
@@ -444,6 +478,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
 
@@ -499,7 +543,7 @@ 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

posthoganalytics/ai/openai/openai_converter.py

@@ -67,6 +67,12 @@ def format_openai_response(response: Any) -> List[FormattedMessage]:
                             }
                         )
 
+                # Handle audio output (gpt-4o-audio-preview)
+                if hasattr(choice.message, "audio") and choice.message.audio:
+                    # Convert Pydantic model to dict to capture all fields from OpenAI
+                    audio_dict = choice.message.audio.model_dump()
+                    content.append({"type": "audio", **audio_dict})
+
         if content:
             output.append(
                 {
@@ -255,6 +261,113 @@ 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).
@@ -312,6 +425,10 @@ def extract_openai_usage_from_response(response: Any) -> TokenUsage:
     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
 
 
@@ -334,6 +451,13 @@ def extract_openai_usage_from_chunk(
     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
 
@@ -386,6 +510,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
 
 

posthoganalytics/ai/sanitization.py

@@ -1,3 +1,4 @@
+import os
 import re
 from typing import Any
 from urllib.parse import urlparse
@@ -5,6 +6,15 @@ from urllib.parse import urlparse
 REDACTED_IMAGE_PLACEHOLDER = "[base64 image redacted]"
 
 
+def _is_multimodal_enabled() -> bool:
+    """Check if multimodal capture is enabled via environment variable."""
+    return os.environ.get("_INTERNAL_LLMA_MULTIMODAL", "").lower() in (
+        "true",
+        "1",
+        "yes",
+    )
+
+
 def is_base64_data_url(text: str) -> bool:
     return re.match(r"^data:([^;]+);base64,", text) is not None
 
@@ -27,6 +37,9 @@ def is_raw_base64(text: str) -> bool:
 
 
 def redact_base64_data_url(value: Any) -> Any:
+    if _is_multimodal_enabled():
+        return value
+
     if not isinstance(value, str):
         return value
 
@@ -83,6 +96,11 @@ def sanitize_openai_image(item: Any) -> Any:
             },
         }
 
+    if item.get("type") == "audio" and "data" in item:
+        if _is_multimodal_enabled():
+            return item
+        return {**item, "data": REDACTED_IMAGE_PLACEHOLDER}
+
     return item
 
 
@@ -100,6 +118,9 @@ def sanitize_openai_response_image(item: Any) -> Any:
 
 
 def sanitize_anthropic_image(item: Any) -> Any:
+    if _is_multimodal_enabled():
+        return item
+
     if not isinstance(item, dict):
         return item
 
@@ -109,8 +130,6 @@ def sanitize_anthropic_image(item: Any) -> Any:
         and item["source"].get("type") == "base64"
         and "data" in item["source"]
     ):
-        # For Anthropic, if the source type is "base64", we should always redact the data
-        # The provider is explicitly telling us this is base64 data
         return {
             **item,
             "source": {
@@ -123,6 +142,9 @@ def sanitize_anthropic_image(item: Any) -> Any:
 
 
 def sanitize_gemini_part(part: Any) -> Any:
+    if _is_multimodal_enabled():
+        return part
+
     if not isinstance(part, dict):
         return part
 
@@ -131,8 +153,6 @@ def sanitize_gemini_part(part: Any) -> Any:
         and isinstance(part["inline_data"], dict)
         and "data" in part["inline_data"]
     ):
-        # For Gemini, the inline_data structure indicates base64 data
-        # We should redact any string data in this context
         return {
             **part,
             "inline_data": {
@@ -185,7 +205,9 @@ def sanitize_langchain_image(item: Any) -> Any:
         and isinstance(item.get("source"), dict)
         and "data" in item["source"]
     ):
-        # Anthropic style - raw base64 in structured format, always redact
+        if _is_multimodal_enabled():
+            return item
+
         return {
             **item,
             "source": {

posthoganalytics/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):