add invoke agent

Author: wrisa

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

@@ -24,8 +24,14 @@
 from opentelemetry.instrumentation.langchain.invocation_manager import (
     _InvocationManager,
 )
+from opentelemetry.instrumentation.langchain.operation_mapping import (
+    OperationName,
+    classify_chain_run,
+    resolve_agent_name,
+)
 from opentelemetry.util.genai.handler import TelemetryHandler
 from opentelemetry.util.genai.invocation import (
+    AgentInvocation,
     InferenceInvocation,
     WorkflowInvocation,
 )
@@ -35,11 +41,6 @@
     OutputMessage,
     Text,
 )
-from opentelemetry.instrumentation.langchain.operation_mapping import (
-    OperationName,
-    classify_chain_run,
-    resolve_agent_name,
-)
 
 
 class OpenTelemetryLangChainCallbackHandler(BaseCallbackHandler):
@@ -72,21 +73,73 @@ def on_chain_start(
             workflow_name_override = (
                 metadata.get("workflow_name") if metadata else None
             )
-            wf = self._telemetry_handler.start_workflow(
+            workflow = self._telemetry_handler.start_workflow(
                 name=workflow_name_override or workflow_name
             )
-            self._invocation_manager.add_invocation_state(run_id, None, wf)
+            self._invocation_manager.add_invocation_state(
+                run_id, parent_run_id, workflow
+            )
         elif operation == OperationName.INVOKE_AGENT:
-            agent_name = resolve_agent_name(serialized, metadata, kwargs)
-
-            if metadata:
-                agent_id = metadata.get("agent_id")
-                agent_desc = metadata.get("agent_description")
+            # agent name passed by the user
+            suggested_agent_name = resolve_agent_name(
+                serialized, metadata, kwargs
+            )
+            # find if there is an agent already
+            agent_invocation = self._find_nearest_agent(parent_run_id)
+            agent_invocation_name = (
+                agent_invocation.agent_name if agent_invocation else None
+            )
+            if suggested_agent_name:
+                suggested_agent_name_lower = suggested_agent_name.lower()
+                agent_invocation_name_lower = (
+                    agent_invocation_name.lower()
+                    if agent_invocation_name
+                    else None
+                )
+                if suggested_agent_name_lower != agent_invocation_name_lower:
+                    agent = self._telemetry_handler.start_invoke_local_agent(
+                        provider=metadata.get("ls_provider", "unknown")
+                        if metadata
+                        else "unknown",
+                    )
+                    agent.agent_name = suggested_agent_name
 
-                for key in ("thread_id", "session_id", "conversation_id"):
-                    conv_id = metadata.get(key)
+                    if metadata:
+                        agent.agent_id = metadata.get("agent_id")
+                        agent.agent_description = metadata.get(
+                            "agent_description"
+                        )
 
-                provider = metadata.get("ls_provider", "unknown")
+                        for key in (
+                            "thread_id",
+                            "session_id",
+                            "conversation_id",
+                        ):
+                            conv_id = metadata.get(key)
+                            if conv_id:
+                                agent.conversation_id = conv_id
+                                break
+
+                    self._invocation_manager.add_invocation_state(
+                        run_id, parent_run_id, agent
+                    )
+                else:
+                    # We create invoke_agent span for the initial chain for agent. All follow-up chains invoked for agent invocation will not create agent span.
+                    self._invocation_manager.add_invocation_state(
+                        run_id, parent_run_id, None
+                    )
+            else:
+                # No agent name could be resolved; still register the run_id so that
+                # parent-child traversal (e.g. _find_nearest_agent) is not broken for
+                # any children of this node.
+                self._invocation_manager.add_invocation_state(
+                    run_id, parent_run_id, None
+                )
+        else:
+            # For unclassified chains, we still want to track them in the invocation manager to maintain the parent-child relationships, even though we won't create spans for them.
+            self._invocation_manager.add_invocation_state(
+                run_id, parent_run_id, None
+            )
 
     def on_chain_end(
         self,
@@ -98,9 +151,10 @@ def on_chain_end(
     ) -> Any:
         invocation = self._invocation_manager.get_invocation(run_id=run_id)
         if invocation is None or not isinstance(
-            invocation, WorkflowInvocation
+            invocation, (WorkflowInvocation, AgentInvocation)
         ):
             # If the invocation does not exist, we cannot set attributes or end it
+            self._invocation_manager.delete_invocation_state(run_id)
             return
 
         invocation.stop()
@@ -118,9 +172,10 @@ def on_chain_error(
     ) -> Any:
         invocation = self._invocation_manager.get_invocation(run_id=run_id)
         if invocation is None or not isinstance(
-            invocation, WorkflowInvocation
+            invocation, (WorkflowInvocation, AgentInvocation)
         ):
             # If the invocation does not exist, we cannot set attributes or end it
+            self._invocation_manager.delete_invocation_state(run_id)
             return
 
         invocation.fail(error)
@@ -349,3 +404,16 @@ def on_llm_error(
         llm_invocation.fail(error)
         if not llm_invocation.span.is_recording():
             self._invocation_manager.delete_invocation_state(run_id=run_id)
+
+    def _find_nearest_agent(
+        self, run_id: Optional[UUID]
+    ) -> Optional[AgentInvocation]:
+        current = run_id
+        visited = set()
+        while current is not None and current not in visited:
+            visited.add(current)
+            entity = self._invocation_manager.get_invocation(current)
+            if isinstance(entity, AgentInvocation):
+                return entity
+            current = self._invocation_manager.get_parent_run_id(current)
+        return None

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

@@ -25,6 +25,8 @@
 class _InvocationState:
     invocation: GenAIInvocation
     children: List[UUID] = field(default_factory=lambda: list())
+    parent_run_id: Optional[UUID] = None
+    ended: bool = False
 
 
 class _InvocationManager:
@@ -39,23 +41,45 @@ def add_invocation_state(
         self,
         run_id: UUID,
         parent_run_id: Optional[UUID],
-        invocation: GenAIInvocation,
+        invocation: Optional[GenAIInvocation],
     ) -> None:
         invocation_state = _InvocationState(invocation=invocation)
-        self._invocations[run_id] = invocation_state
 
         if parent_run_id is not None and parent_run_id in self._invocations:
+            invocation_state.parent_run_id = parent_run_id
+
             parent_invocation_state = self._invocations[parent_run_id]
             parent_invocation_state.children.append(run_id)
 
+        self._invocations[run_id] = invocation_state
+
     def get_invocation(self, run_id: UUID) -> Optional[GenAIInvocation]:
         invocation_state = self._invocations.get(run_id)
         return invocation_state.invocation if invocation_state else None
 
+    def get_parent_run_id(self, run_id: UUID) -> Optional[UUID]:
+        invocation_state = self._invocations.get(run_id)
+        return invocation_state.parent_run_id if invocation_state else None
+
     def delete_invocation_state(self, run_id: UUID) -> None:
         invocation_state = self._invocations.get(run_id)
         if not invocation_state:
             return
-        for child_id in list(invocation_state.children):
-            self._invocations.pop(child_id, None)
+
+        invocation_state.ended = True
+
+        # Defer removal if any children are still live, so upward traversal
+        # (e.g. _find_nearest_agent) can still walk through this node.
+        if any(c in self._invocations for c in invocation_state.children):
+            return
+
         self._invocations.pop(run_id, None)
+
+        # Propagate cleanup upward: if the parent has already ended and has no
+        # more live children, it can now be removed too.
+        if invocation_state.parent_run_id:
+            parent_state = self._invocations.get(
+                invocation_state.parent_run_id
+            )
+            if parent_state is not None and parent_state.ended:
+                self.delete_invocation_state(invocation_state.parent_run_id)

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

@@ -47,10 +47,7 @@
 class OperationName:
     """Canonical GenAI semantic convention operation names."""
 
-    CHAT: str = GenAI.GenAiOperationNameValues.CHAT.value
-    TEXT_COMPLETION: str = GenAI.GenAiOperationNameValues.TEXT_COMPLETION.value
     INVOKE_AGENT: str = GenAI.GenAiOperationNameValues.INVOKE_AGENT.value
-    EXECUTE_TOOL: str = GenAI.GenAiOperationNameValues.EXECUTE_TOOL.value
     # invoke_workflow is not yet in the semconv enum; use the expected
     # string value so the mapping is forward-compatible.
     INVOKE_WORKFLOW: str = "invoke_workflow"
@@ -101,7 +98,7 @@ def resolve_agent_name(
     if name:
         return str(name)
 
-    name = serialized.get("name")
+    name = serialized.get("name") if serialized else None
     if name:
         return str(name)
 
@@ -124,30 +121,6 @@ def _has_agent_signals(metadata: Optional[dict[str, Any]]) -> bool:
     )
 
 
-def _is_langgraph_agent_node(
-    serialized: dict[str, Any],
-    metadata: Optional[dict[str, Any]],
-    kwargs: dict[str, Any],
-) -> bool:
-    """Detect a LangGraph agent node that is not a start/middleware node."""
-    if not metadata:
-        return False
-
-    node = metadata.get(LANGGRAPH_NODE_KEY)
-    if not node:
-        return False
-
-    # Exclude start and middleware nodes.
-    if node == LANGGRAPH_START_NODE:
-        return False
-
-    name = resolve_agent_name(serialized, metadata, kwargs)
-    if name and name.startswith(MIDDLEWARE_PREFIX):
-        return False
-
-    return True
-
-
 def _looks_like_workflow(
     serialized: dict[str, Any],
     metadata: Optional[dict[str, Any]],
@@ -162,13 +135,16 @@ def _looks_like_workflow(
         return True
 
     # Heuristic: check for LangGraph identifier in the serialized repr.
-    name = serialized.get("name", "")
-    graph_id = (
-        serialized.get("graph", {}).get("id", "")
-        if isinstance(serialized.get("graph"), dict)
-        else ""
-    )
-    return LANGGRAPH_IDENTIFIER in name or LANGGRAPH_IDENTIFIER in graph_id
+    if serialized:
+        name = serialized.get("name", "")
+        graph_id = (
+            serialized.get("graph", {}).get("id", "")
+            if isinstance(serialized.get("graph"), dict)
+            else ""
+        )
+        return LANGGRAPH_IDENTIFIER in name or LANGGRAPH_IDENTIFIER in graph_id
+
+    return True
 
 
 # ---------------------------------------------------------------------------
@@ -179,7 +155,6 @@ def _looks_like_workflow(
 def should_ignore_chain(
     metadata: Optional[dict[str, Any]],
     agent_name: Optional[str],
-    parent_run_id: Optional[UUID],
     kwargs: dict[str, Any],
 ) -> bool:
     """Return True if the chain callback should be silently suppressed.
@@ -238,19 +213,16 @@ def classify_chain_run(
     agent_name = resolve_agent_name(serialized, metadata, kwargs)
 
     # 1. Suppress known noise.
-    if should_ignore_chain(metadata, agent_name, parent_run_id, kwargs):
+    if should_ignore_chain(metadata, agent_name, kwargs):
         return None
 
     # 2. Agent detection.
     if _has_agent_signals(metadata):
         return OperationName.INVOKE_AGENT
 
-    if _is_langgraph_agent_node(serialized, metadata, kwargs):
-        return OperationName.INVOKE_AGENT
-
     # 3. Workflow / orchestration detection.
     if _looks_like_workflow(serialized, metadata, parent_run_id):
         return OperationName.INVOKE_WORKFLOW
 
     # 4. Default: suppress unclassified chains.
-    return None
+    return None

instrumentation-genai/opentelemetry-instrumentation-langchain/tests/conftest.py

@@ -208,9 +208,10 @@ def deserialize(cassette_string):
         return yaml.load(cassette_string, Loader=yaml.Loader)
 
 
-@pytest.fixture(scope="module", autouse=True)
+@pytest.fixture(scope="function", autouse=True)
 def fixture_vcr(vcr):
-    vcr.register_serializer("yaml", PrettyPrintJSONBody)
+    if vcr is not None:
+        vcr.register_serializer("yaml", PrettyPrintJSONBody)
     return vcr