chore: cleanup formatting

Author: k11kirky

llm_observability_examples.py

@@ -21,6 +21,7 @@
     posthog_client=posthog,
 )
 
+
 def main_sync():
     trace_id = str(uuid.uuid4())
     print("Trace ID:", trace_id)
@@ -31,6 +32,7 @@ def main_sync():
     except Exception as e:
         print("Error during OpenAI call:", str(e))
 
+
 async def main_async():
     try:
         await basic_async_openai_call()
@@ -42,7 +44,10 @@ async def main_async():
 def basic_openai_call():
     response = openai_client.chat.completions.create(
         model="gpt-4o-mini",
-        messages=[{"role": "system", "content": "You are a complex problem solver."}, {"role": "user", "content": "Explain quantum computing in simple terms."}],
+        messages=[
+            {"role": "system", "content": "You are a complex problem solver."},
+            {"role": "user", "content": "Explain quantum computing in simple terms."},
+        ],
         max_tokens=100,
         temperature=0.7,
     )
@@ -52,10 +57,14 @@ def basic_openai_call():
         print("No response or unexpected format returned.")
     return response
 
+
 async def basic_async_openai_call():
     response = await async_openai_client.chat.completions.create(
         model="gpt-4o-mini",
-        messages=[{"role": "system", "content": "You are a complex problem solver."}, {"role": "user", "content": "Explain quantum computing in simple terms."}],
+        messages=[
+            {"role": "system", "content": "You are a complex problem solver."},
+            {"role": "user", "content": "Explain quantum computing in simple terms."},
+        ],
         max_tokens=100,
         temperature=0.7,
     )
@@ -65,10 +74,14 @@ async def basic_async_openai_call():
         print("No response or unexpected format returned.")
     return response
 
+
 def streaming_openai_call():
     response = openai_client.chat.completions.create(
         model="gpt-4o-mini",
-        messages=[{"role": "system", "content": "You are a complex problem solver."}, {"role": "user", "content": "Explain quantum computing in simple terms."}],
+        messages=[
+            {"role": "system", "content": "You are a complex problem solver."},
+            {"role": "user", "content": "Explain quantum computing in simple terms."},
+        ],
         max_tokens=100,
         temperature=0.7,
         stream=True,
@@ -79,10 +92,14 @@ def streaming_openai_call():
 
     return response
 
+
 async def streaming_async_openai_call():
     response = await async_openai_client.chat.completions.create(
         model="gpt-4o-mini",
-        messages=[{"role": "system", "content": "You are a complex problem solver."}, {"role": "user", "content": "Explain quantum computing in simple terms."}],
+        messages=[
+            {"role": "system", "content": "You are a complex problem solver."},
+            {"role": "user", "content": "Explain quantum computing in simple terms."},
+        ],
         max_tokens=100,
         temperature=0.7,
         stream=True,
@@ -93,10 +110,11 @@ async def streaming_async_openai_call():
 
     return response
 
+
 # HOW TO RUN:
 # comment out one of these to run the other
 
 if __name__ == "__main__":
     main_sync()
 
-# asyncio.run(main_async())
+# asyncio.run(main_async())

posthog/ai/__init__.py

@@ -2,4 +2,4 @@
 from .providers.openai.openai_async import AsyncOpenAI
 
 __all__ = ["OpenAI", "AsyncOpenAI"]
-# TODO: add Azure OpenAI wrapper
+# TODO: add Azure OpenAI wrapper

posthog/ai/providers/openai/openai.py

@@ -4,13 +4,12 @@
 try:
     import openai
 except ImportError:
-    raise ModuleNotFoundError(
-        "Please install OpenAI to use this feature: 'pip install openai'"
-    )
+    raise ModuleNotFoundError("Please install OpenAI to use this feature: 'pip install openai'")
 
 from posthog.client import Client as PostHogClient
 from posthog.ai.utils import get_model_params, format_response, process_sync_streaming_response, track_usage
 
+
 class OpenAI:
     """
     A wrapper around the OpenAI SDK that automatically sends LLM usage events to PostHog.
@@ -83,11 +82,6 @@ def call_method(**call_kwargs):
             return self._openai_client.chat.completions.create(**call_kwargs)
 
         response = track_usage(
-            distinct_id,
-            self._ph_client,
-            posthog_trace_id,
-            posthog_properties,
-            call_method,
-            **kwargs
+            distinct_id, self._ph_client, posthog_trace_id, posthog_properties, call_method, **kwargs
         )
-        return response
+        return response

posthog/ai/providers/openai/openai_async.py

@@ -4,13 +4,12 @@
 try:
     import openai
 except ImportError:
-    raise ModuleNotFoundError(
-        "Please install OpenAI to use this feature: 'pip install openai'"
-    )
+    raise ModuleNotFoundError("Please install OpenAI to use this feature: 'pip install openai'")
 
 from posthog.client import Client as PostHogClient
 from posthog.ai.utils import get_model_params, format_response, process_async_streaming_response, track_usage_async
 
+
 class AsyncOpenAI:
     """
     An async wrapper around the OpenAI SDK that automatically sends LLM usage events to PostHog.
@@ -60,7 +59,7 @@ async def create(
         """
         Wraps openai chat completions (async) and captures a $ai_generation event in PostHog.
 
-        To use streaming in async mode: 
+        To use streaming in async mode:
             async for chunk in async_openai.chat.completions.create(stream=True, ...):
                 ...
         """
@@ -81,11 +80,6 @@ async def call_async_method(**call_kwargs):
             return await self._openai_client.chat.completions.create(**call_kwargs)
 
         response = await track_usage_async(
-            distinct_id,
-            self._ph_client,
-            posthog_trace_id,
-            posthog_properties,
-            call_async_method,
-            **kwargs
+            distinct_id, self._ph_client, posthog_trace_id, posthog_properties, call_async_method, **kwargs
         )
-        return response
+        return response

posthog/ai/utils.py

@@ -2,6 +2,7 @@
 import time
 from posthog.client import Client as PostHogClient
 
+
 def get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]:
     """
     Extracts model parameters from the kwargs dictionary.
@@ -56,7 +57,6 @@ def process_sync_streaming_response(
                 accumulated_content.append(chunk.choices[0].delta.content)
             yield chunk
     finally:
-        # Once we've finished, capture the final content in PostHog
         final_content = "".join(accumulated_content)
         event_properties["$ai_output"] = {
             "choices": [
@@ -107,17 +107,18 @@ async def process_async_streaming_response(
                 properties=event_properties,
             )
 
+
 def track_usage(
     distinct_id: str,
     ph_client: PostHogClient,
     posthog_trace_id: Optional[str],
     posthog_properties: Optional[Dict[str, Any]],
-    call_method: Callable[..., Any],  # This will be the function that actually calls OpenAI
+    call_method: Callable[..., Any],
     **kwargs: Any,
 ) -> Any:
     """
     Common usage-tracking logic for both sync and async calls.
-    call_method: a function or coroutine that actually implements the call to openai
+    call_method: the llm call method (e.g. openai.chat.completions.create)
     """
     start_time = time.time()
     response = None
@@ -134,11 +135,9 @@ def track_usage(
         end_time = time.time()
         latency = end_time - start_time
 
-        # Extract usage if available
         if response and hasattr(response, "usage"):
             usage = response.usage.model_dump()
 
-        # Prepare analytics data
         input_tokens = usage.get("prompt_tokens", 0)
         output_tokens = usage.get("completion_tokens", 0)
         event_properties = {
@@ -155,20 +154,20 @@ def track_usage(
             "$ai_posthog_properties": posthog_properties,
         }
 
-        # Capture usage in PostHog
+        # send the event to posthog
         if hasattr(ph_client, "capture") and callable(ph_client.capture):
             ph_client.capture(
                 distinct_id=distinct_id,
                 event="$ai_generation",
                 properties=event_properties,
             )
 
-    # Re-raise the error if it occurred
     if error:
         raise error
 
     return response
 
+
 async def track_usage_async(
     distinct_id: str,
     ph_client: PostHogClient,
@@ -192,11 +191,9 @@ async def track_usage_async(
         end_time = time.time()
         latency = end_time - start_time
 
-        # Extract usage if available
         if response and hasattr(response, "usage"):
             usage = response.usage.model_dump()
 
-        # Prepare analytics data
         input_tokens = usage.get("prompt_tokens", 0)
         output_tokens = usage.get("completion_tokens", 0)
         event_properties = {
@@ -213,7 +210,7 @@ async def track_usage_async(
             "$ai_posthog_properties": posthog_properties,
         }
 
-        # Capture usage in PostHog
+        # send the event to posthog
         if hasattr(ph_client, "capture") and callable(ph_client.capture):
             ph_client.capture(
                 distinct_id=distinct_id,
@@ -224,4 +221,4 @@ async def track_usage_async(
     if error:
         raise error
 
-    return response
+    return response