Skip to content

Commit

Permalink
impl llmobs agent
Browse files Browse the repository at this point in the history
  • Loading branch information
gary-huang committed Feb 20, 2025
1 parent 1dc457b commit 871d358
Show file tree
Hide file tree
Showing 16 changed files with 572 additions and 15 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,7 @@ private HttpUrl getAgentlessUrl(Intake intake) {

public enum Intake {
API("api", "v2", Config::isCiVisibilityAgentlessEnabled, Config::getCiVisibilityAgentlessUrl),
LLMOBS_API("api", "v2", Config::isLlmObsAgentlessEnabled, Config::getLlMObsAgentlessUrl),
LOGS(
"http-intake.logs",
"v2",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import datadog.trace.api.config.GeneralConfig;
import datadog.trace.api.config.IastConfig;
import datadog.trace.api.config.JmxFetchConfig;
import datadog.trace.api.config.LlmObsConfig;
import datadog.trace.api.config.ProfilingConfig;
import datadog.trace.api.config.RemoteConfigConfig;
import datadog.trace.api.config.TraceInstrumentationConfig;
Expand Down Expand Up @@ -109,7 +110,8 @@ private enum AgentFeature {
false),
DATA_JOBS(propertyNameToSystemPropertyName(GeneralConfig.DATA_JOBS_ENABLED), false),
AGENTLESS_LOG_SUBMISSION(
propertyNameToSystemPropertyName(GeneralConfig.AGENTLESS_LOG_SUBMISSION_ENABLED), false);
propertyNameToSystemPropertyName(GeneralConfig.AGENTLESS_LOG_SUBMISSION_ENABLED), false),
LLMOBS(propertyNameToSystemPropertyName(LlmObsConfig.LLMOBS_ENABLED), false);

private final String systemProp;
private final boolean enabledByDefault;
Expand Down Expand Up @@ -150,6 +152,7 @@ public boolean isEnabledByDefault() {
private static boolean iastFullyDisabled;
private static boolean cwsEnabled = false;
private static boolean ciVisibilityEnabled = false;
private static boolean llmObsEnabled = false;
private static boolean usmEnabled = false;
private static boolean telemetryEnabled = true;
private static boolean debuggerEnabled = false;
Expand Down Expand Up @@ -268,6 +271,7 @@ public static void start(
exceptionDebuggingEnabled = isFeatureEnabled(AgentFeature.EXCEPTION_DEBUGGING);
spanOriginEnabled = isFeatureEnabled(AgentFeature.SPAN_ORIGIN);
agentlessLogSubmissionEnabled = isFeatureEnabled(AgentFeature.AGENTLESS_LOG_SUBMISSION);
llmObsEnabled = isFeatureEnabled(AgentFeature.LLMOBS);

if (profilingEnabled) {
if (!isOracleJDK8()) {
Expand Down Expand Up @@ -537,6 +541,7 @@ public void execute() {
maybeStartAppSec(scoClass, sco);
maybeStartIast(instrumentation, scoClass, sco);
maybeStartCiVisibility(instrumentation, scoClass, sco);
maybeStartLLMObs(instrumentation, scoClass, sco);
// start debugger before remote config to subscribe to it before starting to poll
maybeStartDebugger(instrumentation, scoClass, sco);
maybeStartRemoteConfig(scoClass, sco);
Expand Down Expand Up @@ -883,6 +888,24 @@ private static void maybeStartCiVisibility(Instrumentation inst, Class<?> scoCla
}
}

private static void maybeStartLLMObs(Instrumentation inst, Class<?> scoClass, Object sco) {
if (llmObsEnabled) {
StaticEventLogger.begin("LLM Observability");

try {
final Class<?> llmObsSysClass =
AGENT_CLASSLOADER.loadClass("datadog.trace.llmobs.LLMObsSystem");
final Method llmObsInstallerMethod =
llmObsSysClass.getMethod("start", Instrumentation.class, scoClass);
llmObsInstallerMethod.invoke(null, inst, sco);
} catch (final Throwable e) {
log.warn("Not starting LLM Observability subsystem", e);
}

StaticEventLogger.end("LLM Observability");
}
}

private static void maybeInstallLogsIntake(Class<?> scoClass, Object sco) {
if (agentlessLogSubmissionEnabled) {
StaticEventLogger.begin("Logs Intake");
Expand Down
60 changes: 60 additions & 0 deletions dd-java-agent/agent-llmobs/build.gradle
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
buildscript {
repositories {
mavenCentral()
}

dependencies {
classpath group: 'org.jetbrains.kotlin', name: 'kotlin-gradle-plugin', version: libs.versions.kotlin.get()
}
}

plugins {
id 'com.github.johnrengelman.shadow'
id 'java-test-fixtures'
}

apply from: "$rootDir/gradle/java.gradle"
apply from: "$rootDir/gradle/version.gradle"
apply from: "$rootDir/gradle/test-with-kotlin.gradle"
apply from: "$rootDir/gradle/test-with-scala.gradle"

minimumBranchCoverage = 0.0
minimumInstructionCoverage = 0.0

dependencies {
api libs.slf4j

implementation libs.bundles.asm
implementation group: 'org.jacoco', name: 'org.jacoco.core', version: '0.8.12'
implementation group: 'org.jacoco', name: 'org.jacoco.report', version: '0.8.12'

implementation project(':communication')
implementation project(':components:json')
implementation project(':internal-api')
implementation project(':internal-api:internal-api-9')

testImplementation project(":utils:test-utils")
testImplementation("com.google.jimfs:jimfs:1.1") // an in-memory file system for testing code that works with files

testImplementation libs.scala
testImplementation libs.kotlin

testFixturesApi project(':dd-java-agent:testing')
testFixturesApi project(':utils:test-utils')

testFixturesApi group: 'org.skyscreamer', name: 'jsonassert', version: '1.5.1'
testFixturesApi group: 'org.freemarker', name: 'freemarker', version: '2.3.30'
testFixturesApi group: 'com.jayway.jsonpath', name: 'json-path', version: '2.8.0'
testFixturesApi group: 'com.fasterxml.jackson.core', name: 'jackson-databind', version: '2.16.0'
testFixturesApi group: 'org.msgpack', name: 'jackson-dataformat-msgpack', version: '0.9.6'

testRuntimeOnly("org.junit.platform:junit-platform-launcher:1.9.2") // Required to update dependency lock files
}

shadowJar {
dependencies deps.excludeShared
}

jar {
archiveClassifier = 'unbundled'
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
package datadog.trace.llmobs;

import datadog.communication.BackendApi;
import datadog.communication.BackendApiFactory;
import datadog.communication.ddagent.SharedCommunicationObjects;
import datadog.trace.api.Config;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class LLMObsServices {

private static final Logger logger = LoggerFactory.getLogger(LLMObsServices.class);

final Config config;
final BackendApi backendApi;

LLMObsServices(Config config, SharedCommunicationObjects sco) {
this.config = config;
this.backendApi =
new BackendApiFactory(config, sco).createBackendApi(BackendApiFactory.Intake.LLMOBS_API);
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,100 @@
package datadog.trace.llmobs;

import datadog.communication.ddagent.SharedCommunicationObjects;
import datadog.trace.api.Config;
import datadog.trace.api.llmobs.LLMObs;
import datadog.trace.api.llmobs.LLMObsSpan;
import datadog.trace.api.llmobs.LLMObsTags;
import datadog.trace.bootstrap.instrumentation.api.Tags;
import datadog.trace.llmobs.domain.DDLLMObsSpan;
import datadog.trace.llmobs.domain.LLMObsInternal;
import java.lang.instrument.Instrumentation;
import org.jetbrains.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class LLMObsSystem {

private static final Logger LOGGER = LoggerFactory.getLogger(LLMObsSystem.class);

public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
Config config = Config.get();
if (!config.isLlmObsEnabled()) {
LOGGER.debug("LLM Observability is disabled");
return;
}

sco.createRemaining(config);

LLMObsServices llmObsServices = new LLMObsServices(config, sco);
LLMObsInternal.setLLMObsSpanFactory(
new LLMObsManualSpanFactory(
config.getLlmObsMlApp(), config.getServiceName(), llmObsServices));
}

private static class LLMObsManualSpanFactory implements LLMObs.LLMObsSpanFactory {

private final LLMObsServices llmObsServices;
private final String serviceName;
private final String defaultMLApp;

public LLMObsManualSpanFactory(
String defaultMLApp, String serviceName, LLMObsServices llmObsServices) {
this.defaultMLApp = defaultMLApp;
this.llmObsServices = llmObsServices;
this.serviceName = serviceName;
}

@Override
public LLMObsSpan startLLMSpan(
String spanName,
String modelName,
String modelProvider,
@Nullable String mlApp,
@Nullable String sessionID) {

DDLLMObsSpan span =
new DDLLMObsSpan(
Tags.LLMOBS_LLM_SPAN_KIND, spanName, getMLApp(mlApp), sessionID, serviceName);

span.setTag(LLMObsTags.MODEL_NAME, modelName);
span.setTag(LLMObsTags.MODEL_PROVIDER, modelProvider);
return span;
}

@Override
public LLMObsSpan startAgentSpan(
String spanName, @Nullable String mlApp, @Nullable String sessionID) {
return new DDLLMObsSpan(
Tags.LLMOBS_AGENT_SPAN_KIND, spanName, getMLApp(mlApp), sessionID, serviceName);
}

@Override
public LLMObsSpan startToolSpan(
String spanName, @Nullable String mlApp, @Nullable String sessionID) {
return new DDLLMObsSpan(
Tags.LLMOBS_TOOL_SPAN_KIND, spanName, getMLApp(mlApp), sessionID, serviceName);
}

@Override
public LLMObsSpan startTaskSpan(
String spanName, @Nullable String mlApp, @Nullable String sessionID) {
return new DDLLMObsSpan(
Tags.LLMOBS_TASK_SPAN_KIND, spanName, getMLApp(mlApp), sessionID, serviceName);
}

@Override
public LLMObsSpan startWorkflowSpan(
String spanName, @Nullable String mlApp, @Nullable String sessionID) {
return new DDLLMObsSpan(
Tags.LLMOBS_WORKFLOW_SPAN_KIND, spanName, getMLApp(mlApp), sessionID, serviceName);
}

private String getMLApp(String mlApp) {
if (mlApp == null || mlApp.isEmpty()) {
return defaultMLApp;
}
return mlApp;
}
}
}
Loading

0 comments on commit 871d358

Please sign in to comment.