elastic-opentelemetry-instrumentation-openai: Add support for tracing embeddings calls (#20)

{
        "name": "embeddings all-minilm:33m",
        "context": {
            "trace_id": "0x34b6f145212c169ea39a980ff670dee3",
            "span_id": "0x74b9bd94f2ea6f31",
            "trace_state": "[]"
        },
        "kind": "SpanKind.CLIENT",
        "parent_id": null,
        "start_time": "2024-09-27T14:55:13.104397Z",
        "end_time": "2024-09-27T14:55:13.261547Z",
        "status": {
            "status_code": "UNSET"
        },
        "attributes": {
            "gen_ai.operation.name": "embeddings",
            "gen_ai.request.model": "all-minilm:33m",
            "gen_ai.system": "openai",
            "server.address": "localhost",
            "server.port": 11434,
            "gen_ai.request.encoding_format": "float",
            "gen_ai.response.model": "all-minilm:33m",
            "gen_ai.usage.input_tokens": 9
        },
        "events": [],
        "links": [],
        "resource": {
            "attributes": {
                "telemetry.sdk.language": "python",
                "telemetry.sdk.name": "opentelemetry",
                "telemetry.sdk.version": "1.27.0",
                "telemetry.distro.name": "elastic",
                "telemetry.distro.version": "0.2.0",
                "process.runtime.description": "3.10.12 (main, Sep 11 2024, 15:47:36) [GCC 11.4.0]",
                "process.runtime.name": "cpython",
                "process.runtime.version": "3.10.12",
                "os.type": "linux",
                "os.version": "6.1.0-1034-oem",
                "telemetry.auto.version": "0.48b0",
                "service.name": "unknown_service"
            },
            "schema_url": ""
        }
    }

Author: xrmx

instrumentation/elastic-opentelemetry-instrumentation-openai/README.md

@@ -51,6 +51,11 @@ chat_completion = client.chat.completions.create(model="gpt-4o-mini", messages=m
 - `ELASTIC_OTEL_GENAI_CAPTURE_CONTENT` (default: `false`): when sets to `true` collect more
 informations about prompts and responses by enabling content capture
 
+### Elastic specific semantic conventions
+
+- New `embeddings` value for `gen_ai.operation.name`
+- New `gen_ai.request.encoding_format` attribute with openai specific values `[float, base64]`
+
 ## Development
 
 We use [pytest](https://docs.pytest.org/en/stable/) to execute tests written with the standard

instrumentation/elastic-opentelemetry-instrumentation-openai/src/opentelemetry/instrumentation/openai/__init__.py

@@ -28,11 +28,13 @@
     ELASTIC_OTEL_GENAI_CAPTURE_CONTENT,
 )
 from opentelemetry.instrumentation.openai.helpers import (
+    _get_embeddings_span_attributes_from_wrapper,
     _get_span_attributes_from_wrapper,
     _message_from_choice,
     _record_token_usage_metrics,
     _record_operation_duration_metric,
     _set_span_attributes_from_response,
+    _set_embeddings_span_attributes_from_response,
 )
 from opentelemetry.instrumentation.openai.package import _instruments
 from opentelemetry.instrumentation.openai.version import __version__
@@ -95,6 +97,16 @@ def _patch(self, _module):
             "AsyncCompletions.create",
             self._async_chat_completion_wrapper,
         )
+        wrap_function_wrapper(
+            "openai.resources.embeddings",
+            "Embeddings.create",
+            self._embeddings_wrapper,
+        )
+        wrap_function_wrapper(
+            "openai.resources.embeddings",
+            "AsyncEmbeddings.create",
+            self._async_embeddings_wrapper,
+        )
 
     def _uninstrument(self, **kwargs):
         # unwrap only supports uninstrementing real module references so we
@@ -103,6 +115,8 @@ def _uninstrument(self, **kwargs):
 
         unwrap(openai.resources.chat.completions.Completions, "create")
         unwrap(openai.resources.chat.completions.AsyncCompletions, "create")
+        unwrap(openai.resources.embeddings.Embeddings, "create")
+        unwrap(openai.resources.embeddings.AsyncEmbeddings, "create")
 
     def _chat_completion_wrapper(self, wrapped, instance, args, kwargs):
         logger.debug(f"openai.resources.chat.completions.Completions.create kwargs: {kwargs}")
@@ -226,3 +240,63 @@ async def _async_chat_completion_wrapper(self, wrapped, instance, args, kwargs):
             span.end()
 
             return result
+
+    def _embeddings_wrapper(self, wrapped, instance, args, kwargs):
+        span_attributes = _get_embeddings_span_attributes_from_wrapper(instance, kwargs)
+
+        span_name = f"{span_attributes[GEN_AI_OPERATION_NAME]} {span_attributes[GEN_AI_REQUEST_MODEL]}"
+        with self.tracer.start_as_current_span(
+            name=span_name,
+            kind=SpanKind.CLIENT,
+            attributes=span_attributes,
+            # this is important to avoid having the span closed before ending the stream
+            end_on_exit=False,
+        ) as span:
+            start_time = default_timer()
+            try:
+                result = wrapped(*args, **kwargs)
+            except Exception as exc:
+                span.set_status(StatusCode.ERROR, str(exc))
+                span.set_attribute(ERROR_TYPE, exc.__class__.__qualname__)
+                span.end()
+                _record_operation_duration_metric(self.operation_duration_metric, span, start_time)
+                raise
+
+            _set_embeddings_span_attributes_from_response(span, result.model, result.usage)
+
+            _record_token_usage_metrics(self.token_usage_metric, span, result.usage)
+            _record_operation_duration_metric(self.operation_duration_metric, span, start_time)
+
+            span.end()
+
+            return result
+
+    async def _async_embeddings_wrapper(self, wrapped, instance, args, kwargs):
+        span_attributes = _get_embeddings_span_attributes_from_wrapper(instance, kwargs)
+
+        span_name = f"{span_attributes[GEN_AI_OPERATION_NAME]} {span_attributes[GEN_AI_REQUEST_MODEL]}"
+        with self.tracer.start_as_current_span(
+            name=span_name,
+            kind=SpanKind.CLIENT,
+            attributes=span_attributes,
+            # this is important to avoid having the span closed before ending the stream
+            end_on_exit=False,
+        ) as span:
+            start_time = default_timer()
+            try:
+                result = await wrapped(*args, **kwargs)
+            except Exception as exc:
+                span.set_status(StatusCode.ERROR, str(exc))
+                span.set_attribute(ERROR_TYPE, exc.__class__.__qualname__)
+                span.end()
+                _record_operation_duration_metric(self.operation_duration_metric, span, start_time)
+                raise
+
+            _set_embeddings_span_attributes_from_response(span, result.model, result.usage)
+
+            _record_token_usage_metrics(self.token_usage_metric, span, result.usage)
+            _record_operation_duration_metric(self.operation_duration_metric, span, start_time)
+
+            span.end()
+
+            return result

instrumentation/elastic-opentelemetry-instrumentation-openai/src/opentelemetry/instrumentation/openai/helpers.py

@@ -46,6 +46,8 @@
 else:
     CompletionUsage = None
 
+GEN_AI_REQUEST_ENCODING_FORMAT = "gen_ai.request.encoding_format"
+
 
 def _set_span_attributes_from_response(
     span: Span, response_id: str, model: str, choices, usage: CompletionUsage
@@ -61,6 +63,11 @@ def _set_span_attributes_from_response(
         span.set_attribute(GEN_AI_USAGE_OUTPUT_TOKENS, usage.completion_tokens)
 
 
+def _set_embeddings_span_attributes_from_response(span: Span, model: str, usage: CompletionUsage) -> None:
+    span.set_attribute(GEN_AI_RESPONSE_MODEL, model)
+    span.set_attribute(GEN_AI_USAGE_INPUT_TOKENS, usage.prompt_tokens)
+
+
 def _decode_function_arguments(arguments: str):
     try:
         return json.loads(arguments)
@@ -98,6 +105,23 @@ def _message_from_stream_choices(choices):
     return message
 
 
+def _attributes_from_client(client):
+    span_attributes = {}
+
+    if base_url := getattr(client, "_base_url", None):
+        if host := getattr(base_url, "host", None):
+            span_attributes[SERVER_ADDRESS] = host
+        if port := getattr(base_url, "port", None):
+            span_attributes[SERVER_PORT] = port
+        elif scheme := getattr(base_url, "scheme", None):
+            if scheme == "http":
+                span_attributes[SERVER_PORT] = 80
+            elif scheme == "https":
+                span_attributes[SERVER_PORT] = 443
+
+    return span_attributes
+
+
 def _get_span_attributes_from_wrapper(instance, kwargs):
     span_attributes = {
         GEN_AI_OPERATION_NAME: "chat",
@@ -106,16 +130,7 @@ def _get_span_attributes_from_wrapper(instance, kwargs):
     }
 
     if client := getattr(instance, "_client", None):
-        if base_url := getattr(client, "_base_url", None):
-            if host := getattr(base_url, "host", None):
-                span_attributes[SERVER_ADDRESS] = host
-            if port := getattr(base_url, "port", None):
-                span_attributes[SERVER_PORT] = port
-            elif scheme := getattr(base_url, "scheme", None):
-                if scheme == "http":
-                    span_attributes[SERVER_PORT] = 80
-                elif scheme == "https":
-                    span_attributes[SERVER_PORT] = 443
+        span_attributes.update(_attributes_from_client(client))
 
     if (frequency_penalty := kwargs.get("frequency_penalty")) is not None:
         span_attributes[GEN_AI_REQUEST_FREQUENCY_PENALTY] = frequency_penalty
@@ -135,6 +150,22 @@ def _get_span_attributes_from_wrapper(instance, kwargs):
     return span_attributes
 
 
+def _get_embeddings_span_attributes_from_wrapper(instance, kwargs):
+    span_attributes = {
+        GEN_AI_OPERATION_NAME: "embeddings",
+        GEN_AI_REQUEST_MODEL: kwargs["model"],
+        GEN_AI_SYSTEM: "openai",
+    }
+
+    if client := getattr(instance, "_client", None):
+        span_attributes.update(_attributes_from_client(client))
+
+    if (encoding_format := kwargs.get("encoding_format")) is not None:
+        span_attributes[GEN_AI_REQUEST_ENCODING_FORMAT] = encoding_format
+
+    return span_attributes
+
+
 def _get_attributes_if_set(span: Span, names: Iterable) -> dict:
     """Returns a dict with any attribute found in the span attributes"""
     attributes = span.attributes
@@ -154,7 +185,9 @@ def _record_token_usage_metrics(metric: Histogram, span: Span, usage: Completion
         ),
     )
     metric.record(usage.prompt_tokens, {**token_usage_metric_attrs, GEN_AI_TOKEN_TYPE: "input"})
-    metric.record(usage.completion_tokens, {**token_usage_metric_attrs, GEN_AI_TOKEN_TYPE: "output"})
+    # embeddings responses only have input tokens
+    if hasattr(usage, "completion_tokens"):
+        metric.record(usage.completion_tokens, {**token_usage_metric_attrs, GEN_AI_TOKEN_TYPE: "output"})
 
 
 def _record_operation_duration_metric(metric: Histogram, span: Span, start: float):