posthoganalytics 6.7.14__py3-none-any.whl → 6.9.0__py3-none-any.whl

posthoganalytics/__init__.py

@@ -10,6 +10,9 @@ from posthoganalytics.contexts import (
     tag as inner_tag,
     set_context_session as inner_set_context_session,
     identify_context as inner_identify_context,
+    set_capture_exception_code_variables_context as inner_set_capture_exception_code_variables_context,
+    set_code_variables_mask_patterns_context as inner_set_code_variables_mask_patterns_context,
+    set_code_variables_ignore_patterns_context as inner_set_code_variables_ignore_patterns_context,
 )
 from posthoganalytics.feature_flags import InconclusiveMatchError, RequiresServerEvaluation
 from posthoganalytics.types import FeatureFlag, FlagsAndPayloads, FeatureFlagResult
@@ -20,13 +23,14 @@ __version__ = VERSION
 """Context management."""
 
 
-def new_context(fresh=False, capture_exceptions=True):
+def new_context(fresh=False, capture_exceptions=True, client=None):
     """
     Create a new context scope that will be active for the duration of the with block.
 
     Args:
         fresh: Whether to start with a fresh context (default: False)
         capture_exceptions: Whether to capture exceptions raised within the context (default: True)
+        client: Optional Posthog client instance to use for this context (default: None)
 
     Examples:
         ```python
@@ -39,7 +43,9 @@ def new_context(fresh=False, capture_exceptions=True):
     Category:
         Contexts
     """
-    return inner_new_context(fresh=fresh, capture_exceptions=capture_exceptions)
+    return inner_new_context(
+        fresh=fresh, capture_exceptions=capture_exceptions, client=client
+    )
 
 
 def scoped(fresh=False, capture_exceptions=True):
@@ -103,6 +109,27 @@ def identify_context(distinct_id: str):
     return inner_identify_context(distinct_id)
 
 
+def set_capture_exception_code_variables_context(enabled: bool):
+    """
+    Set whether code variables are captured for the current context.
+    """
+    return inner_set_capture_exception_code_variables_context(enabled)
+
+
+def set_code_variables_mask_patterns_context(mask_patterns: list):
+    """
+    Variable names matching these patterns will be masked with *** when capturing code variables.
+    """
+    return inner_set_code_variables_mask_patterns_context(mask_patterns)
+
+
+def set_code_variables_ignore_patterns_context(ignore_patterns: list):
+    """
+    Variable names matching these patterns will be ignored completely when capturing code variables.
+    """
+    return inner_set_code_variables_ignore_patterns_context(ignore_patterns)
+
+
 def tag(name: str, value: Any):
     """
     Add a tag to the current context.

posthoganalytics/ai/anthropic/anthropic_async.py

@@ -14,14 +14,9 @@ from posthoganalytics import setup
 from posthoganalytics.ai.types import StreamingContentBlock, TokenUsage, ToolInProgress
 from posthoganalytics.ai.utils import (
     call_llm_and_track_usage_async,
-    extract_available_tool_calls,
-    get_model_params,
-    merge_system_prompt,
     merge_usage_stats,
-    with_privacy_mode,
 )
 from posthoganalytics.ai.anthropic.anthropic_converter import (
-    format_anthropic_streaming_content,
     extract_anthropic_usage_from_event,
     handle_anthropic_content_block_start,
     handle_anthropic_text_delta,
@@ -220,66 +215,34 @@ class AsyncWrappedMessages(AsyncMessages):
         content_blocks: List[StreamingContentBlock],
         accumulated_content: str,
     ):
-        if posthog_trace_id is None:
-            posthog_trace_id = str(uuid.uuid4())
-
-        # Format output using converter
-        formatted_content = format_anthropic_streaming_content(content_blocks)
-        formatted_output = []
-
-        if formatted_content:
-            formatted_output = [{"role": "assistant", "content": formatted_content}]
-        else:
-            # Fallback to accumulated content if no blocks
-            formatted_output = [
-                {
-                    "role": "assistant",
-                    "content": [{"type": "text", "text": accumulated_content}],
-                }
-            ]
-
-        event_properties = {
-            "$ai_provider": "anthropic",
-            "$ai_model": kwargs.get("model"),
-            "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": with_privacy_mode(
-                self._client._ph_client,
-                posthog_privacy_mode,
-                sanitize_anthropic(merge_system_prompt(kwargs, "anthropic")),
-            ),
-            "$ai_output_choices": with_privacy_mode(
-                self._client._ph_client,
-                posthog_privacy_mode,
-                formatted_output,
-            ),
-            "$ai_http_status": 200,
-            "$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
-            ),
-            "$ai_cache_creation_input_tokens": usage_stats.get(
-                "cache_creation_input_tokens", 0
+        from posthoganalytics.ai.types import StreamingEventData
+        from posthoganalytics.ai.anthropic.anthropic_converter import (
+            format_anthropic_streaming_input,
+            format_anthropic_streaming_output_complete,
+        )
+        from posthoganalytics.ai.utils import capture_streaming_event
+
+        # Prepare standardized event data
+        formatted_input = format_anthropic_streaming_input(kwargs)
+        sanitized_input = sanitize_anthropic(formatted_input)
+
+        event_data = StreamingEventData(
+            provider="anthropic",
+            model=kwargs.get("model", "unknown"),
+            base_url=str(self._client.base_url),
+            kwargs=kwargs,
+            formatted_input=sanitized_input,
+            formatted_output=format_anthropic_streaming_output_complete(
+                content_blocks, accumulated_content
             ),
-            "$ai_latency": latency,
-            "$ai_trace_id": posthog_trace_id,
-            "$ai_base_url": str(self._client.base_url),
-            **(posthog_properties or {}),
-        }
-
-        # Add tools if available
-        available_tools = extract_available_tool_calls("anthropic", kwargs)
-
-        if available_tools:
-            event_properties["$ai_tools"] = available_tools
-
-        if posthog_distinct_id is None:
-            event_properties["$process_person_profile"] = False
-
-        if hasattr(self._client._ph_client, "capture"):
-            self._client._ph_client.capture(
-                distinct_id=posthog_distinct_id or posthog_trace_id,
-                event="$ai_generation",
-                properties=event_properties,
-                groups=posthog_groups,
-            )
+            usage_stats=usage_stats,
+            latency=latency,
+            distinct_id=posthog_distinct_id,
+            trace_id=posthog_trace_id,
+            properties=posthog_properties,
+            privacy_mode=posthog_privacy_mode,
+            groups=posthog_groups,
+        )
+
+        # Use the common capture function
+        capture_streaming_event(self._client._ph_client, event_data)

posthoganalytics/ai/anthropic/anthropic_converter.py

@@ -163,6 +163,32 @@ def format_anthropic_streaming_content(
     return formatted
 
 
+def extract_anthropic_web_search_count(response: Any) -> int:
+    """
+    Extract web search count from Anthropic response.
+
+    Anthropic provides exact web search counts via usage.server_tool_use.web_search_requests.
+
+    Args:
+        response: The response from Anthropic API
+
+    Returns:
+        Number of web search requests (0 if none)
+    """
+    if not hasattr(response, "usage"):
+        return 0
+
+    if not hasattr(response.usage, "server_tool_use"):
+        return 0
+
+    server_tool_use = response.usage.server_tool_use
+
+    if hasattr(server_tool_use, "web_search_requests"):
+        return max(0, int(getattr(server_tool_use, "web_search_requests", 0)))
+
+    return 0
+
+
 def extract_anthropic_usage_from_response(response: Any) -> TokenUsage:
     """
     Extract usage from a full Anthropic response (non-streaming).
@@ -191,6 +217,10 @@ def extract_anthropic_usage_from_response(response: Any) -> TokenUsage:
         if cache_creation and cache_creation > 0:
             result["cache_creation_input_tokens"] = cache_creation
 
+    web_search_count = extract_anthropic_web_search_count(response)
+    if web_search_count > 0:
+        result["web_search_count"] = web_search_count
+
     return result
 
 
@@ -222,6 +252,16 @@ def extract_anthropic_usage_from_event(event: Any) -> TokenUsage:
     if hasattr(event, "usage") and event.usage:
         usage["output_tokens"] = getattr(event.usage, "output_tokens", 0)
 
+        # Extract web search count from usage
+        if hasattr(event.usage, "server_tool_use"):
+            server_tool_use = event.usage.server_tool_use
+            if hasattr(server_tool_use, "web_search_requests"):
+                web_search_count = int(
+                    getattr(server_tool_use, "web_search_requests", 0)
+                )
+                if web_search_count > 0:
+                    usage["web_search_count"] = web_search_count
+
     return usage
 
 

posthoganalytics/ai/gemini/gemini_converter.py

@@ -338,6 +338,61 @@ def format_gemini_input(contents: Any) -> List[FormattedMessage]:
     return [_format_object_message(contents)]
 
 
+def extract_gemini_web_search_count(response: Any) -> int:
+    """
+    Extract web search count from Gemini response.
+
+    Gemini bills per request that uses grounding, not per query.
+    Returns 1 if grounding_metadata is present with actual search data, 0 otherwise.
+
+    Args:
+        response: The response from Gemini API
+
+    Returns:
+        1 if web search/grounding was used, 0 otherwise
+    """
+
+    # Check for grounding_metadata in candidates
+    if hasattr(response, "candidates"):
+        for candidate in response.candidates:
+            if (
+                hasattr(candidate, "grounding_metadata")
+                and candidate.grounding_metadata
+            ):
+                grounding_metadata = candidate.grounding_metadata
+
+                # Check if web_search_queries exists and is non-empty
+                if hasattr(grounding_metadata, "web_search_queries"):
+                    queries = grounding_metadata.web_search_queries
+
+                    if queries is not None and len(queries) > 0:
+                        return 1
+
+                # Check if grounding_chunks exists and is non-empty
+                if hasattr(grounding_metadata, "grounding_chunks"):
+                    chunks = grounding_metadata.grounding_chunks
+
+                    if chunks is not None and len(chunks) > 0:
+                        return 1
+
+            # Also check for google_search or grounding in function call names
+            if hasattr(candidate, "content") and candidate.content:
+                if hasattr(candidate.content, "parts") and candidate.content.parts:
+                    for part in candidate.content.parts:
+                        if hasattr(part, "function_call") and part.function_call:
+                            function_name = getattr(
+                                part.function_call, "name", ""
+                            ).lower()
+
+                            if (
+                                "google_search" in function_name
+                                or "grounding" in function_name
+                            ):
+                                return 1
+
+    return 0
+
+
 def _extract_usage_from_metadata(metadata: Any) -> TokenUsage:
     """
     Common logic to extract usage from Gemini metadata.
@@ -382,7 +437,14 @@ def extract_gemini_usage_from_response(response: Any) -> TokenUsage:
     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)
+    usage = _extract_usage_from_metadata(response.usage_metadata)
+
+    # Add web search count if present
+    web_search_count = extract_gemini_web_search_count(response)
+    if web_search_count > 0:
+        usage["web_search_count"] = web_search_count
+
+    return usage
 
 
 def extract_gemini_usage_from_chunk(chunk: Any) -> TokenUsage:
@@ -398,11 +460,19 @@ def extract_gemini_usage_from_chunk(chunk: Any) -> TokenUsage:
 
     usage: TokenUsage = TokenUsage()
 
+    # Extract web search count from the chunk before checking for usage_metadata
+    # Web search indicators can appear on any chunk, not just those with usage data
+    web_search_count = extract_gemini_web_search_count(chunk)
+    if web_search_count > 0:
+        usage["web_search_count"] = web_search_count
+
     if not hasattr(chunk, "usage_metadata") or not chunk.usage_metadata:
         return usage
 
-    # Use the shared helper to extract usage
-    usage = _extract_usage_from_metadata(chunk.usage_metadata)
+    usage_from_metadata = _extract_usage_from_metadata(chunk.usage_metadata)
+
+    # Merge the usage from metadata with any web search count we found
+    usage.update(usage_from_metadata)
 
     return usage
 

posthoganalytics/ai/openai/openai_async.py

@@ -213,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
 
@@ -444,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
 

posthoganalytics/ai/openai/openai_converter.py

@@ -255,6 +255,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 +419,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 +445,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 +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
 
 

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

posthoganalytics/ai/utils.py

@@ -53,6 +53,12 @@ def merge_usage_stats(
         if source_reasoning is not None:
             current = target.get("reasoning_tokens") or 0
             target["reasoning_tokens"] = current + source_reasoning
+
+        source_web_search = source.get("web_search_count")
+        if source_web_search is not None:
+            current = target.get("web_search_count") or 0
+            target["web_search_count"] = max(current, source_web_search)
+
     elif mode == "cumulative":
         # Replace with latest values (already cumulative)
         if source.get("input_tokens") is not None:
@@ -67,6 +73,9 @@ def merge_usage_stats(
             ]
         if source.get("reasoning_tokens") is not None:
             target["reasoning_tokens"] = source["reasoning_tokens"]
+        if source.get("web_search_count") is not None:
+            target["web_search_count"] = source["web_search_count"]
+
     else:
         raise ValueError(f"Invalid mode: {mode}. Must be 'incremental' or 'cumulative'")
 
@@ -311,6 +320,10 @@ def call_llm_and_track_usage(
         if reasoning is not None and reasoning > 0:
             event_properties["$ai_reasoning_tokens"] = reasoning
 
+        web_search_count = usage.get("web_search_count")
+        if web_search_count is not None and web_search_count > 0:
+            event_properties["$ai_web_search_count"] = web_search_count
+
         if posthog_distinct_id is None:
             event_properties["$process_person_profile"] = False
 
@@ -414,6 +427,14 @@ async def call_llm_and_track_usage_async(
         if cache_creation is not None and cache_creation > 0:
             event_properties["$ai_cache_creation_input_tokens"] = cache_creation
 
+        reasoning = usage.get("reasoning_tokens")
+        if reasoning is not None and reasoning > 0:
+            event_properties["$ai_reasoning_tokens"] = reasoning
+
+        web_search_count = usage.get("web_search_count")
+        if web_search_count is not None and web_search_count > 0:
+            event_properties["$ai_web_search_count"] = web_search_count
+
         if posthog_distinct_id is None:
             event_properties["$process_person_profile"] = False
 
@@ -535,6 +556,15 @@ def capture_streaming_event(
             if value is not None and isinstance(value, int) and value > 0:
                 event_properties[f"$ai_{field}"] = value
 
+    # Add web search count if present (all providers)
+    web_search_count = event_data["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
+
     # Handle provider-specific fields
     if (
         event_data["provider"] == "openai"

posthoganalytics/client.py

@@ -19,6 +19,9 @@ from posthoganalytics.exception_utils import (
     handle_in_app,
     exception_is_already_captured,
     mark_exception_as_captured,
+    try_attach_code_variables_to_frames,
+    DEFAULT_CODE_VARIABLES_MASK_PATTERNS,
+    DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS,
 )
 from posthoganalytics.feature_flags import (
     InconclusiveMatchError,
@@ -39,6 +42,9 @@ from posthoganalytics.contexts import (
     _get_current_context,
     get_context_distinct_id,
     get_context_session_id,
+    get_capture_exception_code_variables_context,
+    get_code_variables_mask_patterns_context,
+    get_code_variables_ignore_patterns_context,
     new_context,
 )
 from posthoganalytics.types import (
@@ -178,6 +184,9 @@ class Client(object):
         before_send=None,
         flag_fallback_cache_url=None,
         enable_local_evaluation=True,
+        capture_exception_code_variables=False,
+        code_variables_mask_patterns=None,
+        code_variables_ignore_patterns=None,
     ):
         """
         Initialize a new PostHog client instance.
@@ -233,6 +242,18 @@ class Client(object):
         self.privacy_mode = privacy_mode
         self.enable_local_evaluation = enable_local_evaluation
 
+        self.capture_exception_code_variables = capture_exception_code_variables
+        self.code_variables_mask_patterns = (
+            code_variables_mask_patterns
+            if code_variables_mask_patterns is not None
+            else DEFAULT_CODE_VARIABLES_MASK_PATTERNS
+        )
+        self.code_variables_ignore_patterns = (
+            code_variables_ignore_patterns
+            if code_variables_ignore_patterns is not None
+            else DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS
+        )
+
         if project_root is None:
             try:
                 project_root = os.getcwd()
@@ -979,6 +1000,34 @@ class Client(object):
                 **properties,
             }
 
+            context_enabled = get_capture_exception_code_variables_context()
+            context_mask = get_code_variables_mask_patterns_context()
+            context_ignore = get_code_variables_ignore_patterns_context()
+
+            enabled = (
+                context_enabled
+                if context_enabled is not None
+                else self.capture_exception_code_variables
+            )
+            mask_patterns = (
+                context_mask
+                if context_mask is not None
+                else self.code_variables_mask_patterns
+            )
+            ignore_patterns = (
+                context_ignore
+                if context_ignore is not None
+                else self.code_variables_ignore_patterns
+            )
+
+            if enabled:
+                try_attach_code_variables_to_frames(
+                    all_exceptions_with_trace_and_in_app,
+                    exc_info,
+                    mask_patterns=mask_patterns,
+                    ignore_patterns=ignore_patterns,
+                )
+
             if self.log_captured_exceptions:
                 self.log.exception(exception, extra=kwargs)