test(llm-observability): Account for LangGraph 0.3.29 changes

posthog/test/ai/langchain/test_callbacks.py

@@ -618,89 +618,57 @@ def test_graph_state(mock_client):
     assert isinstance(result["messages"][2], AIMessage)
     assert result["messages"][2].content == "It's a type of greeble."
 
-    assert mock_client.capture.call_count == 12
+    assert mock_client.capture.call_count == 6
     calls = [call[1] for call in mock_client.capture.call_args_list]
 
-    trace_args = calls[11]
-    trace_props = calls[11]["properties"]
+    # The trace event is captured at the end
+    trace_args = calls[-1]
+    trace_props = calls[-1]["properties"]
 
     # Events are captured in the reverse order.
     # Check all trace_ids
     for call in calls:
         assert call["properties"]["$ai_trace_id"] == trace_props["$ai_trace_id"]
 
-    # First span, write the state
+    # 1. Span, finish initialization
+    second_state = {
+        "messages": [HumanMessage(content="What's a bar?"), AIMessage(content="Let's explore bar.")],
+        "xyz": "abc",
+    }
+
+    # 1. Span - the fake_plain node, which doesn't do anything
     assert calls[0]["event"] == "$ai_span"
-    assert calls[0]["properties"]["$ai_parent_id"] == calls[2]["properties"]["$ai_span_id"]
+    assert calls[0]["properties"]["$ai_parent_id"] == trace_props["$ai_trace_id"]
     assert "$ai_span_id" in calls[0]["properties"]
+    assert calls[0]["properties"]["$ai_span_name"] == "fake_plain"
     assert calls[0]["properties"]["$ai_input_state"] == initial_state
-    assert calls[0]["properties"]["$ai_output_state"] == initial_state
+    assert calls[0]["properties"]["$ai_output_state"] == second_state
 
-    # Second span, set the START node
+    # 2. Span - the ChatPromptTemplate within fake_llm's FakeMessagesListChatModel
     assert calls[1]["event"] == "$ai_span"
-    assert calls[1]["properties"]["$ai_parent_id"] == calls[2]["properties"]["$ai_span_id"]
+    assert calls[1]["properties"]["$ai_parent_id"] == calls[3]["properties"]["$ai_span_id"]
     assert "$ai_span_id" in calls[1]["properties"]
-    assert calls[1]["properties"]["$ai_input_state"] == initial_state
-    assert calls[1]["properties"]["$ai_output_state"] == initial_state
+    assert calls[1]["properties"]["$ai_span_name"] == "ChatPromptTemplate"
 
-    # Third span, finish initialization
-    assert calls[2]["event"] == "$ai_span"
+    # 3. Generation - the FakeMessagesListChatModel within fake_llm's RunnableSequence
+    assert calls[2]["event"] == "$ai_generation"
+    assert calls[2]["properties"]["$ai_parent_id"] == calls[3]["properties"]["$ai_span_id"]
     assert "$ai_span_id" in calls[2]["properties"]
-    assert calls[2]["properties"]["$ai_span_name"] == START
-    assert calls[2]["properties"]["$ai_parent_id"] == trace_props["$ai_trace_id"]
-    assert calls[2]["properties"]["$ai_input_state"] == initial_state
-    assert calls[2]["properties"]["$ai_output_state"] == initial_state
-
-    # Fourth span, save the value of fake_plain during its execution
-    second_state = {
-        "messages": [HumanMessage(content="What's a bar?"), AIMessage(content="Let's explore bar.")],
-        "xyz": "abc",
-    }
+    assert calls[2]["properties"]["$ai_span_name"] == "FakeMessagesListChatModel"
 
-    assert calls[4]["event"] == "$ai_span"
-    assert calls[4]["properties"]["$ai_parent_id"] == calls[5]["properties"]["$ai_span_id"]
+    # 4. Span - RunnableSequence within fake_llm
+    assert calls[3]["event"] == "$ai_span"
+    assert calls[3]["properties"]["$ai_parent_id"] == calls[4]["properties"]["$ai_span_id"]
     assert "$ai_span_id" in calls[3]["properties"]
-    assert calls[4]["properties"]["$ai_input_state"] == second_state
-    assert calls[4]["properties"]["$ai_output_state"] == second_state
+    assert calls[3]["properties"]["$ai_span_name"] == "RunnableSequence"
 
-    # Fifth span, run the fake_plain node
-    assert calls[5]["event"] == "$ai_span"
+    # 5. Span - the fake_llm node
+    assert calls[4]["event"] == "$ai_span"
+    assert calls[4]["properties"]["$ai_parent_id"] == trace_props["$ai_trace_id"]
     assert "$ai_span_id" in calls[4]["properties"]
-    assert calls[5]["properties"]["$ai_span_name"] == "fake_plain"
-    assert calls[5]["properties"]["$ai_parent_id"] == trace_props["$ai_trace_id"]
-    assert calls[5]["properties"]["$ai_input_state"] == initial_state
-    assert calls[5]["properties"]["$ai_output_state"] == second_state
-
-    # Sixth span, chat prompt template
-    assert calls[6]["event"] == "$ai_span"
-    assert calls[6]["properties"]["$ai_parent_id"] == calls[8]["properties"]["$ai_span_id"]
-    assert "$ai_span_id" in calls[6]["properties"]
-    assert calls[6]["properties"]["$ai_span_name"] == "ChatPromptTemplate"
-
-    # 7. Generation, fake_llm
-    assert calls[7]["event"] == "$ai_generation"
-    assert calls[7]["properties"]["$ai_parent_id"] == calls[8]["properties"]["$ai_span_id"]
-    assert "$ai_span_id" in calls[7]["properties"]
-    assert calls[7]["properties"]["$ai_span_name"] == "FakeMessagesListChatModel"
-
-    # 8. Span, RunnableSequence
-    assert calls[8]["event"] == "$ai_span"
-    assert calls[8]["properties"]["$ai_parent_id"] == calls[10]["properties"]["$ai_span_id"]
-    assert "$ai_span_id" in calls[8]["properties"]
-    assert calls[8]["properties"]["$ai_span_name"] == "RunnableSequence"
-
-    # 9. Span, fake_llm write
-    assert calls[9]["event"] == "$ai_span"
-    assert calls[9]["properties"]["$ai_parent_id"] == calls[10]["properties"]["$ai_span_id"]
-    assert "$ai_span_id" in calls[9]["properties"]
-
-    # 10. Span, fake_llm node
-    assert calls[10]["event"] == "$ai_span"
-    assert calls[10]["properties"]["$ai_parent_id"] == trace_props["$ai_trace_id"]
-    assert "$ai_span_id" in calls[10]["properties"]
-    assert calls[10]["properties"]["$ai_span_name"] == "fake_llm"
-
-    # 11. Trace
+    assert calls[4]["properties"]["$ai_span_name"] == "fake_llm"
+
+    # 6. Trace
     assert trace_args["event"] == "$ai_trace"
     assert trace_props["$ai_span_name"] == "LangGraph"