Remove run_id from GenAI Utils (#228)

* lint updates, test fixes
* remove auto generation of agent_id
* resolve rebase conflicts
* restore context clarification and tests
* clean up dead code in llamaindex, align agent_id references

Author: keith-decker

instrumentation-genai/opentelemetry-instrumentation-langchain/src/opentelemetry/instrumentation/langchain/callback_handler.py

@@ -7,7 +7,7 @@
 from __future__ import annotations
 
 import json
-from typing import Any, Optional, List
+from typing import Any, Optional, List, Dict
 from uuid import UUID
 
 from langchain_core.callbacks import BaseCallbackHandler
@@ -296,13 +296,48 @@ def _extract_tool_details(
     }
 
 
+def _agent_span_id(agent: AgentInvocation) -> Optional[str]:
+    """Return the agent's span ID as a hex string, if available."""
+    if agent.span_id is not None:
+        return f"{agent.span_id:016x}"
+    return None
+
+
+class _InvocationManager:
+    """Local store mapping LangChain run_ids to GenAI invocation objects.
+
+    Keyed by LangChain's own run_id UUID (from callback API), not the
+    GenAI run_id field. Tracks parent-child relationships so the parent
+    chain can be walked without a central registry.
+    """
+
+    def __init__(self) -> None:
+        self._invocations: Dict[UUID, Any] = {}
+        self._parents: Dict[UUID, Optional[UUID]] = {}
+
+    def add(self, run_id: UUID, parent_run_id: Optional[UUID], invocation: Any) -> None:
+        self._invocations[run_id] = invocation
+        self._parents[run_id] = parent_run_id
+
+    def get(self, run_id: UUID) -> Any:
+        return self._invocations.get(run_id)
+
+    def get_parent_id(self, run_id: UUID) -> Optional[UUID]:
+        return self._parents.get(run_id)
+
+    def remove(self, run_id: UUID) -> None:
+        self._invocations.pop(run_id, None)
+        self._parents.pop(run_id, None)
+
+
 class LangchainCallbackHandler(BaseCallbackHandler):
     def __init__(
         self,
         telemetry_handler: Optional[TelemetryHandler] = None,
     ) -> None:
         super().__init__()
         self._handler = telemetry_handler
+        self._invocation_manager = _InvocationManager()
         # Tracks ContextVar state before we push an inferred conversation_id
         # so it can be restored when the root entity finishes.
         self._inferred_context_prev: dict[UUID, GenAIContext | None] = {}
@@ -321,12 +356,12 @@ def _find_nearest_agent(self, run_id: Optional[UUID]) -> Optional[AgentInvocatio
         visited = set()
         while current is not None and current not in visited:
             visited.add(current)
-            entity = self._handler.get_entity(current)
+            entity = self._invocation_manager.get(current)
             if isinstance(entity, AgentInvocation):
                 return entity
             if entity is None:
                 break
-            current = getattr(entity, "parent_run_id", None)
+            current = self._invocation_manager.get_parent_id(current)
         return None
 
     def _start_agent_invocation(
@@ -344,12 +379,10 @@ def _start_agent_invocation(
     ) -> AgentInvocation:
         agent = AgentInvocation(
             name=name,
-            run_id=run_id,
             attributes=attrs,
         )
         agent.input_messages = command_input_messages or _make_input_message(inputs)
         agent.agent_name = _safe_str(agent_name) if agent_name else name
-        agent.parent_run_id = parent_run_id
         agent.framework = "langchain"
         if conversation_id:
             agent.conversation_id = conversation_id
@@ -361,6 +394,7 @@ def _start_agent_invocation(
             if metadata.get("system"):
                 agent.system = _safe_str(metadata["system"])
         self._handler.start_agent(agent)
+        self._invocation_manager.add(run_id, parent_run_id, agent)
         return agent
 
     def on_chain_start(
@@ -436,7 +470,7 @@ def on_chain_start(
                 if is_resume:
                     agent.attributes[GEN_AI_WORKFLOW_COMMAND] = "resume"
             else:
-                wf = Workflow(name=name, run_id=run_id, attributes=attrs)
+                wf = Workflow(name=name, attributes=attrs)
                 wf.input_messages = command_input_messages or _make_input_message(
                     inputs
                 )
@@ -445,12 +479,16 @@ def on_chain_start(
                 if is_resume:
                     wf.attributes[GEN_AI_WORKFLOW_COMMAND] = "resume"
                 self._handler.start_workflow(wf)
+                self._invocation_manager.add(run_id, None, wf)
             return
         else:
             # Skip if parent entity no longer exists (e.g., LangGraph
             # replays the interrupted node during resume — the parent
             # workflow from the previous trace was already ended).
-            if self._handler.get_entity(parent_run_id) is None:
+            # TODO: _invocation_manager is per-handler instance; consider
+            #  making LangchainCallbackHandler a singleton so all
+            #  invocations share the same manager.
+            if self._invocation_manager.get(parent_run_id) is None:
                 return
 
             context_agent = self._find_nearest_agent(parent_run_id)
@@ -477,7 +515,7 @@ def on_chain_start(
                     return
             tool_info = _extract_tool_details(metadata)
             if tool_info is not None:
-                existing = self._handler.get_entity(run_id)
+                existing = self._invocation_manager.get(run_id)
                 if isinstance(existing, ToolCall):
                     tool = existing
                     if context_agent is not None:
@@ -487,7 +525,7 @@ def on_chain_start(
                         if not getattr(tool, "agent_name", None):
                             tool.agent_name = _safe_str(agent_name_value)
                         if not getattr(tool, "agent_id", None):
-                            tool.agent_id = str(context_agent.run_id)
+                            tool.agent_id = _agent_span_id(context_agent)
                 else:
                     # Filter out tool-specific metadata from attributes
                     # since they're stored in dedicated fields
@@ -505,15 +543,14 @@ def on_chain_start(
                         name=tool_info.get("name", name),
                         id=tool_info.get("id"),
                         arguments=arguments,
-                        run_id=run_id,
-                        parent_run_id=parent_run_id,
                         attributes=tool_attrs,
                     )
                     tool.framework = "langchain"
                     if context_agent is not None and context_agent_name is not None:
                         tool.agent_name = context_agent_name
-                        tool.agent_id = str(context_agent.run_id)
+                        tool.agent_id = _agent_span_id(context_agent)
                     self._handler.start_tool_call(tool)
+                    self._invocation_manager.add(run_id, parent_run_id, tool)
                 if inputs is not None and getattr(tool, "arguments", None) is None:
                     tool.arguments = inputs
                 if getattr(tool, "arguments", None) is not None:
@@ -523,16 +560,15 @@ def on_chain_start(
             else:
                 step = Step(
                     name=name,
-                    run_id=run_id,
-                    parent_run_id=parent_run_id,
                     step_type="chain",
                     attributes=attrs,
                 )
                 if context_agent is not None:
                     if context_agent_name is not None:
                         step.agent_name = context_agent_name
-                    step.agent_id = str(context_agent.run_id)
+                    step.agent_id = _agent_span_id(context_agent)
                 self._handler.start_step(step)
+                self._invocation_manager.add(run_id, parent_run_id, step)
 
     def on_chain_end(
         self,
@@ -542,7 +578,7 @@ def on_chain_end(
         parent_run_id: Optional[UUID] = None,
         **_kwargs: Any,
     ) -> None:
-        entity = self._handler.get_entity(run_id)
+        entity = self._invocation_manager.get(run_id)
         if entity is None:
             return
 
@@ -571,6 +607,7 @@ def on_chain_end(
             if serialized is not None:
                 entity.attributes.setdefault("tool.response", serialized)
             self._handler.stop_tool_call(entity)
+        self._invocation_manager.remove(run_id)
 
     def on_chat_model_start(
         self,
@@ -656,8 +693,6 @@ def on_chat_model_start(
             request_model=request_model,
             input_messages=input_messages,
             attributes=attrs,
-            run_id=run_id,
-            parent_run_id=parent_run_id,
         )
         if provider:
             inv.provider = provider
@@ -666,8 +701,9 @@ def on_chat_model_start(
             if context_agent is not None:
                 agent_name_value = context_agent.agent_name or context_agent.name
                 inv.agent_name = _safe_str(agent_name_value)
-                inv.agent_id = str(context_agent.run_id)
+                inv.agent_id = _agent_span_id(context_agent)
         self._handler.start_llm(inv)
+        self._invocation_manager.add(run_id, parent_run_id, inv)
 
     def on_llm_start(
         self,
@@ -690,7 +726,7 @@ def on_llm_start(
             metadata=metadata,
             **extra,
         )
-        inv = self._handler.get_entity(run_id)
+        inv = self._invocation_manager.get(run_id)
         if isinstance(inv, LLMInvocation):
             inv.operation = "generate_text"
 
@@ -702,7 +738,7 @@ def on_llm_end(
         parent_run_id: Optional[UUID] = None,
         **_kwargs: Any,
     ) -> None:
-        inv = self._handler.get_entity(run_id)
+        inv = self._invocation_manager.get(run_id)
         if not isinstance(inv, LLMInvocation):
             return
         generations = getattr(response, "generations", [])
@@ -755,6 +791,7 @@ def on_llm_end(
                     break
 
         self._handler.stop_llm(inv)
+        self._invocation_manager.remove(run_id)
 
     def on_tool_start(
         self,
@@ -794,7 +831,7 @@ def on_tool_start(
         else:
             id_value = None
         arguments: Any = inputs if inputs is not None else input_str
-        existing = self._handler.get_entity(run_id)
+        existing = self._invocation_manager.get(run_id)
         if isinstance(existing, ToolCall):
             if arguments is not None:
                 existing.arguments = arguments
@@ -807,29 +844,28 @@ def on_tool_start(
                 ):
                     existing.agent_name = context_agent_name
                 if not getattr(existing, "agent_id", None):
-                    existing.agent_id = str(context_agent.run_id)
+                    existing.agent_id = _agent_span_id(context_agent)
             if existing.framework is None:
                 existing.framework = "langchain"
             return
         tool = ToolCall(
             name=name,
             id=id_value,
             arguments=arguments,
-            run_id=run_id,
-            parent_run_id=parent_run_id,
             attributes=attrs,
         )
         tool.framework = "langchain"
         if context_agent is not None and context_agent_name is not None:
             tool.agent_name = context_agent_name
-            tool.agent_id = str(context_agent.run_id)
+            tool.agent_id = _agent_span_id(context_agent)
         if arguments is not None:
             serialized_args = _serialize(arguments)
             if serialized_args is not None:
                 tool.attributes.setdefault("tool.arguments", serialized_args)
         if inputs is None and input_str:
             tool.attributes.setdefault("tool.input_str", _safe_str(input_str))
         self._handler.start_tool_call(tool)
+        self._invocation_manager.add(run_id, parent_run_id, tool)
 
     def on_tool_end(
         self,
@@ -839,24 +875,36 @@ def on_tool_end(
         parent_run_id: Optional[UUID] = None,
         **_kwargs: Any,
     ) -> None:
-        tool = self._handler.get_entity(run_id)
+        tool = self._invocation_manager.get(run_id)
         if not isinstance(tool, ToolCall):
             return
         serialized = _serialize(output)
         if serialized is not None:
             tool.attributes.setdefault("tool.response", serialized)
         self._handler.stop_tool_call(tool)
+        self._invocation_manager.remove(run_id)
 
     def _fail(self, run_id: UUID, error: BaseException) -> None:
         classification = _classify_error(error)
-        self._handler.fail_by_run_id(
-            run_id,
-            GenAIError(
-                message=str(error),
-                type=type(error),
-                classification=classification,
-            ),
+        entity = self._invocation_manager.get(run_id)
+        if entity is None:
+            return
+        genai_error = GenAIError(
+            message=str(error),
+            type=type(error),
+            classification=classification,
         )
+        if isinstance(entity, Workflow):
+            self._handler.fail_workflow(entity, genai_error)
+        elif isinstance(entity, AgentInvocation):
+            self._handler.fail_agent(entity, genai_error)
+        elif isinstance(entity, Step):
+            self._handler.fail_step(entity, genai_error)
+        elif isinstance(entity, LLMInvocation):
+            self._handler.fail_llm(entity, genai_error)
+        elif isinstance(entity, ToolCall):
+            self._handler.fail_tool_call(entity, genai_error)
+        self._invocation_manager.remove(run_id)
 
     def on_llm_error(
         self,