Skip to content

Commit 6e3ea48

Browse files
gary-huanggary-huang
committed
impl llmobs agent and llmobs apis
1 parent ec0066e commit 6e3ea48

File tree

14 files changed

+789
-1
lines changed

14 files changed

+789
-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: 46 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;
@@ -38,6 +39,7 @@
3839
import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
3940
import datadog.trace.bootstrap.instrumentation.api.AgentTracer.TracerAPI;
4041
import datadog.trace.bootstrap.instrumentation.api.ProfilingContextIntegration;
42+
import datadog.trace.bootstrap.instrumentation.api.WriterConstants;
4143
import datadog.trace.bootstrap.instrumentation.jfr.InstrumentationBasedProfiling;
4244
import datadog.trace.util.AgentTaskScheduler;
4345
import datadog.trace.util.AgentThreadFactory.AgentThread;
@@ -105,7 +107,10 @@ private enum AgentFeature {
105107
EXCEPTION_REPLAY(DebuggerConfig.EXCEPTION_REPLAY_ENABLED, false),
106108
CODE_ORIGIN(TraceInstrumentationConfig.CODE_ORIGIN_FOR_SPANS_ENABLED, false),
107109
DATA_JOBS(GeneralConfig.DATA_JOBS_ENABLED, false),
108-
AGENTLESS_LOG_SUBMISSION(GeneralConfig.AGENTLESS_LOG_SUBMISSION_ENABLED, false);
110+
AGENTLESS_LOG_SUBMISSION(GeneralConfig.AGENTLESS_LOG_SUBMISSION_ENABLED, false),
111+
LLMOBS(propertyNameToSystemPropertyName(LlmObsConfig.LLMOBS_ENABLED), false),
112+
LLMOBS_AGENTLESS(
113+
propertyNameToSystemPropertyName(LlmObsConfig.LLMOBS_AGENTLESS_ENABLED), false);
109114

110115
private final String configKey;
111116
private final String systemProp;
@@ -152,6 +157,8 @@ public boolean isEnabledByDefault() {
152157
private static boolean iastFullyDisabled;
153158
private static boolean cwsEnabled = false;
154159
private static boolean ciVisibilityEnabled = false;
160+
private static boolean llmObsEnabled = false;
161+
private static boolean llmObsAgentlessEnabled = false;
155162
private static boolean usmEnabled = false;
156163
private static boolean telemetryEnabled = true;
157164
private static boolean dynamicInstrumentationEnabled = false;
@@ -277,6 +284,25 @@ public static void start(
277284
exceptionReplayEnabled = isFeatureEnabled(AgentFeature.EXCEPTION_REPLAY);
278285
codeOriginEnabled = isFeatureEnabled(AgentFeature.CODE_ORIGIN);
279286
agentlessLogSubmissionEnabled = isFeatureEnabled(AgentFeature.AGENTLESS_LOG_SUBMISSION);
287+
llmObsEnabled = isFeatureEnabled(AgentFeature.LLMOBS);
288+
289+
// setup writers when llmobs is enabled to accomodate apm and llmobs
290+
if (llmObsEnabled) {
291+
// for llm obs spans, use agent proxy by default, apm spans will use agent writer
292+
setSystemPropertyDefault(
293+
propertyNameToSystemPropertyName(TracerConfig.WRITER_TYPE),
294+
WriterConstants.MULTI_WRITER_TYPE
295+
+ ":"
296+
+ WriterConstants.DD_INTAKE_WRITER_TYPE
297+
+ ","
298+
+ WriterConstants.DD_AGENT_WRITER_TYPE);
299+
if (llmObsAgentlessEnabled) {
300+
// use API writer only
301+
setSystemPropertyDefault(
302+
propertyNameToSystemPropertyName(TracerConfig.WRITER_TYPE),
303+
WriterConstants.DD_INTAKE_WRITER_TYPE);
304+
}
305+
}
280306

281307
if (profilingEnabled) {
282308
if (!isOracleJDK8()) {
@@ -565,6 +591,7 @@ public void execute() {
565591

566592
maybeStartAppSec(scoClass, sco);
567593
maybeStartCiVisibility(instrumentation, scoClass, sco);
594+
maybeStartLLMObs(instrumentation, scoClass, sco);
568595
// start debugger before remote config to subscribe to it before starting to poll
569596
maybeStartDebugger(instrumentation, scoClass, sco);
570597
maybeStartRemoteConfig(scoClass, sco);
@@ -920,6 +947,24 @@ private static void maybeStartCiVisibility(Instrumentation inst, Class<?> scoCla
920947
}
921948
}
922949

950+
private static void maybeStartLLMObs(Instrumentation inst, Class<?> scoClass, Object sco) {
951+
if (llmObsEnabled) {
952+
StaticEventLogger.begin("LLM Observability");
953+
954+
try {
955+
final Class<?> llmObsSysClass =
956+
AGENT_CLASSLOADER.loadClass("datadog.trace.llmobs.LLMObsSystem");
957+
final Method llmObsInstallerMethod =
958+
llmObsSysClass.getMethod("start", Instrumentation.class, scoClass);
959+
llmObsInstallerMethod.invoke(null, inst, sco);
960+
} catch (final Throwable e) {
961+
log.warn("Not starting LLM Observability subsystem", e);
962+
}
963+
964+
StaticEventLogger.end("LLM Observability");
965+
}
966+
}
967+
923968
private static void maybeInstallLogsIntake(Class<?> scoClass, Object sco) {
924969
if (agentlessLogSubmissionEnabled) {
925970
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.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+
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/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)