feat: anthroipic

Author: k11kirky

CHANGELOG.md

@@ -1,3 +1,8 @@
+## 3.8.4 - 2025-01-17
+
+1. Add Anthropic support for LLM Observability.
+2. Update LLM Observability to use output_choices.
+
 ## 3.8.3 - 2025-01-14
 
 1. Fix setuptools to include the `posthog.ai.openai` and `posthog.ai.langchain` packages for the `posthoganalytics` package.

posthog/ai/anthropic/anthropic.py

@@ -80,13 +80,16 @@ def create(
 
     def stream(
         self,
-        posthog_distinct_id: Optional[str],
-        posthog_trace_id: Optional[str],
-        posthog_properties: Optional[Dict[str, Any]],
-        posthog_privacy_mode: bool,
-        posthog_groups: Optional[Dict[str, Any]],
+        posthog_distinct_id: Optional[str] = None,
+        posthog_trace_id: Optional[str] = None,
+        posthog_properties: Optional[Dict[str, Any]] = None,
+        posthog_privacy_mode: bool = False,
+        posthog_groups: Optional[Dict[str, Any]] = None,
         **kwargs: Any,
     ):
+        if posthog_trace_id is None:
+            posthog_trace_id = uuid.uuid4()
+        
         return self._create_streaming(
             posthog_distinct_id,
             posthog_trace_id,
@@ -106,7 +109,7 @@ def _create_streaming(
         **kwargs: Any,
     ):
         start_time = time.time()
-        usage_stats: Dict[str, int] = {}
+        usage_stats: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
         accumulated_content = []
         response = super().create(**kwargs)
 
@@ -117,8 +120,11 @@ def generator():
                 for event in response:
                     if hasattr(event, "usage") and event.usage:
                         usage_stats = {
-                            "input_tokens": getattr(event.usage, "input_tokens", 0),
-                            "output_tokens": getattr(event.usage, "output_tokens", 0),
+                            k: getattr(event.usage, k, 0)
+                            for k in [
+                                "input_tokens",
+                                "output_tokens",
+                            ]
                         }
 
                     if hasattr(event, "content") and event.content:
@@ -130,6 +136,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,

posthog/ai/utils.py

@@ -28,6 +28,23 @@ def get_model_params(kwargs: Dict[str, Any]) -> Dict[str, Any]:
     return model_params
 
 
+def get_usage(response, provider: str) -> Dict[str, Any]:
+    if provider == "anthropic":
+        return {
+            "input_tokens": response.usage.input_tokens,
+            "output_tokens": response.usage.output_tokens,
+        }
+    elif provider == "openai":
+        return {
+            "input_tokens": response.usage.prompt_tokens,
+            "output_tokens": response.usage.completion_tokens,
+        }
+    return {
+        "input_tokens": 0,
+        "output_tokens": 0,
+    }
+
+
 def format_response(response, provider: str):
     """
     Format a regular (non-streaming) response.
@@ -103,19 +120,17 @@ def call_llm_and_track_usage(
             posthog_trace_id = uuid.uuid4()
 
         if response and hasattr(response, "usage"):
-            usage = response.usage.model_dump()
+            usage = get_usage(response, provider)
 
-        input_tokens = usage.get("prompt_tokens", 0)
-        output_tokens = usage.get("completion_tokens", 0)
         event_properties = {
             "$ai_provider": provider,
             "$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_http_status": http_status,
-            "$ai_input_tokens": input_tokens,
-            "$ai_output_tokens": output_tokens,
+            "$ai_input_tokens": usage.get("input_tokens", 0),
+            "$ai_output_tokens": usage.get("output_tokens", 0),
             "$ai_latency": latency,
             "$ai_trace_id": posthog_trace_id,
             "$ai_base_url": str(base_url),
@@ -171,19 +186,17 @@ async def call_llm_and_track_usage_async(
             posthog_trace_id = uuid.uuid4()
 
         if response and hasattr(response, "usage"):
-            usage = response.usage.model_dump()
+            usage = get_usage(response, provider)
 
-        input_tokens = usage.get("prompt_tokens", 0)
-        output_tokens = usage.get("completion_tokens", 0)
         event_properties = {
             "$ai_provider": provider,
             "$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_http_status": http_status,
-            "$ai_input_tokens": input_tokens,
-            "$ai_output_tokens": output_tokens,
+            "$ai_input_tokens": usage.get("input_tokens", 0),
+            "$ai_output_tokens": usage.get("output_tokens", 0),
             "$ai_latency": latency,
             "$ai_trace_id": posthog_trace_id,
             "$ai_base_url": str(base_url),

posthog/test/ai/anthropic/test_anthropic.py

@@ -2,7 +2,7 @@
 from unittest.mock import patch
 
 import pytest
-from anthropic.types import Message, Usage, MessageStreamEvent
+from anthropic.types import Message, Usage
 
 from posthog.ai.anthropic import Anthropic
 
@@ -33,26 +33,23 @@ def mock_anthropic_response():
 
 @pytest.fixture
 def mock_anthropic_stream():
-    yield from [
-        MessageStreamEvent(
-            id="msg_123",
-            type="message",
-            role="assistant",
-            content="A",
-        ),
-        MessageStreamEvent(
-            id="msg_123",
-            type="message",
-            role="assistant",
-            content="B",
-        ),
-        MessageStreamEvent(
-            id="msg_123",
-            type="message",
-            role="assistant",
-            content="C",
-        ),
-    ]
+    class MockStreamEvent:
+        def __init__(self, content, usage=None):
+            self.content = content
+            self.usage = usage
+
+    def stream_generator():
+        yield MockStreamEvent("A")
+        yield MockStreamEvent("B")
+        yield MockStreamEvent(
+            "C",
+            usage=Usage(
+                input_tokens=20,
+                output_tokens=10,
+            ),
+        )
+
+    return stream_generator()
 
 
 def test_basic_completion(mock_client, mock_anthropic_response):
@@ -101,6 +98,9 @@ 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
 
         call_args = mock_client.capture.call_args[1]
@@ -111,12 +111,11 @@ def test_streaming(mock_client, mock_anthropic_stream):
         assert props["$ai_provider"] == "anthropic"
         assert props["$ai_model"] == "claude-3-opus-20240229"
         assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
-        assert props["$ai_output_choices"] == [{"role": "assistant", "content": "Test response"}]
+        assert props["$ai_output_choices"] == [{"role": "assistant", "content": "ABC"}]
         assert props["$ai_input_tokens"] == 20
         assert props["$ai_output_tokens"] == 10
-        assert props["$ai_http_status"] == 200
-        assert props["foo"] == "bar"
         assert isinstance(props["$ai_latency"], float)
+        assert props["foo"] == "bar"
 
 
 def test_streaming_with_stream_endpoint(mock_client, mock_anthropic_stream):
@@ -135,6 +134,9 @@ 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
 
         call_args = mock_client.capture.call_args[1]
@@ -145,12 +147,11 @@ def test_streaming_with_stream_endpoint(mock_client, mock_anthropic_stream):
         assert props["$ai_provider"] == "anthropic"
         assert props["$ai_model"] == "claude-3-opus-20240229"
         assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
-        assert props["$ai_output_choices"] == [{"role": "assistant", "content": "Test response"}]
+        assert props["$ai_output_choices"] == [{"role": "assistant", "content": "ABC"}]
         assert props["$ai_input_tokens"] == 20
         assert props["$ai_output_tokens"] == 10
-        assert props["$ai_http_status"] == 200
-        assert props["foo"] == "bar"
         assert isinstance(props["$ai_latency"], float)
+        assert props["foo"] == "bar"
 
 
 def test_groups(mock_client, mock_anthropic_response):

posthog/version.py

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