Fix subclassing

Author: Twixes

posthog/ai/providers/openai/openai.py

@@ -2,6 +2,8 @@
 import uuid
 from typing import Any, Dict, Optional
 
+import openai.resources
+
 try:
     import openai
 except ImportError:
@@ -30,29 +32,19 @@ def __init__(self, posthog_client: PostHogClient, **kwargs):
         """
         super().__init__(**kwargs)
         self._ph_client = posthog_client
-
-    @property
-    def chat(self) -> "ChatNamespace":
-        """OpenAI `chat` wrapped with PostHog usage tracking."""
-        return ChatNamespace(self)
+        self.chat = WrappedChat(self)
 
 
-class ChatNamespace:
-    _openai_client: OpenAI
-
-    def __init__(self, openai_client: OpenAI):
-        self._openai_client = openai_client
+class WrappedChat(openai.resources.chat.Chat):
+    _client: OpenAI
 
     @property
     def completions(self):
-        return ChatCompletions(self._openai_client)
-
+        return WrappedCompletions(self._client)
 
-class ChatCompletions:
-    _openai_client: OpenAI
 
-    def __init__(self, openai_client: OpenAI):
-        self._client = openai_client
+class WrappedCompletions(openai.resources.chat.completions.Completions):
+    _client: OpenAI
 
     def create(
         self,
@@ -71,16 +63,13 @@ def create(
                 **kwargs,
             )
 
-        def call_method(**call_kwargs):
-            return self._openai_client.chat.completions.create(**call_kwargs)
-
         return call_llm_and_track_usage(
             distinct_id,
-            self._openai_client._ph_client,
+            self._client._ph_client,
             posthog_trace_id,
             posthog_properties,
-            call_method,
-            self._openai_client.base_url,
+            self._client.base_url,
+            super().create,
             **kwargs,
         )
 
@@ -95,7 +84,7 @@ def _create_streaming(
         usage_stats: Dict[str, int] = {}
         accumulated_content = []
         stream_options = {"include_usage": True}
-        response = self._openai_client.chat.completions.create(
+        response = self._client.chat.completions.create(
             **kwargs, stream_options=stream_options
         )
 
@@ -158,7 +147,9 @@ def _capture_streaming_event(
                     }
                 ]
             },
-            "$ai_request_url": str(self._openai_client.base_url.join("chat/completions")),
+            "$ai_request_url": str(
+                self._client.base_url.join("chat/completions")
+            ),
             "$ai_http_status": 200,
             "$ai_input_tokens": usage_stats.get("prompt_tokens", 0),
             "$ai_output_tokens": usage_stats.get("completion_tokens", 0),
@@ -167,8 +158,8 @@ def _capture_streaming_event(
             "$ai_posthog_properties": posthog_properties,
         }
 
-        if hasattr(self._openai_client._ph_client, "capture"):
-            self._openai_client._ph_client.capture(
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
                 distinct_id=distinct_id,
                 event="$ai_generation",
                 properties=event_properties,

posthog/ai/providers/openai/openai_async.py

@@ -2,6 +2,8 @@
 import uuid
 from typing import Any, Dict, Optional
 
+import openai.resources
+
 try:
     import openai
 except ImportError:
@@ -28,31 +30,20 @@ def __init__(self, posthog_client: PostHogClient, **kwargs):
             **openai_config: Additional keyword args (e.g. organization="xxx").
         """
         super().__init__(**kwargs)
-        super().chat
         self._ph_client = posthog_client
-
-    @property
-    def chat(self) -> "AsyncChatNamespace":
-        """OpenAI `chat` wrapped with PostHog usage tracking."""
-        return AsyncChatNamespace(self)
+        self.chat = WrappedChat(self)
 
 
-class AsyncChatNamespace:
-    _openai_client: AsyncOpenAI
-
-    def __init__(self, openai_client: AsyncOpenAI):
-        self._openai_client = openai_client
+class WrappedChat(openai.resources.chat.AsyncChat):
+    _client: AsyncOpenAI
 
     @property
     def completions(self):
-        return AsyncChatCompletions(self._openai_client)
-
+        return WrappedCompletions(self._client)
 
-class AsyncChatCompletions:
-    _openai_client: AsyncOpenAI
 
-    def __init__(self, openai_client: AsyncOpenAI):
-        self._openai_client = openai_client
+class WrappedCompletions(openai.resources.chat.completions.AsyncCompletions):
+    _client: AsyncOpenAI
 
     async def create(
         self,
@@ -72,17 +63,13 @@ async def create(
                 **kwargs,
             )
 
-        # Non-streaming: let track_usage_async handle request and analytics
-        async def call_async_method(**call_kwargs):
-            return await self._openai_client.chat.completions.create(**call_kwargs)
-
         response = await call_llm_and_track_usage_async(
             distinct_id,
-            self._openai_client._ph_client,
+            self._client._ph_client,
             posthog_trace_id,
             posthog_properties,
-            call_async_method,
-            self._openai_client.base_url,
+            self._client.base_url,
+            super().create,
             **kwargs,
         )
         return response
@@ -98,7 +85,7 @@ async def _create_streaming(
         usage_stats: Dict[str, int] = {}
         accumulated_content = []
         stream_options = {"include_usage": True}
-        response = await self._openai_client.chat.completions.create(
+        response = await self._client.chat.completions.create(
             **kwargs, stream_options=stream_options
         )
 
@@ -166,11 +153,13 @@ def _capture_streaming_event(
             "$ai_latency": latency,
             "$ai_trace_id": posthog_trace_id,
             "$ai_posthog_properties": posthog_properties,
-            "$ai_request_url": str(self._openai_client.base_url.join("chat/completions")),
+            "$ai_request_url": str(
+                self._client.base_url.join("chat/completions")
+            ),
         }
 
-        if hasattr(self._openai_client._ph_client, "capture"):
-            self._openai_client._ph_client.capture(
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
                 distinct_id=distinct_id,
                 event="$ai_generation",
                 properties=event_properties,

posthog/ai/utils.py

@@ -50,8 +50,8 @@ def call_llm_and_track_usage(
     ph_client: PostHogClient,
     posthog_trace_id: Optional[str],
     posthog_properties: Optional[Dict[str, Any]],
-    call_method: Callable[..., Any],
     base_url: Optional[str],
+    call_method: Callable[..., Any],
     **kwargs: Any,
 ) -> Any:
     """
@@ -115,8 +115,8 @@ async def call_llm_and_track_usage_async(
     ph_client: PostHogClient,
     posthog_trace_id: Optional[str],
     posthog_properties: Optional[Dict[str, Any]],
-    call_async_method: Callable[..., Any],
     base_url: URL,
+    call_async_method: Callable[..., Any],
     **kwargs: Any,
 ) -> Any:
     start_time = time.time()