feat: embeddings

Author: k11kirky

llm_observability_examples.py

@@ -24,14 +24,14 @@
 def main_sync():
     trace_id = str(uuid.uuid4())
     print("Trace ID:", trace_id)
-    distinct_id = "test_distinct_id"
+    distinct_id = "test2_distinct_id"
     properties = {"test_property": "test_value"}
 
     try:
         basic_openai_call(distinct_id, trace_id, properties)
         streaming_openai_call(distinct_id, trace_id, properties)
+        embedding_openai_call(distinct_id, trace_id, properties)
         image_openai_call()
-        embedding_openai_call()
     except Exception as e:
         print("Error during OpenAI call:", str(e))
 
@@ -45,7 +45,7 @@ async def main_async():
     try:
         await basic_async_openai_call(distinct_id, trace_id, properties)
         await streaming_async_openai_call(distinct_id, trace_id, properties)
-        await embedding_async_openai_call()
+        await embedding_async_openai_call(distinct_id, trace_id, properties)
         await image_async_openai_call()
     except Exception as e:
         print("Error during OpenAI call:", str(e))
@@ -150,13 +150,13 @@ async def image_async_openai_call():
     return response
 
 
-def embedding_openai_call():
-    response = openai_client.embeddings.create(input="The hedgehog is cute", model="text-embedding-3-small")
+def embedding_openai_call(posthog_distinct_id, posthog_trace_id, posthog_properties):
+    response = openai_client.embeddings.create(input="The hedgehog is cute", model="text-embedding-3-small", posthog_distinct_id=posthog_distinct_id, posthog_trace_id=posthog_trace_id, posthog_properties=posthog_properties)
     print(response)
     return response
 
-async def embedding_async_openai_call():
-    response = await async_openai_client.embeddings.create(input="The hedgehog is cute", model="text-embedding-3-small")
+async def embedding_async_openai_call(posthog_distinct_id, posthog_trace_id, posthog_properties):
+    response = await async_openai_client.embeddings.create(input="The hedgehog is cute", model="text-embedding-3-small", posthog_distinct_id=posthog_distinct_id, posthog_trace_id=posthog_trace_id, posthog_properties=posthog_properties)
     print(response)
     return response
 

posthog/ai/openai/openai.py

@@ -30,6 +30,7 @@ def __init__(self, posthog_client: PostHogClient, **kwargs):
         super().__init__(**kwargs)
         self._ph_client = posthog_client
         self.chat = WrappedChat(self)
+        self.embeddings = WrappedEmbeddings(self)
 
 
 class WrappedChat(openai.resources.chat.Chat):
@@ -167,4 +168,70 @@ def _capture_streaming_event(
                 distinct_id=posthog_distinct_id or posthog_trace_id,
                 event="$ai_generation",
                 properties=event_properties,
-            )
+            )
+
+
+class WrappedEmbeddings(openai.resources.embeddings.Embeddings):
+    _client: OpenAI
+
+    def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Create an embedding using OpenAI's 'embeddings.create' method, but also track usage in PostHog.
+
+        Args:
+            posthog_distinct_id: Optional ID to associate with the usage event.
+            posthog_trace_id: Optional trace UUID for linking events.
+            posthog_properties: Optional dictionary of extra properties to include in the event.
+            **kwargs: Any additional parameters for the OpenAI Embeddings API.
+
+        Returns:
+            The response from OpenAI's embeddings.create call.
+        """
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        start_time = time.time()
+        response = super().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": kwargs.get("input"),
+            "$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,
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # Send capture event for embeddings
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_embedding",
+                properties=event_properties,
+            )
+
+        return response

posthog/ai/openai/openai_async.py

@@ -29,7 +29,7 @@ def __init__(self, posthog_client: PostHogClient, **kwargs):
         super().__init__(**kwargs)
         self._ph_client = posthog_client
         self.chat = WrappedChat(self)
-
+        self.embeddings = WrappedEmbeddings(self)
 
 class WrappedChat(openai.resources.chat.AsyncChat):
     _client: AsyncOpenAI
@@ -167,3 +167,68 @@ def _capture_streaming_event(
                 event="$ai_generation",
                 properties=event_properties,
             )
+
+class WrappedEmbeddings(openai.resources.embeddings.AsyncEmbeddings):
+    _client: AsyncOpenAI
+
+    async def create(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Create an embedding using OpenAI's 'embeddings.create' method, but also track usage in PostHog.
+
+        Args:
+            posthog_distinct_id: Optional ID to associate with the usage event.
+            posthog_trace_id: Optional trace UUID for linking events.
+            posthog_properties: Optional dictionary of extra properties to include in the event.
+            **kwargs: Any additional parameters for the OpenAI Embeddings API.
+
+        Returns:
+            The response from OpenAI's embeddings.create call.
+        """
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+
+        start_time = time.time()
+        response = await super().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": kwargs.get("input"),
+            "$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,
+        }
+
+        if posthog_distinct_id is None:
+            event_properties["$process_person_profile"] = False
+
+        # Send capture event for embeddings
+        if hasattr(self._client._ph_client, "capture"):
+            self._client._ph_client.capture(
+                distinct_id=posthog_distinct_id or posthog_trace_id,
+                event="$ai_embedding",
+                properties=event_properties,
+            )
+
+        return response