fix: tests ai_output

Author: k11kirky

posthog/test/ai/langchain/test_callbacks.py

@@ -120,7 +120,7 @@ def test_basic_chat_chain(mock_client, stream):
         {"role": "system", "content": "You are a helpful assistant."},
         {"role": "user", "content": "Who won the world series in 2020?"},
     ]
-    assert props["$ai_output"] == {
+    assert props["$ai_output_choices"] == {
         "choices": [{"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."}]
     }
     assert props["$ai_input_tokens"] == 10
@@ -165,7 +165,7 @@ async def test_async_basic_chat_chain(mock_client, stream):
         {"role": "system", "content": "You are a helpful assistant."},
         {"role": "user", "content": "Who won the world series in 2020?"},
     ]
-    assert props["$ai_output"] == {
+    assert props["$ai_output_choices"] == {
         "choices": [{"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."}]
     }
     assert props["$ai_input_tokens"] == 10
@@ -200,7 +200,7 @@ def test_basic_llm_chain(mock_client, Model, stream):
     assert "$ai_model" in props
     assert "$ai_provider" in props
     assert props["$ai_input"] == ["Who won the world series in 2020?"]
-    assert props["$ai_output"] == {"choices": ["The Los Angeles Dodgers won the World Series in 2020."]}
+    assert props["$ai_output_choices"] == {"choices": ["The Los Angeles Dodgers won the World Series in 2020."]}
     assert props["$ai_http_status"] == 200
     assert props["$ai_trace_id"] is not None
     assert isinstance(props["$ai_latency"], float)
@@ -231,7 +231,7 @@ async def test_async_basic_llm_chain(mock_client, Model, stream):
     assert "$ai_model" in props
     assert "$ai_provider" in props
     assert props["$ai_input"] == ["Who won the world series in 2020?"]
-    assert props["$ai_output"] == {"choices": ["The Los Angeles Dodgers won the World Series in 2020."]}
+    assert props["$ai_output_choices"] == {"choices": ["The Los Angeles Dodgers won the World Series in 2020."]}
     assert props["$ai_http_status"] == 200
     assert props["$ai_trace_id"] is not None
     assert isinstance(props["$ai_latency"], float)
@@ -258,7 +258,7 @@ def test_trace_id_for_multiple_chains(mock_client):
     assert "$ai_model" in first_call_props
     assert "$ai_provider" in first_call_props
     assert first_call_props["$ai_input"] == [{"role": "user", "content": "Foo"}]
-    assert first_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
+    assert first_call_props["$ai_output_choices"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
     assert first_call_props["$ai_http_status"] == 200
     assert first_call_props["$ai_trace_id"] is not None
     assert isinstance(first_call_props["$ai_latency"], float)
@@ -270,7 +270,7 @@ def test_trace_id_for_multiple_chains(mock_client):
     assert "$ai_model" in second_call_props
     assert "$ai_provider" in second_call_props
     assert second_call_props["$ai_input"] == [{"role": "assistant", "content": "Bar"}]
-    assert second_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
+    assert second_call_props["$ai_output_choices"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
     assert second_call_props["$ai_http_status"] == 200
     assert second_call_props["$ai_trace_id"] is not None
     assert isinstance(second_call_props["$ai_latency"], float)
@@ -338,7 +338,7 @@ def test_metadata(mock_client):
     assert first_call_props["$ai_trace_id"] == "test-trace-id"
     assert first_call_props["foo"] == "bar"
     assert first_call_props["$ai_input"] == [{"role": "user", "content": "Foo"}]
-    assert first_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
+    assert first_call_props["$ai_output_choices"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
     assert first_call_props["$ai_http_status"] == 200
     assert isinstance(first_call_props["$ai_latency"], float)
 
@@ -392,7 +392,7 @@ def test_openai_error(mock_client):
     props = args["properties"]
     assert props["$ai_http_status"] == 401
     assert props["$ai_input"] == [{"role": "user", "content": "Foo"}]
-    assert "$ai_output" not in props
+    assert "$ai_output_choices" not in props
 
 
 @pytest.mark.skipif(not OPENAI_API_KEY, reason="OpenAI API key not set")
@@ -443,7 +443,7 @@ def test_openai_chain(mock_client):
         {"role": "system", "content": 'You must always answer with "Bar".'},
         {"role": "user", "content": "Foo"},
     ]
-    assert first_call_props["$ai_output"] == {
+    assert first_call_props["$ai_output_choices"] == {
         "choices": [
             {
                 "role": "assistant",
@@ -486,7 +486,7 @@ def test_openai_captures_multiple_generations(mock_client):
         {"role": "system", "content": 'You must always answer with "Bar".'},
         {"role": "user", "content": "Foo"},
     ]
-    assert first_call_props["$ai_output"] == {
+    assert first_call_props["$ai_output_choices"] == {
         "choices": [
             {
                 "role": "assistant",
@@ -544,7 +544,7 @@ def test_openai_streaming(mock_client):
         {"role": "system", "content": 'You must always answer with "Bar".'},
         {"role": "user", "content": "Foo"},
     ]
-    assert first_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
+    assert first_call_props["$ai_output_choices"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
     assert first_call_props["$ai_http_status"] == 200
     assert first_call_props["$ai_input_tokens"] == 20
     assert first_call_props["$ai_output_tokens"] == 1
@@ -576,7 +576,7 @@ async def test_async_openai_streaming(mock_client):
         {"role": "system", "content": 'You must always answer with "Bar".'},
         {"role": "user", "content": "Foo"},
     ]
-    assert first_call_props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
+    assert first_call_props["$ai_output_choices"] == {"choices": [{"role": "assistant", "content": "Bar"}]}
     assert first_call_props["$ai_http_status"] == 200
     assert first_call_props["$ai_input_tokens"] == 20
     assert first_call_props["$ai_output_tokens"] == 1
@@ -630,7 +630,7 @@ def test_privacy_mode_local(mock_client):
     assert mock_client.capture.call_count == 1
     call = mock_client.capture.call_args[1]
     assert call["properties"]["$ai_input"] is None
-    assert call["properties"]["$ai_output"] is None
+    assert call["properties"]["$ai_output_choices"] is None
 
 
 def test_privacy_mode_global(mock_client):
@@ -649,4 +649,4 @@ def test_privacy_mode_global(mock_client):
     assert mock_client.capture.call_count == 1
     call = mock_client.capture.call_args[1]
     assert call["properties"]["$ai_input"] is None
-    assert call["properties"]["$ai_output"] is None
+    assert call["properties"]["$ai_output_choices"] is None

posthog/test/ai/openai/test_openai.py

@@ -83,7 +83,7 @@ def test_basic_completion(mock_client, mock_openai_response):
         assert props["$ai_provider"] == "openai"
         assert props["$ai_model"] == "gpt-4"
         assert props["$ai_input"] == [{"role": "user", "content": "Hello"}]
-        assert props["$ai_output"] == {"choices": [{"role": "assistant", "content": "Test response"}]}
+        assert props["$ai_output_choices"] == {"choices": [{"role": "assistant", "content": "Test response"}]}
         assert props["$ai_input_tokens"] == 20
         assert props["$ai_output_tokens"] == 10
         assert props["$ai_http_status"] == 200
@@ -152,7 +152,7 @@ def test_privacy_mode_local(mock_client, mock_openai_response):
         call_args = mock_client.capture.call_args[1]
         props = call_args["properties"]
         assert props["$ai_input"] is None
-        assert props["$ai_output"] is None
+        assert props["$ai_output_choices"] is None
 
 
 def test_privacy_mode_global(mock_client, mock_openai_response):
@@ -172,4 +172,4 @@ def test_privacy_mode_global(mock_client, mock_openai_response):
         call_args = mock_client.capture.call_args[1]
         props = call_args["properties"]
         assert props["$ai_input"] is None
-        assert props["$ai_output"] is None
+        assert props["$ai_output_choices"] is None