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. 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..55bf3885 --- /dev/null +++ b/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) 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 6567241b..e24958db 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, ) @@ -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, 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_metrics.py b/util/opentelemetry-util-genai/tests/test_handler_metrics.py index 37a13cdd..d5eaa9b0 100644 --- a/util/opentelemetry-util-genai/tests/test_handler_metrics.py +++ b/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) 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..32346cb4 --- /dev/null +++ b/util/opentelemetry-util-genai/tests/test_handler_retrieval.py @@ -0,0 +1,379 @@ +# 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 tearDown(self) -> None: + _OpenTelemetrySemanticConventionStability._initialized = False + + def _get_finished_spans(self): + return self.span_exporter.get_finished_spans() + + +class TelemetryHandlerRetrievalTest(_RetrievalTestBase): # pylint: disable=too-many-public-methods + # ------------------------------------------------------------------ + # retrieval + # ------------------------------------------------------------------ + + def test_retrieval_creates_span(self) -> None: + invocation = self.handler.retrieval() + self.assertIsNot(invocation.span, INVALID_SPAN) + invocation.stop() + + 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_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_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_retrieval_records_monotonic_start(self) -> None: + with patch("timeit.default_timer", return_value=42.0): + invocation = self.handler.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.assertIsInstance(attrs[GenAI.GEN_AI_REQUEST_TOP_K], float) + 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.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)