initial commit

Author: JWinermaSplunk

util/opentelemetry-util-genai/examples/retrievals_example.py

@@ -55,6 +55,13 @@
 GEN_AI_EMBEDDINGS_INPUT_TEXTS = "gen_ai.embeddings.input.texts"
 GEN_AI_REQUEST_ENCODING_FORMATS = "gen_ai.request.encoding_formats"
 
+# Retrieval attributes
+GEN_AI_RETRIEVAL_TYPE = "gen_ai.retrieval.type"
+GEN_AI_RETRIEVAL_QUERY_TEXT = "gen_ai.retrieval.query.text"
+GEN_AI_RETRIEVAL_TOP_K = "gen_ai.retrieval.top_k"
+GEN_AI_RETRIEVAL_DOCUMENTS_RETRIEVED = "gen_ai.retrieval.documents_retrieved"
+GEN_AI_RETRIEVAL_DOCUMENTS = "gen_ai.retrieval.documents"
+
 # Server attributes (from semantic conventions)
 SERVER_ADDRESS = "server.address"
 SERVER_PORT = "server.port"

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

@@ -16,6 +16,7 @@
 from ..types import (
     AgentInvocation,
     EmbeddingInvocation,
+    RetrievalInvocation,
     Error,
     LLMInvocation,
     ToolCall,
@@ -50,6 +51,9 @@ def __init__(self, meter: Optional[Meter] = None):
         self._agent_duration_histogram: Histogram = (
             instruments.agent_duration_histogram
         )
+        self._retrieval_duration_histogram: Histogram = (
+            instruments.retrieval_duration_histogram
+        )
 
     def on_start(self, obj: Any) -> None:  # no-op for metrics
         return None
@@ -146,6 +150,9 @@ def on_end(self, obj: Any) -> None:
                 span=getattr(embedding_invocation, "span", None),
             )
 
+        if isinstance(obj, RetrievalInvocation):
+            self._record_retrieval_metrics(obj)
+
     def on_error(self, error: Error, obj: Any) -> None:
         # Handle new agentic types
         if isinstance(obj, Workflow):
@@ -242,6 +249,9 @@ def on_error(self, error: Error, obj: Any) -> None:
                 span=getattr(embedding_invocation, "span", None),
             )
 
+        if isinstance(obj, RetrievalInvocation):
+            self._record_retrieval_metrics(obj, error)
+
     def handles(self, obj: Any) -> bool:
         return isinstance(
             obj,
@@ -251,6 +261,7 @@ def handles(self, obj: Any) -> bool:
                 Workflow,
                 AgentInvocation,
                 EmbeddingInvocation,
+                RetrievalInvocation,
             ),
         )
 
@@ -306,3 +317,40 @@ def _record_agent_metrics(self, agent: AgentInvocation) -> None:
         self._agent_duration_histogram.record(
             duration, attributes=metric_attrs, context=context
         )
+
+    def _record_retrieval_metrics(
+        self, retrieval: RetrievalInvocation, error: Optional[Error] = None
+    ) -> None:
+        """Record metrics for a retrieval operation."""
+        if retrieval.end_time is None:
+            return
+        duration = retrieval.end_time - retrieval.start_time
+        metric_attrs = {
+            GenAI.GEN_AI_OPERATION_NAME: retrieval.operation_name,
+        }
+        if retrieval.retriever_type:
+            metric_attrs["gen_ai.retrieval.type"] = retrieval.retriever_type
+        if retrieval.framework:
+            metric_attrs["gen_ai.framework"] = retrieval.framework
+        if retrieval.provider:
+            metric_attrs[GenAI.GEN_AI_PROVIDER_NAME] = retrieval.provider
+        # Add agent context if available
+        if retrieval.agent_name:
+            metric_attrs[GenAI.GEN_AI_AGENT_NAME] = retrieval.agent_name
+        if retrieval.agent_id:
+            metric_attrs[GenAI.GEN_AI_AGENT_ID] = retrieval.agent_id
+        # Add error type if present
+        if error is not None and getattr(error, "type", None) is not None:
+            metric_attrs[ErrorAttributes.ERROR_TYPE] = error.type.__qualname__
+
+        context = None
+        span = getattr(retrieval, "span", None)
+        if span is not None:
+            try:
+                context = trace.set_span_in_context(span)
+            except (ValueError, RuntimeError):  # pragma: no cover - defensive
+                context = None
+
+        self._retrieval_duration_histogram.record(
+            duration, attributes=metric_attrs, context=context
+        )

util/opentelemetry-util-genai/src/opentelemetry/util/genai/emitters/metrics.py

@@ -26,6 +26,10 @@
     GEN_AI_OUTPUT_MESSAGES,
     GEN_AI_PROVIDER_NAME,
     GEN_AI_REQUEST_ENCODING_FORMATS,
+    GEN_AI_RETRIEVAL_DOCUMENTS_RETRIEVED,
+    GEN_AI_RETRIEVAL_QUERY_TEXT,
+    GEN_AI_RETRIEVAL_TOP_K,
+    GEN_AI_RETRIEVAL_TYPE,
     GEN_AI_STEP_ASSIGNED_AGENT,
     GEN_AI_STEP_NAME,
     GEN_AI_STEP_OBJECTIVE,
@@ -45,6 +49,7 @@
     AgentInvocation,
     ContentCapturingMode,
     EmbeddingInvocation,
+    RetrievalInvocation,
     Error,
     LLMInvocation,
     Step,
@@ -201,9 +206,10 @@ def _apply_start_attrs(self, invocation: GenAIType):
         provider = getattr(invocation, "provider", None)
         if provider:
             span.set_attribute(GEN_AI_PROVIDER_NAME, provider)
-        # framework (named field)
-        if isinstance(invocation, LLMInvocation) and invocation.framework:
-            span.set_attribute("gen_ai.framework", invocation.framework)
+        # framework (named field) - applies to all invocation types
+        framework = getattr(invocation, "framework", None)
+        if framework:
+            span.set_attribute("gen_ai.framework", framework)
         # function definitions (semantic conv derived from structured list)
         if isinstance(invocation, LLMInvocation):
             _apply_function_definitions(span, invocation.request_functions)
@@ -302,6 +308,8 @@ def on_start(
             self._apply_start_attrs(invocation)
         elif isinstance(invocation, EmbeddingInvocation):
             self._start_embedding(invocation)
+        elif isinstance(invocation, RetrievalInvocation):
+            self._start_retrieval(invocation)
         else:
             # Use operation field for span name (defaults to "chat")
             operation = getattr(invocation, "operation", "chat")
@@ -335,6 +343,8 @@ def on_end(self, invocation: LLMInvocation | EmbeddingInvocation) -> None:
             self._finish_step(invocation)
         elif isinstance(invocation, EmbeddingInvocation):
             self._finish_embedding(invocation)
+        elif isinstance(invocation, RetrievalInvocation):
+            self._finish_retrieval(invocation)
         else:
             span = getattr(invocation, "span", None)
             if span is None:
@@ -359,6 +369,8 @@ def on_error(
             self._error_step(error, invocation)
         elif isinstance(invocation, EmbeddingInvocation):
             self._error_embedding(error, invocation)
+        elif isinstance(invocation, RetrievalInvocation):
+            self._error_retrieval(error, invocation)
         else:
             span = getattr(invocation, "span", None)
             if span is None:
@@ -771,3 +783,70 @@ def _error_embedding(
                 token.__exit__(None, None, None)  # type: ignore[misc]
             except Exception:
                 pass
+
+    # ---- Retrieval lifecycle ---------------------------------------------
+    def _start_retrieval(self, retrieval: RetrievalInvocation) -> None:
+        """Start a retrieval span."""
+        span_name = f"{retrieval.operation_name}"
+        if retrieval.provider:
+            span_name = f"{retrieval.operation_name} {retrieval.provider}"
+        parent_span = getattr(retrieval, "parent_span", None)
+        parent_ctx = (
+            trace.set_span_in_context(parent_span)
+            if parent_span is not None
+            else None
+        )
+        cm = self._tracer.start_as_current_span(
+            span_name,
+            kind=SpanKind.CLIENT,
+            end_on_exit=False,
+            context=parent_ctx,
+        )
+        span = cm.__enter__()
+        self._attach_span(retrieval, span, cm)
+        self._apply_start_attrs(retrieval)
+
+        # Set retrieval-specific start attributes
+        if retrieval.top_k is not None:
+            span.set_attribute(GEN_AI_RETRIEVAL_TOP_K, retrieval.top_k)
+        if self._capture_content and retrieval.query:
+            span.set_attribute(GEN_AI_RETRIEVAL_QUERY_TEXT, retrieval.query)
+
+    def _finish_retrieval(self, retrieval: RetrievalInvocation) -> None:
+        """Finish a retrieval span."""
+        span = retrieval.span
+        if span is None:
+            return
+        # Apply finish-time semantic conventions
+        if retrieval.documents_retrieved is not None:
+            span.set_attribute(
+                GEN_AI_RETRIEVAL_DOCUMENTS_RETRIEVED,
+                retrieval.documents_retrieved,
+            )
+        token = retrieval.context_token
+        if token is not None and hasattr(token, "__exit__"):
+            try:
+                token.__exit__(None, None, None)  # type: ignore[misc]
+            except Exception:
+                pass
+        span.end()
+
+    def _error_retrieval(
+        self, error: Error, retrieval: RetrievalInvocation
+    ) -> None:
+        """Fail a retrieval span with error status."""
+        span = retrieval.span
+        if span is None:
+            return
+        span.set_status(Status(StatusCode.ERROR, error.message))
+        if span.is_recording():
+            span.set_attribute(
+                ErrorAttributes.ERROR_TYPE, error.type.__qualname__
+            )
+        token = retrieval.context_token
+        if token is not None and hasattr(token, "__exit__"):
+            try:
+                token.__exit__(None, None, None)  # type: ignore[misc]
+            except Exception:
+                pass
+        span.end()

util/opentelemetry-util-genai/src/opentelemetry/util/genai/emitters/span.py

@@ -87,6 +87,7 @@ def genai_debug_log(*_args: Any, **_kwargs: Any) -> None:  # type: ignore
     AgentInvocation,
     ContentCapturingMode,
     EmbeddingInvocation,
+    RetrievalInvocation,
     Error,
     EvaluationResult,
     GenAI,
@@ -475,6 +476,70 @@ def fail_embedding(
                 pass
         return invocation
 
+    def start_retrieval(
+        self, invocation: RetrievalInvocation
+    ) -> RetrievalInvocation:
+        """Start a retrieval invocation and create a pending span entry."""
+        self._refresh_capture_content()
+        if (
+            not invocation.agent_name or not invocation.agent_id
+        ) and self._agent_context_stack:
+            top_name, top_id = self._agent_context_stack[-1]
+            if not invocation.agent_name:
+                invocation.agent_name = top_name
+            if not invocation.agent_id:
+                invocation.agent_id = top_id
+        invocation.start_time = time.time()
+        self._emitter.on_start(invocation)
+        span = getattr(invocation, "span", None)
+        if span is not None:
+            self._span_registry[str(invocation.run_id)] = span
+        self._entity_registry[str(invocation.run_id)] = invocation
+        return invocation
+
+    def stop_retrieval(
+        self, invocation: RetrievalInvocation
+    ) -> RetrievalInvocation:
+        """Finalize a retrieval invocation successfully and end its span."""
+        invocation.end_time = time.time()
+
+        # Determine if this invocation should be sampled for evaluation
+        invocation.sample_for_evaluation = self._should_sample_for_evaluation(
+            invocation.trace_id
+        )
+
+        self._emitter.on_end(invocation)
+        self._notify_completion(invocation)
+        self._entity_registry.pop(str(invocation.run_id), None)
+        # Force flush metrics if a custom provider with force_flush is present
+        if (
+            hasattr(self, "_meter_provider")
+            and self._meter_provider is not None
+        ):
+            try:  # pragma: no cover
+                self._meter_provider.force_flush()  # type: ignore[attr-defined]
+            except Exception:
+                pass
+        return invocation
+
+    def fail_retrieval(
+        self, invocation: RetrievalInvocation, error: Error
+    ) -> RetrievalInvocation:
+        """Fail a retrieval invocation and end its span with error status."""
+        invocation.end_time = time.time()
+        self._emitter.on_error(error, invocation)
+        self._notify_completion(invocation)
+        self._entity_registry.pop(str(invocation.run_id), None)
+        if (
+            hasattr(self, "_meter_provider")
+            and self._meter_provider is not None
+        ):
+            try:  # pragma: no cover
+                self._meter_provider.force_flush()  # type: ignore[attr-defined]
+            except Exception:
+                pass
+        return invocation
+
     # ToolCall lifecycle --------------------------------------------------
     def start_tool_call(self, invocation: ToolCall) -> ToolCall:
         """Start a tool call invocation and create a pending span entry."""
@@ -880,6 +945,8 @@ def start(self, obj: Any) -> Any:
             return self.start_llm(obj)
         if isinstance(obj, EmbeddingInvocation):
             return self.start_embedding(obj)
+        if isinstance(obj, RetrievalInvocation):
+            return self.start_retrieval(obj)
         if isinstance(obj, ToolCall):
             return self.start_tool_call(obj)
         return obj
@@ -960,6 +1027,8 @@ def finish(self, obj: Any) -> Any:
             return self.stop_llm(obj)
         if isinstance(obj, EmbeddingInvocation):
             return self.stop_embedding(obj)
+        if isinstance(obj, RetrievalInvocation):
+            return self.stop_retrieval(obj)
         if isinstance(obj, ToolCall):
             return self.stop_tool_call(obj)
         return obj
@@ -976,6 +1045,8 @@ def fail(self, obj: Any, error: Error) -> Any:
             return self.fail_llm(obj, error)
         if isinstance(obj, EmbeddingInvocation):
             return self.fail_embedding(obj, error)
+        if isinstance(obj, RetrievalInvocation):
+            return self.fail_retrieval(obj, error)
         if isinstance(obj, ToolCall):
             return self.fail_tool_call(obj, error)
         return obj

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

@@ -42,3 +42,8 @@ def __init__(self, meter: Meter):
             unit="s",
             description="Duration of agent operations",
         )
+        self.retrieval_duration_histogram: Histogram = meter.create_histogram(
+            name="gen_ai.retrieval.duration",
+            unit="s",
+            description="Duration of retrieval operations",
+        )

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

@@ -317,6 +317,42 @@ class EmbeddingInvocation(GenAI):
     )
     error_type: Optional[str] = None
 
+@dataclass
+class RetrievalInvocation(GenAI):
+    """Represents a single retrieval/search invocation."""
+
+    #Required attribute
+    operation_name: str = field(
+        default="retrieval",
+        metadata={"semconv": GenAIAttributes.GEN_AI_OPERATION_NAME},
+    )
+    
+    # Recommended attributes
+    retriever_type: Optional[str] = field(
+        default=None,
+        metadata={"semconv": "gen_ai.retrieval.type"},
+    )
+    query: Optional[str] = field(
+        default=None,
+        metadata={"semconv": "gen_ai.retrieval.query.text"},
+    )
+    top_k: Optional[int] = field(
+        default=None,
+        metadata={"semconv": "gen_ai.retrieval.top_k"},
+    )
+    documents_retrieved: Optional[int] = field(
+        default=None,
+        metadata={"semconv": "gen_ai.retrieval.documents_retrieved"},
+    )
+    
+    # Opt-in attribute
+    results: list[dict[str, Any]] = field(
+        default_factory=list,
+        metadata={"semconv": "gen_ai.retrieval.documents"},
+    )
+    
+    # Additional utility fields (not in semantic conventions)
+    query_vector: Optional[list[float]] = None
 
 @dataclass
 class Workflow(GenAI):