diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/.changelog/698.added b/instrumentation/opentelemetry-instrumentation-genai-llama-index/.changelog/698.added new file mode 100644 index 000000000..f1f83d878 --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/.changelog/698.added @@ -0,0 +1 @@ +Add retrieval span instrumentation for LlamaIndex ``BaseRetriever`` operations. diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/README.rst b/instrumentation/opentelemetry-instrumentation-genai-llama-index/README.rst index 73b769c42..05ce39e1f 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/README.rst +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/README.rst @@ -12,7 +12,8 @@ This package contains OpenTelemetry instrumentation for 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 +executes tools. It also emits ``retrieval`` spans for synchronous and +asynchronous ``BaseRetriever`` operations. Model calls delegated to provider SDKs are intentionally left to those SDKs' OpenTelemetry instrumentations. @@ -34,6 +35,36 @@ Usage LlamaIndexInstrumentor().instrument() +How instrumentation works +-------------------------- + +The instrumentor registers a LlamaIndex span handler with its dispatcher. The +handler observes LlamaIndex-owned operations and delegates span creation and +completion to ``opentelemetry-util-genai``. Provider model calls are left to +the provider's instrumentation, so they can be composed without duplicate +inference spans. + +.. code-block:: mermaid + + flowchart TD + A[LlamaIndexInstrumentor.instrument] --> B[LlamaIndex dispatcher] + B --> C{Span callback} + C -->|AgentWorkflow.run| D[workflow invocation] + C -->|BaseWorkflowAgent.run / run_agent_step| E[agent invocation] + C -->|call_tool / FunctionTool.call| F[tool invocation] + C -->|BaseRetriever.retrieve / aretrieve| G[retrieval invocation] + D --> H[TelemetryHandler.workflow] + E --> I[TelemetryHandler.invoke_local_agent] + F --> J[TelemetryHandler.tool] + G --> K[TelemetryHandler.retrieval] + H --> L[Start OTel span] + I --> L + J --> L + K --> L + L --> M[Dispatcher exit or error callback] + M --> N[Set attributes and stop/fail invocation] + P[Provider SDK instrumentation] -. model calls .-> Q[inference spans] + Configuration ------------- 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 c08e5dbe8..8698a54d5 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 @@ -21,6 +21,7 @@ ToolCall, ToolCallResult, ) +from llama_index.core.base.base_retriever import BaseRetriever from llama_index.core.base.llms.types import ( AudioBlock, ChatMessage, @@ -32,6 +33,7 @@ ) from llama_index.core.instrumentation.span import BaseSpan from llama_index.core.instrumentation.span_handlers import BaseSpanHandler +from llama_index.core.schema import NodeWithScore, QueryBundle from llama_index.core.tools import BaseTool, FunctionTool, ToolOutput from pydantic import PrivateAttr @@ -41,6 +43,7 @@ from opentelemetry.util.genai.invocation import ( GenAIInvocation, LocalAgentInvocation, + RetrievalInvocation, ToolInvocation, WorkflowInvocation, ) @@ -298,6 +301,52 @@ def _request_model(agent: BaseWorkflowAgent) -> str | None: return model_name if isinstance(model_name, str) and model_name else None +def _retrieval_query(bound_args: inspect.BoundArguments) -> str | None: + """Extract text from either accepted LlamaIndex retrieval query form.""" + query = bound_args.arguments.get("str_or_query_bundle") + if isinstance(query, str): + return query + if isinstance(query, QueryBundle): + return query.query_str + return None + + +def _retrieval_top_k(retriever: BaseRetriever) -> int | None: + """Read the common top-k setting without requiring a retriever subtype.""" + top_k = getattr(retriever, "similarity_top_k", None) + if isinstance(top_k, int) and not isinstance(top_k, bool): + return top_k + return None + + +def _retrieval_documents( + result: object, +) -> list[dict[str, Any]] | None: + """Convert retrieved LlamaIndex nodes to semconv document objects.""" + if not isinstance(result, Sequence): + return None + candidates = cast(Sequence[object], result) + documents: list[dict[str, Any]] = [] + for candidate in candidates: + if not isinstance(candidate, NodeWithScore): + continue + try: + document: dict[str, Any] = { + "id": candidate.node_id, + "content": candidate.node.get_content(), + } + if candidate.score is not None: + document["score"] = candidate.score + documents.append(document) + except BaseException: + continue + # Preserve [] for a genuine empty result, but omit the attribute when a + # non-empty result could not be converted into semantic-convention docs. + if documents: + return documents + return [] if len(candidates) == 0 else None + + def _tool_attributes( candidate: object, ) -> tuple[str, str, str | None] | None: @@ -756,7 +805,7 @@ def finalize_workflow_agents( class LlamaIndexSpanHandler(BaseSpanHandler[_LlamaIndexInvocation]): - """Map LlamaIndex-owned agent and tool operations to GenAI spans.""" + """Map LlamaIndex-owned agent, tool, and retrieval operations to spans.""" _handler: TelemetryHandler = PrivateAttr() @@ -784,7 +833,7 @@ def new_span( tags: dict[str, Any] | None = None, **kwargs: Any, ) -> _LlamaIndexInvocation | None: - """Start GenAI invocations for LlamaIndex-owned agents and tools. + """Start GenAI invocations for agents, tools, and retrievers. Provider inference is deliberately ignored so its own instrumentation can emit inference telemetry, and nested tool callbacks are deduplicated. @@ -913,6 +962,15 @@ def new_span( if workflow_run_id is not None: parent.register_workflow_agent(workflow_run_id, agent) workflow_agent = agent + elif isinstance(instance, BaseRetriever) and method_name in { + "retrieve", + "aretrieve", + }: + retrieval_invocation = self._handler.retrieval() + retrieval_invocation.top_k = _retrieval_top_k(instance) + if retrieval_invocation.should_capture_content: + retrieval_invocation.query_text = _retrieval_query(bound_args) + invocation = retrieval_invocation elif method_name == "call_tool" and isinstance( (tool_call := bound_args.arguments.get("ev")), ToolCall ): @@ -1138,6 +1196,9 @@ def prepare_to_exit_span( if not _agent_step_is_complete(result): self._expect_workflow_tools(span, result) return span + elif isinstance(span._invocation, RetrievalInvocation): + if span._invocation.should_capture_content: + span._invocation.documents = _retrieval_documents(result) elif isinstance(span._invocation, ToolInvocation): span.reset_workflow_tool() tool_output: ToolOutput | None = None diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/conformance/retrieval.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/conformance/retrieval.py new file mode 100644 index 000000000..752478776 --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/conformance/retrieval.py @@ -0,0 +1,66 @@ +# Copyright The OpenTelemetry Authors +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Any + +from llama_index.core.base.base_retriever import BaseRetriever +from llama_index.core.schema import NodeWithScore, QueryBundle, TextNode + +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 ( + ExpectedViolation, + Scenario, +) +from opentelemetry.test_util_genai.instrumentor import instrument + + +class _ConformanceRetriever(BaseRetriever): + def __init__(self) -> None: + super().__init__() + self.similarity_top_k = 1 + + def _retrieve(self, query_bundle: QueryBundle) -> list[NodeWithScore]: + return [ + NodeWithScore( + node=TextNode( + id_="capital-france", + text="Paris is the capital of France.", + ), + score=0.99, + ) + ] + + +class RetrievalScenario(Scenario): + expected_spans = {"retrieval": 1} + expected_metrics = ("gen_ai.client.operation.duration",) + expected_violations = ( + ExpectedViolation( + advice_id="genai_expected_attribute_missing", + message_substring="server.address", + ), + ) + + def run( + self, + *, + tracer_provider: TracerProvider, + meter_provider: MeterProvider, + logger_provider: LoggerProvider, + vcr: Any, + ) -> None: + with instrument( + LlamaIndexInstrumentor(), + tracer_provider=tracer_provider, + logger_provider=logger_provider, + meter_provider=meter_provider, + content_capture="SPAN_ONLY", + ): + _ConformanceRetriever().retrieve("What is the capital of France?") 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 c8103fb31..2ec46b42e 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_conformance.py +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_conformance.py @@ -14,12 +14,13 @@ from opentelemetry.test_util_genai.conformance import Scenario, run_conformance from .conformance.agent import AgentScenario +from .conformance.retrieval import RetrievalScenario from .conformance.workflow import WorkflowScenario @pytest.mark.parametrize( "scenario", - [AgentScenario(), WorkflowScenario()], + [AgentScenario(), WorkflowScenario(), RetrievalScenario()], ids=lambda scenario: type(scenario).__name__, ) def test_conformance( diff --git a/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_retrieval.py b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_retrieval.py new file mode 100644 index 000000000..65687f25f --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_retrieval.py @@ -0,0 +1,150 @@ +# Copyright The OpenTelemetry Authors +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json + +import pytest +from llama_index.core.base.base_retriever import BaseRetriever +from llama_index.core.schema import NodeWithScore, QueryBundle, TextNode + +from opentelemetry.instrumentation.genai.llama_index._handler import ( + _retrieval_documents, +) +from opentelemetry.semconv._incubating.attributes import ( + gen_ai_attributes as GenAIAttributes, +) +from opentelemetry.semconv.attributes import ( + error_attributes as ErrorAttributes, +) +from opentelemetry.trace import SpanKind, StatusCode + +_GEN_AI_RETRIEVAL_TOP_K = "gen_ai.retrieval.top_k" + + +def test_unconvertible_retrieval_results_are_omitted() -> None: + assert _retrieval_documents([]) == [] + assert _retrieval_documents([object()]) is None + + +class _Retriever(BaseRetriever): + def __init__(self, error: BaseException | None = None) -> None: + super().__init__() + self.similarity_top_k = 2 + self._error = error + + def _retrieve(self, query_bundle: QueryBundle) -> list[NodeWithScore]: + if self._error: + raise self._error + return [ + NodeWithScore( + node=TextNode(id_="doc-1", text="Paris is in France."), + score=0.9, + ) + ] + + +def _span(exporter): + spans = [s for s in exporter.get_finished_spans() if s.name == "retrieval"] + assert len(spans) == 1 + return spans[0] + + +def test_retrieval_captures_documents_and_query( + span_exporter, instrument_llama_index_with_content +) -> None: + _Retriever().retrieve("Where is Paris?") + span = _span(span_exporter) + assert span.kind == SpanKind.CLIENT + assert ( + span.attributes[GenAIAttributes.GEN_AI_OPERATION_NAME] == "retrieval" + ) + top_k = span.attributes[_GEN_AI_RETRIEVAL_TOP_K] + query_text = span.attributes[GenAIAttributes.GEN_AI_RETRIEVAL_QUERY_TEXT] + documents = span.attributes[GenAIAttributes.GEN_AI_RETRIEVAL_DOCUMENTS] + assert type(top_k) is int + assert type(query_text) is str + assert type(documents) is str + assert top_k == 2 + assert query_text == "Where is Paris?" + assert json.loads(documents) == [ + {"id": "doc-1", "content": "Paris is in France.", "score": 0.9} + ] + + +def test_retrieval_query_bundle_captures_query( + span_exporter, instrument_llama_index_with_content +) -> None: + _Retriever().retrieve(QueryBundle(query_str="Where is Paris?")) + span = _span(span_exporter) + assert span.attributes[GenAIAttributes.GEN_AI_RETRIEVAL_QUERY_TEXT] == ( + "Where is Paris?" + ) + + +def test_retrieval_omits_content_without_capture( + span_exporter, instrument_llama_index +) -> None: + _Retriever().retrieve("Where is Paris?") + attrs = _span(span_exporter).attributes + assert GenAIAttributes.GEN_AI_RETRIEVAL_QUERY_TEXT not in attrs + assert GenAIAttributes.GEN_AI_RETRIEVAL_DOCUMENTS not in attrs + + +def test_retrieval_uses_snapshotted_capture_setting( + span_exporter, instrument_llama_index_with_content, monkeypatch +) -> None: + # Changing the environment after instrumentation must not alter the + # invocation's content-capture decision. + monkeypatch.setenv( + "OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT", "NO_CONTENT" + ) + _Retriever().retrieve("Where is Paris?") + attrs = _span(span_exporter).attributes + assert attrs[GenAIAttributes.GEN_AI_RETRIEVAL_QUERY_TEXT] == ( + "Where is Paris?" + ) + assert GenAIAttributes.GEN_AI_RETRIEVAL_DOCUMENTS in attrs + + +@pytest.mark.asyncio +async def test_async_retrieval_captures_query_and_documents( + span_exporter, instrument_llama_index_with_content +) -> None: + result = await _Retriever().aretrieve("Where is Paris?") + assert result[0].node_id == "doc-1" + span = _span(span_exporter) + query_text = span.attributes[GenAIAttributes.GEN_AI_RETRIEVAL_QUERY_TEXT] + documents = span.attributes[GenAIAttributes.GEN_AI_RETRIEVAL_DOCUMENTS] + assert type(query_text) is str + assert type(documents) is str + assert query_text == "Where is Paris?" + assert json.loads(documents) == [ + {"id": "doc-1", "content": "Paris is in France.", "score": 0.9} + ] + + +def test_sync_retrieval_error_is_unchanged( + span_exporter, instrument_llama_index +) -> None: + error = ValueError("retrieval failed") + with pytest.raises(ValueError) as caught: + _Retriever(error).retrieve("query") + assert caught.value is error + span = _span(span_exporter) + assert span.status.status_code == StatusCode.ERROR + assert span.attributes[ErrorAttributes.ERROR_TYPE] == "ValueError" + + +@pytest.mark.asyncio +async def test_async_retrieval_error_is_unchanged( + span_exporter, instrument_llama_index +) -> None: + error = ValueError("retrieval failed") + with pytest.raises(ValueError) as caught: + await _Retriever(error).aretrieve("query") + assert caught.value is error + span = _span(span_exporter) + assert span.status.status_code == StatusCode.ERROR + assert span.attributes[ErrorAttributes.ERROR_TYPE] == "ValueError"