Skip to content

Commit 1701c89

Browse files
gary-huanggary-huang
committed
impl llmobs agent
working with spans sent to apm
1 parent 3a3f08e commit 1701c89

File tree

14 files changed

+399
-15
lines changed

14 files changed

+399
-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: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
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.bootstrap.instrumentation.api.Tags;
8+
import datadog.trace.llmobs.domain.DDLLMObsSpan;
9+
import datadog.trace.llmobs.domain.LLMObsInternal;
10+
import java.lang.instrument.Instrumentation;
11+
import org.jetbrains.annotations.Nullable;
12+
import org.slf4j.Logger;
13+
import org.slf4j.LoggerFactory;
14+
15+
public class LLMObsSystem {
16+
17+
private static final Logger LOGGER = LoggerFactory.getLogger(LLMObsSystem.class);
18+
19+
public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
20+
Config config = Config.get();
21+
if (!config.isLlmObsEnabled()) {
22+
LOGGER.debug("LLM Observability is disabled");
23+
return;
24+
}
25+
26+
sco.createRemaining(config);
27+
28+
LLMObsServices llmObsServices = new LLMObsServices(config, sco);
29+
LLMObsInternal.setLLMObsSpanFactory(
30+
new LLMObsManualSpanFactory(config.getLlmObsMlApp(), llmObsServices));
31+
}
32+
33+
private static class LLMObsManualSpanFactory implements LLMObs.LLMObsSpanFactory {
34+
35+
private final LLMObsServices llmObsServices;
36+
private final String defaultMLApp;
37+
38+
public LLMObsManualSpanFactory(String defaultMLApp, LLMObsServices llmObsServices) {
39+
this.defaultMLApp = defaultMLApp;
40+
this.llmObsServices = llmObsServices;
41+
}
42+
43+
@Override
44+
public LLMObsSpan startLLMSpan(
45+
String spanName,
46+
String modelName,
47+
String modelProvider,
48+
@Nullable String mlApp,
49+
@Nullable String sessionID) {
50+
if (mlApp == null || mlApp.isEmpty()) {
51+
mlApp = defaultMLApp;
52+
}
53+
54+
return new DDLLMObsSpan(Tags.LLMOBS_LLM_SPAN_KIND, spanName, mlApp, sessionID);
55+
}
56+
57+
@Override
58+
public LLMObsSpan startAgentSpan(
59+
String spanName, @Nullable String mlApp, @Nullable String sessionID) {
60+
return null;
61+
}
62+
63+
@Override
64+
public LLMObsSpan startToolSpan(
65+
String spanName, @Nullable String mlApp, @Nullable String sessionID) {
66+
return null;
67+
}
68+
69+
@Override
70+
public LLMObsSpan startTaskSpan(
71+
String spanName, @Nullable String mlApp, @Nullable String sessionID) {
72+
return null;
73+
}
74+
75+
@Override
76+
public LLMObsSpan startWorkflowSpan(
77+
String spanName, @Nullable String mlApp, @Nullable String sessionID) {
78+
return null;
79+
}
80+
}
81+
}

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

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
package datadog.trace.llmobs.domain;
2+
3+
import datadog.trace.api.DDSpanTypes;
4+
import datadog.trace.api.llmobs.LLMObsSpan;
5+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
6+
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
7+
import datadog.trace.bootstrap.instrumentation.api.Tags;
8+
import java.util.HashMap;
9+
import java.util.List;
10+
import java.util.Map;
11+
import java.util.Optional;
12+
13+
public class DDLLMObsSpan implements LLMObsSpan {
14+
private String sessionID = null;
15+
16+
private String mlApp;
17+
private final AgentSpan span;
18+
private String kind;
19+
private SpanIO spanInput;
20+
private SpanIO spanOutput;
21+
private Map<String, Number> metrics;
22+
23+
private boolean finished = false;
24+
25+
public DDLLMObsSpan(String kind, String spanName, String mlApp, String sessionID) {
26+
AgentTracer.SpanBuilder spanBuilder =
27+
AgentTracer.get()
28+
.buildSpan("LLMOBS", spanName)
29+
.withSpanType(DDSpanTypes.LLMOBS);
30+
31+
this.mlApp = mlApp;
32+
this.kind = kind;
33+
this.sessionID = sessionID;
34+
this.metrics = new HashMap<>();
35+
36+
this.span = spanBuilder.start();
37+
this.span.setTag(Tags.SPAN_KIND, kind);
38+
}
39+
40+
// span_id: str
41+
// trace_id: str
42+
// parent_id: str
43+
// session_id: str
44+
// tags: List[str]
45+
// service: str
46+
// name: str
47+
// start_ns: int
48+
// duration: float
49+
// status: str
50+
// status_message: str
51+
// meta: Dict[str, Any]
52+
// metrics: Dict[str, Any]
53+
// collection_errors: List[str]
54+
55+
@Override
56+
public String toString() {
57+
return super.toString()
58+
+ ", kind="
59+
+ this.kind
60+
+ ", mlApp="
61+
+ this.mlApp
62+
+ ", sessionID="
63+
+ this.sessionID
64+
+ ", context="
65+
+ this.span.context();
66+
}
67+
68+
@Override
69+
public void annotateIO(
70+
List<Map<String, Object>> inputData, List<Map<String, Object>> outputData) {
71+
if (finished) {
72+
return;
73+
}
74+
if (inputData != null && !inputData.isEmpty()) {
75+
this.spanInput = new SpanIO(inputData);
76+
}
77+
if (outputData != null && !outputData.isEmpty()) {
78+
this.spanOutput = new SpanIO(outputData);
79+
}
80+
}
81+
82+
@Override
83+
public void annotateIO(String inputData, String outputData) {
84+
if (finished) {
85+
return;
86+
}
87+
if (inputData != null && !inputData.isEmpty()) {
88+
this.spanInput = new SpanIO(inputData);
89+
}
90+
if (outputData != null && !outputData.isEmpty()) {
91+
this.spanOutput = new SpanIO(outputData);
92+
}
93+
}
94+
95+
@Override
96+
public void setMetadata(Map<String, Object> metadata) {}
97+
98+
@Override
99+
public void setMetrics(Map<String, Number> metrics) {
100+
if (finished) {
101+
return;
102+
}
103+
// TODO handle concurrency? unless using DDSpanContext
104+
this.metrics.putAll(metrics);
105+
}
106+
107+
@Override
108+
public void setTags(Map<String, Object> tags) {
109+
if (finished) {
110+
return;
111+
}
112+
if (tags != null && !tags.isEmpty()) {
113+
for (Map.Entry<String, Object> entry : tags.entrySet()) {
114+
this.span.setTag(entry.getKey(), entry.getValue());
115+
}
116+
}
117+
}
118+
119+
@Override
120+
public void setTag(String key, String value) {
121+
this.span.setTag(key, value);
122+
}
123+
124+
@Override
125+
public void finish() {
126+
this.span.finish();
127+
this.finished = true;
128+
}
129+
130+
private class SpanIO {
131+
Optional<List<Map<String, Object>>> listIO;
132+
Optional<String> strIO;
133+
134+
SpanIO(List<Map<String, Object>> listIO) {
135+
this.listIO = Optional.of(listIO);
136+
}
137+
138+
SpanIO(String strIO) {
139+
this.strIO = Optional.of(strIO);
140+
}
141+
}
142+
}

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
package datadog.trace.llmobs.domain;
2+
3+
import datadog.trace.api.llmobs.LLMObs;
4+
5+
public class LLMObsInternal extends LLMObs {
6+
7+
public static void setLLMObsSpanFactory(final LLMObsSpanFactory factory) {
8+
LLMObs.SPAN_FACTORY = factory;
9+
}
10+
}

0 commit comments

Comments
 (0)