TestOpenAiLlmObs::test_embedding[java-test-ml-app-tcp] PASSED

Author: ygree

dd-java-agent/instrumentation/openai-java/openai-java-1.0/src/main/java/datadog/trace/instrumentation/openai_java/OpenAiDecorator.java

@@ -14,6 +14,7 @@
 import com.openai.models.completions.CompletionCreateParams;
 import com.openai.models.completions.CompletionUsage;
 import com.openai.models.embeddings.CreateEmbeddingResponse;
+import com.openai.models.embeddings.Embedding;
 import com.openai.models.embeddings.EmbeddingCreateParams;
 import com.openai.models.responses.Response;
 import com.openai.models.responses.ResponseCreateParams;
@@ -329,28 +330,39 @@ public void withEmbeddingCreateParams(AgentSpan span, EmbeddingCreateParams para
     }
     params.model()._value().asString().ifPresent(str -> span.setTag(REQUEST_MODEL, str));
 
-    // span.setTag("_ml_obs_tag.input", llmMessages(params.input()));
+    span.setTag("_ml_obs_tag.input", embeddingDocuments(params.input()));
 
-    // "_ml_obs.meta.input.documents"
+    Map<String, Object> metadata = new HashMap<>();
+    Optional<String> encodingFormat = params.encodingFormat().flatMap(v -> v._value().asString());
+    encodingFormat.ifPresent(v -> metadata.put("encoding_format", v));
+    params.dimensions().ifPresent(v -> metadata.put("dimensions", v));
+    span.setTag("_ml_obs_tag.metadata", metadata);
   }
 
-  private List<LLMObs.LLMMessage> llmMessages(EmbeddingCreateParams.Input input) {
+  private List<LLMObs.Document> embeddingDocuments(EmbeddingCreateParams.Input input) {
     List<String> inputs = Collections.emptyList();
     if (input.isString()) {
       inputs = Collections.singletonList(input.asString());
     } else if (input.isArrayOfStrings()) {
       inputs = input.asArrayOfStrings();
     }
-    return inputs.stream()
-        .map(str -> LLMObs.LLMMessage.from(null, str))
-        .collect(Collectors.toList());
+    return inputs.stream().map(LLMObs.Document::from).collect(Collectors.toList());
   }
 
   public void withCreateEmbeddingResponse(AgentSpan span, CreateEmbeddingResponse response) {
     String modelName = response.model();
     span.setTag(RESPONSE_MODEL, modelName);
     span.setTag("_ml_obs_tag.model_name", modelName);
     span.setTag("_ml_obs_tag.model_provider", "openai");
+
+    if (!response.data().isEmpty()) {
+      int embeddingCount = response.data().size();
+      Embedding firstEmbedding = response.data().get(0);
+      int embeddingSize = firstEmbedding.embedding().size();
+      span.setTag(
+          "_ml_obs_tag.output",
+          String.format("[%d embedding(s) returned with size %d]", embeddingCount, embeddingSize));
+    }
   }
 
   public void withResponseCreateParams(AgentSpan span, ResponseCreateParams params) {

dd-java-agent/instrumentation/openai-java/openai-java-1.0/src/test/groovy/EmbeddingServiceTest.groovy

@@ -3,6 +3,7 @@ import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
 import com.openai.core.http.HttpResponseFor
 import com.openai.models.embeddings.CreateEmbeddingResponse
 import datadog.trace.api.DDSpanTypes
+import datadog.trace.api.llmobs.LLMObs
 import datadog.trace.bootstrap.instrumentation.api.Tags
 import datadog.trace.instrumentation.openai_java.OpenAiDecorator
 
@@ -50,8 +51,10 @@ class EmbeddingServiceTest extends OpenAiTest {
           tags {
             "_ml_obs_tag.span.kind" "embedding"
             "_ml_obs_tag.model_provider" "openai"
-            "_ml_obs_tag.model_name" String
-
+            "_ml_obs_tag.model_name" "text-embedding-ada-002-v2"
+            "_ml_obs_tag.input" List<LLMObs.Document>
+            "_ml_obs_tag.metadata" Map
+            "_ml_obs_tag.output" "[1 embedding(s) returned with size 1536]"
             "_ml_obs_tag.parent_id" "undefined"
             "openai.request.method" "POST"
             "openai.request.endpoint" "v1/embeddings"

dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObs.java

@@ -209,4 +209,20 @@ public List<ToolCall> getToolCalls() {
       return toolCalls;
     }
   }
+
+  public static class Document {
+    private String text;
+
+    public static Document from(String text) {
+      return new Document(text);
+    }
+
+    private Document(String text) {
+      this.text = text;
+    }
+
+    public String getText() {
+      return text;
+    }
+  }
 }

dd-trace-core/src/main/java/datadog/trace/llmobs/writer/ddintake/LLMObsSpanMapper.java

@@ -288,10 +288,7 @@ public void accept(Metadata metadata) {
         if (key.equals(INPUT) || key.equals(OUTPUT)) {
           writable.writeString(key, null);
           writable.startMap(1);
-          if (!spanKind.equals(Tags.LLMOBS_LLM_SPAN_KIND)) {
-            writable.writeString("value", null);
-            writable.writeObject(val, null);
-          } else {
+          if (spanKind.equals(Tags.LLMOBS_LLM_SPAN_KIND)) {
             if (!(val instanceof List)) {
               LOGGER.warn(
                   "unexpectedly found incorrect type for LLM span IO {}, expecting list",
@@ -334,6 +331,24 @@ public void accept(Metadata metadata) {
                 }
               }
             }
+          } else if (spanKind.equals(Tags.LLMOBS_EMBEDDING_SPAN_KIND) && key.equals(INPUT)) {
+            if (!(val instanceof List)) {
+              LOGGER.warn(
+                  "unexpectedly found incorrect type for embedding span input {}, expecting list",
+                  val.getClass().getName());
+              continue;
+            }
+            List<LLMObs.Document> documents = (List<LLMObs.Document>) val;
+            writable.writeString("documents", null);
+            writable.startArray(documents.size());
+            for (LLMObs.Document document : documents) {
+              writable.startMap(1);
+              writable.writeString("text", null);
+              writable.writeString(document.getText(), null);
+            }
+          } else {
+            writable.writeString("value", null);
+            writable.writeObject(val, null);
           }
         } else if (key.equals(LLMObsTags.METADATA) && val instanceof Map) {
           Map<String, Object> metadataMap = (Map) val;