feat: add beta parse method support

Author: k11kirky

llm_observability_examples.py

@@ -3,10 +3,10 @@
 
 import posthog
 from posthog.ai.openai import AsyncOpenAI, OpenAI
+from pydantic import BaseModel
 
 # Example credentials - replace these with your own or use environment variables
 posthog.project_api_key = os.getenv("POSTHOG_PROJECT_API_KEY", "your-project-api-key")
-posthog.personal_api_key = os.getenv("POSTHOG_PERSONAL_API_KEY", "your-personal-api-key")
 posthog.host = os.getenv("POSTHOG_HOST", "http://localhost:8000")  # Or https://app.posthog.com
 posthog.debug = True
 # change this to False to see usage events
@@ -22,7 +22,6 @@
     posthog_client=posthog,
 )
 
-
 def main_sync():
     trace_id = str(uuid.uuid4())
     print("Trace ID:", trace_id)
@@ -31,10 +30,11 @@ def main_sync():
     groups = {"company": "test_company"}
 
     try:
-        basic_openai_call(distinct_id, trace_id, properties, groups)
-        streaming_openai_call(distinct_id, trace_id, properties, groups)
-        embedding_openai_call(distinct_id, trace_id, properties, groups)
-        image_openai_call()
+        # basic_openai_call(distinct_id, trace_id, properties, groups)
+        # streaming_openai_call(distinct_id, trace_id, properties, groups)
+        # embedding_openai_call(distinct_id, trace_id, properties, groups)
+        # image_openai_call()
+        beta_openai_call(distinct_id, trace_id, properties, groups)
     except Exception as e:
         print("Error during OpenAI call:", str(e))
 
@@ -187,10 +187,32 @@ async def embedding_async_openai_call(posthog_distinct_id, posthog_trace_id, pos
     return response
 
 
+# TODO: add beta client
+
+class CalendarEvent(BaseModel):
+    name: str
+    date: str
+    participants: list[str]
+
+def beta_openai_call(distinct_id, trace_id, properties, groups):
+    response = openai_client.beta.chat.completions.parse(
+        model="gpt-4o-mini",
+        messages=[
+            {"role": "system", "content": "Extract the event information."},
+            {"role": "user", "content": "Alice and Bob are going to a science fair on Friday."},
+        ],
+        response_format=CalendarEvent,
+        posthog_distinct_id=distinct_id,
+        posthog_trace_id=trace_id,
+        posthog_properties=properties,
+        posthog_groups=groups,
+    )
+    print(response)
+    return response
+
 # HOW TO RUN:
 # comment out one of these to run the other
 
 if __name__ == "__main__":
     main_sync()
-
 # asyncio.run(main_async())

posthog/ai/openai/openai.py

@@ -31,6 +31,7 @@ def __init__(self, posthog_client: PostHogClient, **kwargs):
         self._ph_client = posthog_client
         self.chat = WrappedChat(self)
         self.embeddings = WrappedEmbeddings(self)
+        self.beta = WrappedBeta(self)
 
 
 class WrappedChat(openai.resources.chat.Chat):
@@ -249,3 +250,46 @@ def create(
             )
 
         return response
+
+
+class WrappedBeta(openai.resources.beta.Beta):
+    _client: OpenAI
+
+    @property
+    def chat(self):
+        return WrappedBetaChat(self._client)
+
+
+class WrappedBetaChat(openai.resources.beta.chat.Chat):
+    _client: OpenAI
+
+    @property
+    def completions(self):
+        return WrappedBetaCompletions(self._client)
+
+
+class WrappedBetaCompletions(openai.resources.beta.chat.completions.Completions):
+    _client: OpenAI
+
+    def parse(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        print("kwargs", kwargs)
+        return call_llm_and_track_usage(
+            posthog_distinct_id,
+            self._client._ph_client,
+            "openai",
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            self._client.base_url,
+            super().parse,
+            **kwargs,
+        )

posthog/ai/openai/openai_async.py

@@ -248,3 +248,44 @@ async def create(
             )
 
         return response
+
+class WrappedBeta(openai.resources.beta.Beta):
+    _client: AsyncOpenAI
+
+    @property
+    def chat(self):
+        return WrappedBetaChat(self._client)
+
+
+class WrappedBetaChat(openai.resources.beta.chat.Chat):
+    _client: AsyncOpenAI
+
+    @property
+    def completions(self):
+        return WrappedBetaCompletions(self._client)
+
+
+class WrappedBetaCompletions(openai.resources.beta.chat.completions.Completions):
+    _client: AsyncOpenAI
+
+    def parse(
+        self,
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ):
+        return call_llm_and_track_usage_async(
+            posthog_distinct_id,
+            self._client._ph_client,
+            "openai",
+            posthog_trace_id,
+            posthog_properties,
+            posthog_privacy_mode,
+            posthog_groups,
+            self._client.base_url,
+            super().parse,
+            **kwargs,
+        )