chore(llma): add integration tests for $ai_stop_reason and sampo changeset

Author: carlos-marchal-ph

.sampo/changesets/add-ai-stop-reason.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: minor
+---
+
+feat(ai): add $ai_stop_reason extraction for all providers

posthog/test/ai/anthropic/test_anthropic.py

@@ -1,4 +1,5 @@
 import json
+import os
 from unittest.mock import patch
 
 import pytest
@@ -1401,3 +1402,22 @@ def test_explicit_distinct_id_overrides_outer_context(
 
         call_args = mock_client.capture.call_args[1]
         assert call_args["distinct_id"] == "explicit-user-789"
+
+
+ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
+
+
+@pytest.mark.skipif(not ANTHROPIC_API_KEY, reason="ANTHROPIC_API_KEY is not set")
+def test_integration_stop_reason(mock_client):
+    client = Anthropic(api_key=ANTHROPIC_API_KEY, posthog_client=mock_client)
+    client.messages.create(
+        model="claude-sonnet-4-20250514",
+        messages=[{"role": "user", "content": "Say hi"}],
+        max_tokens=5,
+        posthog_distinct_id="test-id",
+    )
+
+    props = mock_client.capture.call_args[1]["properties"]
+    assert props["$ai_stop_reason"] in ("end_turn", "max_tokens")
+    assert props["$ai_provider"] == "anthropic"
+    assert props["$ai_input_tokens"] > 0

posthog/test/ai/gemini/test_gemini.py

@@ -1,4 +1,5 @@
 import json
+import os
 from unittest.mock import MagicMock, patch
 
 import pytest
@@ -1286,3 +1287,22 @@ def mock_streaming_response():
     assert mock_client.capture.call_count == 1
     props = mock_client.capture.call_args[1]["properties"]
     assert props["$ai_stop_reason"] == "STOP"
+
+
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+
+
+@pytest.mark.skipif(not GEMINI_API_KEY, reason="GEMINI_API_KEY is not set")
+def test_integration_stop_reason(mock_client):
+    client = Client(api_key=GEMINI_API_KEY, posthog_client=mock_client)
+    client.models.generate_content(
+        model="gemini-2.0-flash",
+        contents="Say hi",
+        posthog_distinct_id="test-id",
+    )
+
+    props = mock_client.capture.call_args[1]["properties"]
+    assert props["$ai_stop_reason"] is not None
+    assert isinstance(props["$ai_stop_reason"], str)
+    assert props["$ai_provider"] == "gemini"
+    assert props["$ai_input_tokens"] > 0

posthog/test/ai/langchain/test_callbacks.py

@@ -894,6 +894,7 @@ def test_openai_chain(mock_client):
     )
     assert gen_props["$ai_input_tokens"] == 20
     assert gen_props["$ai_output_tokens"] == 1
+    assert gen_props["$ai_stop_reason"] == "stop"
 
 
 @pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
@@ -1172,6 +1173,7 @@ def test_anthropic_chain(mock_client):
     )
     assert gen_props["$ai_input_tokens"] == 17
     assert gen_props["$ai_output_tokens"] == 4
+    assert gen_props["$ai_stop_reason"] == "end_turn"
 
     assert trace_args["event"] == "$ai_trace"
     assert trace_props["$ai_input_state"] == {}

posthog/test/ai/openai/test_openai.py

@@ -1,4 +1,5 @@
 import json
+import os
 import time
 from unittest.mock import AsyncMock, patch
 
@@ -2154,3 +2155,22 @@ async def chunk_iterable():
     props = call_args["properties"]
 
     assert props["$ai_model"] == "gpt-4o-mini-async-stored"
+
+
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+
+
+@pytest.mark.skipif(not OPENAI_API_KEY, reason="OPENAI_API_KEY is not set")
+def test_integration_stop_reason(mock_client):
+    client = OpenAI(api_key=OPENAI_API_KEY, posthog_client=mock_client)
+    client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[{"role": "user", "content": "Say hi"}],
+        max_tokens=5,
+        posthog_distinct_id="test-id",
+    )
+
+    props = mock_client.capture.call_args[1]["properties"]
+    assert props["$ai_stop_reason"] in ("stop", "length")
+    assert props["$ai_provider"] == "openai"
+    assert props["$ai_input_tokens"] > 0