From c7288f6f022bdfbb6b8355fd005cc83d95dc3c80 Mon Sep 17 00:00:00 2001 From: eternalcuriouslearner Date: Wed, 9 Sep 2026 21:57:23 -0400 Subject: [PATCH 1/8] Instrument LlamaIndex agents and workflows --- .../.changelog/668.added | 1 + .../.changelog/668.fixed | 1 + .../README.rst | 7 +- .../genai/llama_index/_handler.py | 579 ++++++++++++++++- .../tests/conformance/workflow.py | 102 +++ .../tests/test_agent.py | 590 +++++++++++++++++- .../tests/test_composition.py | 65 +- .../tests/test_conformance.py | 7 +- 8 files changed, 1325 insertions(+), 27 deletions(-) create mode 100644 instrumentation/opentelemetry-instrumentation-genai-llama-index/.changelog/668.added create mode 100644 instrumentation/opentelemetry-instrumentation-genai-llama-index/.changelog/668.fixed create mode 100644 instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/conformance/workflow.py diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/.changelog/668.added b/instrumentation/opentelemetry-instrumentation-genai-llama-index/.changelog/668.added new file mode 100644 index 000000000..2637cab07 --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/.changelog/668.added @@ -0,0 +1 @@ +Add tracing for LlamaIndex AgentWorkflow runs and member agent executions. diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/.changelog/668.fixed b/instrumentation/opentelemetry-instrumentation-genai-llama-index/.changelog/668.fixed new file mode 100644 index 000000000..37b4a3a59 --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/.changelog/668.fixed @@ -0,0 +1 @@ +Keep one invoke_agent span open across an AgentWorkflow member's tool loop, parent execute_tool spans to the requesting agent, and keep a handing-off agent's span open until every tool call of that turn ends. diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/README.rst b/instrumentation/opentelemetry-instrumentation-genai-llama-index/README.rst index 9f89244c0..73b769c42 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/README.rst +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/README.rst @@ -9,9 +9,10 @@ OpenTelemetry LlamaIndex Instrumentation This package contains OpenTelemetry instrumentation for `LlamaIndex `_. -It emits ``invoke_agent`` spans for LlamaIndex ``FunctionAgent`` and -``ReActAgent`` runs, and ``execute_tool`` spans when LlamaIndex executes -function tools. Model calls +It emits ``invoke_workflow`` spans for ``AgentWorkflow`` runs, +``invoke_agent`` spans for standalone and workflow-member ``FunctionAgent`` +and ``ReActAgent`` executions, and ``execute_tool`` spans when LlamaIndex +executes tools. Model calls delegated to provider SDKs are intentionally left to those SDKs' OpenTelemetry instrumentations. diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py index 87801c306..b2d0aed07 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py @@ -3,16 +3,21 @@ from __future__ import annotations +import contextvars import inspect from base64 import b64decode from binascii import Error as BinasciiError -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, MutableMapping, Sequence from contextvars import ContextVar, Token from mimetypes import guess_type from typing import Any, cast +from weakref import WeakKeyDictionary from llama_index.core.agent.workflow.base_agent import BaseWorkflowAgent +from llama_index.core.agent.workflow.multi_agent_workflow import AgentWorkflow from llama_index.core.agent.workflow.workflow_events import ( + AgentOutput, + AgentSetup, ToolCall, ToolCallResult, ) @@ -30,11 +35,14 @@ from llama_index.core.tools import BaseTool, FunctionTool, ToolOutput from pydantic import PrivateAttr +from opentelemetry.context import Context, attach, detach +from opentelemetry.trace import set_span_in_context from opentelemetry.util.genai.handler import TelemetryHandler from opentelemetry.util.genai.invocation import ( AgentInvocation, GenAIInvocation, ToolInvocation, + WorkflowInvocation, ) from opentelemetry.util.genai.types import ( BlobPart, @@ -56,6 +64,52 @@ _AGENT_TOOL_ATTRIBUTES: ContextVar[ dict[str, _ToolExecutionAttributes] | None ] = ContextVar("llama_index_agent_tool_attributes", default=None) +_ACTIVE_WORKFLOW_TOOL: ContextVar[tuple[str, ToolInvocation] | None] = ( + ContextVar("llama_index_active_workflow_tool", default=None) +) + + +_MEMBER_AGENT_CONTEXTS: MutableMapping[ + AgentInvocation, contextvars.Context +] = WeakKeyDictionary() + + +def _start_member_agent( + start: Callable[[], AgentInvocation], +) -> AgentInvocation: + """Open a member-agent invocation inside a context this module owns. + + The invocation stays open across workflow steps, and each step runs in its + own asyncio task. ``TelemetryHandler`` attaches the span to whatever context + is current when the invocation starts, and that attachment can only be undone + from the same context -- so the context is kept and reused to finish in. + """ + context = contextvars.copy_context() + invocation = context.run(start) + _MEMBER_AGENT_CONTEXTS[invocation] = context + return invocation + + +def _finish_member_agent( + invocation: AgentInvocation, error: BaseException | None = None +) -> None: + """Finish a member-agent invocation in the context that started it.""" + + def finish() -> None: + if error is None: + invocation.stop() + else: + invocation.fail(error) + + context = _MEMBER_AGENT_CONTEXTS.pop(invocation, None) + if context is None: + finish() + return + try: + context.run(finish) + except RuntimeError: + # Already entered further up the stack; finishing is idempotent. + finish() def _method_name(span_id: str) -> str: @@ -216,6 +270,25 @@ def _agent_input(bound_args: inspect.BoundArguments) -> list[InputMessage]: return messages +def _agent_step_input( + event: AgentSetup, system_prompt: str | None +) -> list[InputMessage]: + """Recover the member agent input from an AgentWorkflow step. + + AgentWorkflow prepends the member's system prompt to ``AgentSetup.input``; + it is captured separately as the agent's system instruction. + """ + messages = list(event.input) + if ( + system_prompt + and messages + and messages[0].role.value == "system" + and messages[0].content == system_prompt + ): + messages.pop(0) + return [_input_message(message) for message in messages] + + def _request_model(agent: BaseWorkflowAgent) -> str | None: """Best-effort extraction of the model name across LLM integrations.""" try: @@ -316,6 +389,37 @@ def _tool_definitions(agent: BaseWorkflowAgent) -> list[ToolDefinition] | None: return definitions or None +def _workflow_tool_definitions( + workflow: AgentWorkflow, agent: BaseWorkflowAgent +) -> list[ToolDefinition] | None: + """Capture static tools plus AgentWorkflow's generated handoff tool. + + ``get_tools`` may perform retrieval, so observe only the workflow's + deterministic handoff resolution here rather than triggering another + retrieval pass solely for telemetry. + """ + tools_value: object = getattr(agent, "tools", None) + candidates: list[object] = ( + list(cast(Sequence[object], tools_value)) + if isinstance(tools_value, Sequence) + else [] + ) + get_handoff_tool = getattr(workflow, "_get_handoff_tool", None) + if callable(get_handoff_tool): + try: + handoff_tool = get_handoff_tool(agent) + except Exception: + handoff_tool = None + if handoff_tool is not None: + candidates.append(handoff_tool) + definitions = [ + definition + for candidate in candidates + if (definition := _tool_definition(candidate)) is not None + ] + return definitions or None + + def _set_agent_output(invocation: AgentInvocation, result: Any) -> None: """Copy the final chat response out of LlamaIndex's workflow result.""" output = getattr(result, "result", None) @@ -324,6 +428,34 @@ def _set_agent_output(invocation: AgentInvocation, result: Any) -> None: invocation.output_messages = [_output_message(response)] +def _set_agent_step_output(invocation: AgentInvocation, result: Any) -> None: + """Copy a member agent's response out of an AgentWorkflow step.""" + if isinstance(result, AgentOutput): + invocation.output_messages = [_output_message(result.response)] + + +def _agent_step_is_complete(result: Any) -> bool: + """Return whether a workflow agent step produced a final response. + + ``response.blocks`` does not reliably contain the tool selections that + drive the next workflow step; ``AgentOutput`` exposes those selections and + retry messages explicitly. + """ + if not isinstance(result, AgentOutput): + return False + if result.retry_messages: + return False + return not result.tool_calls + + +def _set_workflow_output(invocation: WorkflowInvocation, result: Any) -> None: + """Copy the final response out of an AgentWorkflow stop event.""" + output = getattr(result, "result", None) + response = getattr(output, "response", None) + if isinstance(response, ChatMessage): + invocation.output_messages = [_output_message(response)] + + def _tool_arguments( tool: FunctionTool, bound_args: inspect.BoundArguments ) -> dict[str, Any]: @@ -362,6 +494,21 @@ class _LlamaIndexInvocation(BaseSpan): _tool_attributes_token: ( Token[dict[str, _ToolExecutionAttributes] | None] | None ) = PrivateAttr() + _workflow_agents: dict[str, BaseWorkflowAgent] = PrivateAttr() + _workflow_agents_by_run_id: dict[str, BaseWorkflowAgent] = PrivateAttr() + _workflow_invocations_by_run_id: dict[str, AgentInvocation] = PrivateAttr() + _workflow_invocations_by_key: dict[tuple[str, str], AgentInvocation] = ( + PrivateAttr() + ) + _workflow_agent_invocation: AgentInvocation | None = PrivateAttr() + _workflow_handoff: bool = PrivateAttr() + _workflow_agent_context_token: Token[Context] | None = PrivateAttr() + _workflow_tool_token: Token[tuple[str, ToolInvocation] | None] | None = ( + PrivateAttr() + ) + _workflow_run_id: str | None = PrivateAttr() + _workflow_tool_counts: dict[str, int] = PrivateAttr() + _workflow_pending_handoffs: dict[str, AgentInvocation] = PrivateAttr() def __init__( self, @@ -373,11 +520,87 @@ def __init__( dict[str, _ToolExecutionAttributes] | None ] | None = None, + workflow_tool_token: Token[tuple[str, ToolInvocation] | None] + | None = None, + workflow_agents: Mapping[str, BaseWorkflowAgent] | None = None, + workflow_run_id: str | None = None, + workflow_agent: BaseWorkflowAgent | None = None, + workflow_agent_invocation: AgentInvocation | None = None, + workflow_handoff: bool = False, ) -> None: """Create the adapter used by LlamaIndex's span-handler lifecycle.""" super().__init__(id_=id_, parent_id=parent_id) self._invocation = invocation self._tool_attributes_token = tool_attributes_token + self._workflow_tool_token = workflow_tool_token + self._workflow_run_id = workflow_run_id + self._workflow_agents = dict(workflow_agents or {}) + self._workflow_agents_by_run_id = {} + self._workflow_invocations_by_run_id = {} + self._workflow_invocations_by_key = {} + self._workflow_agent_invocation = workflow_agent_invocation + self._workflow_handoff = workflow_handoff + self._workflow_agent_context_token = None + self._workflow_tool_counts = {} + self._workflow_pending_handoffs = {} + if workflow_run_id is not None and workflow_agent is not None: + self.register_workflow_agent(workflow_run_id, workflow_agent) + if ( + workflow_run_id is not None + and workflow_agent_invocation is not None + ): + self._workflow_invocations_by_run_id[workflow_run_id] = ( + workflow_agent_invocation + ) + + def workflow_agent(self, name: str) -> BaseWorkflowAgent | None: + """Return a member agent owned by this workflow invocation.""" + return self._workflow_agents.get(name) + + def register_workflow_agent( + self, run_id: str, agent: BaseWorkflowAgent + ) -> None: + """Associate a workflow run with its currently executing agent.""" + self._workflow_agents_by_run_id[run_id] = agent + + def workflow_agent_for_run_id( + self, run_id: str | None + ) -> BaseWorkflowAgent | None: + """Return the agent executing the current step for a workflow run.""" + if run_id is None: + return None + return self._workflow_agents_by_run_id.get(run_id) + + def workflow_invocation_for_run_id( + self, run_id: str | None, agent_name: str | None = None + ) -> AgentInvocation | None: + """Return the reusable member-agent invocation for a workflow run.""" + if self._workflow_agent_invocation is not None: + return self._workflow_agent_invocation + if run_id is None: + return None + if agent_name is not None: + return self._workflow_invocations_by_key.get((run_id, agent_name)) + return self._workflow_invocations_by_run_id.get(run_id) + + def register_workflow_invocation( + self, run_id: str, agent_name: str, invocation: AgentInvocation + ) -> None: + """Keep one member-agent invocation open across workflow turns.""" + self._workflow_invocations_by_run_id[run_id] = invocation + self._workflow_invocations_by_key[(run_id, agent_name)] = invocation + + def remove_workflow_invocation(self, invocation: AgentInvocation) -> None: + """Forget a completed member invocation so a later turn can restart it.""" + for key, value in list(self._workflow_invocations_by_key.items()): + if value is invocation: + del self._workflow_invocations_by_key[key] + run_id = key[0] + if ( + self._workflow_invocations_by_run_id.get(run_id) + is invocation + ): + del self._workflow_invocations_by_run_id[run_id] def reset_tool_attributes(self) -> None: """Restore task-local tool metadata after an agent run finishes.""" @@ -388,6 +611,74 @@ def reset_tool_attributes(self) -> None: pass self._tool_attributes_token = None + def reset_workflow_tool(self) -> None: + """Stop exposing a workflow tool while its nested SDK call unwinds.""" + if self._workflow_tool_token is not None: + try: + _ACTIVE_WORKFLOW_TOOL.reset(self._workflow_tool_token) + except ValueError: + pass + self._workflow_tool_token = None + + def expect_workflow_tools(self, run_id: str, count: int) -> None: + """Record how many tool calls the agent's current turn requested. + + AgentWorkflow dispatches one ``ToolCall`` event per selection and runs + them as separate steps, so counting the ``call_tool`` spans that have + already opened would miss the ones still queued. + """ + if count: + self._workflow_tool_counts[run_id] = count + else: + self._workflow_tool_counts.pop(run_id, None) + + def release_workflow_tool(self, run_id: str | None) -> bool: + """Release one completed tool and report whether the turn is drained.""" + if run_id is None: + return False + remaining = self._workflow_tool_counts.get(run_id, 0) - 1 + if remaining > 0: + self._workflow_tool_counts[run_id] = remaining + return False + self._workflow_tool_counts.pop(run_id, None) + return True + + def set_pending_handoff( + self, run_id: str, invocation: AgentInvocation + ) -> None: + """Hold a handing-off agent open until its whole turn has drained.""" + self._workflow_pending_handoffs[run_id] = invocation + + def take_pending_handoff(self, run_id: str) -> AgentInvocation | None: + """Claim the handing-off agent owed a close, if there is one.""" + return self._workflow_pending_handoffs.pop(run_id, None) + + def activate_workflow_agent(self) -> None: + """Make a resumed member-agent span current for this workflow step.""" + if self._workflow_agent_context_token is None: + self._workflow_agent_context_token = attach( + set_span_in_context(self._invocation.span) + ) + + def reset_workflow_agent(self) -> None: + """Detach the temporary context used by a resumed agent step.""" + if self._workflow_agent_context_token is not None: + detach(self._workflow_agent_context_token) + self._workflow_agent_context_token = None + + def finalize_workflow_agents( + self, error: BaseException | None = None + ) -> None: + """Finish member-agent spans left open when the workflow terminates.""" + invocations: list[AgentInvocation] = [] + for candidate in self._workflow_invocations_by_key.values(): + if all(candidate is not existing for existing in invocations): + invocations.append(candidate) + for agent_invocation in invocations: + _finish_member_agent(agent_invocation, error) + self._workflow_invocations_by_key.clear() + self._workflow_invocations_by_run_id.clear() + class LlamaIndexSpanHandler(BaseSpanHandler[_LlamaIndexInvocation]): """Map LlamaIndex-owned agent and tool operations to GenAI spans.""" @@ -399,6 +690,16 @@ def __init__(self, handler: TelemetryHandler) -> None: super().__init__() self._handler = handler + def _is_open_tool(self, invocation: ToolInvocation) -> bool: + """Check that a task-local tool still belongs to this handler. + + ``BaseSpanHandler`` mutates ``open_spans`` under its lock from worker + threads, so iterating it unguarded can raise ``RuntimeError``. + """ + with self.lock: + adapters = list(self.open_spans.values()) + return any(adapter._invocation is invocation for adapter in adapters) + def new_span( self, id_: str, @@ -418,8 +719,36 @@ def new_span( tool_attributes_token: ( Token[dict[str, _ToolExecutionAttributes] | None] | None ) = None + workflow_agents: Mapping[str, BaseWorkflowAgent] | None = None + workflow_run_id: str | None = None + workflow_agent: BaseWorkflowAgent | None = None + workflow_agent_invocation: AgentInvocation | None = None + workflow_handoff = False + member_agent_step = False + workflow_tool_token: ( + Token[tuple[str, ToolInvocation] | None] | None + ) = None - if isinstance(instance, BaseWorkflowAgent) and method_name == "run": + if isinstance(instance, AgentWorkflow) and method_name == "run": + capture_content = self._handler.should_capture_content() + input_messages = ( + _agent_input(bound_args) if capture_content else [] + ) + workflow_agents = instance.agents + workflow_name = getattr(instance, "workflow_name", None) + default_workflow_name = ( + f"{type(instance).__module__}.{type(instance).__qualname__}" + ) + if ( + not isinstance(workflow_name, str) + or not workflow_name + or workflow_name == default_workflow_name + ): + workflow_name = type(instance).__name__ + workflow_invocation = self._handler.workflow(name=workflow_name) + workflow_invocation.input_messages = input_messages + invocation = workflow_invocation + elif isinstance(instance, BaseWorkflowAgent) and method_name == "run": capture_content = self._handler.should_capture_content() agent_name = instance.name or type(instance).__name__ request_model = _request_model(instance) @@ -429,7 +758,7 @@ def new_span( ) tool_definitions = _tool_definitions(instance) system_prompt = instance.system_prompt - system_instruction: list[SystemInstructionPart] = ( + agent_system_instruction: list[SystemInstructionPart] = ( [TextPart(content=system_prompt)] if capture_content and system_prompt else [] @@ -441,21 +770,129 @@ def new_span( agent_invocation.agent_description = agent_description agent_invocation.input_messages = input_messages agent_invocation.tool_definitions = tool_definitions - agent_invocation.system_instruction = system_instruction + agent_invocation.system_instruction = agent_system_instruction invocation = agent_invocation tool_attributes_token = _AGENT_TOOL_ATTRIBUTES.set( _agent_tool_attribute_map(instance) ) + elif method_name == "run_agent_step" and isinstance( + (agent_setup := bound_args.arguments.get("ev")), AgentSetup + ): + parent = self.open_spans.get(parent_span_id or "") + agent = ( + parent.workflow_agent(agent_setup.current_agent_name) + if parent is not None + else None + ) + if agent is None: + return None + capture_content = self._handler.should_capture_content() + agent_name = agent.name or type(agent).__name__ + request_model = _request_model(agent) + agent_description = agent.description + input_messages = ( + _agent_step_input(agent_setup, agent.system_prompt) + if capture_content + else [] + ) + tool_definitions = _workflow_tool_definitions( + cast(AgentWorkflow, instance), agent + ) + system_instruction: list[SystemInstructionPart] = ( + [TextPart(content=agent.system_prompt)] + if capture_content and agent.system_prompt + else [] + ) + workflow_run_id = ( + tags.get("llamaindex.run_id") if tags is not None else None + ) + workflow_agent_invocation = ( + parent.workflow_invocation_for_run_id( + workflow_run_id, agent_name + ) + if parent is not None + else None + ) + member_agent_step = True + if workflow_agent_invocation is None: + workflow_agent_invocation = _start_member_agent( + lambda: self._handler.invoke_local_agent( + request_model=request_model, + agent_name=agent_name, + ) + ) + workflow_agent_invocation.agent_description = agent_description + workflow_agent_invocation.input_messages = input_messages + workflow_agent_invocation.tool_definitions = tool_definitions + workflow_agent_invocation.system_instruction = ( + system_instruction + ) + if parent is not None and workflow_run_id is not None: + parent.register_workflow_invocation( + workflow_run_id, agent_name, workflow_agent_invocation + ) + agent_invocation = workflow_agent_invocation + invocation = agent_invocation + if parent is not None: + if workflow_run_id is not None: + parent.register_workflow_agent(workflow_run_id, agent) + workflow_agent = agent elif method_name == "call_tool" and isinstance( (tool_call := bound_args.arguments.get("ev")), ToolCall ): - tool_type, tool_description = _agent_tool_attributes( - instance or bound_args.arguments.get("self"), - tool_call.tool_name, + parent = self.open_spans.get(parent_span_id or "") + active_agent = ( + parent.workflow_agent_for_run_id( + tags.get("llamaindex.run_id") if tags is not None else None + ) + if parent is not None + else None ) - tool_invocation = self._handler.tool( - tool_call.tool_name, - tool_type=tool_type, + active_invocation = ( + parent.workflow_invocation_for_run_id( + tags.get("llamaindex.run_id") if tags is not None else None + ) + if parent is not None + else None + ) + if active_invocation is None and parent is not None: + if isinstance(parent._invocation, AgentInvocation): + active_invocation = parent._invocation + workflow_run_id = ( + tags.get("llamaindex.run_id") if tags is not None else None + ) + tool_type, tool_description = ( + _agent_tool_attributes(active_agent, tool_call.tool_name) + if active_agent is not None + else (None, None) + ) + if tool_type is None: + tool_type, tool_description = _agent_tool_attributes( + instance or bound_args.arguments.get("self"), + tool_call.tool_name, + ) + if tool_type is None and tool_call.tool_name == "handoff": + # AgentWorkflow's built-in handoff is emitted as a ToolCall, + # although its generated tool metadata is not available here. + tool_type = "function" + # The member-agent span stays open across workflow steps, each of + # which runs in its own asyncio task. Pass its context explicitly + # so the tool span nests under the agent that requested the call. + agent_context_token = ( + attach(set_span_in_context(active_invocation.span)) + if active_invocation is not None + else None + ) + try: + tool_invocation = self._handler.tool( + tool_call.tool_name, + tool_type=tool_type, + ) + finally: + if agent_context_token is not None: + detach(agent_context_token) + workflow_tool_token = _ACTIVE_WORKFLOW_TOOL.set( + (tool_call.tool_name, tool_invocation) ) tool_invocation.tool_call_id = tool_call.tool_id tool_invocation.tool_description = tool_description @@ -464,6 +901,11 @@ def new_span( dict[str, Any], cast(Any, tool_call).tool_kwargs ) invocation = tool_invocation + if parent is not None and isinstance( + parent._invocation, WorkflowInvocation + ): + workflow_agent_invocation = active_invocation + workflow_handoff = tool_call.tool_name == "handoff" elif isinstance(instance, FunctionTool) and method_name in { "call", "acall", @@ -471,10 +913,27 @@ def new_span( parent = self.open_spans.get(parent_span_id or "") # LlamaIndex reports an agent tool execution through both call_tool # and the nested FunctionTool.call/acall; the parent records it. - if parent is not None and isinstance( - parent._invocation, ToolInvocation + if ( + parent is not None + and parent._workflow_agent_invocation is not None + and isinstance(parent._invocation, ToolInvocation) ): + # The workflow callback identifies the tool by name only; the + # nested FunctionTool call is the authoritative executing tool. + parent._invocation.tool_description = ( + instance.metadata.description or None + ) return None + active_workflow_tool = _ACTIVE_WORKFLOW_TOOL.get() + if active_workflow_tool is not None: + active_tool_name, active_tool = active_workflow_tool + if active_tool_name == instance.metadata.get_name() and ( + self._is_open_tool(active_tool) + ): + active_tool.tool_description = ( + instance.metadata.description or None + ) + return None metadata = instance.metadata tool_invocation = self._handler.tool( metadata.get_name(), @@ -489,12 +948,66 @@ def new_span( else: return None - return _LlamaIndexInvocation( + adapter = _LlamaIndexInvocation( id_=id_, parent_id=parent_span_id, invocation=invocation, tool_attributes_token=tool_attributes_token, + workflow_tool_token=workflow_tool_token, + workflow_agents=workflow_agents, + workflow_run_id=workflow_run_id, + workflow_agent=workflow_agent, + workflow_agent_invocation=workflow_agent_invocation, + workflow_handoff=workflow_handoff, ) + if method_name == "run_agent_step" and member_agent_step: + adapter.activate_workflow_agent() + return adapter + + def _expect_workflow_tools( + self, span: _LlamaIndexInvocation, result: Any + ) -> None: + """Record the tool calls a member agent's turn just requested.""" + run_id = span._workflow_run_id + if run_id is None or not isinstance(result, AgentOutput): + return + parent = self.open_spans.get(span.parent_id or "") + if parent is not None: + parent.expect_workflow_tools(run_id, len(result.tool_calls)) + + def _release_workflow_invocation( + self, + span: _LlamaIndexInvocation, + invocation: AgentInvocation, + ) -> None: + """Drop a finished member invocation so a later turn opens a new span.""" + parent = self.open_spans.get(span.parent_id or "") + if parent is not None: + parent.remove_workflow_invocation(invocation) + + def _finish_workflow_tool(self, span: _LlamaIndexInvocation) -> None: + """Release one tool of a member agent's turn and close the agent last. + + AgentWorkflow reports a handoff as a tool call made by the agent that is + stepping down, and that turn can request other tools alongside it. The + agent's span has to outlive every one of them, so it is closed only once + the turn's last tool call ends. + """ + run_id = span._workflow_run_id + if run_id is None: + return + parent = self.open_spans.get(span.parent_id or "") + if parent is None: + return + invocation = span._workflow_agent_invocation + if span._workflow_handoff and invocation is not None: + parent.set_pending_handoff(run_id, invocation) + if not parent.release_workflow_tool(run_id): + return + pending = parent.take_pending_handoff(run_id) + if pending is not None: + _finish_member_agent(pending) + parent.remove_workflow_invocation(pending) def prepare_to_exit_span( self, @@ -512,11 +1025,28 @@ def prepare_to_exit_span( span = self.open_spans.get(id_) if span is None: return None - if isinstance(span._invocation, AgentInvocation): + if isinstance(span._invocation, WorkflowInvocation): + if self._handler.should_capture_content(): + _set_workflow_output(span._invocation, result) + span.finalize_workflow_agents() + elif isinstance(span._invocation, AgentInvocation): span.reset_tool_attributes() + span.reset_workflow_tool() if self._handler.should_capture_content(): - _set_agent_output(span._invocation, result) + if isinstance(result, AgentOutput): + _set_agent_step_output(span._invocation, result) + else: + _set_agent_output(span._invocation, result) + if span._workflow_agent_invocation is not None: + span.reset_workflow_agent() + if _agent_step_is_complete(result): + _finish_member_agent(span._invocation) + self._release_workflow_invocation(span, span._invocation) + else: + self._expect_workflow_tools(span, result) + return span elif isinstance(span._invocation, ToolInvocation): + span.reset_workflow_tool() tool_output: ToolOutput | None = None if isinstance(result, ToolCallResult): tool_output = result.tool_output @@ -535,8 +1065,11 @@ def prepare_to_exit_span( else RuntimeError(tool_output.content) ) span._invocation.fail(error) + self._finish_workflow_tool(span) return span span._invocation.stop() + if isinstance(span._invocation, ToolInvocation): + self._finish_workflow_tool(span) return span def prepare_to_drop_span( @@ -552,8 +1085,22 @@ def prepare_to_drop_span( if span is None: return None span.reset_tool_attributes() - if err is None: + span.reset_workflow_tool() + if isinstance(span._invocation, WorkflowInvocation): + span.finalize_workflow_agents(err) + if err is None: + span._invocation.stop() + else: + span._invocation.fail(err) + elif isinstance(span._invocation, AgentInvocation): + _finish_member_agent(span._invocation, err) + elif err is None: span._invocation.stop() else: span._invocation.fail(err) + if isinstance(span._invocation, AgentInvocation): + span.reset_workflow_agent() + self._release_workflow_invocation(span, span._invocation) + elif isinstance(span._invocation, ToolInvocation): + self._finish_workflow_tool(span) return span diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/conformance/workflow.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/conformance/workflow.py new file mode 100644 index 000000000..e443085ab --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/conformance/workflow.py @@ -0,0 +1,102 @@ +# Copyright The OpenTelemetry Authors +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +from typing import Any + +from llama_index.core.agent.workflow import ( + AgentWorkflow, + FunctionAgent, + ReActAgent, +) +from llama_index.core.base.llms.types import ToolCallBlock +from llama_index.core.llms import ChatMessage, MockFunctionCallingLLM + +from opentelemetry.instrumentation.genai.llama_index import ( + LlamaIndexInstrumentor, +) +from opentelemetry.sdk._logs import LoggerProvider +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.test_util_genai.conformance import Scenario +from opentelemetry.test_util_genai.instrumentor import instrument + + +class WorkflowScenario(Scenario): + expected_spans = { + "invoke_workflow": 1, + "invoke_agent": 2, + "execute_tool": 1, + } + expected_metrics = ("gen_ai.client.operation.duration",) + + def run( + self, + *, + tracer_provider: TracerProvider, + meter_provider: MeterProvider, + logger_provider: LoggerProvider, + vcr: Any, + ) -> None: + def function_response( + messages: list[ChatMessage], **kwargs: Any + ) -> ChatMessage: + return ChatMessage( + role="assistant", + blocks=[ + ToolCallBlock( + tool_call_id="handoff-call", + tool_name="handoff", + tool_kwargs={ + "to_agent": "react-member", + "reason": "The ReAct agent should answer.", + }, + ) + ], + ) + + def react_response( + messages: list[ChatMessage], **kwargs: Any + ) -> ChatMessage: + return ChatMessage( + role="assistant", + content="Thought: I can answer.\nAnswer: complete", + ) + + function_agent = FunctionAgent( + name="function-member", + description="Routes the request.", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=function_response, + ), + streaming=False, + ) + react_agent = ReActAgent( + name="react-member", + description="Answers the request.", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=react_response, + ), + streaming=False, + ) + workflow = AgentWorkflow( + agents=[function_agent, react_agent], + root_agent="function-member", + ) + + with instrument( + LlamaIndexInstrumentor(), + tracer_provider=tracer_provider, + logger_provider=logger_provider, + meter_provider=meter_provider, + content_capture="SPAN_ONLY", + ): + + async def run_workflow() -> None: + await workflow.run(user_msg="Complete the request") + + asyncio.run(run_workflow()) diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py index 42519c394..10d5e30ea 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py @@ -4,7 +4,8 @@ from __future__ import annotations import json -from typing import Any +import logging +from typing import Any, cast from unittest.mock import patch from uuid import UUID @@ -36,14 +37,16 @@ from openai import RateLimitError from opentelemetry.instrumentation.genai.llama_index._handler import ( + LlamaIndexSpanHandler, _agent_input, _input_message, + _LlamaIndexInvocation, _method_name, _output_message, _set_agent_output, _tool_definition, ) -from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider from opentelemetry.semconv._incubating.attributes import ( gen_ai_attributes as GenAIAttributes, ) @@ -51,6 +54,7 @@ error_attributes as ErrorAttributes, ) from opentelemetry.trace import SpanKind, StatusCode +from opentelemetry.util.genai.handler import TelemetryHandler from opentelemetry.util.genai.types import ( BlobPart, GenericToolDefinition, @@ -837,7 +841,7 @@ def response_generator(messages, **kwargs): @pytest.mark.asyncio -async def test_agent_workflow_emits_tool_span( +async def test_agent_workflow_emits_span_hierarchy( span_exporter, instrument_llama_index ) -> None: def echo(value: str) -> str: @@ -871,12 +875,315 @@ def response_generator(messages, **kwargs): result = await workflow.run(user_msg="Call echo") assert result.response.content == "done" + + workflow_span = _spans_named( + span_exporter, "invoke_workflow AgentWorkflow" + )[0] + workflow_attrs = dict(workflow_span.attributes or {}) + assert workflow_span.kind == SpanKind.INTERNAL + assert workflow_span.parent is None + assert ( + workflow_attrs[GenAIAttributes.GEN_AI_OPERATION_NAME] + == "invoke_workflow" + ) + assert ( + workflow_attrs[GenAIAttributes.GEN_AI_WORKFLOW_NAME] == "AgentWorkflow" + ) + + agent_spans = _spans_named(span_exporter, "invoke_agent workflow-agent") + assert len(agent_spans) == 1 + assert all( + span.parent is not None + and span.parent.span_id == workflow_span.context.span_id + and span.context.trace_id == workflow_span.context.trace_id + for span in agent_spans + ) + tool_spans = _spans_named(span_exporter, "execute_tool echo") assert len(tool_spans) == 1 - assert tool_spans[0].parent is None + assert tool_spans[0].parent is not None + assert tool_spans[0].parent.span_id == agent_spans[0].context.span_id + assert tool_spans[0].context.trace_id == workflow_span.context.trace_id FunctionTool.from_defaults(echo)(value="after workflow") - assert len(_spans_named(span_exporter, "execute_tool echo")) == 2 + tool_spans = _spans_named(span_exporter, "execute_tool echo") + assert len(tool_spans) == 2 + assert tool_spans[1].parent is None + + +@pytest.mark.asyncio +async def test_agent_workflow_uses_configured_workflow_name( + span_exporter, instrument_llama_index +) -> None: + def response_generator(messages, **kwargs): + return ChatMessage(role="assistant", content="workflow complete") + + agent = FunctionAgent( + name="named-workflow-agent", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=response_generator, + ), + streaming=False, + ) + workflow = AgentWorkflow( + agents=[agent], + workflow_name="customer-support-workflow", + ) + + await workflow.run(user_msg="Run the workflow") + + workflow_span = _spans_named( + span_exporter, "invoke_workflow customer-support-workflow" + )[0] + assert workflow_span.attributes[GenAIAttributes.GEN_AI_WORKFLOW_NAME] == ( + "customer-support-workflow" + ) + + +@pytest.mark.asyncio +async def test_agent_workflow_uses_executing_agent_tool_metadata( + span_exporter, instrument_llama_index +) -> None: + def response_generator(messages, **kwargs): + if any(message.role.value == "tool" for message in messages): + return ChatMessage(role="assistant", content="workflow complete") + return ChatMessage( + role="assistant", + blocks=[ + ToolCallBlock( + tool_call_id="duplicate-tool-call", + tool_name="lookup", + tool_kwargs={"value": "hello"}, + ) + ], + ) + + def first_lookup(value: str) -> str: + return f"first: {value}" + + class GenericLookupTool(AsyncBaseTool): + @property + def metadata(self) -> ToolMetadata: + return ToolMetadata( + name="lookup", description="Second lookup description." + ) + + def call(self, value: str) -> ToolOutput: + return ToolOutput( + tool_name="lookup", + content=f"second: {value}", + raw_input={"value": value}, + raw_output=value, + ) + + async def acall(self, value: str) -> ToolOutput: + return self.call(value) + + first_agent = FunctionAgent( + name="first-agent", + description="First agent.", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=lambda messages, **kwargs: ChatMessage( + role="assistant", content="first complete" + ), + ), + tools=[ + FunctionTool.from_defaults( + first_lookup, + name="lookup", + description="First lookup description.", + ) + ], + streaming=False, + ) + second_agent = FunctionAgent( + name="second-agent", + description="Second agent.", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=response_generator, + ), + tools=[GenericLookupTool()], + streaming=False, + ) + workflow = AgentWorkflow( + agents=[first_agent, second_agent], + root_agent="second-agent", + ) + + await workflow.run(user_msg="Use lookup") + + tool_span = _spans_named(span_exporter, "execute_tool lookup")[0] + assert tool_span.attributes[GenAIAttributes.GEN_AI_TOOL_DESCRIPTION] == ( + "Second lookup description." + ) + assert tool_span.attributes[GenAIAttributes.GEN_AI_TOOL_TYPE] == ( + "GenericLookupTool" + ) + + +@pytest.mark.asyncio +async def test_agent_workflow_captures_content( + span_exporter, instrument_llama_index_with_content +) -> None: + def response_generator(messages, **kwargs): + return ChatMessage(role="assistant", content="workflow complete") + + agent = FunctionAgent( + name="content-agent", + system_prompt="Answer briefly.", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=response_generator, + ), + streaming=False, + ) + workflow = AgentWorkflow(agents=[agent]) + + await workflow.run(user_msg="Run the workflow") + + workflow_span = _spans_named( + span_exporter, "invoke_workflow AgentWorkflow" + )[0] + agent_span = _spans_named(span_exporter, "invoke_agent content-agent")[0] + workflow_attrs = dict(workflow_span.attributes or {}) + agent_attrs = dict(agent_span.attributes or {}) + assert json.loads( + workflow_attrs[GenAIAttributes.GEN_AI_INPUT_MESSAGES] + ) == [ + { + "role": "user", + "parts": [{"type": "text", "content": "Run the workflow"}], + "name": None, + } + ] + assert json.loads(workflow_attrs[GenAIAttributes.GEN_AI_OUTPUT_MESSAGES])[ + 0 + ]["parts"] == [{"type": "text", "content": "workflow complete"}] + assert json.loads(agent_attrs[GenAIAttributes.GEN_AI_INPUT_MESSAGES]) == [ + { + "role": "user", + "parts": [{"type": "text", "content": "Run the workflow"}], + "name": None, + } + ] + assert json.loads( + agent_attrs[GenAIAttributes.GEN_AI_SYSTEM_INSTRUCTIONS] + ) == [{"type": "text", "content": "Answer briefly."}] + assert json.loads(agent_attrs[GenAIAttributes.GEN_AI_OUTPUT_MESSAGES])[0][ + "parts" + ] == [{"type": "text", "content": "workflow complete"}] + + +@pytest.mark.asyncio +async def test_agent_workflow_error_marks_workflow_and_agent_spans( + span_exporter, instrument_llama_index +) -> None: + error = RuntimeError("workflow agent failed") + + def response_generator(messages, **kwargs): + raise error + + agent = FunctionAgent( + name="failing-workflow-agent", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=response_generator, + ), + streaming=False, + ) + workflow = AgentWorkflow(agents=[agent]) + + with pytest.raises(RuntimeError) as caught: + await workflow.run(user_msg="Fail") + assert caught.value is error + + workflow_span = _spans_named( + span_exporter, "invoke_workflow AgentWorkflow" + )[0] + agent_span = _spans_named( + span_exporter, "invoke_agent failing-workflow-agent" + )[0] + for span in (workflow_span, agent_span): + assert span.status.status_code == StatusCode.ERROR + assert span.attributes[ErrorAttributes.ERROR_TYPE] == "RuntimeError" + + +@pytest.mark.asyncio +async def test_agent_workflow_instruments_function_and_react_members( + span_exporter, instrument_llama_index +) -> None: + def function_response(messages, **kwargs): + return ChatMessage( + role="assistant", + blocks=[ + ToolCallBlock( + tool_call_id="handoff-call", + tool_name="handoff", + tool_kwargs={ + "to_agent": "react-member", + "reason": "The ReAct agent should answer.", + }, + ) + ], + ) + + def react_response(messages, **kwargs): + return ChatMessage( + role="assistant", + content="Thought: I can answer.\nAnswer: complete", + ) + + function_agent = FunctionAgent( + name="function-member", + description="Routes the request.", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=function_response, + ), + streaming=False, + ) + react_agent = ReActAgent( + name="react-member", + description="Answers the request.", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=react_response, + ), + streaming=False, + ) + workflow = AgentWorkflow( + agents=[function_agent, react_agent], + root_agent="function-member", + ) + + result = await workflow.run(user_msg="Complete the request") + assert result.response.content == "complete" + + workflow_span = _spans_named( + span_exporter, "invoke_workflow AgentWorkflow" + )[0] + function_span = _spans_named( + span_exporter, "invoke_agent function-member" + )[0] + react_span = _spans_named(span_exporter, "invoke_agent react-member")[0] + handoff_span = _spans_named(span_exporter, "execute_tool handoff")[0] + assert ( + handoff_span.attributes[GenAIAttributes.GEN_AI_TOOL_TYPE] == "function" + ) + for span in (function_span, react_span): + assert span.parent is not None + assert span.parent.span_id == workflow_span.context.span_id + assert span.context.trace_id == workflow_span.context.trace_id + assert handoff_span.parent is not None + assert handoff_span.parent.span_id == function_span.context.span_id + assert handoff_span.context.trace_id == workflow_span.context.trace_id + # The handing-off agent stays open until its handoff tool call ends, so + # the tool span falls inside its parent rather than after it. + assert handoff_span.start_time >= function_span.start_time + assert handoff_span.end_time <= function_span.end_time def test_sync_tool_span( @@ -974,3 +1281,276 @@ def broken() -> None: span = spans[0] assert span.status.status_code == StatusCode.ERROR assert span.attributes[ErrorAttributes.ERROR_TYPE] == "RuntimeError" + + +@pytest.mark.asyncio +async def test_agent_workflow_releases_member_invocation_on_failure( + span_exporter, instrument_llama_index +) -> None: + """A dropped agent step must not leave its ended invocation reusable.""" + + def response_generator(messages, **kwargs): + if any(message.role.value == "tool" for message in messages): + raise RuntimeError("llm failed") + return ChatMessage( + role="assistant", + blocks=[ + ToolCallBlock( + tool_call_id="call-1", + tool_name="echo", + tool_kwargs={"value": "hello"}, + ) + ], + ) + + agent = FunctionAgent( + name="failing-agent", + description="Fails after a tool call.", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=response_generator, + ), + tools=[ + FunctionTool.from_defaults( + lambda value: value, name="echo", description="Echoes." + ) + ], + streaming=False, + ) + workflow = AgentWorkflow(agents=[agent]) + + registries: list[_LlamaIndexInvocation] = [] + register = _LlamaIndexInvocation.register_workflow_invocation + + def capture(self, run_id, agent_name, invocation): + registries.append(self) + return register(self, run_id, agent_name, invocation) + + with patch.object( + _LlamaIndexInvocation, "register_workflow_invocation", capture + ): + with pytest.raises(RuntimeError, match="llm failed"): + await workflow.run(user_msg="Use echo") + + assert registries, "no member invocation was ever registered" + for parent in registries: + assert parent._workflow_invocations_by_key == {} + assert parent._workflow_invocations_by_run_id == {} + + agent_span = _spans_named(span_exporter, "invoke_agent failing-agent")[0] + assert agent_span.status.status_code == StatusCode.ERROR + + +@pytest.mark.asyncio +async def test_agent_workflow_handoff_turn_contains_its_sibling_tools( + span_exporter, instrument_llama_index +) -> None: + """A turn may request a handoff alongside other tools; all nest in the agent.""" + + def handoff_and_echo(to_agent: str): + def _generate(messages, **kwargs): + return ChatMessage( + role="assistant", + blocks=[ + ToolCallBlock( + tool_call_id=f"handoff-{to_agent}", + tool_name="handoff", + tool_kwargs={"to_agent": to_agent, "reason": "next"}, + ), + ToolCallBlock( + tool_call_id="echo-1", + tool_name="echo", + tool_kwargs={"value": "hi"}, + ), + ], + ) + + return _generate + + def only_handoff(to_agent: str): + def _generate(messages, **kwargs): + return ChatMessage( + role="assistant", + blocks=[ + ToolCallBlock( + tool_call_id=f"handoff-{to_agent}", + tool_name="handoff", + tool_kwargs={"to_agent": to_agent, "reason": "next"}, + ) + ], + ) + + return _generate + + first = FunctionAgent( + name="first", + description="Hands off and echoes.", + can_handoff_to=["second"], + llm=MockFunctionCallingLLM( + is_chat_model=True, response_generator=handoff_and_echo("second") + ), + tools=[ + FunctionTool.from_defaults( + lambda value: value, name="echo", description="Echoes." + ) + ], + streaming=False, + ) + second = FunctionAgent( + name="second", + description="Hands off again.", + can_handoff_to=["third"], + llm=MockFunctionCallingLLM( + is_chat_model=True, response_generator=only_handoff("third") + ), + streaming=False, + ) + third = FunctionAgent( + name="third", + description="Answers.", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=lambda messages, **kwargs: ChatMessage( + role="assistant", content="done" + ), + ), + streaming=False, + ) + workflow = AgentWorkflow(agents=[first, second, third], root_agent="first") + + await workflow.run(user_msg="Start") + + first_span = _spans_named(span_exporter, "invoke_agent first")[0] + second_span = _spans_named(span_exporter, "invoke_agent second")[0] + third_span = _spans_named(span_exporter, "invoke_agent third")[0] + echo_span = _spans_named(span_exporter, "execute_tool echo")[0] + + # The sibling tool belongs to the agent that requested it, not the workflow. + assert echo_span.parent is not None + assert echo_span.parent.span_id == first_span.context.span_id + assert echo_span.start_time >= first_span.start_time + assert echo_span.end_time <= first_span.end_time + + # Each member closes before the next one starts. + assert first_span.end_time <= second_span.start_time + assert second_span.end_time <= third_span.start_time + + +@pytest.mark.asyncio +async def test_agent_workflow_does_not_log_context_detach_failures( + span_exporter, instrument_llama_index +) -> None: + """Member-agent spans outlive the task that opened them. + + ``TelemetryHandler`` attaches the span to the context that was current when + the invocation started, and that attachment can only be undone from the same + context. Finishing from a later step's task would make + ``opentelemetry.context.detach`` raise, which it logs with a traceback. + """ + + class _DetachFailures(logging.Handler): + def __init__(self) -> None: + super().__init__() + self.records: list[logging.LogRecord] = [] + + def emit(self, record: logging.LogRecord) -> None: + if "Failed to detach context" in record.getMessage(): + self.records.append(record) + + def handoff_then_answer(to_agent: str): + def _generate(messages, **kwargs): + if any(message.role.value == "tool" for message in messages): + return ChatMessage(role="assistant", content="done") + return ChatMessage( + role="assistant", + blocks=[ + ToolCallBlock( + tool_call_id="handoff-1", + tool_name="handoff", + tool_kwargs={"to_agent": to_agent, "reason": "next"}, + ) + ], + ) + + return _generate + + router = FunctionAgent( + name="router", + description="Routes.", + can_handoff_to=["worker"], + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=handoff_then_answer("worker"), + ), + streaming=False, + ) + worker = FunctionAgent( + name="worker", + description="Answers using a tool.", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=_tool_then_answer("echo", "hello"), + ), + tools=[ + FunctionTool.from_defaults( + lambda value: value, name="echo", description="Echoes." + ) + ], + streaming=False, + ) + workflow = AgentWorkflow(agents=[router, worker], root_agent="router") + + failures = _DetachFailures() + context_logger = logging.getLogger("opentelemetry.context") + context_logger.addHandler(failures) + try: + await workflow.run(user_msg="Start") + finally: + context_logger.removeHandler(failures) + + assert [record.getMessage() for record in failures.records] == [] + # The spans the cross-task finishes produce are still correct. + assert len(_spans_named(span_exporter, "invoke_agent router")) == 1 + assert len(_spans_named(span_exporter, "invoke_agent worker")) == 1 + + +def _tool_then_answer(tool_name: str, value: str): + def _generate(messages, **kwargs): + if any(message.role.value == "tool" for message in messages): + return ChatMessage(role="assistant", content="done") + return ChatMessage( + role="assistant", + blocks=[ + ToolCallBlock( + tool_call_id="tool-1", + tool_name=tool_name, + tool_kwargs={"value": value}, + ) + ], + ) + + return _generate + + +def test_open_span_lookup_holds_the_handler_lock() -> None: + """``open_spans`` is mutated under the lock from LlamaIndex worker threads. + + Iterating it unguarded raises ``RuntimeError: dictionary changed size during + iteration``, which the dispatcher swallows -- losing the tool span silently. + """ + handler = LlamaIndexSpanHandler( + handler=TelemetryHandler(tracer_provider=TracerProvider()) + ) + + class _RecordingSpans(dict): + locked_during_iteration: bool | None = None + + def values(self): # type: ignore[override] + self.locked_during_iteration = handler.lock.locked() + return super().values() + + handler.open_spans = _RecordingSpans() + + assert handler._is_open_tool(cast(Any, object())) is False + assert handler.open_spans.locked_during_iteration is True + assert handler.lock.locked() is False diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_composition.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_composition.py index dffb45cf8..ded3fbefa 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_composition.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_composition.py @@ -4,7 +4,11 @@ from __future__ import annotations import pytest -from llama_index.core.agent.workflow import FunctionAgent, ReActAgent +from llama_index.core.agent.workflow import ( + AgentWorkflow, + FunctionAgent, + ReActAgent, +) from llama_index.core.base.llms.types import ToolCallBlock from llama_index.core.llms import ChatMessage, MockFunctionCallingLLM from llama_index.core.tools import FunctionTool @@ -84,6 +88,7 @@ def react_response(messages, **kwargs): llm=openai_llm, streaming=False, ) + provider_workflow = AgentWorkflow(agents=[provider_agent]) providers = { "tracer_provider": tracer_provider, @@ -95,7 +100,7 @@ def react_response(messages, **kwargs): await function_agent.run(user_msg="What is the weather in Paris?") await react_agent.run(user_msg="What is two plus two?") with vcr.use_cassette("inference.yaml"): - await provider_agent.run(user_msg="Hello!") + await provider_workflow.run(user_msg="Hello!") spans = span_exporter.get_finished_spans() operations = [ @@ -103,6 +108,7 @@ def react_response(messages, **kwargs): for span in spans ] assert operations.count("invoke_agent") == 3 + assert operations.count("invoke_workflow") == 1 assert operations.count("execute_tool") == 1 assert operations.count("chat") == 1 assert all(isinstance(operation, str) for operation in operations) @@ -112,6 +118,7 @@ def react_response(messages, **kwargs): function_span = spans_by_name["invoke_agent weather-agent"] inference_span = spans_by_name["chat gpt-4o-mini"] provider_span = spans_by_name["invoke_agent provider-agent"] + workflow_span = spans_by_name["invoke_workflow AgentWorkflow"] assert tool_span.context.trace_id == function_span.context.trace_id assert tool_span.parent is not None @@ -119,6 +126,9 @@ def react_response(messages, **kwargs): assert inference_span.context.trace_id == provider_span.context.trace_id assert inference_span.parent is not None assert inference_span.parent.span_id == provider_span.context.span_id + assert provider_span.context.trace_id == workflow_span.context.trace_id + assert provider_span.parent is not None + assert provider_span.parent.span_id == workflow_span.context.span_id @pytest.mark.asyncio @@ -169,3 +179,54 @@ async def test_agent_and_inference_provider_errors_compose( assert inference_span.context.trace_id == agent_span.context.trace_id assert inference_span.parent is not None assert inference_span.parent.span_id == agent_span.context.span_id + + +@pytest.mark.asyncio +async def test_standalone_agent_nests_provider_inference( + span_exporter, + tracer_provider, + logger_provider, + meter_provider, + openai_llm, + vcr, +) -> None: + """A provider chat span nests under a standalone agent, not just a member. + + ``BaseWorkflowAgent.run`` takes a different path through the span handler + than an ``AgentWorkflow`` member step, so both need their own coverage. + """ + OpenAIInstrumentor = pytest.importorskip( + "opentelemetry.instrumentation.genai.openai" + ).OpenAIInstrumentor + + agent = FunctionAgent( + name="standalone-agent", + llm=openai_llm, + streaming=False, + ) + providers = { + "tracer_provider": tracer_provider, + "logger_provider": logger_provider, + "meter_provider": meter_provider, + } + with instrument(LlamaIndexInstrumentor(), **providers): + with instrument(OpenAIInstrumentor(), **providers): + with vcr.use_cassette("inference.yaml"): + await agent.run(user_msg="Hello!") + + spans = span_exporter.get_finished_spans() + operations = [ + span.attributes[GenAIAttributes.GEN_AI_OPERATION_NAME] + for span in spans + ] + # A standalone run emits no workflow span. + assert operations.count("invoke_workflow") == 0 + assert operations.count("invoke_agent") == 1 + + spans_by_name = {span.name: span for span in spans} + agent_span = spans_by_name["invoke_agent standalone-agent"] + inference_span = spans_by_name["chat gpt-4o-mini"] + + assert inference_span.context.trace_id == agent_span.context.trace_id + assert inference_span.parent is not None + assert inference_span.parent.span_id == agent_span.context.span_id diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_conformance.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_conformance.py index 0d2280819..c8103fb31 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_conformance.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_conformance.py @@ -14,9 +14,14 @@ from opentelemetry.test_util_genai.conformance import Scenario, run_conformance from .conformance.agent import AgentScenario +from .conformance.workflow import WorkflowScenario -@pytest.mark.parametrize("scenario", [AgentScenario()]) +@pytest.mark.parametrize( + "scenario", + [AgentScenario(), WorkflowScenario()], + ids=lambda scenario: type(scenario).__name__, +) def test_conformance( scenario: Scenario, vcr: Any, weaver_live_check: WeaverLiveCheck ) -> None: From 77d9d07f64bea07bbfcd0264178fef6e5cbb10cc Mon Sep 17 00:00:00 2001 From: eternalcuriouslearner Date: Wed, 9 Sep 2026 22:24:52 -0400 Subject: [PATCH 2/8] Fix return-direct workflow output selection --- .../genai/llama_index/_handler.py | 71 +++++++++++- .../tests/test_agent.py | 108 ++++++++++++++++++ 2 files changed, 174 insertions(+), 5 deletions(-) diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py index b2d0aed07..73cfb7ac2 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py @@ -434,6 +434,17 @@ def _set_agent_step_output(invocation: AgentInvocation, result: Any) -> None: invocation.output_messages = [_output_message(result.response)] +def _set_return_direct_agent_output( + invocation: AgentInvocation, tool_output: ToolOutput +) -> None: + """Record the response synthesized by a return-direct tool execution.""" + invocation.output_messages = [ + _output_message( + ChatMessage(role="assistant", content=tool_output.content) + ) + ] + + def _agent_step_is_complete(result: Any) -> bool: """Return whether a workflow agent step produced a final response. @@ -503,11 +514,13 @@ class _LlamaIndexInvocation(BaseSpan): _workflow_agent_invocation: AgentInvocation | None = PrivateAttr() _workflow_handoff: bool = PrivateAttr() _workflow_agent_context_token: Token[Context] | None = PrivateAttr() + _tool_parent_context_token: Token[Context] | None = PrivateAttr() _workflow_tool_token: Token[tuple[str, ToolInvocation] | None] | None = ( PrivateAttr() ) _workflow_run_id: str | None = PrivateAttr() _workflow_tool_counts: dict[str, int] = PrivateAttr() + _workflow_return_direct_runs: set[str] = PrivateAttr() _workflow_pending_handoffs: dict[str, AgentInvocation] = PrivateAttr() def __init__( @@ -527,6 +540,7 @@ def __init__( workflow_agent: BaseWorkflowAgent | None = None, workflow_agent_invocation: AgentInvocation | None = None, workflow_handoff: bool = False, + tool_parent_context_token: Token[Context] | None = None, ) -> None: """Create the adapter used by LlamaIndex's span-handler lifecycle.""" super().__init__(id_=id_, parent_id=parent_id) @@ -541,7 +555,9 @@ def __init__( self._workflow_agent_invocation = workflow_agent_invocation self._workflow_handoff = workflow_handoff self._workflow_agent_context_token = None + self._tool_parent_context_token = tool_parent_context_token self._workflow_tool_counts = {} + self._workflow_return_direct_runs = set() self._workflow_pending_handoffs = {} if workflow_run_id is not None and workflow_agent is not None: self.register_workflow_agent(workflow_run_id, workflow_agent) @@ -620,6 +636,15 @@ def reset_workflow_tool(self) -> None: pass self._workflow_tool_token = None + def reset_tool_parent_context(self) -> None: + """Detach the agent context after the tool span has finished.""" + if self._tool_parent_context_token is not None: + try: + detach(self._tool_parent_context_token) + except ValueError: + pass + self._tool_parent_context_token = None + def expect_workflow_tools(self, run_id: str, count: int) -> None: """Record how many tool calls the agent's current turn requested. @@ -629,8 +654,22 @@ def expect_workflow_tools(self, run_id: str, count: int) -> None: """ if count: self._workflow_tool_counts[run_id] = count + self._workflow_return_direct_runs.discard(run_id) else: self._workflow_tool_counts.pop(run_id, None) + self._workflow_return_direct_runs.discard(run_id) + + def set_return_direct_agent_output( + self, + run_id: str, + invocation: AgentInvocation, + tool_output: ToolOutput, + ) -> None: + """Keep the first successful return-direct result to arrive.""" + if run_id in self._workflow_return_direct_runs: + return + self._workflow_return_direct_runs.add(run_id) + _set_return_direct_agent_output(invocation, tool_output) def release_workflow_tool(self, run_id: str | None) -> bool: """Release one completed tool and report whether the turn is drained.""" @@ -641,6 +680,7 @@ def release_workflow_tool(self, run_id: str | None) -> bool: self._workflow_tool_counts[run_id] = remaining return False self._workflow_tool_counts.pop(run_id, None) + self._workflow_return_direct_runs.discard(run_id) return True def set_pending_handoff( @@ -728,6 +768,7 @@ def new_span( workflow_tool_token: ( Token[tuple[str, ToolInvocation] | None] | None ) = None + tool_parent_context_token: Token[Context] | None = None if isinstance(instance, AgentWorkflow) and method_name == "run": capture_content = self._handler.should_capture_content() @@ -888,9 +929,11 @@ def new_span( tool_call.tool_name, tool_type=tool_type, ) - finally: + except BaseException: if agent_context_token is not None: detach(agent_context_token) + raise + tool_parent_context_token = agent_context_token workflow_tool_token = _ACTIVE_WORKFLOW_TOOL.set( (tool_call.tool_name, tool_invocation) ) @@ -954,6 +997,7 @@ def new_span( invocation=invocation, tool_attributes_token=tool_attributes_token, workflow_tool_token=workflow_tool_token, + tool_parent_context_token=tool_parent_context_token, workflow_agents=workflow_agents, workflow_run_id=workflow_run_id, workflow_agent=workflow_agent, @@ -1039,10 +1083,7 @@ def prepare_to_exit_span( _set_agent_output(span._invocation, result) if span._workflow_agent_invocation is not None: span.reset_workflow_agent() - if _agent_step_is_complete(result): - _finish_member_agent(span._invocation) - self._release_workflow_invocation(span, span._invocation) - else: + if not _agent_step_is_complete(result): self._expect_workflow_tools(span, result) return span elif isinstance(span._invocation, ToolInvocation): @@ -1053,6 +1094,23 @@ def prepare_to_exit_span( elif isinstance(result, ToolOutput): tool_output = result if tool_output is not None: + if ( + isinstance(result, ToolCallResult) + and result.return_direct + and not tool_output.is_error + and span._workflow_agent_invocation is not None + and self._handler.should_capture_content() + ): + parent = self.open_spans.get(span.parent_id or "") + if ( + parent is not None + and span._workflow_run_id is not None + ): + parent.set_return_direct_agent_output( + span._workflow_run_id, + span._workflow_agent_invocation, + tool_output, + ) if span._invocation.should_capture_content: span._invocation.tool_result = tool_output.raw_output if tool_output.is_error: @@ -1065,9 +1123,11 @@ def prepare_to_exit_span( else RuntimeError(tool_output.content) ) span._invocation.fail(error) + span.reset_tool_parent_context() self._finish_workflow_tool(span) return span span._invocation.stop() + span.reset_tool_parent_context() if isinstance(span._invocation, ToolInvocation): self._finish_workflow_tool(span) return span @@ -1098,6 +1158,7 @@ def prepare_to_drop_span( span._invocation.stop() else: span._invocation.fail(err) + span.reset_tool_parent_context() if isinstance(span._invocation, AgentInvocation): span.reset_workflow_agent() self._release_workflow_invocation(span, span._invocation) diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py index 10d5e30ea..d8d90c8c7 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py @@ -3,6 +3,7 @@ from __future__ import annotations +import asyncio import json import logging from typing import Any, cast @@ -911,6 +912,113 @@ def response_generator(messages, **kwargs): assert tool_spans[1].parent is None +@pytest.mark.asyncio +async def test_agent_workflow_captures_first_return_direct_result( + span_exporter, instrument_llama_index_with_content +) -> None: + def first() -> str: + return "FIRST" + + def second() -> str: + return "SECOND" + + def response_generator(messages, **kwargs): + return ChatMessage( + role="assistant", + blocks=[ + ToolCallBlock( + tool_call_id="first-call", + tool_name="first", + tool_kwargs={}, + ), + ToolCallBlock( + tool_call_id="second-call", + tool_name="second", + tool_kwargs={}, + ), + ], + ) + + agent = FunctionAgent( + name="return-direct-agent", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=response_generator, + ), + tools=[ + FunctionTool.from_defaults(first, return_direct=True), + FunctionTool.from_defaults(second, return_direct=True), + ], + streaming=False, + ) + workflow = AgentWorkflow(agents=[agent]) + + result = await workflow.run(user_msg="Call both tools") + assert result.response.content == "FIRST" + + agent_span = _spans_named( + span_exporter, "invoke_agent return-direct-agent" + )[0] + agent_output = json.loads( + agent_span.attributes[GenAIAttributes.GEN_AI_OUTPUT_MESSAGES] + ) + assert agent_output[0]["parts"] == [{"type": "text", "content": "FIRST"}] + + +@pytest.mark.asyncio +async def test_agent_workflow_captures_first_arriving_return_direct_result( + span_exporter, instrument_llama_index_with_content +) -> None: + async def first() -> str: + await asyncio.sleep(0.02) + return "FIRST" + + async def second() -> str: + return "SECOND" + + def response_generator(messages, **kwargs): + return ChatMessage( + role="assistant", + blocks=[ + ToolCallBlock( + tool_call_id="first-call", + tool_name="first", + tool_kwargs={}, + ), + ToolCallBlock( + tool_call_id="second-call", + tool_name="second", + tool_kwargs={}, + ), + ], + ) + + agent = FunctionAgent( + name="arrival-order-agent", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=response_generator, + ), + tools=[ + FunctionTool.from_defaults(async_fn=first, return_direct=True), + FunctionTool.from_defaults(async_fn=second, return_direct=True), + ], + streaming=False, + ) + workflow = AgentWorkflow(agents=[agent]) + + result = await workflow.run(user_msg="Call both tools") + assert result.response.content == "SECOND" + + agent_span = _spans_named( + span_exporter, "invoke_agent arrival-order-agent" + )[0] + agent_output = json.loads( + agent_span.attributes[GenAIAttributes.GEN_AI_OUTPUT_MESSAGES] + ) + assert agent_output[0]["parts"] == [{"type": "text", "content": "SECOND"}] + + @pytest.mark.asyncio async def test_agent_workflow_uses_configured_workflow_name( span_exporter, instrument_llama_index From abcff90eac22bedcb727d6cd84b9e47a2b5531a0 Mon Sep 17 00:00:00 2001 From: eternalcuriouslearner Date: Wed, 9 Sep 2026 22:29:24 -0400 Subject: [PATCH 3/8] Deduplicate standalone nested tool spans --- .../instrumentation/genai/llama_index/_handler.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py index 73cfb7ac2..d4cc00d39 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py @@ -956,10 +956,8 @@ def new_span( parent = self.open_spans.get(parent_span_id or "") # LlamaIndex reports an agent tool execution through both call_tool # and the nested FunctionTool.call/acall; the parent records it. - if ( - parent is not None - and parent._workflow_agent_invocation is not None - and isinstance(parent._invocation, ToolInvocation) + if parent is not None and isinstance( + parent._invocation, ToolInvocation ): # The workflow callback identifies the tool by name only; the # nested FunctionTool call is the authoritative executing tool. From be5518e1039691b093043ca0f51ecb081ca621a9 Mon Sep 17 00:00:00 2001 From: eternalcuriouslearner Date: Thu, 10 Sep 2026 08:19:50 -0400 Subject: [PATCH 4/8] Fix return-direct regression test --- .../tests/test_agent.py | 53 ------------------- 1 file changed, 53 deletions(-) diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py index d8d90c8c7..98f20cc86 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py @@ -912,59 +912,6 @@ def response_generator(messages, **kwargs): assert tool_spans[1].parent is None -@pytest.mark.asyncio -async def test_agent_workflow_captures_first_return_direct_result( - span_exporter, instrument_llama_index_with_content -) -> None: - def first() -> str: - return "FIRST" - - def second() -> str: - return "SECOND" - - def response_generator(messages, **kwargs): - return ChatMessage( - role="assistant", - blocks=[ - ToolCallBlock( - tool_call_id="first-call", - tool_name="first", - tool_kwargs={}, - ), - ToolCallBlock( - tool_call_id="second-call", - tool_name="second", - tool_kwargs={}, - ), - ], - ) - - agent = FunctionAgent( - name="return-direct-agent", - llm=MockFunctionCallingLLM( - is_chat_model=True, - response_generator=response_generator, - ), - tools=[ - FunctionTool.from_defaults(first, return_direct=True), - FunctionTool.from_defaults(second, return_direct=True), - ], - streaming=False, - ) - workflow = AgentWorkflow(agents=[agent]) - - result = await workflow.run(user_msg="Call both tools") - assert result.response.content == "FIRST" - - agent_span = _spans_named( - span_exporter, "invoke_agent return-direct-agent" - )[0] - agent_output = json.loads( - agent_span.attributes[GenAIAttributes.GEN_AI_OUTPUT_MESSAGES] - ) - assert agent_output[0]["parts"] == [{"type": "text", "content": "FIRST"}] - - @pytest.mark.asyncio async def test_agent_workflow_captures_first_arriving_return_direct_result( span_exporter, instrument_llama_index_with_content From 1543b7c94e1646fba13c32efec896ae5492b8f94 Mon Sep 17 00:00:00 2001 From: eternalcuriouslearner Date: Thu, 10 Sep 2026 09:15:26 -0400 Subject: [PATCH 5/8] Cover workflow streaming and early stopping --- .../genai/llama_index/_handler.py | 24 +++- .../tests/conformance/workflow.py | 5 +- .../tests/test_agent.py | 123 ++++++++++++++++++ 3 files changed, 148 insertions(+), 4 deletions(-) diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py index d4cc00d39..bbd5cf9f9 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py @@ -707,9 +707,24 @@ def reset_workflow_agent(self) -> None: self._workflow_agent_context_token = None def finalize_workflow_agents( - self, error: BaseException | None = None + self, + error: BaseException | None = None, + result: Any | None = None, ) -> None: """Finish member-agent spans left open when the workflow terminates.""" + if error is None and isinstance( + (output := getattr(result, "result", None)), AgentOutput + ): + # ``early_stopping_method="generate"`` creates the final response + # in ``parse_agent_output`` rather than another agent step. + agent_name = output.current_agent_name + for ( + _, + name, + ), invocation in self._workflow_invocations_by_key.items(): + if name == agent_name: + _set_agent_step_output(invocation, output) + break invocations: list[AgentInvocation] = [] for candidate in self._workflow_invocations_by_key.values(): if all(candidate is not existing for existing in invocations): @@ -1068,9 +1083,12 @@ def prepare_to_exit_span( if span is None: return None if isinstance(span._invocation, WorkflowInvocation): - if self._handler.should_capture_content(): + capture_content = self._handler.should_capture_content() + if capture_content: _set_workflow_output(span._invocation, result) - span.finalize_workflow_agents() + span.finalize_workflow_agents( + result=result if capture_content else None + ) elif isinstance(span._invocation, AgentInvocation): span.reset_tool_attributes() span.reset_workflow_tool() diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/conformance/workflow.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/conformance/workflow.py index e443085ab..0d60c3603 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/conformance/workflow.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/conformance/workflow.py @@ -30,7 +30,10 @@ class WorkflowScenario(Scenario): "invoke_agent": 2, "execute_tool": 1, } - expected_metrics = ("gen_ai.client.operation.duration",) + expected_metrics = ( + "gen_ai.client.operation.duration", + "gen_ai.invoke_workflow.duration", + ) def run( self, diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py index 98f20cc86..0397666b9 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py @@ -912,6 +912,129 @@ def response_generator(messages, **kwargs): assert tool_spans[1].parent is None +@pytest.mark.asyncio +async def test_agent_workflow_streaming_member_span_hierarchy( + span_exporter, instrument_llama_index +) -> None: + agent = FunctionAgent( + name="streaming-workflow-agent", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=lambda messages, **kwargs: ChatMessage( + role="assistant", content="streamed answer" + ), + ), + streaming=True, + ) + workflow = AgentWorkflow(agents=[agent]) + + result = await workflow.run(user_msg="Answer by streaming") + assert result.response.content == "streamed answer" + + workflow_span = _spans_named( + span_exporter, "invoke_workflow AgentWorkflow" + )[0] + agent_span = _spans_named( + span_exporter, "invoke_agent streaming-workflow-agent" + )[0] + assert agent_span.parent is not None + assert agent_span.parent.span_id == workflow_span.context.span_id + assert agent_span.context.trace_id == workflow_span.context.trace_id + assert workflow_span.status.status_code == StatusCode.UNSET + assert agent_span.status.status_code == StatusCode.UNSET + + +@pytest.mark.asyncio +async def test_agent_workflow_streaming_member_error_finalization( + span_exporter, instrument_llama_index +) -> None: + error = ConnectionError("workflow stream disconnected") + + async def response_stream(): + yield ChatResponse( + message=ChatMessage(role="assistant", content="partial response") + ) + raise error + + async def failing_stream(*args, **kwargs): + return response_stream() + + agent = ReActAgent( + name="streaming-react-workflow-agent", + llm=MockFunctionCallingLLM(is_chat_model=True), + streaming=True, + ) + workflow = AgentWorkflow(agents=[agent]) + + with patch.object( + MockFunctionCallingLLM, "astream_chat", side_effect=failing_stream + ): + with pytest.raises(ConnectionError) as exc_info: + await workflow.run(user_msg="Answer by streaming") + + assert exc_info.value is error + workflow_span = _spans_named( + span_exporter, "invoke_workflow AgentWorkflow" + )[0] + agent_span = _spans_named( + span_exporter, "invoke_agent streaming-react-workflow-agent" + )[0] + assert workflow_span.status.status_code == StatusCode.ERROR + assert agent_span.status.status_code == StatusCode.ERROR + _assert_error_type(workflow_span, "ConnectionError") + _assert_error_type(agent_span, "ConnectionError") + + +@pytest.mark.asyncio +async def test_agent_workflow_captures_early_stopping_response( + span_exporter, instrument_llama_index_with_content +) -> None: + responses = iter( + [ + ChatMessage( + role="assistant", + blocks=[ + ToolCallBlock( + tool_call_id="echo-call", + tool_name="echo", + tool_kwargs={"value": "partial"}, + ) + ], + ), + ChatMessage(role="assistant", content="generated answer"), + ] + ) + + agent = FunctionAgent( + name="early-stopping-agent", + llm=MockFunctionCallingLLM( + is_chat_model=True, + response_generator=lambda messages, **kwargs: next(responses), + ), + tools=[FunctionTool.from_defaults(lambda value: value, name="echo")], + streaming=False, + early_stopping_method="generate", + ) + workflow = AgentWorkflow(agents=[agent]) + + result = await workflow.run( + user_msg="Answer this", + max_iterations=1, + early_stopping_method="generate", + ) + assert result.response.content == "generated answer" + + agent_span = _spans_named( + span_exporter, "invoke_agent early-stopping-agent" + )[0] + agent_output = json.loads( + agent_span.attributes[GenAIAttributes.GEN_AI_OUTPUT_MESSAGES] + ) + assert agent_output[0]["parts"] == [ + {"type": "text", "content": "generated answer"} + ] + + @pytest.mark.asyncio async def test_agent_workflow_captures_first_arriving_return_direct_result( span_exporter, instrument_llama_index_with_content From 7a15c44abe823a0eaa7e1f9cf13c5565df2bd5c0 Mon Sep 17 00:00:00 2001 From: eternalcuriouslearner Date: Thu, 10 Sep 2026 10:42:05 -0400 Subject: [PATCH 6/8] Handle failed workflow handoffs --- .../genai/llama_index/_handler.py | 16 ++++++++++++---- .../tests/test_agent.py | 5 +++-- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py index bbd5cf9f9..049b5bd07 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py @@ -1042,7 +1042,11 @@ def _release_workflow_invocation( if parent is not None: parent.remove_workflow_invocation(invocation) - def _finish_workflow_tool(self, span: _LlamaIndexInvocation) -> None: + def _finish_workflow_tool( + self, + span: _LlamaIndexInvocation, + handoff_succeeded: bool = True, + ) -> None: """Release one tool of a member agent's turn and close the agent last. AgentWorkflow reports a handoff as a tool call made by the agent that is @@ -1057,7 +1061,11 @@ def _finish_workflow_tool(self, span: _LlamaIndexInvocation) -> None: if parent is None: return invocation = span._workflow_agent_invocation - if span._workflow_handoff and invocation is not None: + if ( + handoff_succeeded + and span._workflow_handoff + and invocation is not None + ): parent.set_pending_handoff(run_id, invocation) if not parent.release_workflow_tool(run_id): return @@ -1140,7 +1148,7 @@ def prepare_to_exit_span( ) span._invocation.fail(error) span.reset_tool_parent_context() - self._finish_workflow_tool(span) + self._finish_workflow_tool(span, handoff_succeeded=False) return span span._invocation.stop() span.reset_tool_parent_context() @@ -1179,5 +1187,5 @@ def prepare_to_drop_span( span.reset_workflow_agent() self._release_workflow_invocation(span, span._invocation) elif isinstance(span._invocation, ToolInvocation): - self._finish_workflow_tool(span) + self._finish_workflow_tool(span, handoff_succeeded=err is None) return span diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py index 0397666b9..aa874d7db 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_agent.py @@ -6,6 +6,7 @@ import asyncio import json import logging +from collections.abc import ValuesView from typing import Any, cast from unittest.mock import patch from uuid import UUID @@ -1720,10 +1721,10 @@ def test_open_span_lookup_holds_the_handler_lock() -> None: handler=TelemetryHandler(tracer_provider=TracerProvider()) ) - class _RecordingSpans(dict): + class _RecordingSpans(dict[str, _LlamaIndexInvocation]): locked_during_iteration: bool | None = None - def values(self): # type: ignore[override] + def values(self) -> ValuesView[_LlamaIndexInvocation]: self.locked_during_iteration = handler.lock.locked() return super().values() From 419eb0516e1506af869f31b94061fafa165662c2 Mon Sep 17 00:00:00 2001 From: eternalcuriouslearner Date: Thu, 10 Sep 2026 12:11:49 -0400 Subject: [PATCH 7/8] Preserve mixed workflow tool failures --- .../genai/llama_index/_handler.py | 35 +++++++++++++++++-- 1 file changed, 32 insertions(+), 3 deletions(-) diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py index 049b5bd07..9eeab7162 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py @@ -521,6 +521,7 @@ class _LlamaIndexInvocation(BaseSpan): _workflow_run_id: str | None = PrivateAttr() _workflow_tool_counts: dict[str, int] = PrivateAttr() _workflow_return_direct_runs: set[str] = PrivateAttr() + _workflow_tool_errors: dict[str, BaseException] = PrivateAttr() _workflow_pending_handoffs: dict[str, AgentInvocation] = PrivateAttr() def __init__( @@ -558,6 +559,7 @@ def __init__( self._tool_parent_context_token = tool_parent_context_token self._workflow_tool_counts = {} self._workflow_return_direct_runs = set() + self._workflow_tool_errors = {} self._workflow_pending_handoffs = {} if workflow_run_id is not None and workflow_agent is not None: self.register_workflow_agent(workflow_run_id, workflow_agent) @@ -655,9 +657,11 @@ def expect_workflow_tools(self, run_id: str, count: int) -> None: if count: self._workflow_tool_counts[run_id] = count self._workflow_return_direct_runs.discard(run_id) + self._workflow_tool_errors.pop(run_id, None) else: self._workflow_tool_counts.pop(run_id, None) self._workflow_return_direct_runs.discard(run_id) + self._workflow_tool_errors.pop(run_id, None) def set_return_direct_agent_output( self, @@ -671,6 +675,16 @@ def set_return_direct_agent_output( self._workflow_return_direct_runs.add(run_id) _set_return_direct_agent_output(invocation, tool_output) + def record_workflow_tool_error( + self, run_id: str, error: BaseException + ) -> None: + """Preserve the first failure across a concurrent tool turn.""" + self._workflow_tool_errors.setdefault(run_id, error) + + def take_workflow_tool_error(self, run_id: str) -> BaseException | None: + """Return and clear the failure recorded for a completed tool turn.""" + return self._workflow_tool_errors.pop(run_id, None) + def release_workflow_tool(self, run_id: str | None) -> bool: """Release one completed tool and report whether the turn is drained.""" if run_id is None: @@ -1046,6 +1060,7 @@ def _finish_workflow_tool( self, span: _LlamaIndexInvocation, handoff_succeeded: bool = True, + error: BaseException | None = None, ) -> None: """Release one tool of a member agent's turn and close the agent last. @@ -1060,6 +1075,8 @@ def _finish_workflow_tool( parent = self.open_spans.get(span.parent_id or "") if parent is None: return + if error is not None: + parent.record_workflow_tool_error(run_id, error) invocation = span._workflow_agent_invocation if ( handoff_succeeded @@ -1071,8 +1088,12 @@ def _finish_workflow_tool( return pending = parent.take_pending_handoff(run_id) if pending is not None: - _finish_member_agent(pending) + _finish_member_agent( + pending, parent.take_workflow_tool_error(run_id) + ) parent.remove_workflow_invocation(pending) + else: + parent.take_workflow_tool_error(run_id) def prepare_to_exit_span( self, @@ -1148,7 +1169,11 @@ def prepare_to_exit_span( ) span._invocation.fail(error) span.reset_tool_parent_context() - self._finish_workflow_tool(span, handoff_succeeded=False) + self._finish_workflow_tool( + span, + handoff_succeeded=False, + error=error, + ) return span span._invocation.stop() span.reset_tool_parent_context() @@ -1187,5 +1212,9 @@ def prepare_to_drop_span( span.reset_workflow_agent() self._release_workflow_invocation(span, span._invocation) elif isinstance(span._invocation, ToolInvocation): - self._finish_workflow_tool(span, handoff_succeeded=err is None) + self._finish_workflow_tool( + span, + handoff_succeeded=err is None, + error=err, + ) return span From 501bfd5636b2248991bc0eae5e4bf1bf7c32dcdb Mon Sep 17 00:00:00 2001 From: Surya Date: Thu, 10 Sep 2026 19:07:53 -0400 Subject: [PATCH 8/8] wip: adding agent name to tool execution span. Co-authored-by: Liudmila Molkova --- .../opentelemetry/instrumentation/genai/llama_index/_handler.py | 1 + 1 file changed, 1 insertion(+) diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py index 9eeab7162..c5da051e3 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py @@ -957,6 +957,7 @@ def new_span( tool_invocation = self._handler.tool( tool_call.tool_name, tool_type=tool_type, + agent_name=getattr(active_invocation, "_agent_name", None), ) except BaseException: if agent_context_token is not None: