feat: azure export w/ async

Author: k11kirky

llm_observability_examples.py

@@ -33,11 +33,11 @@ def main_sync():
     groups = {"company": "test_company"}
 
     try:
-        # basic_openai_call(distinct_id, trace_id, properties, groups)
+        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)
+        # beta_openai_call(distinct_id, trace_id, properties, groups)
     except Exception as e:
         print("Error during OpenAI call:", str(e))
 
@@ -216,6 +216,6 @@ def beta_openai_call(distinct_id, trace_id, properties, groups):
 # HOW TO RUN:
 # comment out one of these to run the other
 
-# if __name__ == "__main__":
-# main_sync()
-asyncio.run(main_async())
+if __name__ == "__main__":
+    main_sync()
+# asyncio.run(main_async())

posthog/ai/openai/openai_providers.py

@@ -0,0 +1,39 @@
+try:
+    import openai
+    import openai.resources
+except ImportError:
+    raise ModuleNotFoundError("Please install the Open AI SDK to use this feature: 'pip install anthropic'")
+
+from posthog.ai.openai.openai import WrappedBeta, WrappedChat, WrappedEmbeddings
+from posthog.ai.openai.openai_async import AsyncWrappedBeta, AsyncWrappedChat, AsyncWrappedEmbeddings
+from posthog.client import Client as PostHogClient
+
+
+class AzureOpenAI(openai.AzureOpenAI):
+    """
+    A wrapper around the Azure OpenAI SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        super().__init__(**kwargs)
+        self._ph_client = posthog_client
+        self.chat = WrappedChat(self)
+        self.embeddings = WrappedEmbeddings(self)
+        self.beta = WrappedBeta(self)
+
+    
+class AsyncAzureOpenAI(openai.AsyncAzureOpenAI):
+    """
+    A wrapper around the Azure OpenAI SDK that automatically sends LLM usage events to PostHog.
+    """
+
+    _ph_client: PostHogClient
+
+    def __init__(self, posthog_client: PostHogClient, **kwargs):
+        super().__init__(**kwargs)
+        self._ph_client = posthog_client
+        self.chat = AsyncWrappedChat(self)
+        self.embeddings = AsyncWrappedEmbeddings(self)
+        self.beta = AsyncWrappedBeta(self)