refactor(openai): share resource and embedding policy

posthog/ai/openai/_embeddings.py

@@ -0,0 +1,49 @@
+from typing import Any, Dict, Optional
+
+from ..utils import _capture_ai_event, finalize_ai_content, with_privacy_mode
+
+
+def _capture_embedding_event(
+    *,
+    posthog_client: Any,
+    base_url: Any,
+    response: Any,
+    request_kwargs: Dict[str, Any],
+    latency: float,
+    distinct_id: Optional[str],
+    trace_id: str,
+    properties: Optional[Dict[str, Any]],
+    privacy_mode: bool,
+    groups: Optional[Dict[str, Any]],
+) -> None:
+    """Build and capture telemetry shared by sync and async embedding wrappers."""
+    usage = getattr(response, "usage", None)
+    input_tokens = getattr(usage, "prompt_tokens", 0) if usage else 0
+
+    event_properties = {
+        "$ai_provider": "openai",
+        "$ai_model": request_kwargs.get("model"),
+        "$ai_input": with_privacy_mode(
+            posthog_client,
+            privacy_mode,
+            finalize_ai_content(request_kwargs.get("input"), posthog_client),
+        ),
+        "$ai_http_status": 200,
+        "$ai_input_tokens": input_tokens,
+        "$ai_latency": latency,
+        "$ai_trace_id": trace_id,
+        "$ai_base_url": str(base_url),
+        **(properties or {}),
+    }
+
+    if distinct_id is None:
+        event_properties["$process_person_profile"] = False
+
+    if hasattr(posthog_client, "capture"):
+        _capture_ai_event(
+            posthog_client,
+            "$ai_embedding",
+            distinct_id=distinct_id or trace_id,
+            properties=event_properties,
+            groups=groups,
+        )

posthog/ai/openai/openai.py

@@ -1,6 +1,6 @@
 import time
 import uuid
-from typing import Any, Dict, Optional
+from typing import TYPE_CHECKING as _TYPE_CHECKING, Any, Dict, Optional
 
 from posthog.ai.types import TokenUsage as TokenUsage
 
@@ -13,11 +13,10 @@
 
 from posthog.ai.utils import (
     call_llm_and_track_usage,
-    _capture_ai_event,
     extract_available_tool_calls as extract_available_tool_calls,
-    finalize_ai_content,
+    finalize_ai_content as finalize_ai_content,
     merge_usage_stats as merge_usage_stats,
-    with_privacy_mode,
+    with_privacy_mode as with_privacy_mode,
 )
 from posthog.ai.openai.openai_converter import (
     accumulate_openai_tool_calls as accumulate_openai_tool_calls,
@@ -34,8 +33,10 @@
     _ResponsesStreamState,
     _build_streaming_event_data,
 )
-from posthog.ai.openai.wrapper_utils import (
+from ._embeddings import _capture_embedding_event
+from .wrapper_utils import (
     _OpenAIWrapperResource,
+    _wrap_openai_resources,
     merge_provider_override,
 )
 
@@ -47,6 +48,12 @@ class OpenAI(openai.OpenAI):
 
     _ph_client: PostHogClient
 
+    if _TYPE_CHECKING:
+        chat: "WrappedChat"
+        embeddings: "WrappedEmbeddings"
+        beta: "WrappedBeta"
+        responses: "WrappedResponses"
+
     def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
         """
         Args:
@@ -59,24 +66,7 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
         super().__init__(**kwargs)
         self._ph_client = posthog_client or setup()
 
-        # Store original objects after parent initialization (only if they exist)
-        self._original_chat = getattr(self, "chat", None)
-        self._original_embeddings = getattr(self, "embeddings", None)
-        self._original_beta = getattr(self, "beta", None)
-        self._original_responses = getattr(self, "responses", None)
-
-        # Replace with wrapped versions (only if originals exist)
-        if self._original_chat is not None:
-            self.chat = WrappedChat(self, self._original_chat)
-
-        if self._original_embeddings is not None:
-            self.embeddings = WrappedEmbeddings(self, self._original_embeddings)
-
-        if self._original_beta is not None:
-            self.beta = WrappedBeta(self, self._original_beta)
-
-        if self._original_responses is not None:
-            self.responses = WrappedResponses(self, self._original_responses)
+        _wrap_openai_resources(self, _SYNC_RESOURCE_WRAPPERS)
 
 
 def _parse_and_track(
@@ -491,46 +481,18 @@ def create(
         response = self._original.create(**kwargs)
         end_time = time.time()
 
-        # Extract usage statistics if available
-        usage_stats = {}
-        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),
-            }
-
-        latency = end_time - start_time
-
-        # Build the event properties
-        event_properties = {
-            "$ai_provider": "openai",
-            "$ai_model": kwargs.get("model"),
-            "$ai_input": with_privacy_mode(
-                self._client._ph_client,
-                posthog_privacy_mode,
-                finalize_ai_content(kwargs.get("input"), self._client._ph_client),
-            ),
-            "$ai_http_status": 200,
-            "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
-            "$ai_latency": latency,
-            "$ai_trace_id": posthog_trace_id,
-            "$ai_base_url": str(self._client.base_url),
-            **(posthog_properties or {}),
-        }
-
-        if posthog_distinct_id is None:
-            event_properties["$process_person_profile"] = False
-
-        # Send capture event for embeddings
-        if hasattr(self._client._ph_client, "capture"):
-            _capture_ai_event(
-                self._client._ph_client,
-                "$ai_embedding",
-                distinct_id=posthog_distinct_id or posthog_trace_id,
-                properties=event_properties,
-                groups=posthog_groups,
-            )
-
+        _capture_embedding_event(
+            posthog_client=self._client._ph_client,
+            base_url=self._client.base_url,
+            response=response,
+            request_kwargs=kwargs,
+            latency=end_time - start_time,
+            distinct_id=posthog_distinct_id,
+            trace_id=posthog_trace_id,
+            properties=posthog_properties,
+            privacy_mode=posthog_privacy_mode,
+            groups=posthog_groups,
+        )
         return response
 
 
@@ -594,3 +556,11 @@ def parse(
             posthog_provider_override,
             **kwargs,
         )
+
+
+_SYNC_RESOURCE_WRAPPERS = {
+    "chat": WrappedChat,
+    "embeddings": WrappedEmbeddings,
+    "beta": WrappedBeta,
+    "responses": WrappedResponses,
+}

posthog/ai/openai/openai_async.py

@@ -1,9 +1,9 @@
 import time
 import uuid
-from typing import Any, Dict, Optional
+from typing import TYPE_CHECKING as _TYPE_CHECKING, Any, Dict, Optional
 
 from posthog.ai.stream import AsyncStreamWrapper
-from posthog.ai.types import TokenUsage
+from posthog.ai.types import TokenUsage as TokenUsage
 
 try:
     import openai
@@ -15,12 +15,11 @@
 from posthog import setup
 from posthog.ai.utils import (
     call_llm_and_track_usage_async,
-    _capture_ai_event,
     extract_available_tool_calls as extract_available_tool_calls,
-    finalize_ai_content,
+    finalize_ai_content as finalize_ai_content,
     get_model_params as get_model_params,
     merge_usage_stats as merge_usage_stats,
-    with_privacy_mode,
+    with_privacy_mode as with_privacy_mode,
 )
 from posthog.ai.openai.openai_converter import (
     accumulate_openai_tool_calls as accumulate_openai_tool_calls,
@@ -36,8 +35,10 @@
     _ResponsesStreamState,
     _build_streaming_event_data,
 )
-from posthog.ai.openai.wrapper_utils import (
+from ._embeddings import _capture_embedding_event
+from .wrapper_utils import (
     _OpenAIWrapperResource,
+    _wrap_openai_resources,
     merge_provider_override,
 )
 
@@ -49,6 +50,12 @@ class AsyncOpenAI(openai.AsyncOpenAI):
 
     _ph_client: PostHogClient
 
+    if _TYPE_CHECKING:
+        chat: "WrappedChat"
+        embeddings: "WrappedEmbeddings"
+        beta: "WrappedBeta"
+        responses: "WrappedResponses"
+
     def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
         """
         Args:
@@ -61,24 +68,7 @@ def __init__(self, posthog_client: Optional[PostHogClient] = None, **kwargs):
         super().__init__(**kwargs)
         self._ph_client = posthog_client or setup()
 
-        # Store original objects after parent initialization (only if they exist)
-        self._original_chat = getattr(self, "chat", None)
-        self._original_embeddings = getattr(self, "embeddings", None)
-        self._original_beta = getattr(self, "beta", None)
-        self._original_responses = getattr(self, "responses", None)
-
-        # Replace with wrapped versions (only if originals exist)
-        if self._original_chat is not None:
-            self.chat = WrappedChat(self, self._original_chat)
-
-        if self._original_embeddings is not None:
-            self.embeddings = WrappedEmbeddings(self, self._original_embeddings)
-
-        if self._original_beta is not None:
-            self.beta = WrappedBeta(self, self._original_beta)
-
-        if self._original_responses is not None:
-            self.responses = WrappedResponses(self, self._original_responses)
+        _wrap_openai_resources(self, _ASYNC_RESOURCE_WRAPPERS)
 
 
 async def _parse_and_track(
@@ -496,47 +486,18 @@ async def create(
         response = await self._original.create(**kwargs)
         end_time = time.time()
 
-        # Extract usage statistics if available
-        usage_stats: TokenUsage = TokenUsage()
-
-        if hasattr(response, "usage") and response.usage:
-            usage_stats = TokenUsage(
-                input_tokens=getattr(response.usage, "prompt_tokens", 0),
-                output_tokens=getattr(response.usage, "completion_tokens", 0),
-            )
-
-        latency = end_time - start_time
-
-        # Build the event properties
-        event_properties = {
-            "$ai_provider": "openai",
-            "$ai_model": kwargs.get("model"),
-            "$ai_input": with_privacy_mode(
-                self._client._ph_client,
-                posthog_privacy_mode,
-                finalize_ai_content(kwargs.get("input"), self._client._ph_client),
-            ),
-            "$ai_http_status": 200,
-            "$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_properties or {}),
-        }
-
-        if posthog_distinct_id is None:
-            event_properties["$process_person_profile"] = False
-
-        # Send capture event for embeddings
-        if hasattr(self._client._ph_client, "capture"):
-            _capture_ai_event(
-                self._client._ph_client,
-                "$ai_embedding",
-                distinct_id=posthog_distinct_id or posthog_trace_id,
-                properties=event_properties,
-                groups=posthog_groups,
-            )
-
+        _capture_embedding_event(
+            posthog_client=self._client._ph_client,
+            base_url=self._client.base_url,
+            response=response,
+            request_kwargs=kwargs,
+            latency=end_time - start_time,
+            distinct_id=posthog_distinct_id,
+            trace_id=posthog_trace_id,
+            properties=posthog_properties,
+            privacy_mode=posthog_privacy_mode,
+            groups=posthog_groups,
+        )
         return response
 
 
@@ -600,3 +561,11 @@ async def parse(
             posthog_provider_override,
             **kwargs,
         )
+
+
+_ASYNC_RESOURCE_WRAPPERS = {
+    "chat": WrappedChat,
+    "embeddings": WrappedEmbeddings,
+    "beta": WrappedBeta,
+    "responses": WrappedResponses,
+}