From 35c949c94aefb42643920d7d1013eb19f27dc89b Mon Sep 17 00:00:00 2001 From: Keith Decker Date: Wed, 13 May 2026 14:12:10 -0600 Subject: [PATCH 1/8] retrieval invocation --- util/opentelemetry-util-genai/AGENTS.md | 5 +- .../util/genai/_retrieval_invocation.py | 139 +++++++ .../src/opentelemetry/util/genai/handler.py | 30 ++ .../opentelemetry/util/genai/invocation.py | 3 + .../tests/test_handler_retrieval.py | 375 ++++++++++++++++++ 5 files changed, 550 insertions(+), 2 deletions(-) create mode 100644 util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py create mode 100644 util/opentelemetry-util-genai/tests/test_handler_retrieval.py diff --git a/util/opentelemetry-util-genai/AGENTS.md b/util/opentelemetry-util-genai/AGENTS.md index 4eaa4fc8..6199652e 100644 --- a/util/opentelemetry-util-genai/AGENTS.md +++ b/util/opentelemetry-util-genai/AGENTS.md @@ -41,18 +41,19 @@ Factory methods on `TelemetryHandler` (`handler.py`): - `inference(provider, request_model, *, server_address, server_port)` → `InferenceInvocation` - `embedding(provider, request_model, *, server_address, server_port)` → `EmbeddingInvocation` +- `retrieval(*, data_source_id, provider, request_model, server_address, server_port)` → `RetrievalInvocation` - `tool(name, *, arguments, tool_call_id, tool_type, tool_description)` → `ToolInvocation` - `workflow(name)` → `WorkflowInvocation` The returned object can also be used as a context manager (`with ... as invocation:`) when the span lifetime maps cleanly to a `with` block. The above factories must map 1:1 to distinct semconv operation types (inference, embeddings, -tool execution, agent invocation, workflow invocation). Names must match the operation +retrieval, tool execution, agent invocation, workflow invocation). Names must match the operation unambiguously — for example, `create_agent` and `invoke_agent` are different operations, so a single `agent()` would be ambiguous and is not acceptable. Add a new factory per operation instead. -Factory names are Python-style singular verbs (`inference`, `embedding`, `tool`, `workflow`); the op names +Factory names are Python-style singular verbs (`inference`, `embedding`, `retrieval`, `tool`, `workflow`); the op names they map to follow semconv operations. Factory methods must accept all attributes that semconv marks as important for sampling diff --git a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py new file mode 100644 index 00000000..c33aea06 --- /dev/null +++ b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py @@ -0,0 +1,139 @@ +# Copyright The OpenTelemetry Authors +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import Any + +from opentelemetry._logs import Logger +from opentelemetry.semconv._incubating.attributes import ( + gen_ai_attributes as GenAI, +) +from opentelemetry.semconv.attributes import server_attributes +from opentelemetry.trace import SpanKind, Tracer +from opentelemetry.util.genai._invocation import Error, GenAIInvocation +from opentelemetry.util.genai.completion_hook import CompletionHook +from opentelemetry.util.genai.metrics import InvocationMetricsRecorder +from opentelemetry.util.genai.utils import ( + ContentCapturingMode, + gen_ai_json_dumps, + get_content_capturing_mode, + is_experimental_mode, +) + + +class RetrievalInvocation(GenAIInvocation): + """Represents a single retrieval invocation (retrieval span). + + Use handler.start_retrieval() or the handler.retrieval() context manager + rather than constructing this directly. + + Reference: https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-spans.md#retrievals + + Semantic convention attributes for retrieval spans: + - gen_ai.operation.name: "retrieval" (Required) + - error.type: Error type if operation failed (Conditionally Required) + - gen_ai.data_source.id: Data source identifier (Conditionally Required, when applicable) + - gen_ai.provider.name: Provider name (Conditionally Required, when applicable) + - gen_ai.request.model: Model name if applicable (Conditionally Required, if available) + - server.port: Server port (Conditionally Required, if server.address is set) + - gen_ai.request.top_k: Top-k sampling setting (Recommended) + - server.address: Server address (Recommended) + - gen_ai.retrieval.documents: Retrieved documents (Opt-In, may contain sensitive data) + - gen_ai.retrieval.query.text: Query text (Opt-In, may contain sensitive data) + """ + + def __init__( + self, + tracer: Tracer, + metrics_recorder: InvocationMetricsRecorder, + logger: Logger, + completion_hook: CompletionHook, + *, + data_source_id: str | None = None, + provider: str | None = None, + request_model: str | None = None, + server_address: str | None = None, + server_port: int | None = None, + ) -> None: + """Use handler.start_retrieval() or handler.retrieval() instead of calling this directly.""" + _operation_name = GenAI.GenAiOperationNameValues.RETRIEVAL.value + super().__init__( + tracer, + metrics_recorder, + logger, + completion_hook, + operation_name=_operation_name, + span_name=f"{_operation_name} {data_source_id}" + if data_source_id + else _operation_name, + span_kind=SpanKind.CLIENT, + ) + self.data_source_id = data_source_id + self.provider = provider + self.request_model = request_model + self.server_address = server_address + self.server_port = server_port + self.top_k: int | None = None + self.query_text: str | None = None + self.documents: Any = None + self._start(self._get_base_attributes()) + + def _get_base_attributes(self) -> dict[str, Any]: + """Return sampling-relevant attributes available at span creation time.""" + optional_attrs = ( + (GenAI.GEN_AI_DATA_SOURCE_ID, self.data_source_id), + (GenAI.GEN_AI_PROVIDER_NAME, self.provider), + (GenAI.GEN_AI_REQUEST_MODEL, self.request_model), + (server_attributes.SERVER_ADDRESS, self.server_address), + (server_attributes.SERVER_PORT, self.server_port), + ) + return { + GenAI.GEN_AI_OPERATION_NAME: self._operation_name, + **{k: v for k, v in optional_attrs if v is not None}, + } + + def _get_metric_attributes(self) -> dict[str, Any]: + # data_source_id intentionally excluded — high cardinality + optional_attrs = ( + (GenAI.GEN_AI_PROVIDER_NAME, self.provider), + (GenAI.GEN_AI_REQUEST_MODEL, self.request_model), + (server_attributes.SERVER_ADDRESS, self.server_address), + (server_attributes.SERVER_PORT, self.server_port), + ) + attrs: dict[str, Any] = { + GenAI.GEN_AI_OPERATION_NAME: self._operation_name, + **{k: v for k, v in optional_attrs if v is not None}, + } + attrs.update(self.metric_attributes) + return attrs + + def _get_content_attributes_for_span(self) -> dict[str, Any]: + if not self.span.is_recording(): + return {} + if not is_experimental_mode() or get_content_capturing_mode() not in ( + ContentCapturingMode.SPAN_ONLY, + ContentCapturingMode.SPAN_AND_EVENT, + ): + return {} + optional_attrs = ( + (GenAI.GEN_AI_RETRIEVAL_QUERY_TEXT, self.query_text), + ( + GenAI.GEN_AI_RETRIEVAL_DOCUMENTS, + gen_ai_json_dumps(self.documents) + if self.documents is not None + else None, + ), + ) + return {k: v for k, v in optional_attrs if v is not None} + + def _apply_finish(self, error: Error | None = None) -> None: + if error is not None: + self._apply_error_attributes(error) + attributes: dict[str, Any] = {} + if self.top_k is not None: + attributes[GenAI.GEN_AI_REQUEST_TOP_K] = self.top_k + attributes.update(self._get_content_attributes_for_span()) + attributes.update(self.attributes) + self.span.set_attributes(attributes) + self._metrics_recorder.record(self) diff --git a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/handler.py b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/handler.py index 3ab7335c..7d2ad70a 100644 --- a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/handler.py +++ b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/handler.py @@ -63,6 +63,7 @@ from opentelemetry.util.genai.invocation import ( EmbeddingInvocation, InferenceInvocation, + RetrievalInvocation, ToolInvocation, WorkflowInvocation, ) @@ -204,6 +205,35 @@ def start_embedding( server_port=server_port, ) + def retrieval( + self, + *, + data_source_id: str | None = None, + provider: str | None = None, + request_model: str | None = None, + server_address: str | None = None, + server_port: int | None = None, + ) -> RetrievalInvocation: + """Returns a Retrieval invocation. Starts span when called. + + Returned object can be used as a ContextManager which automatically calls `stop` or `fail` + to finalize the span upon exiting. If not used as a ContextManager, the caller is + responsible for calling `stop` or `fail` to finalize the span. + + Only set data attributes on the invocation object, do not modify the span or context. + """ + return RetrievalInvocation( + self._tracer, + self._metrics_recorder, + self._logger, + self._completion_hook, + data_source_id=data_source_id, + provider=provider, + request_model=request_model, + server_address=server_address, + server_port=server_port, + ) + def start_tool( self, name: str, diff --git a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/invocation.py b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/invocation.py index 4cb4c298..2d779910 100644 --- a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/invocation.py +++ b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/invocation.py @@ -10,6 +10,7 @@ GenAIInvocation, InferenceInvocation, EmbeddingInvocation, + RetrievalInvocation, ToolInvocation, WorkflowInvocation, ) @@ -23,6 +24,7 @@ Error, GenAIInvocation, ) +from opentelemetry.util.genai._retrieval_invocation import RetrievalInvocation from opentelemetry.util.genai._tool_invocation import ToolInvocation from opentelemetry.util.genai._workflow_invocation import WorkflowInvocation @@ -33,6 +35,7 @@ "GenAIInvocation", "InferenceInvocation", "EmbeddingInvocation", + "RetrievalInvocation", "ToolInvocation", "WorkflowInvocation", ] diff --git a/util/opentelemetry-util-genai/tests/test_handler_retrieval.py b/util/opentelemetry-util-genai/tests/test_handler_retrieval.py new file mode 100644 index 00000000..be0b92e6 --- /dev/null +++ b/util/opentelemetry-util-genai/tests/test_handler_retrieval.py @@ -0,0 +1,375 @@ +# Copyright The OpenTelemetry Authors +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import json +import os +from unittest import TestCase +from unittest.mock import patch + +import pytest + +from opentelemetry.instrumentation._semconv import ( + OTEL_SEMCONV_STABILITY_OPT_IN, + _OpenTelemetrySemanticConventionStability, +) +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from opentelemetry.sdk.trace.sampling import Decision, SamplingResult +from opentelemetry.semconv._incubating.attributes import ( + gen_ai_attributes as GenAI, +) +from opentelemetry.trace import INVALID_SPAN, SpanKind +from opentelemetry.trace.status import StatusCode +from opentelemetry.util.genai.environment_variables import ( + OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT, +) +from opentelemetry.util.genai.handler import TelemetryHandler +from opentelemetry.util.genai.invocation import RetrievalInvocation +from opentelemetry.util.genai.types import Error + + +class _RetrievalTestBase(TestCase): + def setUp(self) -> None: + self.span_exporter = InMemorySpanExporter() + self.tracer_provider = TracerProvider() + self.tracer_provider.add_span_processor( + SimpleSpanProcessor(self.span_exporter) + ) + self.handler = TelemetryHandler( + tracer_provider=self.tracer_provider, + ) + + def _get_finished_spans(self): + return self.span_exporter.get_finished_spans() + + +class TelemetryHandlerRetrievalTest(_RetrievalTestBase): # pylint: disable=too-many-public-methods + # ------------------------------------------------------------------ + # start_retrieval + # ------------------------------------------------------------------ + + def test_start_retrieval_creates_span(self) -> None: + invocation = self.handler.start_retrieval() + self.assertIsNot(invocation.span, INVALID_SPAN) + invocation.stop() + + def test_start_retrieval_span_name_with_data_source_id(self) -> None: + invocation = self.handler.start_retrieval(data_source_id="H7STPQYOND") + invocation.stop() + + spans = self._get_finished_spans() + self.assertEqual(len(spans), 1) + self.assertEqual(spans[0].name, "retrieval H7STPQYOND") + + def test_start_retrieval_span_name_without_data_source_id(self) -> None: + invocation = self.handler.start_retrieval() + invocation.stop() + + spans = self._get_finished_spans() + self.assertEqual(len(spans), 1) + self.assertEqual(spans[0].name, "retrieval") + + def test_start_retrieval_span_kind_is_client(self) -> None: + invocation = self.handler.start_retrieval() + invocation.stop() + + spans = self._get_finished_spans() + self.assertEqual(spans[0].kind, SpanKind.CLIENT) + + def test_start_retrieval_records_monotonic_start(self) -> None: + with patch("timeit.default_timer", return_value=42.0): + invocation = self.handler.start_retrieval() + self.assertEqual(invocation._monotonic_start_s, 42.0) + invocation.stop() + + # ------------------------------------------------------------------ + # stop (required + conditionally required attributes) + # ------------------------------------------------------------------ + + def test_stop_sets_operation_name(self) -> None: + invocation = self.handler.start_retrieval() + invocation.stop() + + spans = self._get_finished_spans() + self.assertEqual( + spans[0].attributes[GenAI.GEN_AI_OPERATION_NAME], "retrieval" + ) + + def test_stop_sets_data_source_id(self) -> None: + invocation = self.handler.start_retrieval(data_source_id="DS123") + invocation.stop() + + spans = self._get_finished_spans() + self.assertEqual( + spans[0].attributes[GenAI.GEN_AI_DATA_SOURCE_ID], "DS123" + ) + + def test_stop_sets_provider_name(self) -> None: + invocation = self.handler.start_retrieval(provider="pinecone") + invocation.stop() + + spans = self._get_finished_spans() + self.assertEqual( + spans[0].attributes[GenAI.GEN_AI_PROVIDER_NAME], "pinecone" + ) + + def test_stop_sets_request_model(self) -> None: + invocation = self.handler.start_retrieval( + request_model="text-embedding-ada-002" + ) + invocation.stop() + + spans = self._get_finished_spans() + self.assertEqual( + spans[0].attributes[GenAI.GEN_AI_REQUEST_MODEL], + "text-embedding-ada-002", + ) + + def test_stop_sets_server_address_and_port(self) -> None: + invocation = self.handler.start_retrieval( + server_address="db.example.com", server_port=443 + ) + invocation.stop() + + spans = self._get_finished_spans() + attrs = spans[0].attributes + self.assertEqual(attrs["server.address"], "db.example.com") + self.assertEqual(attrs["server.port"], 443) + + # ------------------------------------------------------------------ + # stop (recommended + opt-in attributes set after construction) + # ------------------------------------------------------------------ + + def test_stop_sets_top_k(self) -> None: + invocation = self.handler.start_retrieval() + invocation.top_k = 10.0 + invocation.stop() + + spans = self._get_finished_spans() + value = spans[0].attributes[GenAI.GEN_AI_REQUEST_TOP_K] + self.assertIsInstance(value, float) + self.assertEqual(value, 10.0) + + @patch.dict( + os.environ, + { + OTEL_SEMCONV_STABILITY_OPT_IN: "gen_ai_latest_experimental", + OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT: "SPAN_ONLY", + }, + ) + def test_stop_sets_query_text_when_content_capture_enabled(self) -> None: + _OpenTelemetrySemanticConventionStability._initialized = False + _OpenTelemetrySemanticConventionStability._initialize() + + invocation = self.handler.start_retrieval() + invocation.query_text = "What is the capital of France?" + invocation.stop() + + spans = self._get_finished_spans() + value = spans[0].attributes[GenAI.GEN_AI_RETRIEVAL_QUERY_TEXT] + self.assertIsInstance(value, str) + self.assertEqual(value, "What is the capital of France?") + + def test_stop_suppresses_query_text_when_content_capture_disabled( + self, + ) -> None: + invocation = self.handler.start_retrieval() + invocation.query_text = "What is the capital of France?" + invocation.stop() + + spans = self._get_finished_spans() + self.assertNotIn( + GenAI.GEN_AI_RETRIEVAL_QUERY_TEXT, spans[0].attributes + ) + + @patch.dict( + os.environ, + { + OTEL_SEMCONV_STABILITY_OPT_IN: "gen_ai_latest_experimental", + OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT: "SPAN_ONLY", + }, + ) + def test_stop_sets_documents_when_content_capture_enabled(self) -> None: + _OpenTelemetrySemanticConventionStability._initialized = False + _OpenTelemetrySemanticConventionStability._initialize() + + docs = [{"id": "doc_1", "score": 0.95}, {"id": "doc_2", "score": 0.87}] + invocation = self.handler.start_retrieval() + invocation.documents = docs + invocation.stop() + + spans = self._get_finished_spans() + raw = spans[0].attributes[GenAI.GEN_AI_RETRIEVAL_DOCUMENTS] + self.assertIsInstance(raw, str) + self.assertEqual(json.loads(raw), docs) + + def test_stop_suppresses_documents_when_content_capture_disabled( + self, + ) -> None: + docs = [{"id": "doc_1", "score": 0.95}] + invocation = self.handler.start_retrieval() + invocation.documents = docs + invocation.stop() + + spans = self._get_finished_spans() + self.assertNotIn(GenAI.GEN_AI_RETRIEVAL_DOCUMENTS, spans[0].attributes) + + def test_stop_sets_custom_attributes(self) -> None: + invocation = self.handler.start_retrieval() + invocation.attributes["custom.key"] = "value" + invocation.stop() + + spans = self._get_finished_spans() + self.assertEqual(spans[0].attributes["custom.key"], "value") + + def test_stop_omits_none_attributes(self) -> None: + invocation = self.handler.start_retrieval() + invocation.stop() + + spans = self._get_finished_spans() + attrs = spans[0].attributes + self.assertNotIn(GenAI.GEN_AI_DATA_SOURCE_ID, attrs) + self.assertNotIn(GenAI.GEN_AI_PROVIDER_NAME, attrs) + self.assertNotIn(GenAI.GEN_AI_REQUEST_MODEL, attrs) + self.assertNotIn(GenAI.GEN_AI_REQUEST_TOP_K, attrs) + + # ------------------------------------------------------------------ + # fail + # ------------------------------------------------------------------ + + def test_fail_sets_error_status(self) -> None: + invocation = self.handler.start_retrieval() + invocation.fail(Error(message="timeout", type=TimeoutError)) + + spans = self._get_finished_spans() + self.assertEqual(spans[0].status.status_code, StatusCode.ERROR) + self.assertEqual(spans[0].status.description, "timeout") + + def test_fail_sets_error_type_attribute(self) -> None: + invocation = self.handler.start_retrieval() + invocation.fail(Error(message="bad", type=ConnectionError)) + + spans = self._get_finished_spans() + self.assertEqual(spans[0].attributes["error.type"], "ConnectionError") + + def test_fail_sets_operation_name(self) -> None: + invocation = self.handler.start_retrieval() + invocation.fail(Error(message="err", type=RuntimeError)) + + spans = self._get_finished_spans() + self.assertEqual( + spans[0].attributes[GenAI.GEN_AI_OPERATION_NAME], "retrieval" + ) + + def test_fail_with_exception_instance(self) -> None: + invocation = self.handler.start_retrieval() + invocation.fail(ValueError("oops")) + + spans = self._get_finished_spans() + self.assertEqual(spans[0].status.status_code, StatusCode.ERROR) + self.assertEqual(spans[0].attributes["error.type"], "ValueError") + + +class TelemetryHandlerRetrievalContextManagerTest(_RetrievalTestBase): + # ------------------------------------------------------------------ + # retrieval context manager + # ------------------------------------------------------------------ + + def test_context_manager_creates_and_ends_span(self) -> None: + with self.handler.retrieval(data_source_id="DS1") as inv: + self.assertIsNot(inv.span, INVALID_SPAN) + + spans = self._get_finished_spans() + self.assertEqual(len(spans), 1) + self.assertEqual(spans[0].name, "retrieval DS1") + + def test_context_manager_default_invocation(self) -> None: + with self.handler.retrieval() as inv: + self.assertIsInstance(inv, RetrievalInvocation) + self.assertIsNone(inv.data_source_id) + self.assertEqual(inv._operation_name, "retrieval") + + def test_context_manager_success_has_unset_status(self) -> None: + with self.handler.retrieval(): + pass + + spans = self._get_finished_spans() + self.assertEqual(spans[0].status.status_code, StatusCode.UNSET) + + def test_context_manager_reraises_exception(self) -> None: + with pytest.raises(ValueError, match="lookup failed"): + with self.handler.retrieval(): + raise ValueError("lookup failed") + + def test_context_manager_marks_error_on_exception(self) -> None: + with pytest.raises(RuntimeError): + with self.handler.retrieval(): + raise RuntimeError("store down") + + spans = self._get_finished_spans() + self.assertEqual(spans[0].status.status_code, StatusCode.ERROR) + self.assertEqual(spans[0].attributes["error.type"], "RuntimeError") + + def test_context_manager_sets_attributes_on_span(self) -> None: + with self.handler.retrieval(provider="weaviate") as inv: + inv.top_k = 5.0 + + spans = self._get_finished_spans() + attrs = spans[0].attributes + self.assertEqual(attrs[GenAI.GEN_AI_PROVIDER_NAME], "weaviate") + self.assertEqual(attrs[GenAI.GEN_AI_REQUEST_TOP_K], 5.0) + + +class TelemetryHandlerRetrievalSamplingTest(_RetrievalTestBase): + def test_sampling_attributes_available_at_span_creation(self) -> None: + """Sampling-relevant attributes must be present at start_span() time.""" + captured_attributes: dict = {} + + class AttributeCapturingSampler: # pylint: disable=no-self-use + def should_sample( + self, + parent_context, + trace_id, + name, + kind=None, + attributes=None, + links=None, + ): + captured_attributes.update(attributes or {}) + return SamplingResult(Decision.RECORD_AND_SAMPLE, attributes) + + def get_description(self): + return "AttributeCapturingSampler" + + sampler_provider = TracerProvider(sampler=AttributeCapturingSampler()) + sampler_provider.add_span_processor( + SimpleSpanProcessor(self.span_exporter) + ) + handler = TelemetryHandler(tracer_provider=sampler_provider) + + invocation = handler.start_retrieval( + data_source_id="DS42", + provider="pinecone", + server_address="db.example.com", + server_port=443, + ) + invocation.stop() + + self.assertEqual( + captured_attributes[GenAI.GEN_AI_OPERATION_NAME], "retrieval" + ) + self.assertEqual( + captured_attributes[GenAI.GEN_AI_DATA_SOURCE_ID], "DS42" + ) + self.assertEqual( + captured_attributes[GenAI.GEN_AI_PROVIDER_NAME], "pinecone" + ) + self.assertEqual( + captured_attributes["server.address"], "db.example.com" + ) + self.assertEqual(captured_attributes["server.port"], 443) From 12d733c3482e30055e51e36bf9dd0ce071980ae1 Mon Sep 17 00:00:00 2001 From: Keith Decker Date: Mon, 18 May 2026 10:25:11 -0600 Subject: [PATCH 2/8] copilot feedback --- .../util/genai/_retrieval_invocation.py | 4 +- .../tests/test_handler_metrics.py | 125 ++++++++++++++++++ .../tests/test_handler_retrieval.py | 67 +++++----- 3 files changed, 161 insertions(+), 35 deletions(-) diff --git a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py index c33aea06..db9a2b34 100644 --- a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py +++ b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py @@ -25,7 +25,7 @@ class RetrievalInvocation(GenAIInvocation): """Represents a single retrieval invocation (retrieval span). - Use handler.start_retrieval() or the handler.retrieval() context manager + Use handler.retrieval() rather than constructing this directly. rather than constructing this directly. Reference: https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-spans.md#retrievals @@ -56,7 +56,7 @@ def __init__( server_address: str | None = None, server_port: int | None = None, ) -> None: - """Use handler.start_retrieval() or handler.retrieval() instead of calling this directly.""" + """Use handler.retrieval() instead of calling this directly.""" _operation_name = GenAI.GenAiOperationNameValues.RETRIEVAL.value super().__init__( tracer, diff --git a/util/opentelemetry-util-genai/tests/test_handler_metrics.py b/util/opentelemetry-util-genai/tests/test_handler_metrics.py index 37a13cdd..c18a7dc2 100644 --- a/util/opentelemetry-util-genai/tests/test_handler_metrics.py +++ b/util/opentelemetry-util-genai/tests/test_handler_metrics.py @@ -392,3 +392,128 @@ def test_fail_tool_records_duration_with_error(self) -> None: ) self.assertAlmostEqual(duration_point.sum, 1.5, places=3) self.assertNotIn("gen_ai.client.token.usage", metrics) + + +class TelemetryHandlerRetrievalMetricsTest(TestBase): + def _harvest_metrics(self) -> Dict[str, List[Any]]: + metrics = self.get_sorted_metrics() + metrics_by_name: Dict[str, List[Any]] = {} + for metric in metrics or []: + points = metric.data.data_points or [] + metrics_by_name.setdefault(metric.name, []).extend(points) + return metrics_by_name + + def test_stop_retrieval_records_duration(self) -> None: + handler = TelemetryHandler( + tracer_provider=self.tracer_provider, + meter_provider=self.meter_provider, + ) + with patch("timeit.default_timer", return_value=1000.0): + invocation = handler.retrieval( + provider="pinecone", request_model="text-embedding-ada-002" + ) + + with patch("timeit.default_timer", return_value=1001.5): + invocation.stop() + + metrics = self._harvest_metrics() + self.assertIn("gen_ai.client.operation.duration", metrics) + duration_points = metrics["gen_ai.client.operation.duration"] + self.assertEqual(len(duration_points), 1) + duration_point = duration_points[0] + + self.assertEqual( + duration_point.attributes[GenAI.GEN_AI_OPERATION_NAME], + GenAI.GenAiOperationNameValues.RETRIEVAL.value, + ) + self.assertEqual( + duration_point.attributes[GenAI.GEN_AI_PROVIDER_NAME], "pinecone" + ) + self.assertEqual( + duration_point.attributes[GenAI.GEN_AI_REQUEST_MODEL], + "text-embedding-ada-002", + ) + self.assertAlmostEqual(duration_point.sum, 1.5, places=3) + self.assertNotIn("gen_ai.client.token.usage", metrics) + + def test_stop_retrieval_excludes_data_source_id_from_metrics(self) -> None: + handler = TelemetryHandler( + tracer_provider=self.tracer_provider, + meter_provider=self.meter_provider, + ) + invocation = handler.retrieval( + data_source_id="DS_HIGH_CARDINALITY", provider="weaviate" + ) + invocation.stop() + + metrics = self._harvest_metrics() + self.assertIn("gen_ai.client.operation.duration", metrics) + duration_points = metrics["gen_ai.client.operation.duration"] + self.assertEqual(len(duration_points), 1) + duration_point = duration_points[0] + + self.assertNotIn(GenAI.GEN_AI_DATA_SOURCE_ID, duration_point.attributes) + self.assertEqual( + duration_point.attributes[GenAI.GEN_AI_PROVIDER_NAME], "weaviate" + ) + + def test_stop_retrieval_records_duration_with_additional_attributes( + self, + ) -> None: + handler = TelemetryHandler( + tracer_provider=self.tracer_provider, + meter_provider=self.meter_provider, + ) + invocation = handler.retrieval( + provider="pinecone", + server_address="db.example.com", + server_port=443, + ) + invocation.metric_attributes = {"custom.retrieval.attr": "val"} + invocation.attributes = {"should not be on metrics": "value"} + invocation.stop() + + metrics = self._harvest_metrics() + self.assertIn("gen_ai.client.operation.duration", metrics) + duration_points = metrics["gen_ai.client.operation.duration"] + self.assertEqual(len(duration_points), 1) + duration_point = duration_points[0] + + self.assertEqual( + duration_point.attributes["server.address"], "db.example.com" + ) + self.assertEqual(duration_point.attributes["server.port"], 443) + self.assertEqual( + duration_point.attributes["custom.retrieval.attr"], "val" + ) + self.assertIsNone( + duration_point.attributes.get("should not be on metrics") + ) + + def test_fail_retrieval_records_duration_with_error(self) -> None: + handler = TelemetryHandler( + tracer_provider=self.tracer_provider, + meter_provider=self.meter_provider, + ) + with patch("timeit.default_timer", return_value=2000.0): + invocation = handler.retrieval(provider="pinecone") + + error = Error(message="retrieval failed", type=ConnectionError) + with patch("timeit.default_timer", return_value=2003.0): + invocation.fail(error) + + metrics = self._harvest_metrics() + self.assertIn("gen_ai.client.operation.duration", metrics) + duration_points = metrics["gen_ai.client.operation.duration"] + self.assertEqual(len(duration_points), 1) + duration_point = duration_points[0] + + self.assertEqual( + duration_point.attributes["error.type"], "ConnectionError" + ) + self.assertEqual( + duration_point.attributes[GenAI.GEN_AI_OPERATION_NAME], + GenAI.GenAiOperationNameValues.RETRIEVAL.value, + ) + self.assertAlmostEqual(duration_point.sum, 3.0, places=3) + self.assertNotIn("gen_ai.client.token.usage", metrics) diff --git a/util/opentelemetry-util-genai/tests/test_handler_retrieval.py b/util/opentelemetry-util-genai/tests/test_handler_retrieval.py index be0b92e6..b399a274 100644 --- a/util/opentelemetry-util-genai/tests/test_handler_retrieval.py +++ b/util/opentelemetry-util-genai/tests/test_handler_retrieval.py @@ -50,40 +50,40 @@ def _get_finished_spans(self): class TelemetryHandlerRetrievalTest(_RetrievalTestBase): # pylint: disable=too-many-public-methods # ------------------------------------------------------------------ - # start_retrieval + # retrieval # ------------------------------------------------------------------ - def test_start_retrieval_creates_span(self) -> None: - invocation = self.handler.start_retrieval() + def test_retrieval_creates_span(self) -> None: + invocation = self.handler.retrieval() self.assertIsNot(invocation.span, INVALID_SPAN) invocation.stop() - def test_start_retrieval_span_name_with_data_source_id(self) -> None: - invocation = self.handler.start_retrieval(data_source_id="H7STPQYOND") + def test_retrieval_span_name_with_data_source_id(self) -> None: + invocation = self.handler.retrieval(data_source_id="H7STPQYOND") invocation.stop() spans = self._get_finished_spans() self.assertEqual(len(spans), 1) self.assertEqual(spans[0].name, "retrieval H7STPQYOND") - def test_start_retrieval_span_name_without_data_source_id(self) -> None: - invocation = self.handler.start_retrieval() + def test_retrieval_span_name_without_data_source_id(self) -> None: + invocation = self.handler.retrieval() invocation.stop() spans = self._get_finished_spans() self.assertEqual(len(spans), 1) self.assertEqual(spans[0].name, "retrieval") - def test_start_retrieval_span_kind_is_client(self) -> None: - invocation = self.handler.start_retrieval() + def test_retrieval_span_kind_is_client(self) -> None: + invocation = self.handler.retrieval() invocation.stop() spans = self._get_finished_spans() self.assertEqual(spans[0].kind, SpanKind.CLIENT) - def test_start_retrieval_records_monotonic_start(self) -> None: + def test_retrieval_records_monotonic_start(self) -> None: with patch("timeit.default_timer", return_value=42.0): - invocation = self.handler.start_retrieval() + invocation = self.handler.retrieval() self.assertEqual(invocation._monotonic_start_s, 42.0) invocation.stop() @@ -92,7 +92,7 @@ def test_start_retrieval_records_monotonic_start(self) -> None: # ------------------------------------------------------------------ def test_stop_sets_operation_name(self) -> None: - invocation = self.handler.start_retrieval() + invocation = self.handler.retrieval() invocation.stop() spans = self._get_finished_spans() @@ -101,7 +101,7 @@ def test_stop_sets_operation_name(self) -> None: ) def test_stop_sets_data_source_id(self) -> None: - invocation = self.handler.start_retrieval(data_source_id="DS123") + invocation = self.handler.retrieval(data_source_id="DS123") invocation.stop() spans = self._get_finished_spans() @@ -110,7 +110,7 @@ def test_stop_sets_data_source_id(self) -> None: ) def test_stop_sets_provider_name(self) -> None: - invocation = self.handler.start_retrieval(provider="pinecone") + invocation = self.handler.retrieval(provider="pinecone") invocation.stop() spans = self._get_finished_spans() @@ -119,7 +119,7 @@ def test_stop_sets_provider_name(self) -> None: ) def test_stop_sets_request_model(self) -> None: - invocation = self.handler.start_retrieval( + invocation = self.handler.retrieval( request_model="text-embedding-ada-002" ) invocation.stop() @@ -131,7 +131,7 @@ def test_stop_sets_request_model(self) -> None: ) def test_stop_sets_server_address_and_port(self) -> None: - invocation = self.handler.start_retrieval( + invocation = self.handler.retrieval( server_address="db.example.com", server_port=443 ) invocation.stop() @@ -146,14 +146,14 @@ def test_stop_sets_server_address_and_port(self) -> None: # ------------------------------------------------------------------ def test_stop_sets_top_k(self) -> None: - invocation = self.handler.start_retrieval() - invocation.top_k = 10.0 + invocation = self.handler.retrieval() + invocation.top_k = 10 invocation.stop() spans = self._get_finished_spans() value = spans[0].attributes[GenAI.GEN_AI_REQUEST_TOP_K] - self.assertIsInstance(value, float) - self.assertEqual(value, 10.0) + self.assertIsInstance(value, int) + self.assertEqual(value, 10) @patch.dict( os.environ, @@ -166,7 +166,7 @@ def test_stop_sets_query_text_when_content_capture_enabled(self) -> None: _OpenTelemetrySemanticConventionStability._initialized = False _OpenTelemetrySemanticConventionStability._initialize() - invocation = self.handler.start_retrieval() + invocation = self.handler.retrieval() invocation.query_text = "What is the capital of France?" invocation.stop() @@ -178,7 +178,7 @@ def test_stop_sets_query_text_when_content_capture_enabled(self) -> None: def test_stop_suppresses_query_text_when_content_capture_disabled( self, ) -> None: - invocation = self.handler.start_retrieval() + invocation = self.handler.retrieval() invocation.query_text = "What is the capital of France?" invocation.stop() @@ -199,7 +199,7 @@ def test_stop_sets_documents_when_content_capture_enabled(self) -> None: _OpenTelemetrySemanticConventionStability._initialize() docs = [{"id": "doc_1", "score": 0.95}, {"id": "doc_2", "score": 0.87}] - invocation = self.handler.start_retrieval() + invocation = self.handler.retrieval() invocation.documents = docs invocation.stop() @@ -212,7 +212,7 @@ def test_stop_suppresses_documents_when_content_capture_disabled( self, ) -> None: docs = [{"id": "doc_1", "score": 0.95}] - invocation = self.handler.start_retrieval() + invocation = self.handler.retrieval() invocation.documents = docs invocation.stop() @@ -220,7 +220,7 @@ def test_stop_suppresses_documents_when_content_capture_disabled( self.assertNotIn(GenAI.GEN_AI_RETRIEVAL_DOCUMENTS, spans[0].attributes) def test_stop_sets_custom_attributes(self) -> None: - invocation = self.handler.start_retrieval() + invocation = self.handler.retrieval() invocation.attributes["custom.key"] = "value" invocation.stop() @@ -228,7 +228,7 @@ def test_stop_sets_custom_attributes(self) -> None: self.assertEqual(spans[0].attributes["custom.key"], "value") def test_stop_omits_none_attributes(self) -> None: - invocation = self.handler.start_retrieval() + invocation = self.handler.retrieval() invocation.stop() spans = self._get_finished_spans() @@ -243,7 +243,7 @@ def test_stop_omits_none_attributes(self) -> None: # ------------------------------------------------------------------ def test_fail_sets_error_status(self) -> None: - invocation = self.handler.start_retrieval() + invocation = self.handler.retrieval() invocation.fail(Error(message="timeout", type=TimeoutError)) spans = self._get_finished_spans() @@ -251,14 +251,14 @@ def test_fail_sets_error_status(self) -> None: self.assertEqual(spans[0].status.description, "timeout") def test_fail_sets_error_type_attribute(self) -> None: - invocation = self.handler.start_retrieval() + invocation = self.handler.retrieval() invocation.fail(Error(message="bad", type=ConnectionError)) spans = self._get_finished_spans() self.assertEqual(spans[0].attributes["error.type"], "ConnectionError") def test_fail_sets_operation_name(self) -> None: - invocation = self.handler.start_retrieval() + invocation = self.handler.retrieval() invocation.fail(Error(message="err", type=RuntimeError)) spans = self._get_finished_spans() @@ -267,7 +267,7 @@ def test_fail_sets_operation_name(self) -> None: ) def test_fail_with_exception_instance(self) -> None: - invocation = self.handler.start_retrieval() + invocation = self.handler.retrieval() invocation.fail(ValueError("oops")) spans = self._get_finished_spans() @@ -317,12 +317,13 @@ def test_context_manager_marks_error_on_exception(self) -> None: def test_context_manager_sets_attributes_on_span(self) -> None: with self.handler.retrieval(provider="weaviate") as inv: - inv.top_k = 5.0 + inv.top_k = 5 spans = self._get_finished_spans() attrs = spans[0].attributes self.assertEqual(attrs[GenAI.GEN_AI_PROVIDER_NAME], "weaviate") - self.assertEqual(attrs[GenAI.GEN_AI_REQUEST_TOP_K], 5.0) + self.assertIsInstance(attrs[GenAI.GEN_AI_REQUEST_TOP_K], int) + self.assertEqual(attrs[GenAI.GEN_AI_REQUEST_TOP_K], 5) class TelemetryHandlerRetrievalSamplingTest(_RetrievalTestBase): @@ -352,7 +353,7 @@ def get_description(self): ) handler = TelemetryHandler(tracer_provider=sampler_provider) - invocation = handler.start_retrieval( + invocation = handler.retrieval( data_source_id="DS42", provider="pinecone", server_address="db.example.com", From 06689a9c66c81dbcec3763f58dcf62bed4080788 Mon Sep 17 00:00:00 2001 From: Keith Decker Date: Mon, 18 May 2026 12:02:02 -0600 Subject: [PATCH 3/8] addressing copilot feedback --- util/opentelemetry-util-genai/tests/test_handler_retrieval.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/util/opentelemetry-util-genai/tests/test_handler_retrieval.py b/util/opentelemetry-util-genai/tests/test_handler_retrieval.py index b399a274..ffebd28b 100644 --- a/util/opentelemetry-util-genai/tests/test_handler_retrieval.py +++ b/util/opentelemetry-util-genai/tests/test_handler_retrieval.py @@ -44,6 +44,9 @@ def setUp(self) -> None: tracer_provider=self.tracer_provider, ) + def tearDown(self) -> None: + _OpenTelemetrySemanticConventionStability._initialized = False + def _get_finished_spans(self): return self.span_exporter.get_finished_spans() From 2029e371376e2c0683c9905cc95dcbc8949f2b59 Mon Sep 17 00:00:00 2001 From: Keith Decker Date: Mon, 18 May 2026 12:31:16 -0600 Subject: [PATCH 4/8] lint update --- util/opentelemetry-util-genai/tests/test_handler_metrics.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/util/opentelemetry-util-genai/tests/test_handler_metrics.py b/util/opentelemetry-util-genai/tests/test_handler_metrics.py index c18a7dc2..d5eaa9b0 100644 --- a/util/opentelemetry-util-genai/tests/test_handler_metrics.py +++ b/util/opentelemetry-util-genai/tests/test_handler_metrics.py @@ -452,7 +452,9 @@ def test_stop_retrieval_excludes_data_source_id_from_metrics(self) -> None: self.assertEqual(len(duration_points), 1) duration_point = duration_points[0] - self.assertNotIn(GenAI.GEN_AI_DATA_SOURCE_ID, duration_point.attributes) + self.assertNotIn( + GenAI.GEN_AI_DATA_SOURCE_ID, duration_point.attributes + ) self.assertEqual( duration_point.attributes[GenAI.GEN_AI_PROVIDER_NAME], "weaviate" ) From a810ae24cbf97ea638800a2400fdebbb51d91ff2 Mon Sep 17 00:00:00 2001 From: Keith Decker Date: Mon, 18 May 2026 12:41:17 -0600 Subject: [PATCH 5/8] towncrier fix --- .../opentelemetry-instrumentation-langchain/.changelog/36.added | 1 + util/opentelemetry-util-genai/.changelog/36.added | 1 + 2 files changed, 2 insertions(+) create mode 100644 instrumentation/opentelemetry-instrumentation-langchain/.changelog/36.added create mode 100644 util/opentelemetry-util-genai/.changelog/36.added diff --git a/instrumentation/opentelemetry-instrumentation-langchain/.changelog/36.added b/instrumentation/opentelemetry-instrumentation-langchain/.changelog/36.added new file mode 100644 index 00000000..eac55e68 --- /dev/null +++ b/instrumentation/opentelemetry-instrumentation-langchain/.changelog/36.added @@ -0,0 +1 @@ +Add retrieval span support via `on_retriever_start`, `on_retriever_end`, and `on_retriever_error` callbacks using `RetrievalInvocation`. diff --git a/util/opentelemetry-util-genai/.changelog/36.added b/util/opentelemetry-util-genai/.changelog/36.added new file mode 100644 index 00000000..8b9890cd --- /dev/null +++ b/util/opentelemetry-util-genai/.changelog/36.added @@ -0,0 +1 @@ +Add `RetrievalInvocation` type with `start_retrieval` / `retrieval` span lifecycle, supporting `gen_ai.operation.name=retrieval` spans per the GenAI semantic conventions. From 39c8e7f4f266c5ebdbd7f55b2d268a8513d58396 Mon Sep 17 00:00:00 2001 From: Keith Decker Date: Tue, 19 May 2026 09:31:46 -0600 Subject: [PATCH 6/8] clean up old changelog and docstring --- .../opentelemetry-instrumentation-langchain/.changelog/36.added | 1 - .../src/opentelemetry/util/genai/_retrieval_invocation.py | 1 - 2 files changed, 2 deletions(-) delete mode 100644 instrumentation/opentelemetry-instrumentation-langchain/.changelog/36.added diff --git a/instrumentation/opentelemetry-instrumentation-langchain/.changelog/36.added b/instrumentation/opentelemetry-instrumentation-langchain/.changelog/36.added deleted file mode 100644 index eac55e68..00000000 --- a/instrumentation/opentelemetry-instrumentation-langchain/.changelog/36.added +++ /dev/null @@ -1 +0,0 @@ -Add retrieval span support via `on_retriever_start`, `on_retriever_end`, and `on_retriever_error` callbacks using `RetrievalInvocation`. diff --git a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py index db9a2b34..99deba9e 100644 --- a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py +++ b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py @@ -26,7 +26,6 @@ class RetrievalInvocation(GenAIInvocation): """Represents a single retrieval invocation (retrieval span). Use handler.retrieval() rather than constructing this directly. - rather than constructing this directly. Reference: https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-spans.md#retrievals From db685d07c7f321b4075e87678da7ceb83bb8b64d Mon Sep 17 00:00:00 2001 From: Keith Decker Date: Tue, 19 May 2026 11:21:03 -0600 Subject: [PATCH 7/8] use float for top_k per semconvs https://github.com/open-telemetry/semantic-conventions-genai/blob/main/docs/registry/attributes/gen-ai.md#gen-ai-request-top-k --- .../util/genai/_retrieval_invocation.py | 2 +- .../tests/test_handler_retrieval.py | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py index 99deba9e..1da1fb35 100644 --- a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py +++ b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py @@ -73,7 +73,7 @@ def __init__( self.request_model = request_model self.server_address = server_address self.server_port = server_port - self.top_k: int | None = None + self.top_k: float | None = None self.query_text: str | None = None self.documents: Any = None self._start(self._get_base_attributes()) diff --git a/util/opentelemetry-util-genai/tests/test_handler_retrieval.py b/util/opentelemetry-util-genai/tests/test_handler_retrieval.py index ffebd28b..32346cb4 100644 --- a/util/opentelemetry-util-genai/tests/test_handler_retrieval.py +++ b/util/opentelemetry-util-genai/tests/test_handler_retrieval.py @@ -150,13 +150,13 @@ def test_stop_sets_server_address_and_port(self) -> None: def test_stop_sets_top_k(self) -> None: invocation = self.handler.retrieval() - invocation.top_k = 10 + invocation.top_k = 10.0 invocation.stop() spans = self._get_finished_spans() value = spans[0].attributes[GenAI.GEN_AI_REQUEST_TOP_K] - self.assertIsInstance(value, int) - self.assertEqual(value, 10) + self.assertIsInstance(value, float) + self.assertEqual(value, 10.0) @patch.dict( os.environ, @@ -320,13 +320,13 @@ def test_context_manager_marks_error_on_exception(self) -> None: def test_context_manager_sets_attributes_on_span(self) -> None: with self.handler.retrieval(provider="weaviate") as inv: - inv.top_k = 5 + inv.top_k = 5.0 spans = self._get_finished_spans() attrs = spans[0].attributes self.assertEqual(attrs[GenAI.GEN_AI_PROVIDER_NAME], "weaviate") - self.assertIsInstance(attrs[GenAI.GEN_AI_REQUEST_TOP_K], int) - self.assertEqual(attrs[GenAI.GEN_AI_REQUEST_TOP_K], 5) + self.assertIsInstance(attrs[GenAI.GEN_AI_REQUEST_TOP_K], float) + self.assertEqual(attrs[GenAI.GEN_AI_REQUEST_TOP_K], 5.0) class TelemetryHandlerRetrievalSamplingTest(_RetrievalTestBase): From bf52d3bc4215c604af4f4370875f55b29c523ea6 Mon Sep 17 00:00:00 2001 From: Keith Decker Date: Tue, 26 May 2026 14:18:55 -0600 Subject: [PATCH 8/8] migrate away from using Any on attributes --- .../util/genai/_retrieval_invocation.py | 28 +++++++++++-------- 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py index 1da1fb35..55bf3885 100644 --- a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py +++ b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py @@ -3,7 +3,8 @@ from __future__ import annotations -from typing import Any +from collections.abc import Mapping, Sequence +from typing import Any, cast from opentelemetry._logs import Logger from opentelemetry.semconv._incubating.attributes import ( @@ -20,6 +21,7 @@ get_content_capturing_mode, is_experimental_mode, ) +from opentelemetry.util.types import AttributeValue class RetrievalInvocation(GenAIInvocation): @@ -75,12 +77,12 @@ def __init__( self.server_port = server_port self.top_k: float | None = None self.query_text: str | None = None - self.documents: Any = None + self.documents: Sequence[Mapping[str, Any]] | None = None self._start(self._get_base_attributes()) - def _get_base_attributes(self) -> dict[str, Any]: + def _get_base_attributes(self) -> dict[str, AttributeValue]: """Return sampling-relevant attributes available at span creation time.""" - optional_attrs = ( + optional_attrs: tuple[tuple[str, AttributeValue | None], ...] = ( (GenAI.GEN_AI_DATA_SOURCE_ID, self.data_source_id), (GenAI.GEN_AI_PROVIDER_NAME, self.provider), (GenAI.GEN_AI_REQUEST_MODEL, self.request_model), @@ -92,22 +94,23 @@ def _get_base_attributes(self) -> dict[str, Any]: **{k: v for k, v in optional_attrs if v is not None}, } - def _get_metric_attributes(self) -> dict[str, Any]: + def _get_metric_attributes(self) -> dict[str, AttributeValue]: # data_source_id intentionally excluded — high cardinality - optional_attrs = ( + optional_attrs: tuple[tuple[str, AttributeValue | None], ...] = ( (GenAI.GEN_AI_PROVIDER_NAME, self.provider), (GenAI.GEN_AI_REQUEST_MODEL, self.request_model), (server_attributes.SERVER_ADDRESS, self.server_address), (server_attributes.SERVER_PORT, self.server_port), ) - attrs: dict[str, Any] = { + attrs: dict[str, AttributeValue] = { GenAI.GEN_AI_OPERATION_NAME: self._operation_name, **{k: v for k, v in optional_attrs if v is not None}, } - attrs.update(self.metric_attributes) + # TODO: remove cast once base class metric_attributes is typed as dict[str, AttributeValue] + attrs.update(cast(dict[str, AttributeValue], self.metric_attributes)) return attrs - def _get_content_attributes_for_span(self) -> dict[str, Any]: + def _get_content_attributes_for_span(self) -> dict[str, AttributeValue]: if not self.span.is_recording(): return {} if not is_experimental_mode() or get_content_capturing_mode() not in ( @@ -115,7 +118,7 @@ def _get_content_attributes_for_span(self) -> dict[str, Any]: ContentCapturingMode.SPAN_AND_EVENT, ): return {} - optional_attrs = ( + optional_attrs: tuple[tuple[str, AttributeValue | None], ...] = ( (GenAI.GEN_AI_RETRIEVAL_QUERY_TEXT, self.query_text), ( GenAI.GEN_AI_RETRIEVAL_DOCUMENTS, @@ -129,10 +132,11 @@ def _get_content_attributes_for_span(self) -> dict[str, Any]: def _apply_finish(self, error: Error | None = None) -> None: if error is not None: self._apply_error_attributes(error) - attributes: dict[str, Any] = {} + attributes: dict[str, AttributeValue] = {} if self.top_k is not None: attributes[GenAI.GEN_AI_REQUEST_TOP_K] = self.top_k attributes.update(self._get_content_attributes_for_span()) - attributes.update(self.attributes) + # TODO: remove cast once base class self.attributes is typed as dict[str, AttributeValue] + attributes.update(cast(dict[str, AttributeValue], self.attributes)) self.span.set_attributes(attributes) self._metrics_recorder.record(self)