Merge branch 'master' into feat/modal_param_extraction

Author: k11kirky

posthog/ai/anthropic/anthropic.py

@@ -8,7 +8,13 @@
 import uuid
 from typing import Any, Dict, Optional
 
-from posthog.ai.utils import call_llm_and_track_usage, extract_core_model_params, get_model_params, with_privacy_mode
+from posthog.ai.utils import (
+    call_llm_and_track_usage,
+    extract_core_model_params,
+    get_model_params,
+    merge_system_prompt,
+    with_privacy_mode,
+)
 from posthog.client import Client as PostHogClient
 
 
@@ -171,7 +177,11 @@ def _capture_streaming_event(
             "$ai_provider": "anthropic",
             "$ai_model": kwargs.get("model"),
             "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": with_privacy_mode(self._client._ph_client, posthog_privacy_mode, kwargs.get("messages")),
+            "$ai_input": with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                merge_system_prompt(kwargs, "anthropic"),
+            ),
             "$ai_output_choices": with_privacy_mode(
                 self._client._ph_client,
                 posthog_privacy_mode,

posthog/ai/anthropic/anthropic_async.py

@@ -8,7 +8,7 @@
 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.ai.utils import call_llm_and_track_usage_async, get_model_params, merge_system_prompt, with_privacy_mode
 from posthog.client import Client as PostHogClient
 
 
@@ -171,7 +171,11 @@ async def _capture_streaming_event(
             "$ai_provider": "anthropic",
             "$ai_model": kwargs.get("model"),
             "$ai_model_parameters": get_model_params(kwargs),
-            "$ai_input": with_privacy_mode(self._client._ph_client, posthog_privacy_mode, kwargs.get("messages")),
+            "$ai_input": with_privacy_mode(
+                self._client._ph_client,
+                posthog_privacy_mode,
+                merge_system_prompt(kwargs, "anthropic"),
+            ),
             "$ai_output_choices": with_privacy_mode(
                 self._client._ph_client,
                 posthog_privacy_mode,

posthog/ai/utils.py

@@ -115,6 +115,15 @@ def format_response_openai(response):
     return output
 
 
+def merge_system_prompt(kwargs: Dict[str, Any], provider: str):
+    if provider != "anthropic":
+        return kwargs.get("messages")
+    messages = kwargs.get("messages") or []
+    if kwargs.get("system") is None:
+        return messages
+    return [{"role": "system", "content": kwargs.get("system")}] + messages
+
+
 def call_llm_and_track_usage(
     posthog_distinct_id: Optional[str],
     ph_client: PostHogClient,
@@ -152,11 +161,13 @@ def call_llm_and_track_usage(
         if response and hasattr(response, "usage"):
             usage = get_usage(response, provider)
 
+        messages = merge_system_prompt(kwargs, provider)
+
         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_input": with_privacy_mode(ph_client, posthog_privacy_mode, messages),
             "$ai_output_choices": with_privacy_mode(
                 ph_client, posthog_privacy_mode, format_response(response, provider)
             ),
@@ -221,11 +232,13 @@ async def call_llm_and_track_usage_async(
         if response and hasattr(response, "usage"):
             usage = get_usage(response, provider)
 
+        messages = merge_system_prompt(kwargs, provider)
+
         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_input": with_privacy_mode(ph_client, posthog_privacy_mode, messages),
             "$ai_output_choices": with_privacy_mode(
                 ph_client, posthog_privacy_mode, format_response(response, provider)
             ),

posthog/test/ai/anthropic/test_anthropic.py

@@ -218,25 +218,29 @@ def test_basic_integration(mock_client):
     client = Anthropic(posthog_client=mock_client)
     client.messages.create(
         model="claude-3-opus-20240229",
-        messages=[{"role": "user", "content": "You must always answer with 'Bar'."}],
+        messages=[{"role": "user", "content": "Foo"}],
         max_tokens=1,
         temperature=0,
         posthog_distinct_id="test-id",
         posthog_properties={"foo": "bar"},
+        system="You must always answer with 'Bar'.",
     )
 
     assert mock_client.capture.call_count == 1
 
     call_args = mock_client.capture.call_args[1]
     props = call_args["properties"]
-
     assert call_args["distinct_id"] == "test-id"
     assert call_args["event"] == "$ai_generation"
     assert props["$ai_provider"] == "anthropic"
     assert props["$ai_model"] == "claude-3-opus-20240229"
-    assert props["$ai_input"] == [{"role": "user", "content": "You must always answer with 'Bar'."}]
+    assert props["$ai_input"] == [
+        {"role": "system", "content": "You must always answer with 'Bar'."},
+        {"role": "user", "content": "Foo"},
+    ]
     assert props["$ai_output_choices"][0]["role"] == "assistant"
-    assert props["$ai_input_tokens"] == 16
+    assert props["$ai_output_choices"][0]["content"] == "Bar"
+    assert props["$ai_input_tokens"] == 18
     assert props["$ai_output_tokens"] == 1
     assert props["$ai_http_status"] == 200
     assert props["foo"] == "bar"
@@ -273,6 +277,54 @@ async def test_basic_async_integration(mock_client):
     assert isinstance(props["$ai_latency"], float)
 
 
+def test_streaming_system_prompt(mock_client, mock_anthropic_stream):
+    with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_stream):
+        client = Anthropic(api_key="test-key", posthog_client=mock_client)
+        response = client.messages.create(
+            model="claude-3-opus-20240229",
+            system="Foo",
+            messages=[{"role": "user", "content": "Bar"}],
+            stream=True,
+        )
+
+        # Consume the stream
+        list(response)
+
+        # 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]
+        props = call_args["properties"]
+        assert props["$ai_input"] == [{"role": "system", "content": "Foo"}, {"role": "user", "content": "Bar"}]
+
+
+@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set")
+async def test_async_streaming_system_prompt(mock_client, mock_anthropic_stream):
+    client = AsyncAnthropic(posthog_client=mock_client)
+    response = await client.messages.create(
+        model="claude-3-opus-20240229",
+        system="You must always answer with 'Bar'.",
+        messages=[{"role": "user", "content": "Foo"}],
+        stream=True,
+        max_tokens=1,
+    )
+
+    # Consume the stream
+    [c async for c in response]
+
+    # 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]
+    props = call_args["properties"]
+
+    assert props["$ai_input"] == [
+        {"role": "system", "content": "You must always answer with 'Bar'."},
+        {"role": "user", "content": "Foo"},
+    ]
+
 def test_core_model_params(mock_client, mock_anthropic_response):
     with patch("anthropic.resources.Messages.create", return_value=mock_anthropic_response):
         client = Anthropic(api_key="test-key", posthog_client=mock_client)
@@ -287,12 +339,9 @@ def test_core_model_params(mock_client, mock_anthropic_response):
         )
 
         assert response == mock_anthropic_response
-        assert mock_client.capture.call_count == 1
-
-        call_args = mock_client.capture.call_args[1]
-        props = call_args["properties"]
+        props = mock_client.capture.call_args[1]["properties"]
         assert props["$ai_model_parameters"] == {"temperature": 0.5, "max_tokens": 100, "stream": False}
         assert props["$ai_temperature"] == 0.5
         assert props["$ai_max_tokens"] == 100
         assert props["$ai_stream"] == False
-        assert props["foo"] == "bar"
+        assert props["foo"] == "bar"