util-genai | Add RetrievalInvocation and lifecycle API (#36)

* retrieval invocation

* copilot feedback

* addressing copilot feedback

* lint update

* towncrier fix

* clean up old changelog and docstring

* 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

* migrate away from using Any on attributes

---------

Co-authored-by: Liudmila Molkova <neskazu@gmail.com>

Author: keith-decker

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.

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

util/opentelemetry-util-genai/src/opentelemetry/util/genai/_retrieval_invocation.py

@@ -0,0 +1,142 @@
+# Copyright The OpenTelemetry Authors
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any, cast
+
+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,
+)
+from opentelemetry.util.types import AttributeValue
+
+
+class RetrievalInvocation(GenAIInvocation):
+    """Represents a single retrieval invocation (retrieval span).
+
+    Use handler.retrieval() 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.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: float | None = None
+        self.query_text: str | None = None
+        self.documents: Sequence[Mapping[str, Any]] | None = None
+        self._start(self._get_base_attributes())
+
+    def _get_base_attributes(self) -> dict[str, AttributeValue]:
+        """Return sampling-relevant attributes available at span creation time."""
+        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),
+            (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, AttributeValue]:
+        # data_source_id intentionally excluded — high cardinality
+        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, AttributeValue] = {
+            GenAI.GEN_AI_OPERATION_NAME: self._operation_name,
+            **{k: v for k, v in optional_attrs if v is not None},
+        }
+        # 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, AttributeValue]:
+        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: tuple[tuple[str, AttributeValue | None], ...] = (
+            (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, 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())
+        # 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)

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,
 )
@@ -202,6 +203,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,

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",
 ]

util/opentelemetry-util-genai/tests/test_handler_metrics.py

@@ -392,3 +392,130 @@ 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)