Added embedding and retrieval spans (#9297)

Author: nayeem-kamal

dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java

@@ -12,7 +12,7 @@
 import java.lang.instrument.Instrumentation;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
-import org.jetbrains.annotations.Nullable;
+import javax.annotation.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -189,6 +189,30 @@ public LLMObsSpan startWorkflowSpan(
           Tags.LLMOBS_WORKFLOW_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName);
     }
 
+    @Override
+    public LLMObsSpan startEmbeddingSpan(
+        String spanName,
+        @Nullable String mlApp,
+        @Nullable String modelProvider,
+        @Nullable String modelName,
+        @Nullable String sessionId) {
+      if (modelProvider == null) {
+        modelProvider = "custom";
+      }
+      DDLLMObsSpan embeddingSpan =
+          new DDLLMObsSpan(
+              Tags.LLMOBS_EMBEDDING_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName);
+      embeddingSpan.setTag(LLMObsTags.MODEL_PROVIDER, modelProvider);
+      embeddingSpan.setTag(LLMObsTags.MODEL_NAME, modelName);
+      return embeddingSpan;
+    }
+
+    public LLMObsSpan startRetrievalSpan(
+        String spanName, @Nullable String mlApp, @Nullable String sessionId) {
+      return new DDLLMObsSpan(
+          Tags.LLMOBS_RETRIEVAL_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName);
+    }
+
     private String getMLApp(String mlApp) {
       if (mlApp == null || mlApp.isEmpty()) {
         return defaultMLApp;

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

@@ -46,6 +46,20 @@ public static LLMObsSpan startWorkflowSpan(
     return SPAN_FACTORY.startWorkflowSpan(spanName, mlApp, sessionId);
   }
 
+  public LLMObsSpan startEmbeddingSpan(
+      String spanName,
+      @Nullable String mlApp,
+      @Nullable String modelProvider,
+      @Nullable String modelName,
+      @Nullable String sessionId) {
+    return SPAN_FACTORY.startEmbeddingSpan(spanName, mlApp, modelProvider, modelName, sessionId);
+  }
+
+  public LLMObsSpan startRetrievalSpan(
+      String spanName, @Nullable String mlApp, @Nullable String sessionId) {
+    return SPAN_FACTORY.startRetrievalSpan(spanName, mlApp, sessionId);
+  }
+
   public static void SubmitEvaluation(
       LLMObsSpan llmObsSpan, String label, String categoricalValue, Map<String, Object> tags) {
     EVAL_PROCESSOR.SubmitEvaluation(llmObsSpan, label, categoricalValue, tags);
@@ -90,6 +104,16 @@ LLMObsSpan startLLMSpan(
 
     LLMObsSpan startWorkflowSpan(
         String spanName, @Nullable String mlApp, @Nullable String sessionId);
+
+    LLMObsSpan startEmbeddingSpan(
+        String spanName,
+        @Nullable String mlApp,
+        @Nullable String modelProvider,
+        @Nullable String modelName,
+        @Nullable String sessionId);
+
+    LLMObsSpan startRetrievalSpan(
+        String spanName, @Nullable String mlApp, @Nullable String sessionId);
   }
 
   public interface LLMObsEvalProcessor {

dd-trace-api/src/main/java/datadog/trace/api/llmobs/noop/NoOpLLMObsSpanFactory.java

@@ -35,4 +35,18 @@ public LLMObsSpan startWorkflowSpan(
       String spanName, @Nullable String mlApp, @Nullable String sessionId) {
     return NoOpLLMObsSpan.INSTANCE;
   }
+
+  public LLMObsSpan startEmbeddingSpan(
+      String spanName,
+      @Nullable String mlApp,
+      @Nullable String modelProvider,
+      @Nullable String modelName,
+      @Nullable String sessionId) {
+    return NoOpLLMObsSpan.INSTANCE;
+  }
+
+  public LLMObsSpan startRetrievalSpan(
+      String spanName, @Nullable String mlApp, @Nullable String sessionId) {
+    return NoOpLLMObsSpan.INSTANCE;
+  }
 }

internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java

@@ -169,4 +169,6 @@ public class Tags {
   public static final String LLMOBS_TASK_SPAN_KIND = "task";
   public static final String LLMOBS_AGENT_SPAN_KIND = "agent";
   public static final String LLMOBS_TOOL_SPAN_KIND = "tool";
+  public static final String LLMOBS_EMBEDDING_SPAN_KIND = "embedding";
+  public static final String LLMOBS_RETRIEVAL_SPAN_KIND = "retrieval";
 }