Add Support for Annotating Embedding and Retrieval Documents

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/DDLLMObsSpan.java

@@ -156,35 +156,81 @@ public void annotateIO(List<LLMObs.LLMMessage> inputData, List<LLMObs.LLMMessage
   }
 
   @Override
-  public void annotateIO(String inputData, String outputData) {
+  public void annotateEmbeddingIO(List<LLMObs.Document> inputDocuments, String outputData) {
+    if (finished) {
+      return;
+    }
+    if (inputDocuments != null && !inputDocuments.isEmpty()) {
+      span.setTag(INPUT, inputDocuments);
+    }
+    if (outputData != null && !outputData.isEmpty()) {
+      span.setTag(OUTPUT, outputData);
+    }
+  }
+
+  @Override
+  public void annotateRetrievalIO(String inputData, List<LLMObs.Document> outputDocuments) {
     if (finished) {
       return;
     }
-    boolean wrongSpanKind = false;
     if (inputData != null && !inputData.isEmpty()) {
-      if (Tags.LLMOBS_LLM_SPAN_KIND.equals(spanKind)) {
-        wrongSpanKind = true;
-        annotateIO(
-            Collections.singletonList(LLMObs.LLMMessage.from(LLM_MESSAGE_UNKNOWN_ROLE, inputData)),
-            null);
-      } else {
-        span.setTag(INPUT, inputData);
+      span.setTag(INPUT, inputData);
+    }
+    if (outputDocuments != null && !outputDocuments.isEmpty()) {
+      span.setTag(OUTPUT, outputDocuments);
+    }
+  }
+
+  @Override
+  public void annotateIO(String inputData, String outputData) {
+    if (finished) {
+      return;
+    }
+    boolean hasInput = inputData != null && !inputData.isEmpty();
+    boolean hasOutput = outputData != null && !outputData.isEmpty();
+    if (Tags.LLMOBS_LLM_SPAN_KIND.equals(spanKind)) {
+      List<LLMObs.LLMMessage> inputMessages =
+          hasInput
+              ? Collections.singletonList(
+                  LLMObs.LLMMessage.from(LLM_MESSAGE_UNKNOWN_ROLE, inputData))
+              : null;
+      List<LLMObs.LLMMessage> outputMessages =
+          hasOutput
+              ? Collections.singletonList(
+                  LLMObs.LLMMessage.from(LLM_MESSAGE_UNKNOWN_ROLE, outputData))
+              : null;
+      annotateIO(inputMessages, outputMessages);
+      if (hasInput || hasOutput) {
+        LOGGER.warn(
+            "the span being annotated is an LLM span, it is recommended to use the overload with List<LLMObs.LLMMessage> as arguments");
       }
+      return;
     }
-    if (outputData != null && !outputData.isEmpty()) {
-      if (Tags.LLMOBS_LLM_SPAN_KIND.equals(spanKind)) {
-        wrongSpanKind = true;
-        annotateIO(
-            null,
-            Collections.singletonList(
-                LLMObs.LLMMessage.from(LLM_MESSAGE_UNKNOWN_ROLE, outputData)));
-      } else {
-        span.setTag(OUTPUT, outputData);
+    if (Tags.LLMOBS_EMBEDDING_SPAN_KIND.equals(spanKind)) {
+      List<LLMObs.Document> inputDocuments =
+          hasInput ? Collections.singletonList(LLMObs.Document.from(inputData)) : null;
+      annotateEmbeddingIO(inputDocuments, outputData);
+      if (hasInput) {
+        LOGGER.warn(
+            "the span being annotated is an embedding span, it is recommended to use annotateEmbeddingIO");
+      }
+      return;
+    }
+    if (Tags.LLMOBS_RETRIEVAL_SPAN_KIND.equals(spanKind)) {
+      List<LLMObs.Document> outputDocuments =
+          hasOutput ? Collections.singletonList(LLMObs.Document.from(outputData)) : null;
+      annotateRetrievalIO(inputData, outputDocuments);
+      if (hasOutput) {
+        LOGGER.warn(
+            "the span being annotated is a retrieval span, it is recommended to use annotateRetrievalIO");
       }
+      return;
+    }
+    if (hasInput) {
+      span.setTag(INPUT, inputData);
     }
-    if (wrongSpanKind) {
-      LOGGER.warn(
-          "the span being annotated is an LLM span, it is recommended to use the overload with List<LLMObs.LLMMessage> as arguments");
+    if (hasOutput) {
+      span.setTag(OUTPUT, outputData);
     }
   }
 

dd-java-agent/agent-llmobs/src/test/java/datadog/trace/llmobs/domain/DDLLMObsSpanDocumentIOTest.java

@@ -0,0 +1,141 @@
+package datadog.trace.llmobs.domain;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+import datadog.trace.agent.tooling.TracerInstaller;
+import datadog.trace.api.WellKnownTags;
+import datadog.trace.api.llmobs.LLMObs;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import datadog.trace.core.CoreTracer;
+import java.lang.reflect.Field;
+import java.util.Arrays;
+import java.util.List;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+class DDLLMObsSpanDocumentIOTest {
+  private static final String INPUT_TAG = "_ml_obs_tag.input";
+  private static final String OUTPUT_TAG = "_ml_obs_tag.output";
+  private static final Field SPAN_FIELD;
+
+  private static CoreTracer tracer;
+
+  static {
+    try {
+      SPAN_FIELD = DDLLMObsSpan.class.getDeclaredField("span");
+      SPAN_FIELD.setAccessible(true);
+    } catch (ReflectiveOperationException error) {
+      throw new ExceptionInInitializerError(error);
+    }
+  }
+
+  @BeforeAll
+  static void installTracer() {
+    tracer = CoreTracer.builder().build();
+    TracerInstaller.forceInstallGlobalTracer(tracer);
+  }
+
+  @AfterAll
+  static void closeTracer() {
+    TracerInstaller.forceInstallGlobalTracer(null);
+    tracer.close();
+  }
+
+  @Test
+  void formatsEmbeddingStringInputAsDocument() throws IllegalAccessException {
+    DDLLMObsSpan llmObsSpan = newSpan(Tags.LLMOBS_EMBEDDING_SPAN_KIND);
+    try {
+      llmObsSpan.annotateIO("embedding input", "embedding output");
+
+      AgentSpan span = (AgentSpan) SPAN_FIELD.get(llmObsSpan);
+      assertDocument(span.getTag(INPUT_TAG), "embedding input");
+      assertEquals("embedding output", span.getTag(OUTPUT_TAG));
+    } finally {
+      llmObsSpan.finish();
+    }
+  }
+
+  @Test
+  void formatsRetrievalStringOutputAsDocument() throws IllegalAccessException {
+    DDLLMObsSpan llmObsSpan = newSpan(Tags.LLMOBS_RETRIEVAL_SPAN_KIND);
+    try {
+      llmObsSpan.annotateIO("retrieval input", "retrieval output");
+
+      AgentSpan span = (AgentSpan) SPAN_FIELD.get(llmObsSpan);
+      assertEquals("retrieval input", span.getTag(INPUT_TAG));
+      assertDocument(span.getTag(OUTPUT_TAG), "retrieval output");
+    } finally {
+      llmObsSpan.finish();
+    }
+  }
+
+  @Test
+  void acceptsEmbeddingDocumentInputs() throws IllegalAccessException {
+    DDLLMObsSpan llmObsSpan = newSpan(Tags.LLMOBS_EMBEDDING_SPAN_KIND);
+    List<LLMObs.Document> documents =
+        Arrays.asList(
+            LLMObs.Document.from("first input", "first.txt", "input-1", 0.5),
+            LLMObs.Document.from("second input"));
+    try {
+      llmObsSpan.annotateEmbeddingIO(documents, "embedding output");
+
+      AgentSpan span = (AgentSpan) SPAN_FIELD.get(llmObsSpan);
+      assertEquals(documents, span.getTag(INPUT_TAG));
+      assertDocument(documents.get(0), "first input", "first.txt", "input-1", 0.5);
+      assertEquals("embedding output", span.getTag(OUTPUT_TAG));
+    } finally {
+      llmObsSpan.finish();
+    }
+  }
+
+  @Test
+  void acceptsRetrievalDocumentOutputs() throws IllegalAccessException {
+    DDLLMObsSpan llmObsSpan = newSpan(Tags.LLMOBS_RETRIEVAL_SPAN_KIND);
+    List<LLMObs.Document> documents =
+        Arrays.asList(
+            LLMObs.Document.from("first output", "result.txt", "output-1", 0.95),
+            LLMObs.Document.from("second output"));
+    try {
+      llmObsSpan.annotateRetrievalIO("retrieval input", documents);
+
+      AgentSpan span = (AgentSpan) SPAN_FIELD.get(llmObsSpan);
+      assertEquals("retrieval input", span.getTag(INPUT_TAG));
+      assertEquals(documents, span.getTag(OUTPUT_TAG));
+      assertDocument(documents.get(0), "first output", "result.txt", "output-1", 0.95);
+    } finally {
+      llmObsSpan.finish();
+    }
+  }
+
+  private static DDLLMObsSpan newSpan(String kind) {
+    WellKnownTags tags =
+        new WellKnownTags("runtime-id", "hostname", "test", "service", "version", "java");
+    return new DDLLMObsSpan(kind, "span", "ml-app", null, "service", tags);
+  }
+
+  private static void assertDocument(Object value, String expectedText) {
+    List<?> documents = assertInstanceOf(List.class, value);
+    assertEquals(1, documents.size());
+    LLMObs.Document document = assertInstanceOf(LLMObs.Document.class, documents.get(0));
+    assertEquals(expectedText, document.getText());
+    assertNull(document.getName());
+    assertNull(document.getId());
+    assertNull(document.getScore());
+  }
+
+  private static void assertDocument(
+      LLMObs.Document document,
+      String expectedText,
+      String expectedName,
+      String expectedId,
+      double expectedScore) {
+    assertEquals(expectedText, document.getText());
+    assertEquals(expectedName, document.getName());
+    assertEquals(expectedId, document.getId());
+    assertEquals(expectedScore, document.getScore());
+  }
+}

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

@@ -305,17 +305,40 @@ public List<ToolResult> getToolResults() {
 
   public static class Document {
     private String text;
+    private String name;
+    private String id;
+    private Double score;
 
     public static Document from(String text) {
-      return new Document(text);
+      return new Document(text, null, null, null);
+    }
+
+    public static Document from(
+        String text, @Nullable String name, @Nullable String id, @Nullable Double score) {
+      return new Document(text, name, id, score);
     }
 
-    private Document(String text) {
+    private Document(String text, String name, String id, Double score) {
       this.text = text;
+      this.name = name;
+      this.id = id;
+      this.score = score;
     }
 
     public String getText() {
       return text;
     }
+
+    public String getName() {
+      return name;
+    }
+
+    public String getId() {
+      return id;
+    }
+
+    public Double getScore() {
+      return score;
+    }
   }
 }

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

@@ -15,6 +15,26 @@ public interface LLMObsSpan {
    */
   void annotateIO(List<LLMObs.LLMMessage> inputMessages, List<LLMObs.LLMMessage> outputMessages);
 
+  /**
+   * Annotate an embedding span with document inputs and a string output.
+   *
+   * @param inputDocuments The input documents of the span
+   * @param outputData The output data of the span in the form of a string
+   */
+  default void annotateEmbeddingIO(List<LLMObs.Document> inputDocuments, String outputData) {
+    annotateIO((String) null, outputData);
+  }
+
+  /**
+   * Annotate a retrieval span with a string input and document outputs.
+   *
+   * @param inputData The input data of the span in the form of a string
+   * @param outputDocuments The output documents of the span
+   */
+  default void annotateRetrievalIO(String inputData, List<LLMObs.Document> outputDocuments) {
+    annotateIO(inputData, (String) null);
+  }
+
   /**
    * Annotate the span with inputs and outputs
    *

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

@@ -357,6 +357,9 @@ public void accept(Metadata metadata) {
         String key = tag.getKey().substring(LLMOBS_TAG_PREFIX.length());
         Object val = tag.getValue();
         if (key.equals(INPUT) || key.equals(OUTPUT)) {
+          boolean isDocumentIO =
+              (spanKind.equals(Tags.LLMOBS_EMBEDDING_SPAN_KIND) && key.equals(INPUT))
+                  || (spanKind.equals(Tags.LLMOBS_RETRIEVAL_SPAN_KIND) && key.equals(OUTPUT));
           if (spanKind.equals(Tags.LLMOBS_LLM_SPAN_KIND)) {
             writable.writeString(key, null);
             if (val instanceof List) {
@@ -371,24 +374,40 @@ public void accept(Metadata metadata) {
                   val.getClass().getName());
               continue;
             }
-          } 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;
-            }
+          } else if (isDocumentIO && isDocumentList(val)) {
             writable.writeString(key, null);
             writable.startMap(1);
             List<LLMObs.Document> documents = (List<LLMObs.Document>) val;
             writable.writeString("documents", null);
             writable.startArray(documents.size());
             for (LLMObs.Document document : documents) {
-              writable.startMap(1);
+              int documentSize = 1;
+              if (document.getName() != null) documentSize++;
+              if (document.getId() != null) documentSize++;
+              if (document.getScore() != null) documentSize++;
+              writable.startMap(documentSize);
               writable.writeString("text", null);
               writable.writeString(document.getText(), null);
+              if (document.getName() != null) {
+                writable.writeString("name", null);
+                writable.writeString(document.getName(), null);
+              }
+              if (document.getId() != null) {
+                writable.writeString("id", null);
+                writable.writeString(document.getId(), null);
+              }
+              if (document.getScore() != null) {
+                writable.writeString("score", null);
+                writable.writeObject(document.getScore(), null);
+              }
             }
           } else {
+            if (isDocumentIO) {
+              LOGGER.warn(
+                  "unexpectedly found invalid document data for {} span {}, serializing as value",
+                  spanKind,
+                  key);
+            }
             writable.writeString(key, null);
             writable.startMap(1);
             writable.writeString("value", null);
@@ -444,6 +463,18 @@ private void writeToolDefinitions(List<?> toolDefinitions) {
       }
     }
 
+    private static boolean isDocumentList(Object value) {
+      if (!(value instanceof List)) {
+        return false;
+      }
+      for (Object item : (List<?>) value) {
+        if (!(item instanceof LLMObs.Document)) {
+          return false;
+        }
+      }
+      return true;
+    }
+
     private void writeLlmInputMap(Map<?, ?> inputMap) {
       writable.startMap(inputMap.size());
       for (Map.Entry<?, ?> entry : inputMap.entrySet()) {