diff --git a/llm_observability_examples.py b/llm_observability_examples.py index 8328780..9b9a42d 100644 --- a/llm_observability_examples.py +++ b/llm_observability_examples.py @@ -31,6 +31,7 @@ def main_sync(): try: basic_openai_call(distinct_id, trace_id, properties) # streaming_openai_call(distinct_id, trace_id, properties) + non_instrumented_openai_call() except Exception as e: print("Error during OpenAI call:", str(e)) @@ -130,6 +131,12 @@ async def streaming_async_openai_call(distinct_id, trace_id, properties): return response +def non_instrumented_openai_call(): + response = openai_client.images.generate(model="dall-e-3", prompt="A cute baby sea otter", n=1, size="1024x1024") + print(response) + return response + + # HOW TO RUN: # comment out one of these to run the other diff --git a/posthog/ai/__init__.py b/posthog/ai/__init__.py index e3d4311..6fbccd1 100644 --- a/posthog/ai/__init__.py +++ b/posthog/ai/__init__.py @@ -2,4 +2,3 @@ from .providers.openai.openai_async import AsyncOpenAI __all__ = ["OpenAI", "AsyncOpenAI"] -# TODO: add Azure OpenAI wrapper diff --git a/posthog/ai/providers/openai/openai.py b/posthog/ai/providers/openai/openai.py index 7bbae23..58accda 100644 --- a/posthog/ai/providers/openai/openai.py +++ b/posthog/ai/providers/openai/openai.py @@ -8,7 +8,7 @@ from posthog.client import Client as PostHogClient from posthog.ai.utils import ( - track_usage, + call_llm_and_track_usage, get_model_params, ) @@ -33,6 +33,15 @@ def __init__( self._openai_client = openai.OpenAI(**openai_config) self._posthog_client = posthog_client + def __getattr__(self, name: str) -> Any: + """ + Expose all attributes of the underlying openai.OpenAI instance except for the 'chat' property, + which is replaced with a custom ChatNamespace for usage tracking. + """ + if name == "chat": + return self.chat + return getattr(self._openai_client, name) + @property def chat(self) -> "ChatNamespace": return ChatNamespace(self._posthog_client, self._openai_client) @@ -61,7 +70,7 @@ def create( **kwargs: Any, ): distinct_id = posthog_distinct_id or "anonymous_ai_user" - + if kwargs.get("stream", False): return self._create_streaming( distinct_id, @@ -70,11 +79,10 @@ def create( **kwargs, ) - def call_method(**call_kwargs): return self._openai_client.chat.completions.create(**call_kwargs) - return track_usage( + return call_llm_and_track_usage( distinct_id, self._ph_client, posthog_trace_id, @@ -113,7 +121,9 @@ def generator(): end_time = time.time() latency = end_time - start_time output = "".join(accumulated_content) - self._capture_streaming_event(distinct_id, posthog_trace_id, posthog_properties, kwargs, usage_stats, latency, output) + self._capture_streaming_event( + distinct_id, posthog_trace_id, posthog_properties, kwargs, usage_stats, latency, output + ) return generator() @@ -127,7 +137,7 @@ def _capture_streaming_event( latency: float, output: str, ): - + event_properties = { "$ai_provider": "openai", "$ai_model": kwargs.get("model"), @@ -140,7 +150,7 @@ def _capture_streaming_event( "role": "assistant", } ] - }, + }, "$ai_http_status": 200, "$ai_input_tokens": usage_stats.get("prompt_tokens", 0), "$ai_output_tokens": usage_stats.get("completion_tokens", 0), diff --git a/posthog/ai/providers/openai/openai_async.py b/posthog/ai/providers/openai/openai_async.py index 93b8db2..ddce1ac 100644 --- a/posthog/ai/providers/openai/openai_async.py +++ b/posthog/ai/providers/openai/openai_async.py @@ -7,7 +7,7 @@ raise ModuleNotFoundError("Please install OpenAI to use this feature: 'pip install openai'") from posthog.client import Client as PostHogClient -from posthog.ai.utils import track_usage_async, get_model_params +from posthog.ai.utils import call_llm_and_track_usage_async, get_model_params class AsyncOpenAI: @@ -29,6 +29,15 @@ def __init__( self._openai_client = openai.AsyncOpenAI(**openai_config) self._posthog_client = posthog_client + def __getattr__(self, name: str) -> Any: + """ + Expose all attributes of the underlying openai.AsyncOpenAI instance except for the 'chat' property, + which is replaced with a custom AsyncChatNamespace for usage tracking. + """ + if name == "chat": + return self.chat + return getattr(self._openai_client, name) + @property def chat(self) -> "AsyncChatNamespace": return AsyncChatNamespace(self._posthog_client, self._openai_client) @@ -71,11 +80,10 @@ async def create( async def call_async_method(**call_kwargs): return await self._openai_client.chat.completions.create(**call_kwargs) - response = await track_usage_async( + response = await call_llm_and_track_usage_async( distinct_id, self._ph_client, posthog_trace_id, posthog_properties, call_async_method, **kwargs ) return response - async def _create_streaming( self, @@ -106,7 +114,9 @@ async def async_generator(): end_time = time.time() latency = end_time - start_time output = "".join(accumulated_content) - self._capture_streaming_event(distinct_id, posthog_trace_id, posthog_properties, kwargs, usage_stats, latency, output) + self._capture_streaming_event( + distinct_id, posthog_trace_id, posthog_properties, kwargs, usage_stats, latency, output + ) return async_generator() @@ -120,7 +130,7 @@ def _capture_streaming_event( latency: float, output: str, ): - + event_properties = { "$ai_provider": "openai", "$ai_model": kwargs.get("model"), @@ -133,7 +143,7 @@ def _capture_streaming_event( "role": "assistant", } ] - }, + }, "$ai_http_status": 200, "$ai_input_tokens": usage_stats.get("prompt_tokens", 0), "$ai_output_tokens": usage_stats.get("completion_tokens", 0), @@ -148,4 +158,3 @@ def _capture_streaming_event( event="$ai_generation", properties=event_properties, ) - diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index d28629d..5be6483 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, AsyncGenerator, Callable, Optional +from typing import Any, Dict, Callable, Optional import time from posthog.client import Client as PostHogClient @@ -39,7 +39,7 @@ def format_response(response): return output -def track_usage( +def call_llm_and_track_usage( distinct_id: str, ph_client: PostHogClient, posthog_trace_id: Optional[str], @@ -99,7 +99,7 @@ def track_usage( return response -async def track_usage_async( +async def call_llm_and_track_usage_async( distinct_id: str, ph_client: PostHogClient, posthog_trace_id: Optional[str],