diff --git a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java index e184a645f71..00a04e5e59d 100644 --- a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java +++ b/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; diff --git a/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObs.java b/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObs.java index 392e76a4a82..6f3224dfcaa 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/llmobs/LLMObs.java +++ b/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 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 { diff --git a/dd-trace-api/src/main/java/datadog/trace/api/llmobs/noop/NoOpLLMObsSpanFactory.java b/dd-trace-api/src/main/java/datadog/trace/api/llmobs/noop/NoOpLLMObsSpanFactory.java index 080aa41bd82..b142ee43313 100644 --- a/dd-trace-api/src/main/java/datadog/trace/api/llmobs/noop/NoOpLLMObsSpanFactory.java +++ b/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; + } } diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java index 63aec166076..8a75bb12068 100644 --- a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java +++ b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java @@ -166,4 +166,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"; }