41 lines
1.4 KiB
Python
41 lines
1.4 KiB
Python
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)
|