add LLM obs configs

DataDog · Dec 19, 2024 · fd136e0 · fd136e0
1 parent 2b24697
commit fd136e0
Show file tree

Hide file tree

Showing 7 changed files with 192 additions and 1 deletion.
diff --git a/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java b/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java
@@ -311,6 +311,9 @@ public static Set<InstrumenterModule.TargetSystem> getEnabledSystems() {
     if (cfg.isUsmEnabled()) {
       enabledSystems.add(InstrumenterModule.TargetSystem.USM);
     }
+    if (cfg.isLlmObsEnabled()) {
+      enabledSystems.add(InstrumenterModule.TargetSystem.LLMOBS);
+    }
     return enabledSystems;
   }
 

diff --git a/...ava-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java b/...ava-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/InstrumenterModule.java
@@ -49,7 +49,8 @@ public enum TargetSystem {
     APPSEC,
     IAST,
     CIVISIBILITY,
-    USM
+    USM,
+    LLMOBS,
   }
 
   private static final Logger log = LoggerFactory.getLogger(InstrumenterModule.class);

diff --git a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java
@@ -135,6 +135,9 @@ public final class ConfigDefaults {
 
   static final boolean DEFAULT_IAST_STACK_TRACE_ENABLED = true;
 
+  static final boolean DEFAULT_LLM_OBS_ENABLED = false;
+  static final boolean DEFAULT_LLM_OBS_AGENTLESS_ENABLED = false;
+
   static final boolean DEFAULT_USM_ENABLED = false;
 
   static final boolean DEFAULT_CIVISIBILITY_ENABLED = false;

diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/LlmObsConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/LlmObsConfig.java
@@ -0,0 +1,16 @@
+package datadog.trace.api.config;
+
+/**
+ * Constant with names of configuration options for LLM Observability. (EXPERIMENTAL AND SUBJECT TO
+ * CHANGE)
+ */
+public final class LlmObsConfig {
+
+  public static final String LLM_OBS_ENABLED = "llmobs.enabled";
+
+  public static final String LLM_OBS_ML_APP = "llmobs.ml.app";
+
+  public static final String LLM_OBS_AGENTLESS_ENABLED = "llmobs.agentless.enabled";
+
+  private LlmObsConfig() {}
+}
diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java
@@ -13,6 +13,7 @@
 import static datadog.trace.api.config.GeneralConfig.SERVICE_NAME;
 import static datadog.trace.api.config.IastConfig.*;
 import static datadog.trace.api.config.JmxFetchConfig.*;
+import static datadog.trace.api.config.LlmObsConfig.*;
 import static datadog.trace.api.config.ProfilingConfig.*;
 import static datadog.trace.api.config.RemoteConfigConfig.*;
 import static datadog.trace.api.config.TraceInstrumentationConfig.*;
@@ -307,6 +308,9 @@ public static String getHostName() {
   private final boolean iastStackTraceEnabled;
   private final boolean iastExperimentalPropagationEnabled;
 
+  private final boolean llmObsAgentlessEnabled;
+  private final String llmObsMlApp;
+
   private final boolean ciVisibilityTraceSanitationEnabled;
   private final boolean ciVisibilityAgentlessEnabled;
   private final String ciVisibilityAgentlessUrl;
@@ -1319,6 +1323,10 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment())
     iastExperimentalPropagationEnabled =
         configProvider.getBoolean(IAST_EXPERIMENTAL_PROPAGATION_ENABLED, false);
 
+    llmObsAgentlessEnabled =
+        configProvider.getBoolean(LLM_OBS_AGENTLESS_ENABLED, DEFAULT_LLM_OBS_AGENTLESS_ENABLED);
+    llmObsMlApp = configProvider.getString(LLM_OBS_ML_APP);
+
     ciVisibilityTraceSanitationEnabled =
         configProvider.getBoolean(CIVISIBILITY_TRACE_SANITATION_ENABLED, true);
 
@@ -1754,6 +1762,21 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment())
         configProvider.getLong(
             TRACE_POST_PROCESSING_TIMEOUT, ConfigDefaults.DEFAULT_TRACE_POST_PROCESSING_TIMEOUT);
 
+    if (isLlmObsEnabled()) {
+      log.debug("Attempting to enable LLM Observability");
+      if (llmObsMlApp == null || llmObsMlApp.isEmpty()) {
+        throw new IllegalArgumentException(
+            "Attempt to enable LLM Observability without ML app defined."
+                + "Please ensure that the name of the ML app is provided through properties or env variable");
+      }
+      if (llmObsAgentlessEnabled && (apiKey == null || apiKey.isEmpty())) {
+        throw new FatalAgentMisconfigurationError(
+            "Attempt to start LLM Observability in Agentless mode without API key. "
+                + "Please ensure that either an API key is configured, or the tracer is set up to work with the Agent");
+      }
+      log.debug("LLM Observability enabled for ML app {}, agentless mode {}", llmObsMlApp, llmObsAgentlessEnabled);
+    }
+
     if (isCiVisibilityEnabled()
         && ciVisibilityAgentlessEnabled
         && (apiKey == null || apiKey.isEmpty())) {
@@ -2606,6 +2629,18 @@ public boolean isIastExperimentalPropagationEnabled() {
     return iastExperimentalPropagationEnabled;
   }
 
+  public boolean isLlmObsEnabled() {
+    return instrumenterConfig.isLlmObsEnabled();
+  }
+
+  public boolean isLlmObsAgentlessEnabled() {
+    return llmObsAgentlessEnabled;
+  }
+
+  public String getLlmObsMlApp() {
+    return llmObsMlApp;
+  }
+
   public boolean isCiVisibilityEnabled() {
     return instrumenterConfig.isCiVisibilityEnabled();
   }

diff --git a/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java b/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java
@@ -5,6 +5,7 @@
 import static datadog.trace.api.ConfigDefaults.DEFAULT_CODE_ORIGIN_FOR_SPANS_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_IAST_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_INTEGRATIONS_ENABLED;
+import static datadog.trace.api.ConfigDefaults.DEFAULT_LLM_OBS_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_MEASURE_METHODS;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_RESOLVER_RESET_INTERVAL;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_RUNTIME_CONTEXT_FIELD_INJECTION;
@@ -26,6 +27,7 @@
 import static datadog.trace.api.config.GeneralConfig.TRACE_TRIAGE;
 import static datadog.trace.api.config.GeneralConfig.TRIAGE_REPORT_TRIGGER;
 import static datadog.trace.api.config.IastConfig.IAST_ENABLED;
+import static datadog.trace.api.config.LlmObsConfig.LLM_OBS_ENABLED;
 import static datadog.trace.api.config.ProfilingConfig.PROFILING_DIRECT_ALLOCATION_ENABLED;
 import static datadog.trace.api.config.ProfilingConfig.PROFILING_DIRECT_ALLOCATION_ENABLED_DEFAULT;
 import static datadog.trace.api.config.ProfilingConfig.PROFILING_ENABLED;
@@ -111,6 +113,7 @@ public class InstrumenterConfig {
   private final boolean iastFullyDisabled;
   private final boolean usmEnabled;
   private final boolean telemetryEnabled;
+  private final boolean llmObsEnabled;
 
   private final String traceExtensionsPath;
 
@@ -199,6 +202,7 @@ private InstrumenterConfig() {
       iastFullyDisabled = iastEnabled != null && !iastEnabled;
       usmEnabled = configProvider.getBoolean(USM_ENABLED, DEFAULT_USM_ENABLED);
       telemetryEnabled = configProvider.getBoolean(TELEMETRY_ENABLED, DEFAULT_TELEMETRY_ENABLED);
+      llmObsEnabled = configProvider.getBoolean(LLM_OBS_ENABLED, DEFAULT_LLM_OBS_ENABLED);
     } else {
       // disable these features in native-image
       ciVisibilityEnabled = false;
@@ -207,6 +211,7 @@ private InstrumenterConfig() {
       iastFullyDisabled = true;
       telemetryEnabled = false;
       usmEnabled = false;
+      llmObsEnabled = false;
     }
 
     traceExtensionsPath = configProvider.getString(TRACE_EXTENSIONS_PATH);
@@ -355,6 +360,10 @@ public boolean isIastFullyDisabled() {
     return iastFullyDisabled;
   }
 
+  public boolean isLlmObsEnabled() {
+    return llmObsEnabled;
+  }
+
   public boolean isUsmEnabled() {
     return usmEnabled;
   }

diff --git a/internal-api/src/test/groovy/datadog/trace/api/ConfigTest.groovy b/internal-api/src/test/groovy/datadog/trace/api/ConfigTest.groovy
@@ -63,6 +63,9 @@ import static datadog.trace.api.config.JmxFetchConfig.JMX_FETCH_REFRESH_BEANS_PE
 import static datadog.trace.api.config.JmxFetchConfig.JMX_FETCH_STATSD_HOST
 import static datadog.trace.api.config.JmxFetchConfig.JMX_FETCH_STATSD_PORT
 import static datadog.trace.api.config.JmxFetchConfig.JMX_TAGS
+import static datadog.trace.api.config.LlmObsConfig.LLM_OBS_AGENTLESS_ENABLED
+import static datadog.trace.api.config.LlmObsConfig.LLM_OBS_ML_APP
+import static datadog.trace.api.config.LlmObsConfig.LLM_OBS_ENABLED
 import static datadog.trace.api.config.ProfilingConfig.PROFILING_AGENTLESS
 import static datadog.trace.api.config.ProfilingConfig.PROFILING_API_KEY_FILE_OLD
 import static datadog.trace.api.config.ProfilingConfig.PROFILING_API_KEY_FILE_VERY_OLD
@@ -163,6 +166,9 @@ class ConfigTest extends DDSpecification {
   private static final DD_PROFILING_TAGS_ENV = "DD_PROFILING_TAGS"
   private static final DD_PROFILING_PROXY_PASSWORD_ENV = "DD_PROFILING_PROXY_PASSWORD"
   private static final DD_TRACE_X_DATADOG_TAGS_MAX_LENGTH = "DD_TRACE_X_DATADOG_TAGS_MAX_LENGTH"
+  private static final DD_LLMOBS_ENABLED_ENV = "DD_LLMOBS_ENABLED"
+  private static final DD_LLMOBS_ML_APP_ENV = "DD_LLMOBS_ML_APP"
+  private static final DD_LLMOBS_AGENTLESS_ENABLED_ENV = "DD_LLMOBS_AGENTLESS_ENABLED"
 
   def "specify overrides via properties"() {
     setup:
@@ -2208,6 +2214,124 @@ class ConfigTest extends DDSpecification {
     !hostname.trim().isEmpty()
   }
 
+  def "config instantiation should fail if llm obs is enabled via sys prop and ml app is not set"() {
+    setup:
+    Properties properties = new Properties()
+    properties.setProperty(LLM_OBS_ENABLED, "true")
+
+    when:
+    new Config(ConfigProvider.withPropertiesOverride(properties))
+
+    then:
+    thrown IllegalArgumentException
+  }
+
+  def "config instantiation should fail if llm obs is enabled via env var and ml app is not set"() {
+    setup:
+    environmentVariables.set(DD_LLMOBS_ENABLED_ENV, "true")
+
+    when:
+    new Config()
+
+    then:
+    thrown IllegalArgumentException
+  }
+
+
+  def "config instantiation should NOT fail if llm obs is enabled (agentless disabled) via sys prop and ml app is set"() {
+    setup:
+    Properties properties = new Properties()
+    properties.setProperty(LLM_OBS_ENABLED, "true")
+    properties.setProperty(LLM_OBS_AGENTLESS_ENABLED, "false")
+    properties.setProperty(LLM_OBS_ML_APP, "test-ml-app")
+
+    when:
+    def config = new Config(ConfigProvider.withPropertiesOverride(properties))
+
+    then:
+    noExceptionThrown()
+    config.isLlmObsEnabled()
+    !config.isLlmObsAgentlessEnabled()
+    config.llmObsMlApp == "test-ml-app"
+  }
+
+  def "config instantiation should NOT fail if llm obs is enabled (agentless disabled) via env var and ml app is set"() {
+    setup:
+    environmentVariables.set(DD_LLMOBS_ENABLED_ENV, "true")
+    environmentVariables.set(DD_LLMOBS_ML_APP_ENV, "test-ml-app")
+
+    when:
+    def config = new Config()
+
+    then:
+    noExceptionThrown()
+    config.isLlmObsEnabled()
+    !config.isLlmObsAgentlessEnabled()
+    config.llmObsMlApp == "test-ml-app"
+  }
+
+  def "config instantiation should fail if llm obs is in agentless mode via sys prop and API key is not set"() {
+    setup:
+    Properties properties = new Properties()
+    properties.setProperty(LLM_OBS_ENABLED, "true")
+    properties.setProperty(LLM_OBS_AGENTLESS_ENABLED, "true")
+    properties.setProperty(LLM_OBS_ML_APP, "test-ml-app")
+
+    when:
+    new Config(ConfigProvider.withPropertiesOverride(properties))
+
+    then:
+    thrown FatalAgentMisconfigurationError
+  }
+
+  def "config instantiation should fail if llm obs is in agentless mode via env var and API key is not set"() {
+    setup:
+    environmentVariables.set(DD_LLMOBS_ENABLED_ENV, "true")
+    environmentVariables.set(DD_LLMOBS_ML_APP_ENV, "a")
+    environmentVariables.set(DD_LLMOBS_AGENTLESS_ENABLED_ENV, "true")
+
+    when:
+    new Config()
+
+    then:
+    thrown FatalAgentMisconfigurationError
+  }
+
+  def "config instantiation should NOT fail if llm obs is enabled (agentless enabled) and API key & ml app are set via sys prop"() {
+    setup:
+    Properties properties = new Properties()
+    properties.setProperty(LLM_OBS_ENABLED, "true")
+    properties.setProperty(LLM_OBS_AGENTLESS_ENABLED, "true")
+    properties.setProperty(LLM_OBS_ML_APP, "test-ml-app")
+    properties.setProperty(API_KEY, "123456789")
+
+    when:
+    def config = new Config(ConfigProvider.withPropertiesOverride(properties))
+
+    then:
+    noExceptionThrown()
+    config.isLlmObsEnabled()
+    config.isLlmObsAgentlessEnabled()
+    config.llmObsMlApp == "test-ml-app"
+  }
+
+  def "config instantiation should NOT fail if llm obs is enabled (agentless enabled) and API key & ml app are set via env var"() {
+    setup:
+    environmentVariables.set(DD_LLMOBS_ENABLED_ENV, "true")
+    environmentVariables.set(DD_LLMOBS_ML_APP_ENV, "a")
+    environmentVariables.set(DD_LLMOBS_AGENTLESS_ENABLED_ENV, "true")
+    environmentVariables.set(DD_API_KEY_ENV, "8663294466")
+
+    when:
+    def config = new Config()
+
+    then:
+    noExceptionThrown()
+    config.isLlmObsEnabled()
+    config.isLlmObsAgentlessEnabled()
+    config.llmObsMlApp == "a"
+  }
+
   def "config instantiation should fail if CI visibility agentless mode is enabled and API key is not set"() {
     setup:
     Properties properties = new Properties()
-Original file line number
+Diff line change
@@ Expand Up / @@ -49,7 +49,8 @@ public enum TargetSystem { @@
         APPSEC,
         IAST,
         CIVISIBILITY,
-        USM
+        USM,
+        LLMOBS,
       }
       private static final Logger log = LoggerFactory.getLogger(InstrumenterModule.class);
@@ Expand Down @@