feat: azure export w/ async (#200)

* feat: azure export w/ async

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

Author: k11kirky

CHANGELOG.md

@@ -1,3 +1,8 @@
+
+## 3.18.0 - 2025-02-28
+
+1. Add support for Azure OpenAI.
+
 ## 3.17.0 - 2025-02-27
 
 1. The LangChain handler now captures tools in `$ai_generation` events, in property `$ai_tools`. This allows for displaying tools provided to the LLM call in PostHog UI. Note that support for `$ai_tools` in OpenAI and Anthropic SDKs is coming soon.

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/__init__.py

@@ -1,4 +1,5 @@
 from .openai import OpenAI
 from .openai_async import AsyncOpenAI
+from .openai_providers import AsyncAzureOpenAI, AzureOpenAI
 
-__all__ = ["OpenAI", "AsyncOpenAI"]
+__all__ = ["OpenAI", "AsyncOpenAI", "AzureOpenAI", "AsyncAzureOpenAI"]

posthog/ai/openai/openai_providers.py

@@ -0,0 +1,41 @@
+try:
+    import openai
+    import openai.resources
+except ImportError:
+    raise ModuleNotFoundError("Please install the Open AI SDK to use this feature: 'pip install openai'")
+
+from posthog.ai.openai.openai import WrappedBeta, WrappedChat, WrappedEmbeddings
+from posthog.ai.openai.openai_async import WrappedBeta as AsyncWrappedBeta
+from posthog.ai.openai.openai_async import WrappedChat as AsyncWrappedChat
+from posthog.ai.openai.openai_async import WrappedEmbeddings as 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)

posthog/version.py

@@ -1,4 +1,4 @@
-VERSION = "3.17.0"
+VERSION = "3.18.0"
 
 if __name__ == "__main__":
     print(VERSION, end="")  # noqa: T201