Skip to content

Commit 77291d8

Browse files
gary-huanggary-huang
authored
implement LLM Obs SDK spans APIs (#8390)
* add APIs for llm obs * add llm message class to support llm spans * add llm message class to support llm spans * impl llmobs agent and llmobs apis * support llm messages with tool calls * handle default model name and provider * rm unneeded file * spotless * add APIs for llm obs sdk (#8135) * add APIs for llm obs * add llm message class to support llm spans * follow java convention of naming Id instead of ID * add codeowners * rename ID to Id according to java naming conventions * Undo change to integrations-core submodule * fix build gradle * rm empty line * fix test
1 parent 2be1b73 commit 77291d8

File tree

13 files changed

+800
-1
lines changed

13 files changed

+800
-1
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: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import datadog.trace.api.config.GeneralConfig;
2626
import datadog.trace.api.config.IastConfig;
2727
import datadog.trace.api.config.JmxFetchConfig;
28+
import datadog.trace.api.config.LlmObsConfig;
2829
import datadog.trace.api.config.ProfilingConfig;
2930
import datadog.trace.api.config.RemoteConfigConfig;
3031
import datadog.trace.api.config.TraceInstrumentationConfig;
@@ -41,6 +42,7 @@
4142
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
4243
import datadog.trace.bootstrap.instrumentation.api.AgentTracer.TracerAPI;
4344
import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
45+
import datadog.trace.bootstrap.instrumentation.api.WriterConstants;
4446
import datadog.trace.bootstrap.instrumentation.jfr.InstrumentationBasedProfiling;
4547
import datadog.trace.util.AgentTaskScheduler;
4648
import datadog.trace.util.AgentThreadFactory.AgentThread;
@@ -109,7 +111,9 @@ private enum AgentFeature {
109111
EXCEPTION_REPLAY(DebuggerConfig.EXCEPTION_REPLAY_ENABLED, false),
110112
CODE_ORIGIN(TraceInstrumentationConfig.CODE_ORIGIN_FOR_SPANS_ENABLED, false),
111113
DATA_JOBS(GeneralConfig.DATA_JOBS_ENABLED, false),
112-
AGENTLESS_LOG_SUBMISSION(GeneralConfig.AGENTLESS_LOG_SUBMISSION_ENABLED, false);
114+
AGENTLESS_LOG_SUBMISSION(GeneralConfig.AGENTLESS_LOG_SUBMISSION_ENABLED, false),
115+
LLMOBS(LlmObsConfig.LLMOBS_ENABLED, false),
116+
LLMOBS_AGENTLESS(LlmObsConfig.LLMOBS_AGENTLESS_ENABLED, false);
113117

114118
private final String configKey;
115119
private final String systemProp;
@@ -156,6 +160,8 @@ public boolean isEnabledByDefault() {
156160
private static boolean iastFullyDisabled;
157161
private static boolean cwsEnabled = false;
158162
private static boolean ciVisibilityEnabled = false;
163+
private static boolean llmObsEnabled = false;
164+
private static boolean llmObsAgentlessEnabled = false;
159165
private static boolean usmEnabled = false;
160166
private static boolean telemetryEnabled = true;
161167
private static boolean dynamicInstrumentationEnabled = false;
@@ -290,6 +296,25 @@ public static void start(
290296
exceptionReplayEnabled = isFeatureEnabled(AgentFeature.EXCEPTION_REPLAY);
291297
codeOriginEnabled = isFeatureEnabled(AgentFeature.CODE_ORIGIN);
292298
agentlessLogSubmissionEnabled = isFeatureEnabled(AgentFeature.AGENTLESS_LOG_SUBMISSION);
299+
llmObsEnabled = isFeatureEnabled(AgentFeature.LLMOBS);
300+
301+
// setup writers when llmobs is enabled to accomodate apm and llmobs
302+
if (llmObsEnabled) {
303+
// for llm obs spans, use agent proxy by default, apm spans will use agent writer
304+
setSystemPropertyDefault(
305+
propertyNameToSystemPropertyName(TracerConfig.WRITER_TYPE),
306+
WriterConstants.MULTI_WRITER_TYPE
307+
+ ":"
308+
+ WriterConstants.DD_INTAKE_WRITER_TYPE
309+
+ ","
310+
+ WriterConstants.DD_AGENT_WRITER_TYPE);
311+
if (llmObsAgentlessEnabled) {
312+
// use API writer only
313+
setSystemPropertyDefault(
314+
propertyNameToSystemPropertyName(TracerConfig.WRITER_TYPE),
315+
WriterConstants.DD_INTAKE_WRITER_TYPE);
316+
}
317+
}
293318

294319
patchJPSAccess(inst);
295320

@@ -597,6 +622,7 @@ public void execute() {
597622

598623
maybeStartAppSec(scoClass, sco);
599624
maybeStartCiVisibility(instrumentation, scoClass, sco);
625+
maybeStartLLMObs(instrumentation, scoClass, sco);
600626
// start debugger before remote config to subscribe to it before starting to poll
601627
maybeStartDebugger(instrumentation, scoClass, sco);
602628
maybeStartRemoteConfig(scoClass, sco);
@@ -952,6 +978,24 @@ private static void maybeStartCiVisibility(Instrumentation inst, Class<?> scoCla
952978
}
953979
}
954980

981+
private static void maybeStartLLMObs(Instrumentation inst, Class<?> scoClass, Object sco) {
982+
if (llmObsEnabled) {
983+
StaticEventLogger.begin("LLM Observability");
984+
985+
try {
986+
final Class<?> llmObsSysClass =
987+
AGENT_CLASSLOADER.loadClass("datadog.trace.llmobs.LLMObsSystem");
988+
final Method llmObsInstallerMethod =
989+
llmObsSysClass.getMethod("start", Instrumentation.class, scoClass);
990+
llmObsInstallerMethod.invoke(null, inst, sco);
991+
} catch (final Throwable e) {
992+
log.warn("Not starting LLM Observability subsystem", e);
993+
}
994+
995+
StaticEventLogger.end("LLM Observability");
996+
}
997+
}
998+
955999
private static void maybeInstallLogsIntake(Class<?> scoClass, Object sco) {
9561000
if (agentlessLogSubmissionEnabled) {
9571001
StaticEventLogger.begin("Logs Intake");

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

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
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.gradleup.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+
20+
minimumBranchCoverage = 0.0
21+
minimumInstructionCoverage = 0.0
22+
23+
dependencies {
24+
api libs.slf4j
25+
26+
implementation project(':communication')
27+
implementation project(':components:json')
28+
implementation project(':internal-api')
29+
30+
testImplementation project(":utils:test-utils")
31+
32+
testFixturesApi project(':dd-java-agent:testing')
33+
testFixturesApi project(':utils:test-utils')
34+
}
35+
36+
shadowJar {
37+
dependencies deps.excludeShared
38+
}
39+
40+
jar {
41+
archiveClassifier = 'unbundled'
42+
}

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

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
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+
private static final String CUSTOM_MODEL_VAL = "custom";
21+
22+
public static void start(Instrumentation inst, SharedCommunicationObjects sco) {
23+
Config config = Config.get();
24+
if (!config.isLlmObsEnabled()) {
25+
LOGGER.debug("LLM Observability is disabled");
26+
return;
27+
}
28+
29+
sco.createRemaining(config);
30+
31+
LLMObsInternal.setLLMObsSpanFactory(
32+
new LLMObsManualSpanFactory(config.getLlmObsMlApp(), config.getServiceName()));
33+
}
34+
35+
private static class LLMObsManualSpanFactory implements LLMObs.LLMObsSpanFactory {
36+
37+
private final String serviceName;
38+
private final String defaultMLApp;
39+
40+
public LLMObsManualSpanFactory(String defaultMLApp, String serviceName) {
41+
this.defaultMLApp = defaultMLApp;
42+
this.serviceName = serviceName;
43+
}
44+
45+
@Override
46+
public LLMObsSpan startLLMSpan(
47+
String spanName,
48+
String modelName,
49+
String modelProvider,
50+
@Nullable String mlApp,
51+
@Nullable String sessionId) {
52+
53+
DDLLMObsSpan span =
54+
new DDLLMObsSpan(
55+
Tags.LLMOBS_LLM_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName);
56+
57+
if (modelName == null || modelName.isEmpty()) {
58+
modelName = CUSTOM_MODEL_VAL;
59+
}
60+
span.setTag(LLMObsTags.MODEL_NAME, modelName);
61+
62+
if (modelProvider == null || modelProvider.isEmpty()) {
63+
modelProvider = CUSTOM_MODEL_VAL;
64+
}
65+
span.setTag(LLMObsTags.MODEL_PROVIDER, modelProvider);
66+
return span;
67+
}
68+
69+
@Override
70+
public LLMObsSpan startAgentSpan(
71+
String spanName, @Nullable String mlApp, @Nullable String sessionId) {
72+
return new DDLLMObsSpan(
73+
Tags.LLMOBS_AGENT_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName);
74+
}
75+
76+
@Override
77+
public LLMObsSpan startToolSpan(
78+
String spanName, @Nullable String mlApp, @Nullable String sessionId) {
79+
return new DDLLMObsSpan(
80+
Tags.LLMOBS_TOOL_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName);
81+
}
82+
83+
@Override
84+
public LLMObsSpan startTaskSpan(
85+
String spanName, @Nullable String mlApp, @Nullable String sessionId) {
86+
return new DDLLMObsSpan(
87+
Tags.LLMOBS_TASK_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName);
88+
}
89+
90+
@Override
91+
public LLMObsSpan startWorkflowSpan(
92+
String spanName, @Nullable String mlApp, @Nullable String sessionId) {
93+
return new DDLLMObsSpan(
94+
Tags.LLMOBS_WORKFLOW_SPAN_KIND, spanName, getMLApp(mlApp), sessionId, serviceName);
95+
}
96+
97+
private String getMLApp(String mlApp) {
98+
if (mlApp == null || mlApp.isEmpty()) {
99+
return defaultMLApp;
100+
}
101+
return mlApp;
102+
}
103+
}
104+
}

0 commit comments

Comments
 (0)