Skip to content

Commit c437b9c

Browse files
gary-huanggary-huang
committed
impl llmobs agent
1 parent 1dc457b commit c437b9c

File tree

16 files changed

+568
-15
lines changed

16 files changed

+568
-15
lines changed

‎communication/src/main/java/datadog/communication/BackendApiFactory.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ private HttpUrl getAgentlessUrl(Intake intake) {
7272

7373
public enum Intake {
7474
API("api", "v2", Config::isCiVisibilityAgentlessEnabled, Config::getCiVisibilityAgentlessUrl),
75+
LLMOBS_API("api", "v2", Config::isLlmObsAgentlessEnabled, Config::getLlMObsAgentlessUrl),
7576
LOGS(
7677
"http-intake.logs",
7778
"v2",

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import datadog.trace.api.config.GeneralConfig;
2525
import datadog.trace.api.config.IastConfig;
2626
import datadog.trace.api.config.JmxFetchConfig;
27+
import datadog.trace.api.config.LlmObsConfig;
2728
import datadog.trace.api.config.ProfilingConfig;
2829
import datadog.trace.api.config.RemoteConfigConfig;
2930
import datadog.trace.api.config.TraceInstrumentationConfig;
@@ -109,7 +110,8 @@ private enum AgentFeature {
109110
false),
110111
DATA_JOBS(propertyNameToSystemPropertyName(GeneralConfig.DATA_JOBS_ENABLED), false),
111112
AGENTLESS_LOG_SUBMISSION(
112-
propertyNameToSystemPropertyName(GeneralConfig.AGENTLESS_LOG_SUBMISSION_ENABLED), false);
113+
propertyNameToSystemPropertyName(GeneralConfig.AGENTLESS_LOG_SUBMISSION_ENABLED), false),
114+
LLMOBS(propertyNameToSystemPropertyName(LlmObsConfig.LLMOBS_ENABLED), false);
113115

114116
private final String systemProp;
115117
private final boolean enabledByDefault;
@@ -150,6 +152,7 @@ public boolean isEnabledByDefault() {
150152
private static boolean iastFullyDisabled;
151153
private static boolean cwsEnabled = false;
152154
private static boolean ciVisibilityEnabled = false;
155+
private static boolean llmObsEnabled = false;
153156
private static boolean usmEnabled = false;
154157
private static boolean telemetryEnabled = true;
155158
private static boolean debuggerEnabled = false;
@@ -268,6 +271,7 @@ public static void start(
268271
exceptionDebuggingEnabled = isFeatureEnabled(AgentFeature.EXCEPTION_DEBUGGING);
269272
spanOriginEnabled = isFeatureEnabled(AgentFeature.SPAN_ORIGIN);
270273
agentlessLogSubmissionEnabled = isFeatureEnabled(AgentFeature.AGENTLESS_LOG_SUBMISSION);
274+
llmObsEnabled = isFeatureEnabled(AgentFeature.LLMOBS);
271275

272276
if (profilingEnabled) {
273277
if (!isOracleJDK8()) {
@@ -537,6 +541,7 @@ public void execute() {
537541
maybeStartAppSec(scoClass, sco);
538542
maybeStartIast(instrumentation, scoClass, sco);
539543
maybeStartCiVisibility(instrumentation, scoClass, sco);
544+
maybeStartLLMObs(instrumentation, scoClass, sco);
540545
// start debugger before remote config to subscribe to it before starting to poll
541546
maybeStartDebugger(instrumentation, scoClass, sco);
542547
maybeStartRemoteConfig(scoClass, sco);
@@ -883,6 +888,24 @@ private static void maybeStartCiVisibility(Instrumentation inst, Class<?> scoCla
883888
}
884889
}
885890

891+
private static void maybeStartLLMObs(Instrumentation inst, Class<?> scoClass, Object sco) {
892+
if (llmObsEnabled) {
893+
StaticEventLogger.begin("LLM Observability");
894+
895+
try {
896+
final Class<?> llmObsSysClass =
897+
AGENT_CLASSLOADER.loadClass("datadog.trace.llmobs.LLMObsSystem");
898+
final Method llmObsInstallerMethod =
899+
llmObsSysClass.getMethod("start", Instrumentation.class, scoClass);
900+
llmObsInstallerMethod.invoke(null, inst, sco);
901+
} catch (final Throwable e) {
902+
log.warn("Not starting LLM Observability subsystem", e);
903+
}
904+
905+
StaticEventLogger.end("LLM Observability");
906+
}
907+
}
908+
886909
private static void maybeInstallLogsIntake(Class<?> scoClass, Object sco) {
887910
if (agentlessLogSubmissionEnabled) {
888911
StaticEventLogger.begin("Logs Intake");

‎dd-java-agent/agent-llmobs/build.gradle

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
buildscript {
2+
repositories {
3+
mavenCentral()
4+
}
5+
6+
dependencies {
7+
classpath group: 'org.jetbrains.kotlin', name: 'kotlin-gradle-plugin', version: libs.versions.kotlin.get()
8+
}
9+
}
10+
11+
plugins {
12+
id 'com.github.johnrengelman.shadow'
13+
id 'java-test-fixtures'
14+
}
15+
16+
apply from: "$rootDir/gradle/java.gradle"
17+
apply from: "$rootDir/gradle/version.gradle"
18+
apply from: "$rootDir/gradle/test-with-kotlin.gradle"
19+
apply from: "$rootDir/gradle/test-with-scala.gradle"
20+
21+
minimumBranchCoverage = 0.0
22+
minimumInstructionCoverage = 0.0
23+
24+
dependencies {
25+
api libs.slf4j
26+
27+
implementation libs.bundles.asm
28+
implementation group: 'org.jacoco', name: 'org.jacoco.core', version: '0.8.12'
29+
implementation group: 'org.jacoco', name: 'org.jacoco.report', version: '0.8.12'
30+
31+
implementation project(':communication')
32+
implementation project(':components:json')
33+
implementation project(':internal-api')
34+
implementation project(':internal-api:internal-api-9')
35+
36+
testImplementation project(":utils:test-utils")
37+
testImplementation("com.google.jimfs:jimfs:1.1") // an in-memory file system for testing code that works with files
38+
39+
testImplementation libs.scala
40+
testImplementation libs.kotlin
41+
42+
testFixturesApi project(':dd-java-agent:testing')
43+
testFixturesApi project(':utils:test-utils')
44+
45+
testFixturesApi group: 'org.skyscreamer', name: 'jsonassert', version: '1.5.1'
46+
testFixturesApi group: 'org.freemarker', name: 'freemarker', version: '2.3.30'
47+
testFixturesApi group: 'com.jayway.jsonpath', name: 'json-path', version: '2.8.0'
48+
testFixturesApi group: 'com.fasterxml.jackson.core', name: 'jackson-databind', version: '2.16.0'
49+
testFixturesApi group: 'org.msgpack', name: 'jackson-dataformat-msgpack', version: '0.9.6'
50+
51+
testRuntimeOnly("org.junit.platform:junit-platform-launcher:1.9.2") // Required to update dependency lock files
52+
}
53+
54+
shadowJar {
55+
dependencies deps.excludeShared
56+
}
57+
58+
jar {
59+
archiveClassifier = 'unbundled'
60+
}

‎dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsServices.java

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
package datadog.trace.llmobs;
2+
3+
import datadog.communication.BackendApi;
4+
import datadog.communication.BackendApiFactory;
5+
import datadog.communication.ddagent.SharedCommunicationObjects;
6+
import datadog.trace.api.Config;
7+
import org.slf4j.Logger;
8+
import org.slf4j.LoggerFactory;
9+
10+
public class LLMObsServices {
11+
12+
private static final Logger logger = LoggerFactory.getLogger(LLMObsServices.class);
13+
14+
final Config config;
15+
final BackendApi backendApi;
16+
17+
LLMObsServices(Config config, SharedCommunicationObjects sco) {
18+
this.config = config;
19+
this.backendApi =
20+
new BackendApiFactory(config, sco).createBackendApi(BackendApiFactory.Intake.LLMOBS_API);
21+
}
22+
}

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

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
package datadog.trace.llmobs;
2+
3+
import datadog.communication.ddagent.SharedCommunicationObjects;
4+
import datadog.trace.api.Config;
5+
import datadog.trace.api.llmobs.LLMObs;
6+
import datadog.trace.api.llmobs.LLMObsSpan;
7+
import datadog.trace.api.llmobs.LLMObsTags;
8+
import datadog.trace.bootstrap.instrumentation.api.Tags;
9+
import datadog.trace.llmobs.domain.DDLLMObsSpan;
10+
import datadog.trace.llmobs.domain.LLMObsInternal;
11+
import java.lang.instrument.Instrumentation;
12+
import org.jetbrains.annotations.Nullable;
13+
import org.slf4j.Logger;
14+
import org.slf4j.LoggerFactory;
15+
16+
public class LLMObsSystem {
17+
18+
private static final Logger LOGGER = LoggerFactory.getLogger(LLMObsSystem.class);
19+
20+
public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
21+
Config config = Config.get();
22+
if (!config.isLlmObsEnabled()) {
23+
LOGGER.debug("LLM Observability is disabled");
24+
return;
25+
}
26+
27+
sco.createRemaining(config);
28+
29+
LLMObsServices llmObsServices = new LLMObsServices(config, sco);
30+
LLMObsInternal.setLLMObsSpanFactory(
31+
new LLMObsManualSpanFactory(
32+
config.getLlmObsMlApp(), config.getServiceName(), llmObsServices));
33+
}
34+
35+
private static class LLMObsManualSpanFactory implements LLMObs.LLMObsSpanFactory {
36+
37+
private final LLMObsServices llmObsServices;
38+
private final String serviceName;
39+
private final String defaultMLApp;
40+
41+
public LLMObsManualSpanFactory(
42+
String defaultMLApp, String serviceName, LLMObsServices llmObsServices) {
43+
this.defaultMLApp = defaultMLApp;
44+
this.llmObsServices = llmObsServices;
45+
this.serviceName = serviceName;
46+
}
47+
48+
@Override
49+
public LLMObsSpan startLLMSpan(
50+
String spanName,
51+
String modelName,
52+
String modelProvider,
53+
@Nullable String mlApp,
54+
@Nullable String sessionID) {
55+
56+
DDLLMObsSpan span =
57+
new DDLLMObsSpan(
58+
Tags.LLMOBS_LLM_SPAN_KIND, spanName, getMLApp(mlApp), sessionID, serviceName);
59+
60+
span.setTag(LLMObsTags.MODEL_NAME, modelName);
61+
span.setTag(LLMObsTags.MODEL_PROVIDER, modelProvider);
62+
return span;
63+
}
64+
65+
@Override
66+
public LLMObsSpan startAgentSpan(
67+
String spanName, @Nullable String mlApp, @Nullable String sessionID) {
68+
return new DDLLMObsSpan(
69+
Tags.LLMOBS_AGENT_SPAN_KIND, spanName, getMLApp(mlApp), sessionID, serviceName);
70+
}
71+
72+
@Override
73+
public LLMObsSpan startToolSpan(
74+
String spanName, @Nullable String mlApp, @Nullable String sessionID) {
75+
return new DDLLMObsSpan(
76+
Tags.LLMOBS_TOOL_SPAN_KIND, spanName, getMLApp(mlApp), sessionID, serviceName);
77+
}
78+
79+
@Override
80+
public LLMObsSpan startTaskSpan(
81+
String spanName, @Nullable String mlApp, @Nullable String sessionID) {
82+
return new DDLLMObsSpan(
83+
Tags.LLMOBS_TASK_SPAN_KIND, spanName, getMLApp(mlApp), sessionID, serviceName);
84+
}
85+
86+
@Override
87+
public LLMObsSpan startWorkflowSpan(
88+
String spanName, @Nullable String mlApp, @Nullable String sessionID) {
89+
return new DDLLMObsSpan(
90+
Tags.LLMOBS_WORKFLOW_SPAN_KIND, spanName, getMLApp(mlApp), sessionID, serviceName);
91+
}
92+
93+
private String getMLApp(String mlApp) {
94+
if (mlApp == null || mlApp.isEmpty()) {
95+
return defaultMLApp;
96+
}
97+
return mlApp;
98+
}
99+
}
100+
}

0 commit comments

Comments
 (0)