Add embedding invocation type

Author: shuningc

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

@@ -82,9 +82,15 @@
 from opentelemetry.util.genai.span_utils import (
     _apply_error_attributes,
     _apply_llm_finish_attributes,
+    _apply_embedding_finish_attributes,
     _maybe_emit_llm_event,
+    _maybe_emit_embedding_event,
+)
+from opentelemetry.util.genai.types import (
+    Error,
+    LLMInvocation,
+    EmbeddingInvocation,
 )
-from opentelemetry.util.genai.types import Error, LLMInvocation
 from opentelemetry.util.genai.version import __version__
 
 
@@ -131,6 +137,18 @@ def _record_llm_metrics(
             error_type=error_type,
         )
 
+    def _record_embedding_metrics(
+        self,
+        invocation: EmbeddingInvocation,
+        span: Span | None = None,
+        *,
+        error_type: str | None = None,
+    ) -> None:
+        # Metrics recorder currently supports LLMInvocation fields only.
+        # Keep embedding metrics as a no-op until dedicated embedding
+        # metric support is added.
+        return
+
     def start_llm(
         self,
         invocation: LLMInvocation,
@@ -208,6 +226,62 @@ def llm(
             raise
         self.stop_llm(invocation)
 
+    def start_embedding(
+        self, invocation: EmbeddingInvocation
+    ) -> EmbeddingInvocation:
+        """Start an embedding invocation and create a pending span entry."""
+
+        span = self._tracer.start_span(
+            name=f"{invocation.operation_name} {invocation.request_model}",
+            kind=SpanKind.CLIENT,
+        )
+        # Record a monotonic start timestamp (seconds) for duration
+        # calculation using timeit.default_timer.
+        invocation.monotonic_start_s = timeit.default_timer()
+        invocation.span = span
+        invocation.context_token = otel_context.attach(
+            set_span_in_context(span)
+        )
+        return invocation
+
+    def stop_embedding(
+        self, invocation: EmbeddingInvocation
+    ) -> EmbeddingInvocation:
+        """Finalize an embedding invocation successfully and end its span."""
+        if invocation.context_token is None or invocation.span is None:
+            # TODO: Provide feedback that this invocation was not started
+            return invocation
+
+        span = invocation.span
+        _apply_embedding_finish_attributes(span, invocation)
+        self._record_embedding_metrics(invocation, span)
+        _maybe_emit_embedding_event(self._logger, span, invocation)
+        # Detach context and end span
+        otel_context.detach(invocation.context_token)
+        span.end()
+        return invocation
+
+    def fail_embedding(
+        self, invocation: EmbeddingInvocation, error: Error
+    ) -> EmbeddingInvocation:
+        """Fail an embedding invocation and end its span with error status."""
+        if invocation.context_token is None or invocation.span is None:
+            # TODO: Provide feedback that this invocation was not started
+            return invocation
+
+        span = invocation.span
+        _apply_embedding_finish_attributes(invocation.span, invocation)
+        _apply_error_attributes(invocation.span, error)
+        error_type = getattr(error.type, "__qualname__", None)
+        self._record_embedding_metrics(
+            invocation, span, error_type=error_type
+        )
+        _maybe_emit_embedding_event(self._logger, span, invocation, error)
+        # Detach context and end span
+        otel_context.detach(invocation.context_token)
+        span.end()
+        return invocation
+
 
 def get_telemetry_handler(
     tracer_provider: TracerProvider | None = None,

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

@@ -35,6 +35,7 @@
     Error,
     InputMessage,
     LLMInvocation,
+    EmbeddingInvocation,
     MessagePart,
     OutputMessage,
 )
@@ -68,11 +69,35 @@ def _get_llm_common_attributes(
     }
 
 
+def _get_embedding_common_attributes(
+    invocation: EmbeddingInvocation,
+) -> dict[str, Any]:
+    """Get common Embedding attributes shared by finish() and error() paths.
+
+    Returns a dictionary of attributes.
+    """
+    optional_attrs = (
+        (server_attributes.SERVER_ADDRESS, invocation.server_address),
+        (server_attributes.SERVER_PORT, invocation.server_port),
+    )
+
+    return {
+        GenAI.GEN_AI_OPERATION_NAME: invocation.operation_name,
+        GenAI.GEN_AI_PROVIDER_NAME: invocation.provider,
+        **{key: value for key, value in optional_attrs if value is not None},
+    }
+
+
 def _get_llm_span_name(invocation: LLMInvocation) -> str:
     """Get the span name for an LLM invocation."""
     return f"{invocation.operation_name} {invocation.request_model}".strip()
 
 
+def _get_embedding_span_name(invocation: EmbeddingInvocation) -> str:
+    """Get the span name for an Embedding invocation."""
+    return f"{invocation.operation_name} {invocation.request_model}".strip()
+
+
 def _get_llm_messages_attributes_for_span(
     input_messages: list[InputMessage],
     output_messages: list[OutputMessage],
@@ -192,6 +217,44 @@ def _maybe_emit_llm_event(
     logger.emit(event)
 
 
+def _maybe_emit_embedding_event(
+    logger: Logger | None,
+    span: Span,
+    invocation: EmbeddingInvocation,
+    error: Error | None = None,
+) -> None:
+    """Emit a gen_ai.client.inference.operation.details event to the logger.
+
+    This function creates a LogRecord event following the semantic convention
+    for gen_ai.client.inference.operation.details as specified in the GenAI
+    event semantic conventions.
+
+    For more details, see the semantic convention documentation:
+    https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-events.md#event-eventgen_aiclientinferenceoperationdetails
+    """
+    if not is_experimental_mode() or not should_emit_event() or logger is None:
+        return
+
+    # Build event attributes by reusing the attribute getter functions
+    attributes: dict[str, Any] = {}
+    attributes.update(_get_embedding_common_attributes(invocation))
+    attributes.update(_get_embedding_request_attributes(invocation))
+    attributes.update(_get_embedding_response_attributes(invocation))
+
+    # Add error.type if operation ended in error
+    if error is not None:
+        attributes[error_attributes.ERROR_TYPE] = error.type.__qualname__
+
+    # Create and emit the event
+    context = set_span_in_context(span, get_current())
+    event = LogRecord(
+        event_name="gen_ai.client.embedding.operation.details",
+        attributes=attributes,
+        context=context,
+    )
+    logger.emit(event)
+
+
 def _apply_llm_finish_attributes(
     span: Span, invocation: LLMInvocation
 ) -> None:
@@ -218,6 +281,26 @@ def _apply_llm_finish_attributes(
         span.set_attributes(attributes)
 
 
+def _apply_embedding_finish_attributes(
+    span: Span, invocation: EmbeddingInvocation
+) -> None:
+    """Apply attributes common to embedding finish() paths."""
+    # Update span name
+    span.update_name(_get_embedding_span_name(invocation))
+
+    # Build all attributes by reusing the attribute getter functions
+    attributes: dict[str, Any] = {}
+    attributes.update(_get_embedding_common_attributes(invocation))
+    attributes.update(_get_embedding_request_attributes(invocation))
+    attributes.update(_get_embedding_response_attributes(invocation))
+
+    attributes.update(invocation.attributes)
+
+    # Set all attributes on the span
+    if attributes:
+        span.set_attributes(attributes)
+
+
 def _apply_error_attributes(span: Span, error: Error) -> None:
     """Apply status and error attributes common to error() paths."""
     span.set_status(Status(StatusCode.ERROR, error.message))
@@ -244,6 +327,19 @@ def _get_llm_request_attributes(
     return {key: value for key, value in optional_attrs if value is not None}
 
 
+def _get_embedding_request_attributes(
+    invocation: EmbeddingInvocation,
+) -> dict[str, Any]:
+    """Get GenAI request semantic convention attributes."""
+    optional_attrs = (
+        (GenAI.GEN_AI_REQUEST_MODEL, invocation.request_model),
+        (GenAI.GEN_AI_EMBEDDING_DIMENSION_COUNT, invocation.dimension_count),
+        (GenAI.GEN_AI_REQUEST_ENCODING_FORMATS, invocation.encoding_formats),
+    )
+
+    return {key: value for key, value in optional_attrs if value is not None}
+
+
 def _get_llm_response_attributes(
     invocation: LLMInvocation,
 ) -> dict[str, Any]:
@@ -279,6 +375,17 @@ def _get_llm_response_attributes(
     return {key: value for key, value in optional_attrs if value is not None}
 
 
+def _get_embedding_response_attributes(
+    invocation: EmbeddingInvocation,
+) -> dict[str, Any]:
+    """Get GenAI response semantic convention attributes."""
+    optional_attrs = (
+        (GenAI.GEN_AI_USAGE_INPUT_TOKENS, invocation.input_tokens),
+    )
+
+    return {key: value for key, value in optional_attrs if value is not None}
+
+
 __all__ = [
     "_apply_llm_finish_attributes",
     "_apply_error_attributes",
@@ -287,4 +394,10 @@ def _get_llm_response_attributes(
     "_get_llm_response_attributes",
     "_get_llm_span_name",
     "_maybe_emit_llm_event",
+    "_apply_embedding_finish_attributes",
+    "_get_embedding_common_attributes",
+    "_get_embedding_request_attributes",
+    "_get_embedding_response_attributes",
+    "_get_embedding_span_name",
+    "_maybe_emit_embedding_event",
 ]

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

@@ -236,6 +236,61 @@ class LLMInvocation(GenAIInvocation):
     monotonic_start_s: float | None = None
 
 
+@dataclass
+class EmbeddingInvocation(GenAIInvocation):
+    """
+    Represents a single embedding model invocation. When creating an
+    EmbeddingInvocation object, only update the data attributes. The span
+    and context_token attributes are set by the TelemetryHandler.
+    """
+
+    operation_name: str = field(
+        default=GenAI.GenAiOperationNameValues.EMBEDDINGS.value,
+        metadata={"semconv": GenAI.GEN_AI_OPERATION_NAME},
+    )
+    provider: str | None = None  # e.g., azure.ai.openai, openai, aws.bedrock
+
+    request_model: str | None = field(
+        default=None,
+        metadata={"semconv": GenAI.GEN_AI_REQUEST_MODEL},
+    )
+
+    server_address: str | None = None
+    server_port: int | None = None
+    error_type: str | None = None
+
+    # encoding_formats can be multi-value -> combinational cardinality risk.
+    # Keep on spans/events only.
+    encoding_formats: list[str] = field(
+        default_factory=list,
+        metadata={"semconv": GenAI.GEN_AI_REQUEST_ENCODING_FORMATS},
+    )
+
+    input_tokens: int | None = field(
+        default=None,
+        metadata={"semconv": GenAI.GEN_AI_USAGE_INPUT_TOKENS},
+    )
+    dimension_count: int | None = None
+
+    attributes: dict[str, Any] = field(default_factory=_new_str_any_dict)
+    """
+    Additional attributes to set on spans and/or events. These attributes
+    will not be set on metrics.
+    """
+
+    metric_attributes: dict[str, Any] = field(
+        default_factory=_new_str_any_dict
+    )
+    """
+    Additional attributes to set on metrics. Must be of a low cardinality.
+    These attributes will not be set on spans or events.
+    """
+    # Monotonic start time in seconds (from timeit.default_timer) used
+    # for duration calculations to avoid mixing clock sources. This is
+    # populated by the TelemetryHandler when starting an invocation.
+    monotonic_start_s: float | None = None
+
+
 @dataclass
 class Error:
     message: str