Skip to content

Commit 4bce391

Browse files
Add LlamaIndex retrieval instrumentation (#698)
* Add LlamaIndex retrieval instrumentation * Add LlamaIndex retrieval changelog * Use shared retrieval document model * Update retrieval changelog fragments * Refactor import statements and improve code formatting in retrieval tests. * Fix LlamaIndex retrieval document content * Test LlamaIndex QueryBundle retrieval * Keep retrieval documents instrumentation-local * Change exception handling to BaseException Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> * Harden retrieval metadata extraction * Assert retrieval attribute types * Refactor _retrieval_documents to improve type casting and return logic * Match retrieval metadata access pattern * Add missing import for json in test_retrieval.py --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
1 parent 627049d commit 4bce391

6 files changed

Lines changed: 314 additions & 4 deletions

File tree

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Add retrieval span instrumentation for LlamaIndex ``BaseRetriever`` operations.

‎instrumentation/opentelemetry-instrumentation-genai-llama-index/README.rst‎

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,8 @@ This package contains OpenTelemetry instrumentation for
1212
It emits ``invoke_workflow`` spans for ``AgentWorkflow`` runs,
1313
``invoke_agent`` spans for standalone and workflow-member ``FunctionAgent``
1414
and ``ReActAgent`` executions, and ``execute_tool`` spans when LlamaIndex
15-
executes tools. Model calls
15+
executes tools. It also emits ``retrieval`` spans for synchronous and
16+
asynchronous ``BaseRetriever`` operations. Model calls
1617
delegated to provider SDKs are intentionally left to those SDKs' OpenTelemetry
1718
instrumentations.
1819

@@ -34,6 +35,36 @@ Usage
3435
3536
LlamaIndexInstrumentor().instrument()
3637
38+
How instrumentation works
39+
--------------------------
40+
41+
The instrumentor registers a LlamaIndex span handler with its dispatcher. The
42+
handler observes LlamaIndex-owned operations and delegates span creation and
43+
completion to ``opentelemetry-util-genai``. Provider model calls are left to
44+
the provider's instrumentation, so they can be composed without duplicate
45+
inference spans.
46+
47+
.. code-block:: mermaid
48+
49+
flowchart TD
50+
A[LlamaIndexInstrumentor.instrument] --> B[LlamaIndex dispatcher]
51+
B --> C{Span callback}
52+
C -->|AgentWorkflow.run| D[workflow invocation]
53+
C -->|BaseWorkflowAgent.run / run_agent_step| E[agent invocation]
54+
C -->|call_tool / FunctionTool.call| F[tool invocation]
55+
C -->|BaseRetriever.retrieve / aretrieve| G[retrieval invocation]
56+
D --> H[TelemetryHandler.workflow]
57+
E --> I[TelemetryHandler.invoke_local_agent]
58+
F --> J[TelemetryHandler.tool]
59+
G --> K[TelemetryHandler.retrieval]
60+
H --> L[Start OTel span]
61+
I --> L
62+
J --> L
63+
K --> L
64+
L --> M[Dispatcher exit or error callback]
65+
M --> N[Set attributes and stop/fail invocation]
66+
P[Provider SDK instrumentation] -. model calls .-> Q[inference spans]
67+
3768
Configuration
3869
-------------
3970

‎instrumentation/opentelemetry-instrumentation-genai-llama-index/src/opentelemetry/instrumentation/genai/llama_index/_handler.py‎

Lines changed: 63 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
ToolCall,
2222
ToolCallResult,
2323
)
24+
from llama_index.core.base.base_retriever import BaseRetriever
2425
from llama_index.core.base.llms.types import (
2526
AudioBlock,
2627
ChatMessage,
@@ -32,6 +33,7 @@
3233
)
3334
from llama_index.core.instrumentation.span import BaseSpan
3435
from llama_index.core.instrumentation.span_handlers import BaseSpanHandler
36+
from llama_index.core.schema import NodeWithScore, QueryBundle
3537
from llama_index.core.tools import BaseTool, FunctionTool, ToolOutput
3638
from pydantic import PrivateAttr
3739

@@ -41,6 +43,7 @@
4143
from opentelemetry.util.genai.invocation import (
4244
GenAIInvocation,
4345
LocalAgentInvocation,
46+
RetrievalInvocation,
4447
ToolInvocation,
4548
WorkflowInvocation,
4649
)
@@ -298,6 +301,52 @@ def _request_model(agent: BaseWorkflowAgent) -> str | None:
298301
return model_name if isinstance(model_name, str) and model_name else None
299302

300303

304+
def _retrieval_query(bound_args: inspect.BoundArguments) -> str | None:
305+
"""Extract text from either accepted LlamaIndex retrieval query form."""
306+
query = bound_args.arguments.get("str_or_query_bundle")
307+
if isinstance(query, str):
308+
return query
309+
if isinstance(query, QueryBundle):
310+
return query.query_str
311+
return None
312+
313+
314+
def _retrieval_top_k(retriever: BaseRetriever) -> int | None:
315+
"""Read the common top-k setting without requiring a retriever subtype."""
316+
top_k = getattr(retriever, "similarity_top_k", None)
317+
if isinstance(top_k, int) and not isinstance(top_k, bool):
318+
return top_k
319+
return None
320+
321+
322+
def _retrieval_documents(
323+
result: object,
324+
) -> list[dict[str, Any]] | None:
325+
"""Convert retrieved LlamaIndex nodes to semconv document objects."""
326+
if not isinstance(result, Sequence):
327+
return None
328+
candidates = cast(Sequence[object], result)
329+
documents: list[dict[str, Any]] = []
330+
for candidate in candidates:
331+
if not isinstance(candidate, NodeWithScore):
332+
continue
333+
try:
334+
document: dict[str, Any] = {
335+
"id": candidate.node_id,
336+
"content": candidate.node.get_content(),
337+
}
338+
if candidate.score is not None:
339+
document["score"] = candidate.score
340+
documents.append(document)
341+
except BaseException:
342+
continue
343+
# Preserve [] for a genuine empty result, but omit the attribute when a
344+
# non-empty result could not be converted into semantic-convention docs.
345+
if documents:
346+
return documents
347+
return [] if len(candidates) == 0 else None
348+
349+
301350
def _tool_attributes(
302351
candidate: object,
303352
) -> tuple[str, str, str | None] | None:
@@ -756,7 +805,7 @@ def finalize_workflow_agents(
756805

757806

758807
class LlamaIndexSpanHandler(BaseSpanHandler[_LlamaIndexInvocation]):
759-
"""Map LlamaIndex-owned agent and tool operations to GenAI spans."""
808+
"""Map LlamaIndex-owned agent, tool, and retrieval operations to spans."""
760809

761810
_handler: TelemetryHandler = PrivateAttr()
762811

@@ -784,7 +833,7 @@ def new_span(
784833
tags: dict[str, Any] | None = None,
785834
**kwargs: Any,
786835
) -> _LlamaIndexInvocation | None:
787-
"""Start GenAI invocations for LlamaIndex-owned agents and tools.
836+
"""Start GenAI invocations for agents, tools, and retrievers.
788837
789838
Provider inference is deliberately ignored so its own instrumentation
790839
can emit inference telemetry, and nested tool callbacks are deduplicated.
@@ -913,6 +962,15 @@ def new_span(
913962
if workflow_run_id is not None:
914963
parent.register_workflow_agent(workflow_run_id, agent)
915964
workflow_agent = agent
965+
elif isinstance(instance, BaseRetriever) and method_name in {
966+
"retrieve",
967+
"aretrieve",
968+
}:
969+
retrieval_invocation = self._handler.retrieval()
970+
retrieval_invocation.top_k = _retrieval_top_k(instance)
971+
if retrieval_invocation.should_capture_content:
972+
retrieval_invocation.query_text = _retrieval_query(bound_args)
973+
invocation = retrieval_invocation
916974
elif method_name == "call_tool" and isinstance(
917975
(tool_call := bound_args.arguments.get("ev")), ToolCall
918976
):
@@ -1138,6 +1196,9 @@ def prepare_to_exit_span(
11381196
if not _agent_step_is_complete(result):
11391197
self._expect_workflow_tools(span, result)
11401198
return span
1199+
elif isinstance(span._invocation, RetrievalInvocation):
1200+
if span._invocation.should_capture_content:
1201+
span._invocation.documents = _retrieval_documents(result)
11411202
elif isinstance(span._invocation, ToolInvocation):
11421203
span.reset_workflow_tool()
11431204
tool_output: ToolOutput | None = None
Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
# Copyright The OpenTelemetry Authors
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
from __future__ import annotations
5+
6+
from typing import Any
7+
8+
from llama_index.core.base.base_retriever import BaseRetriever
9+
from llama_index.core.schema import NodeWithScore, QueryBundle, TextNode
10+
11+
from opentelemetry.instrumentation.genai.llama_index import (
12+
LlamaIndexInstrumentor,
13+
)
14+
from opentelemetry.sdk._logs import LoggerProvider
15+
from opentelemetry.sdk.metrics import MeterProvider
16+
from opentelemetry.sdk.trace import TracerProvider
17+
from opentelemetry.test_util_genai.conformance import (
18+
ExpectedViolation,
19+
Scenario,
20+
)
21+
from opentelemetry.test_util_genai.instrumentor import instrument
22+
23+
24+
class _ConformanceRetriever(BaseRetriever):
25+
def __init__(self) -> None:
26+
super().__init__()
27+
self.similarity_top_k = 1
28+
29+
def _retrieve(self, query_bundle: QueryBundle) -> list[NodeWithScore]:
30+
return [
31+
NodeWithScore(
32+
node=TextNode(
33+
id_="capital-france",
34+
text="Paris is the capital of France.",
35+
),
36+
score=0.99,
37+
)
38+
]
39+
40+
41+
class RetrievalScenario(Scenario):
42+
expected_spans = {"retrieval": 1}
43+
expected_metrics = ("gen_ai.client.operation.duration",)
44+
expected_violations = (
45+
ExpectedViolation(
46+
advice_id="genai_expected_attribute_missing",
47+
message_substring="server.address",
48+
),
49+
)
50+
51+
def run(
52+
self,
53+
*,
54+
tracer_provider: TracerProvider,
55+
meter_provider: MeterProvider,
56+
logger_provider: LoggerProvider,
57+
vcr: Any,
58+
) -> None:
59+
with instrument(
60+
LlamaIndexInstrumentor(),
61+
tracer_provider=tracer_provider,
62+
logger_provider=logger_provider,
63+
meter_provider=meter_provider,
64+
content_capture="SPAN_ONLY",
65+
):
66+
_ConformanceRetriever().retrieve("What is the capital of France?")

‎instrumentation/opentelemetry-instrumentation-genai-llama-index/tests/test_conformance.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,13 @@
1414
from opentelemetry.test_util_genai.conformance import Scenario, run_conformance
1515

1616
from .conformance.agent import AgentScenario
17+
from .conformance.retrieval import RetrievalScenario
1718
from .conformance.workflow import WorkflowScenario
1819

1920

2021
@pytest.mark.parametrize(
2122
"scenario",
22-
[AgentScenario(), WorkflowScenario()],
23+
[AgentScenario(), WorkflowScenario(), RetrievalScenario()],
2324
ids=lambda scenario: type(scenario).__name__,
2425
)
2526
def test_conformance(

0 commit comments

Comments
 (0)