|
21 | 21 | ToolCall, |
22 | 22 | ToolCallResult, |
23 | 23 | ) |
| 24 | +from llama_index.core.base.base_retriever import BaseRetriever |
24 | 25 | from llama_index.core.base.llms.types import ( |
25 | 26 | AudioBlock, |
26 | 27 | ChatMessage, |
|
32 | 33 | ) |
33 | 34 | from llama_index.core.instrumentation.span import BaseSpan |
34 | 35 | from llama_index.core.instrumentation.span_handlers import BaseSpanHandler |
| 36 | +from llama_index.core.schema import NodeWithScore, QueryBundle |
35 | 37 | from llama_index.core.tools import BaseTool, FunctionTool, ToolOutput |
36 | 38 | from pydantic import PrivateAttr |
37 | 39 |
|
|
41 | 43 | from opentelemetry.util.genai.invocation import ( |
42 | 44 | GenAIInvocation, |
43 | 45 | LocalAgentInvocation, |
| 46 | + RetrievalInvocation, |
44 | 47 | ToolInvocation, |
45 | 48 | WorkflowInvocation, |
46 | 49 | ) |
@@ -298,6 +301,52 @@ def _request_model(agent: BaseWorkflowAgent) -> str | None: |
298 | 301 | return model_name if isinstance(model_name, str) and model_name else None |
299 | 302 |
|
300 | 303 |
|
| 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 | + |
301 | 350 | def _tool_attributes( |
302 | 351 | candidate: object, |
303 | 352 | ) -> tuple[str, str, str | None] | None: |
@@ -756,7 +805,7 @@ def finalize_workflow_agents( |
756 | 805 |
|
757 | 806 |
|
758 | 807 | 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.""" |
760 | 809 |
|
761 | 810 | _handler: TelemetryHandler = PrivateAttr() |
762 | 811 |
|
@@ -784,7 +833,7 @@ def new_span( |
784 | 833 | tags: dict[str, Any] | None = None, |
785 | 834 | **kwargs: Any, |
786 | 835 | ) -> _LlamaIndexInvocation | None: |
787 | | - """Start GenAI invocations for LlamaIndex-owned agents and tools. |
| 836 | + """Start GenAI invocations for agents, tools, and retrievers. |
788 | 837 |
|
789 | 838 | Provider inference is deliberately ignored so its own instrumentation |
790 | 839 | can emit inference telemetry, and nested tool callbacks are deduplicated. |
@@ -913,6 +962,15 @@ def new_span( |
913 | 962 | if workflow_run_id is not None: |
914 | 963 | parent.register_workflow_agent(workflow_run_id, agent) |
915 | 964 | 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 |
916 | 974 | elif method_name == "call_tool" and isinstance( |
917 | 975 | (tool_call := bound_args.arguments.get("ev")), ToolCall |
918 | 976 | ): |
@@ -1138,6 +1196,9 @@ def prepare_to_exit_span( |
1138 | 1196 | if not _agent_step_is_complete(result): |
1139 | 1197 | self._expect_workflow_tools(span, result) |
1140 | 1198 | return span |
| 1199 | + elif isinstance(span._invocation, RetrievalInvocation): |
| 1200 | + if span._invocation.should_capture_content: |
| 1201 | + span._invocation.documents = _retrieval_documents(result) |
1141 | 1202 | elif isinstance(span._invocation, ToolInvocation): |
1142 | 1203 | span.reset_workflow_tool() |
1143 | 1204 | tool_output: ToolOutput | None = None |
|
0 commit comments