fi9x: isort + black

Author: k11kirky

posthog/ai/anthropic/__init__.py

@@ -1,5 +1,12 @@
 from .anthropic import Anthropic
 from .anthropic_async import AsyncAnthropic
-from .anthropic_providers import AnthropicBedrock, AsyncAnthropicBedrock, AnthropicVertex, AsyncAnthropicVertex
+from .anthropic_providers import AnthropicBedrock, AnthropicVertex, AsyncAnthropicBedrock, AsyncAnthropicVertex
 
-__all__ = ["Anthropic", "AsyncAnthropic", "AnthropicBedrock", "AsyncAnthropicBedrock", "AnthropicVertex", "AsyncAnthropicVertex"]
+__all__ = [
+    "Anthropic",
+    "AsyncAnthropic",
+    "AnthropicBedrock",
+    "AsyncAnthropicBedrock",
+    "AnthropicVertex",
+    "AsyncAnthropicVertex",
+]

posthog/ai/anthropic/anthropic.py

@@ -4,11 +4,12 @@
 except ImportError:
     raise ModuleNotFoundError("Please install the Anthropic SDK to use this feature: 'pip install anthropic'")
 
+import time
+import uuid
+from typing import Any, Dict, Optional
+
 from posthog.ai.utils import call_llm_and_track_usage, get_model_params, with_privacy_mode
 from posthog.client import Client as PostHogClient
-from typing import Any, Dict, Optional
-import uuid
-import time
 
 
 class Anthropic(anthropic.Anthropic):
@@ -43,7 +44,7 @@ def create(
     ):
         """
         Create a message using Anthropic's API while tracking usage in PostHog.
-        
+
         Args:
             posthog_distinct_id: Optional ID to associate with the usage event
             posthog_trace_id: Optional trace UUID for linking events
@@ -77,7 +78,7 @@ def create(
             super().create,
             **kwargs,
         )
-    
+
     def stream(
         self,
         posthog_distinct_id: Optional[str] = None,
@@ -89,7 +90,7 @@ def stream(
     ):
         if posthog_trace_id is None:
             posthog_trace_id = uuid.uuid4()
-        
+
         return self._create_streaming(
             posthog_distinct_id,
             posthog_trace_id,
@@ -136,7 +137,7 @@ def generator():
                 end_time = time.time()
                 latency = end_time - start_time
                 output = "".join(accumulated_content)
-                
+
                 self._capture_streaming_event(
                     posthog_distinct_id,
                     posthog_trace_id,
@@ -195,4 +196,3 @@ def _capture_streaming_event(
                 properties=event_properties,
                 groups=posthog_groups,
             )
-

posthog/ai/anthropic/anthropic_async.py

@@ -4,11 +4,12 @@
 except ImportError:
     raise ModuleNotFoundError("Please install the Anthropic SDK to use this feature: 'pip install anthropic'")
 
+import time
+import uuid
+from typing import Any, Dict, Optional
+
 from posthog.ai.utils import call_llm_and_track_usage_async, get_model_params, with_privacy_mode
 from posthog.client import Client as PostHogClient
-from typing import Any, Dict, Optional
-import uuid
-import time
 
 
 class AsyncAnthropic(anthropic.AsyncAnthropic):
@@ -43,7 +44,7 @@ async def create(
     ):
         """
         Create a message using Anthropic's API while tracking usage in PostHog.
-        
+
         Args:
             posthog_distinct_id: Optional ID to associate with the usage event
             posthog_trace_id: Optional trace UUID for linking events
@@ -77,7 +78,7 @@ async def create(
             super().create,
             **kwargs,
         )
-    
+
     async def stream(
         self,
         posthog_distinct_id: Optional[str] = None,
@@ -89,7 +90,7 @@ async def stream(
     ):
         if posthog_trace_id is None:
             posthog_trace_id = uuid.uuid4()
-        
+
         return await self._create_streaming(
             posthog_distinct_id,
             posthog_trace_id,
@@ -136,7 +137,7 @@ async def generator():
                 end_time = time.time()
                 latency = end_time - start_time
                 output = "".join(accumulated_content)
-                
+
                 await self._capture_streaming_event(
                     posthog_distinct_id,
                     posthog_trace_id,
@@ -195,4 +196,3 @@ async def _capture_streaming_event(
                 properties=event_properties,
                 groups=posthog_groups,
             )
-

posthog/ai/anthropic/anthropic_providers.py

@@ -7,6 +7,7 @@
 from posthog.ai.anthropic.anthropic_async import AsyncWrappedMessages
 from posthog.client import Client as PostHogClient
 
+
 class AnthropicBedrock(anthropic.AnthropicBedrock):
     """
     A wrapper around the Anthropic Bedrock SDK that automatically sends LLM usage events to PostHog.

posthog/ai/utils.py

@@ -127,7 +127,9 @@ def call_llm_and_track_usage(
             "$ai_model": kwargs.get("model"),
             "$ai_model_parameters": get_model_params(kwargs),
             "$ai_input": with_privacy_mode(ph_client, posthog_privacy_mode, kwargs.get("messages")),
-            "$ai_output_choices": with_privacy_mode(ph_client, posthog_privacy_mode, format_response(response, provider)),
+            "$ai_output_choices": with_privacy_mode(
+                ph_client, posthog_privacy_mode, format_response(response, provider)
+            ),
             "$ai_http_status": http_status,
             "$ai_input_tokens": usage.get("input_tokens", 0),
             "$ai_output_tokens": usage.get("output_tokens", 0),
@@ -193,7 +195,9 @@ async def call_llm_and_track_usage_async(
             "$ai_model": kwargs.get("model"),
             "$ai_model_parameters": get_model_params(kwargs),
             "$ai_input": with_privacy_mode(ph_client, posthog_privacy_mode, kwargs.get("messages")),
-            "$ai_output_choices": with_privacy_mode(ph_client, posthog_privacy_mode, format_response(response, provider)),
+            "$ai_output_choices": with_privacy_mode(
+                ph_client, posthog_privacy_mode, format_response(response, provider)
+            ),
             "$ai_http_status": http_status,
             "$ai_input_tokens": usage.get("input_tokens", 0),
             "$ai_output_tokens": usage.get("output_tokens", 0),

posthog/test/ai/anthropic/test_anthropic.py

@@ -98,7 +98,7 @@ def test_streaming(mock_client, mock_anthropic_stream):
         assert chunks[0].content == "A"
         assert chunks[1].content == "B"
         assert chunks[2].content == "C"
-        
+
         # Wait a bit to ensure the capture is called
         time.sleep(0.1)
         assert mock_client.capture.call_count == 1
@@ -134,7 +134,7 @@ def test_streaming_with_stream_endpoint(mock_client, mock_anthropic_stream):
         assert chunks[0].content == "A"
         assert chunks[1].content == "B"
         assert chunks[2].content == "C"
-        
+
         # Wait a bit to ensure the capture is called
         time.sleep(0.1)
         assert mock_client.capture.call_count == 1