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

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:
@@ -14,8 +16,17 @@ from posthoganalytics.ai.utils import (
     call_llm_and_track_usage_async,
     extract_available_tool_calls,
     get_model_params,
+    merge_usage_stats,
     with_privacy_mode,
 )
+from posthoganalytics.ai.openai.openai_converter import (
+    extract_openai_usage_from_chunk,
+    extract_openai_content_from_chunk,
+    extract_openai_tool_calls_from_chunk,
+    accumulate_openai_tool_calls,
+    format_openai_streaming_output,
+)
+from posthoganalytics.ai.sanitization import sanitize_openai, sanitize_openai_response
 from posthoganalytics.client import Client as PostHogClient
 
 
@@ -34,6 +45,7 @@ class AsyncOpenAI(openai.AsyncOpenAI):
                             of the global posthog.
             **openai_config: Any additional keyword args to set on openai (e.g. organization="xxx").
         """
+
         super().__init__(**kwargs)
         self._ph_client = posthog_client or setup()
 
@@ -66,6 +78,7 @@ class WrappedResponses:
 
     def __getattr__(self, name):
         """Fallback to original responses object for any methods we don't explicitly handle."""
+
         return getattr(self._original, name)
 
     async def create(
@@ -113,45 +126,36 @@ class WrappedResponses:
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {}
+        usage_stats: TokenUsage = TokenUsage()
         final_content = []
+        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:
-                    if hasattr(chunk, "type") and chunk.type == "response.completed":
-                        res = chunk.response
-                        if res.output and len(res.output) > 0:
-                            final_content.append(res.output[0])
-
-                    if hasattr(chunk, "usage") and chunk.usage:
-                        usage_stats = {
-                            k: getattr(chunk.usage, k, 0)
-                            for k in [
-                                "input_tokens",
-                                "output_tokens",
-                                "total_tokens",
-                            ]
-                        }
-
-                        # Add support for cached tokens
-                        if hasattr(chunk.usage, "output_tokens_details") and hasattr(
-                            chunk.usage.output_tokens_details, "reasoning_tokens"
+                    # 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"
                         ):
-                            usage_stats["reasoning_tokens"] = (
-                                chunk.usage.output_tokens_details.reasoning_tokens
-                            )
+                            model_from_response = chunk.response.model
 
-                        if hasattr(chunk.usage, "input_tokens_details") and hasattr(
-                            chunk.usage.input_tokens_details, "cached_tokens"
-                        ):
-                            usage_stats["cache_read_input_tokens"] = (
-                                chunk.usage.input_tokens_details.cached_tokens
-                            )
+                    # Extract usage stats from chunk
+                    chunk_usage = extract_openai_usage_from_chunk(chunk, "responses")
+
+                    if chunk_usage:
+                        merge_usage_stats(usage_stats, chunk_usage)
+
+                    # Extract content from chunk
+                    content = extract_openai_content_from_chunk(chunk, "responses")
+
+                    if content is not None:
+                        final_content.append(content)
 
                     yield chunk
 
@@ -159,6 +163,7 @@ class WrappedResponses:
                 end_time = time.time()
                 latency = end_time - start_time
                 output = final_content
+
                 await self._capture_streaming_event(
                     posthog_distinct_id,
                     posthog_trace_id,
@@ -170,6 +175,7 @@ class WrappedResponses:
                     latency,
                     output,
                     extract_available_tool_calls("openai", kwargs),
+                    model_from_response,
                 )
 
         return async_generator()
@@ -182,25 +188,31 @@ 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,
+        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, posthog_privacy_mode, kwargs.get("input")
+                self._client._ph_client,
+                posthog_privacy_mode,
+                sanitize_openai_response(kwargs.get("input")),
             ),
             "$ai_output_choices": with_privacy_mode(
                 self._client._ph_client,
                 posthog_privacy_mode,
-                output,
+                format_openai_streaming_output(output, "responses"),
             ),
             "$ai_http_status": 200,
             "$ai_input_tokens": usage_stats.get("input_tokens", 0),
@@ -215,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
 
@@ -340,8 +361,10 @@ 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]] = {}
+        model_from_response: Optional[str] = None
 
         if "stream_options" not in kwargs:
             kwargs["stream_options"] = {}
@@ -351,50 +374,45 @@ class WrappedCompletions:
         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:
-                    if hasattr(chunk, "usage") and chunk.usage:
-                        usage_stats = {
-                            k: getattr(chunk.usage, k, 0)
-                            for k in [
-                                "prompt_tokens",
-                                "completion_tokens",
-                                "total_tokens",
-                            ]
-                        }
-
-                        # Add support for cached tokens
-                        if hasattr(chunk.usage, "prompt_tokens_details") and hasattr(
-                            chunk.usage.prompt_tokens_details, "cached_tokens"
-                        ):
-                            usage_stats["cache_read_input_tokens"] = (
-                                chunk.usage.prompt_tokens_details.cached_tokens
-                            )
-
-                        if hasattr(chunk.usage, "output_tokens_details") and hasattr(
-                            chunk.usage.output_tokens_details, "reasoning_tokens"
-                        ):
-                            usage_stats["reasoning_tokens"] = (
-                                chunk.usage.output_tokens_details.reasoning_tokens
-                            )
-
-                    if (
-                        hasattr(chunk, "choices")
-                        and chunk.choices
-                        and len(chunk.choices) > 0
-                    ):
-                        if chunk.choices[0].delta and chunk.choices[0].delta.content:
-                            content = chunk.choices[0].delta.content
-                            if content:
-                                accumulated_content.append(content)
+                    # 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:
+                        merge_usage_stats(usage_stats, chunk_usage)
+
+                    # Extract content from chunk
+                    content = extract_openai_content_from_chunk(chunk, "chat")
+                    if content is not None:
+                        accumulated_content.append(content)
+
+                    # Extract and accumulate tool calls from chunk
+                    chunk_tool_calls = extract_openai_tool_calls_from_chunk(chunk)
+                    if chunk_tool_calls:
+                        accumulate_openai_tool_calls(
+                            accumulated_tool_calls, chunk_tool_calls
+                        )
 
                     yield chunk
 
             finally:
                 end_time = time.time()
                 latency = end_time - start_time
-                output = "".join(accumulated_content)
+
+                # Convert accumulated tool calls dict to list
+                tool_calls_list = (
+                    list(accumulated_tool_calls.values())
+                    if accumulated_tool_calls
+                    else None
+                )
+
                 await self._capture_streaming_event(
                     posthog_distinct_id,
                     posthog_trace_id,
@@ -404,8 +422,10 @@ class WrappedCompletions:
                     kwargs,
                     usage_stats,
                     latency,
-                    output,
+                    accumulated_content,
+                    tool_calls_list,
                     extract_available_tool_calls("openai", kwargs),
+                    model_from_response,
                 )
 
         return async_generator()
@@ -418,29 +438,36 @@ 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,
         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, posthog_privacy_mode, kwargs.get("messages")
+                self._client._ph_client,
+                posthog_privacy_mode,
+                sanitize_openai(kwargs.get("messages")),
             ),
             "$ai_output_choices": with_privacy_mode(
                 self._client._ph_client,
                 posthog_privacy_mode,
-                [{"content": output, "role": "assistant"}],
+                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
             ),
@@ -451,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
 
@@ -475,6 +512,7 @@ class WrappedEmbeddings:
 
     def __getattr__(self, name):
         """Fallback to original embeddings object for any methods we don't explicitly handle."""
+
         return getattr(self._original, name)
 
     async def create(
@@ -500,6 +538,7 @@ class WrappedEmbeddings:
         Returns:
             The response from OpenAI's embeddings.create call.
         """
+
         if posthog_trace_id is None:
             posthog_trace_id = str(uuid.uuid4())
 
@@ -508,12 +547,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
 
@@ -522,10 +562,12 @@ class WrappedEmbeddings:
             "$ai_provider": "openai",
             "$ai_model": kwargs.get("model"),
             "$ai_input": with_privacy_mode(
-                self._client._ph_client, posthog_privacy_mode, kwargs.get("input")
+                self._client._ph_client,
+                posthog_privacy_mode,
+                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),
@@ -556,6 +598,7 @@ class WrappedBeta:
 
     def __getattr__(self, name):
         """Fallback to original beta object for any methods we don't explicitly handle."""
+
         return getattr(self._original, name)
 
     @property
@@ -572,6 +615,7 @@ class WrappedBetaChat:
 
     def __getattr__(self, name):
         """Fallback to original beta chat object for any methods we don't explicitly handle."""
+
         return getattr(self._original, name)
 
     @property
@@ -588,6 +632,7 @@ class WrappedBetaCompletions:
 
     def __getattr__(self, name):
         """Fallback to original beta completions object for any methods we don't explicitly handle."""
+
         return getattr(self._original, name)
 
     async def parse(