fix: correctly capture a parent id when a custom trace_id is set

Author: skoob13

posthog/ai/langchain/callbacks.py

@@ -5,7 +5,6 @@
 
 import logging
 import time
-import uuid
 from dataclasses import dataclass
 from typing import (
     Any,
@@ -402,9 +401,17 @@ def _pop_run_metadata(self, run_id: UUID) -> Optional[RunMetadataUnion]:
     def _get_trace_id(self, run_id: UUID):
         trace_id = self._trace_id or self._find_root_run(run_id)
         if not trace_id:
-            trace_id = uuid.uuid4()
+            return run_id
         return trace_id
 
+    def _get_parent_run_id(self, trace_id: Any, run_id: UUID, parent_run_id: Optional[UUID]):
+        """
+        Replace the parent run ID with the trace ID for second level runs when a custom trace ID is set.
+        """
+        if parent_run_id is not None and parent_run_id not in self._parent_tree:
+            return trace_id
+        return parent_run_id
+
     def _pop_run_and_capture_trace_or_span(self, run_id: UUID, parent_run_id: Optional[UUID], outputs: Any):
         trace_id = self._get_trace_id(run_id)
         self._pop_parent_of_run(run_id)
@@ -414,7 +421,9 @@ def _pop_run_and_capture_trace_or_span(self, run_id: UUID, parent_run_id: Option
         if isinstance(run, GenerationMetadata):
             log.warning(f"Run {run_id} is not a generation but attempted to be captured as a trace or span.")
             return
-        self._capture_trace_or_span(trace_id, run_id, run, outputs, parent_run_id)
+        self._capture_trace_or_span(
+            trace_id, run_id, run, outputs, self._get_parent_run_id(trace_id, run_id, parent_run_id)
+        )
 
     def _capture_trace_or_span(
         self,
@@ -462,7 +471,9 @@ def _pop_run_and_capture_generation(
         if not isinstance(run, GenerationMetadata):
             log.warning(f"Run {run_id} is not a generation but attempted to be captured as a generation.")
             return
-        self._capture_generation(trace_id, run_id, run, response, parent_run_id)
+        self._capture_generation(
+            trace_id, run_id, run, response, self._get_parent_run_id(trace_id, run_id, parent_run_id)
+        )
 
     def _capture_generation(
         self,

posthog/test/ai/langchain/test_callbacks.py

@@ -1265,3 +1265,27 @@ def get_weather(city: Literal["nyc", "sf"]):
     assert len([call for call in calls if call["event"] == "$ai_generation"]) == 2
     assert len([call for call in calls if call["event"] == "$ai_span"]) == 18
     assert len([call for call in calls if call["event"] == "$ai_trace"]) == 1
+
+
+@pytest.mark.parametrize("trace_id", ["test-trace-id", None])
+def test_span_set_parent_ids(mock_client, trace_id):
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", "You are a helpful assistant."),
+            ("user", "Who won the world series in 2020?"),
+        ]
+    )
+    model = FakeMessagesListChatModel(
+        responses=[AIMessage(content="The Los Angeles Dodgers won the World Series in 2020.")]
+    )
+    callbacks = [CallbackHandler(mock_client, trace_id=trace_id)]
+    chain = prompt | model
+    chain.invoke({}, config={"callbacks": callbacks})
+
+    assert mock_client.capture.call_count == 3
+
+    span_props = mock_client.capture.call_args_list[0][1]
+    assert span_props["properties"]["$ai_trace_id"] == span_props["properties"]["$ai_parent_id"]
+
+    generation_props = mock_client.capture.call_args_list[1][1]
+    assert generation_props["properties"]["$ai_trace_id"] == generation_props["properties"]["$ai_parent_id"]