diff --git a/dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfilingIntegration.java b/dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfilingIntegration.java index 00a0358d346..3b143919bbb 100644 --- a/dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfilingIntegration.java +++ b/dd-java-agent/agent-profiling/profiling-ddprof/src/main/java/com/datadog/profiling/ddprof/DatadogProfilingIntegration.java @@ -82,6 +82,14 @@ public void clearContext() { DDPROF.clearContextValue(RESOURCE_NAME_INDEX); } + public static void setLlmPhase(String phaseToken) { + DDPROF.setAgentPhase(phaseToken); + } + + public static void clearLlmPhase() { + DDPROF.clearAgentPhase(); + } + @Override public ProfilingContextAttribute createContextAttribute(String attribute) { return new DatadogProfilerContextSetter(attribute, DDPROF); diff --git a/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/build.gradle b/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/build.gradle index a5a2793836c..1fe2fbe48d6 100644 --- a/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/build.gradle +++ b/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/build.gradle @@ -13,7 +13,32 @@ muzzle { addTestSuiteForDir('latestDepTest', 'test') dependencies { + compileOnly project(':dd-java-agent:agent-profiling:profiling-ddprof') compileOnly group: 'dev.langchain4j', name: 'langchain4j-core', version: minVer testImplementation group: 'dev.langchain4j', name: 'langchain4j', version: minVer + testImplementation group: 'dev.langchain4j', name: 'langchain4j-ollama', version: '1.2.0' latestDepTestImplementation group: 'dev.langchain4j', name: 'langchain4j', version: '+' + latestDepTestCompileOnly group: 'dev.langchain4j', name: 'langchain4j-ollama', version: '+' +} + +tasks.register('runOllamaDemo', JavaExec) { + dependsOn testClasses, ':dd-java-agent:shadowJar' + classpath = sourceSets.test.runtimeClasspath + mainClass = 'datadog.trace.instrumentation.langchain4j.demo.OllamaLlmPipelineDemo' + // Pick the highest semver agent jar in build/libs, ignoring classifier jars. + def agentJar = fileTree("$rootDir/dd-java-agent/build/libs") + .include("dd-java-agent-*.jar") + .exclude("*-sources.jar", "*-javadoc.jar") + .files + .sort { f -> + def m = f.name =~ /dd-java-agent-(\d+)\.(\d+)\.(\d+)/ + m ? (m[0][1] as int) * 1_000_000 + (m[0][2] as int) * 1_000 + (m[0][3] as int) : 0 + }.last() + jvmArgs "-javaagent:${agentJar}", + "-Ddd.profiling.enabled=true", + "-Ddd.profiling.context.attributes.llm.phase.enabled=true", + "-Ddd.trace.enabled=false", + "-Ddd.profiling.start-force-first=true", + "-Ddd.profiling.upload.period=10", + "-Ddd.profiling.debug.dump_path=/tmp/llm-demo-profiles" } diff --git a/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/main/java/datadog/trace/instrumentation/langchain4j/AiServicesInstrumentation.java b/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/main/java/datadog/trace/instrumentation/langchain4j/AiServicesInstrumentation.java index 7033aa04cd6..00a8874cf21 100644 --- a/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/main/java/datadog/trace/instrumentation/langchain4j/AiServicesInstrumentation.java +++ b/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/main/java/datadog/trace/instrumentation/langchain4j/AiServicesInstrumentation.java @@ -3,9 +3,8 @@ import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.named; +import com.datadog.profiling.ddprof.DatadogProfilingIntegration; import datadog.trace.agent.tooling.Instrumenter; -import datadog.trace.api.profiling.Profiling; -import datadog.trace.api.profiling.ProfilingScope; import net.bytebuddy.asm.Advice; public class AiServicesInstrumentation @@ -25,16 +24,13 @@ public void methodAdvice(MethodTransformer transformer) { public static final class InvokeAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static ProfilingScope enter() { - ProfilingScope scope = Profiling.get().newScope(); - scope.setContextValue("llm.agent.phase", "context_build"); - return scope; + public static void enter() { + DatadogProfilingIntegration.setLlmPhase("context_build"); } @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) - public static void exit(@Advice.Enter final ProfilingScope scope) { - scope.clearContextValue("llm.agent.phase"); - scope.close(); + public static void exit() { + DatadogProfilingIntegration.clearLlmPhase(); } } } diff --git a/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/main/java/datadog/trace/instrumentation/langchain4j/ChatModelInstrumentation.java b/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/main/java/datadog/trace/instrumentation/langchain4j/ChatModelInstrumentation.java index 20f286d6589..bead8a76547 100644 --- a/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/main/java/datadog/trace/instrumentation/langchain4j/ChatModelInstrumentation.java +++ b/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/main/java/datadog/trace/instrumentation/langchain4j/ChatModelInstrumentation.java @@ -5,9 +5,8 @@ import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesArgument; +import com.datadog.profiling.ddprof.DatadogProfilingIntegration; import datadog.trace.agent.tooling.Instrumenter; -import datadog.trace.api.profiling.Profiling; -import datadog.trace.api.profiling.ProfilingScope; import net.bytebuddy.asm.Advice; import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.matcher.ElementMatcher; @@ -36,16 +35,13 @@ public void methodAdvice(MethodTransformer transformer) { public static final class ChatAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static ProfilingScope enter() { - ProfilingScope scope = Profiling.get().newScope(); - scope.setContextValue("llm.agent.phase", "awaiting_inference"); - return scope; + public static void enter() { + DatadogProfilingIntegration.setLlmPhase("awaiting_inference"); } @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) - public static void exit(@Advice.Enter final ProfilingScope scope) { - scope.clearContextValue("llm.agent.phase"); - scope.close(); + public static void exit() { + DatadogProfilingIntegration.clearLlmPhase(); } } } diff --git a/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/main/java/datadog/trace/instrumentation/langchain4j/ToolExecutorInstrumentation.java b/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/main/java/datadog/trace/instrumentation/langchain4j/ToolExecutorInstrumentation.java index 2ac69441885..0090b27ed4e 100644 --- a/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/main/java/datadog/trace/instrumentation/langchain4j/ToolExecutorInstrumentation.java +++ b/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/main/java/datadog/trace/instrumentation/langchain4j/ToolExecutorInstrumentation.java @@ -5,9 +5,8 @@ import static net.bytebuddy.matcher.ElementMatchers.isMethod; import static net.bytebuddy.matcher.ElementMatchers.takesArguments; +import com.datadog.profiling.ddprof.DatadogProfilingIntegration; import datadog.trace.agent.tooling.Instrumenter; -import datadog.trace.api.profiling.Profiling; -import datadog.trace.api.profiling.ProfilingScope; import net.bytebuddy.asm.Advice; import net.bytebuddy.description.type.TypeDescription; import net.bytebuddy.matcher.ElementMatcher; @@ -36,16 +35,13 @@ public void methodAdvice(MethodTransformer transformer) { public static final class ExecuteAdvice { @Advice.OnMethodEnter(suppress = Throwable.class) - public static ProfilingScope enter() { - ProfilingScope scope = Profiling.get().newScope(); - scope.setContextValue("llm.agent.phase", "tool_execution"); - return scope; + public static void enter() { + DatadogProfilingIntegration.setLlmPhase("tool_execution"); } @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) - public static void exit(@Advice.Enter final ProfilingScope scope) { - scope.clearContextValue("llm.agent.phase"); - scope.close(); + public static void exit() { + DatadogProfilingIntegration.clearLlmPhase(); } } } diff --git a/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/test/java/datadog/trace/instrumentation/langchain4j/demo/MockLlmPipelineTest.java b/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/test/java/datadog/trace/instrumentation/langchain4j/demo/MockLlmPipelineTest.java new file mode 100644 index 00000000000..5ec16c0bff9 --- /dev/null +++ b/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/test/java/datadog/trace/instrumentation/langchain4j/demo/MockLlmPipelineTest.java @@ -0,0 +1,62 @@ +package datadog.trace.instrumentation.langchain4j.demo; + +import dev.langchain4j.agent.tool.Tool; +import dev.langchain4j.agent.tool.ToolExecutionRequest; +import dev.langchain4j.data.message.AiMessage; +import dev.langchain4j.memory.chat.MessageWindowChatMemory; +import dev.langchain4j.model.chat.ChatModel; +import dev.langchain4j.model.chat.request.ChatRequest; +import dev.langchain4j.model.chat.response.ChatResponse; +import dev.langchain4j.service.AiServices; +import org.junit.jupiter.api.Test; + +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.assertNotNull; + +public class MockLlmPipelineTest { + + interface WeatherAssistant { + String chat(String message); + } + + static class MockWeatherTool { + @Tool("Get current weather for a location") + public String getWeather(String location) { + return "Sunny, 22°C in " + location; + } + } + + static class TwoTurnMockModel implements ChatModel { + private final AtomicInteger calls = new AtomicInteger(); + + @Override + public ChatResponse chat(ChatRequest request) { + try { Thread.sleep(30); } catch (InterruptedException ignored) {} + if (calls.getAndIncrement() == 0) { + return ChatResponse.builder() + .aiMessage(AiMessage.from( + ToolExecutionRequest.builder() + .name("getWeather") + .arguments("{\"location\": \"Amsterdam\"}") + .build())) + .build(); + } + return ChatResponse.builder() + .aiMessage(AiMessage.from("The weather in Amsterdam is Sunny, 22°C.")) + .build(); + } + } + + @Test + public void pipelineExercisesAllThreeInstrumentationPoints() { + WeatherAssistant assistant = AiServices.builder(WeatherAssistant.class) + .chatModel(new TwoTurnMockModel()) + .chatMemory(MessageWindowChatMemory.withMaxMessages(10)) + .tools(new MockWeatherTool()) + .build(); + + String response = assistant.chat("What is the weather in Amsterdam?"); + assertNotNull(response, "Expected a non-null response from the mock pipeline"); + } +} diff --git a/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/test/java/datadog/trace/instrumentation/langchain4j/demo/OllamaLlmPipelineDemo.java b/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/test/java/datadog/trace/instrumentation/langchain4j/demo/OllamaLlmPipelineDemo.java new file mode 100644 index 00000000000..455e918f51c --- /dev/null +++ b/dd-java-agent/instrumentation/langchain4j/langchain4j-1.0/src/test/java/datadog/trace/instrumentation/langchain4j/demo/OllamaLlmPipelineDemo.java @@ -0,0 +1,58 @@ +package datadog.trace.instrumentation.langchain4j.demo; + +import dev.langchain4j.memory.chat.MessageWindowChatMemory; +import dev.langchain4j.model.ollama.OllamaChatModel; +import dev.langchain4j.service.AiServices; + +import java.time.Duration; + +/** + * Generates a JFR recording with llm.agent.phase tags against a local Ollama server. + * + * Run: + * java -javaagent:dd-java-agent.jar + * -Ddd.profiling.enabled=true + * -Ddd.profiling.context.attributes.llm.phase.enabled=true + * -cp + * datadog.trace.instrumentation.langchain4j.demo.OllamaLlmPipelineDemo + * + * Prerequisites: ollama serve && ollama pull llama3 + */ +public class OllamaLlmPipelineDemo { + + interface Assistant { + String chat(String message); + } + + public static void main(String[] args) { + OllamaChatModel model = OllamaChatModel.builder() + .baseUrl("http://localhost:11434") + .modelName("llama3") + .timeout(Duration.ofMinutes(2)) + .build(); + + Assistant assistant = AiServices.builder(Assistant.class) + .chatModel(model) + .chatMemory(MessageWindowChatMemory.withMaxMessages(10)) + .build(); + + String[] questions = { + "Explain Java garbage collection in one sentence.", + "What is a Java virtual thread?", + "Name one advantage of the G1 garbage collector." + }; + for (String q : questions) { + System.out.println("Q: " + q); + System.out.println("A: " + assistant.chat(q)); + } + + // Hold the JVM alive so the profiler flushes at least two recording cycles + // (dd.profiling.upload.period=10 → flush at ~10 s and ~20 s). + System.out.println("Waiting 25 s for profiling data to flush..."); + try { + Thread.sleep(25_000); + } catch (InterruptedException ignored) { + } + System.out.println("Done."); + } +}