diff --git a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java index e184a645f71..42404944c2b 100644 --- a/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java +++ b/dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/LLMObsSystem.java @@ -12,7 +12,7 @@ import java.lang.instrument.Instrumentation; import java.util.Map; import java.util.concurrent.TimeUnit; -import org.jetbrains.annotations.Nullable; +import javax.annotation.Nullable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -22,7 +22,7 @@ public class LLMObsSystem { private static final String CUSTOM_MODEL_VAL = "custom"; - public static void start(Instrumentation inst, SharedCommunicationObjects sco) { + public static void start(@Nullable Instrumentation inst, SharedCommunicationObjects sco) { Config config = Config.get(); if (!config.isLlmObsEnabled()) { LOGGER.debug("LLM Observability is disabled"); diff --git a/dd-java-agent/instrumentation/openai-java-2.8/build.gradle b/dd-java-agent/instrumentation/openai-java-2.8/build.gradle new file mode 100644 index 00000000000..ab67cf1d288 --- /dev/null +++ b/dd-java-agent/instrumentation/openai-java-2.8/build.gradle @@ -0,0 +1,27 @@ +muzzle { + pass { + group = "com.openai" + module = "openai-java" + versions = "[2.8.0,)" + assertInverse = true + } +} + +// Instrumentation is currently in preview. Additional testing will be implemented before enabling this by default. + +apply from: "$rootDir/gradle/java.gradle" + +addTestSuiteForDir('latestDepTest', 'test') + +dependencies { + compileOnly(group: 'com.openai', name: 'openai-java', version: '2.8.0') + + testImplementation(group: 'com.openai', name: 'openai-java') { + version { + strictly '[2.8.0,)' + } + } + + latestDepTestImplementation group: 'com.openai', name: 'openai-java', version: '+' + +} diff --git a/dd-java-agent/instrumentation/openai-java-2.8/gradle.lockfile b/dd-java-agent/instrumentation/openai-java-2.8/gradle.lockfile new file mode 100644 index 00000000000..09cf7528a00 --- /dev/null +++ b/dd-java-agent/instrumentation/openai-java-2.8/gradle.lockfile @@ -0,0 +1,4 @@ +# This is a Gradle generated file for dependency locking. +# This file is expected to be part of source control. +# Manual entries are allowed. +empty= \ No newline at end of file diff --git a/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/ChatCompletionServiceInstrumentation.java b/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/ChatCompletionServiceInstrumentation.java new file mode 100644 index 00000000000..9dad1c236cf --- /dev/null +++ b/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/ChatCompletionServiceInstrumentation.java @@ -0,0 +1,106 @@ +package datadog.trace.instrumentation.openaiclient; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.*; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; + +import com.openai.models.chat.completions.ChatCompletion; +import com.openai.models.chat.completions.ChatCompletionCreateParams; +import com.openai.services.blocking.CompletionService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.InstrumenterConfig; +import datadog.trace.api.llmobs.LLMObsSpan; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +public class ChatCompletionServiceInstrumentation extends InstrumenterModule.Tracing + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + public ChatCompletionServiceInstrumentation() { + super("openai-java", "openai-java-2.8", "openai-client"); + } + + @Override + protected boolean defaultEnabled() { + return super.defaultEnabled() + && InstrumenterConfig.get().isIntegrationEnabled(Collections.singleton("openai"), false); + } + + @Override + public String[] helperClassNames() { + return new String[] { + packageName + ".OpenAIClientInfo", + }; + } + + @Override + public Map contextStore() { + Map contextStores = new HashMap<>(1); + contextStores.put( + "com.openai.services.blocking.ChatCompletionService", OpenAIClientInfo.class.getName()); + return contextStores; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isMethod() + .and(isPublic()) + .and(named("create")) + .and( + takesArgument( + 0, named("com.openai.models.completions.ChatCompletionCreateParams"))), + getClass().getName() + "$ChatCompletionServiceAdvice"); + } + + @Override + public String hierarchyMarkerType() { + return "com.openai.services.blocking.chat.ChatCompletionService"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(named(hierarchyMarkerType())); + } + + public static class ChatCompletionServiceAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static Map methodEnter( + @Advice.Argument(0) final ChatCompletionCreateParams params) { + Map spans = new HashMap<>(); + spans.put( + "datadog.trace.api.llmobs.LLMObsSpan", + OpenAIClientDecorator.DECORATE.startLLMObsChatCompletionSpan(params)); + spans.put( + "datadog.trace.bootstrap.instrumentation.api.AgentScope", + OpenAIClientDecorator.DECORATE.startChatCompletionSpan(params)); + return spans; + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void methodExit( + @Advice.Enter final Map spans, + @Advice.This final CompletionService completionService, + @Advice.Return final Object result, + @Advice.Thrown final Throwable throwable) { + OpenAIClientInfo info = + InstrumentationContext.get(CompletionService.class, OpenAIClientInfo.class) + .get(completionService); + OpenAIClientDecorator.DECORATE.finishLLMObsChatCompletionSpan( + (LLMObsSpan) spans.get("datadog.trace.api.llmobs.LLMObsSpan"), + (ChatCompletion) result, + throwable); + OpenAIClientDecorator.DECORATE.finishSpan( + (AgentScope) spans.get("datadog.trace.bootstrap.instrumentation.api.AgentScope"), + result, + throwable); + } + } +} diff --git a/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/CompletionServiceInstrumentation.java b/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/CompletionServiceInstrumentation.java new file mode 100644 index 00000000000..edad2d0f071 --- /dev/null +++ b/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/CompletionServiceInstrumentation.java @@ -0,0 +1,88 @@ +package datadog.trace.instrumentation.openaiclient; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.*; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; + +import com.openai.models.completions.CompletionCreateParams; +import com.openai.services.blocking.CompletionService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.InstrumenterConfig; +import datadog.trace.bootstrap.InstrumentationContext; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +public class CompletionServiceInstrumentation extends InstrumenterModule.Tracing + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + public CompletionServiceInstrumentation() { + super("openai-java", "openai-java-2.8", "openai-client"); + } + + @Override + protected boolean defaultEnabled() { + return super.defaultEnabled() + && InstrumenterConfig.get().isIntegrationEnabled(Collections.singleton("openai"), false); + } + + @Override + public String[] helperClassNames() { + return new String[] { + packageName + ".OpenAIClientInfo", + }; + } + + @Override + public Map contextStore() { + Map contextStores = new HashMap<>(1); + contextStores.put( + "com.openai.services.blocking.CompletionService", OpenAIClientInfo.class.getName()); + return contextStores; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isMethod() + .and(isPublic()) + .and(named("create")) + .and(takesArgument(0, named("com.openai.models.completions.CompletionCreateParams"))), + getClass().getName() + "$CompletionServiceAdvice"); + } + + @Override + public String hierarchyMarkerType() { + return "com.openai.services.blocking.CompletionService"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(named(hierarchyMarkerType())); + } + + public static class CompletionServiceAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope methodEnter(@Advice.Argument(0) final CompletionCreateParams params) { + + return OpenAIClientDecorator.DECORATE.startCompletionSpan(params); + } + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void methodExit( + @Advice.Enter final AgentScope span, + @Advice.This final CompletionService completionService, + @Advice.Return final Object result, + @Advice.Thrown final Throwable throwable) { + OpenAIClientInfo info = + InstrumentationContext.get(CompletionService.class, OpenAIClientInfo.class) + .get(completionService); + OpenAIClientDecorator.DECORATE.finishSpan(span, result, throwable); + } + } +} diff --git a/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/EmbeddingServiceInstrumentation.java b/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/EmbeddingServiceInstrumentation.java new file mode 100644 index 00000000000..6e313adf14a --- /dev/null +++ b/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/EmbeddingServiceInstrumentation.java @@ -0,0 +1,79 @@ +package datadog.trace.instrumentation.openaiclient; + +import static datadog.trace.agent.tooling.bytebuddy.matcher.HierarchyMatchers.implementsInterface; +import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named; +import static net.bytebuddy.matcher.ElementMatchers.isMethod; +import static net.bytebuddy.matcher.ElementMatchers.isPublic; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; + +import com.google.auto.service.AutoService; +import com.openai.models.embeddings.EmbeddingCreateParams; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.InstrumenterConfig; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import java.util.Collections; +import net.bytebuddy.asm.Advice; +import net.bytebuddy.description.type.TypeDescription; +import net.bytebuddy.matcher.ElementMatcher; + +@AutoService(InstrumenterModule.class) +public class EmbeddingServiceInstrumentation extends InstrumenterModule.Tracing + implements Instrumenter.ForTypeHierarchy, Instrumenter.HasMethodAdvice { + + public EmbeddingServiceInstrumentation() { + super("openai-java", "openai-java-2.8", "openai-client"); + } + + @Override + protected boolean defaultEnabled() { + return super.defaultEnabled() + && InstrumenterConfig.get().isIntegrationEnabled(Collections.singleton("openai"), false); + } + + @Override + public String[] helperClassNames() { + return new String[] { + packageName + ".OpenAIDecorator", + }; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + // Instrument embedding creation methods + transformer.applyAdvice( + isMethod() + .and(named("create")) + .and(isPublic()) + .and(takesArgument(0, named("com.openai.models.embeddings.EmbeddingCreateParams"))), + EmbeddingServiceInstrumentation.class.getName() + "$EmbeddingAdvice"); + } + + @Override + public String hierarchyMarkerType() { + return "com.openai.services.embedding.EmbeddingService"; + } + + @Override + public ElementMatcher hierarchyMatcher() { + return implementsInterface(named(hierarchyMarkerType())); + } + + public static class EmbeddingAdvice { + @Advice.OnMethodEnter(suppress = Throwable.class) + public static AgentScope onEnter( + @Advice.Argument(0) final EmbeddingCreateParams embeddingParams) { + + return OpenAIClientDecorator.DECORATE.startEmbeddingSpan(embeddingParams); + } + + @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class) + public static void onExit( + @Advice.Enter final AgentScope span, + @Advice.Return final Object result, + @Advice.Thrown final Throwable throwable) { + + OpenAIClientDecorator.DECORATE.finishSpan(span, result, throwable); + } + } +} diff --git a/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/OpenAIClientDecorator.java b/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/OpenAIClientDecorator.java new file mode 100644 index 00000000000..a3e6762bccd --- /dev/null +++ b/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/OpenAIClientDecorator.java @@ -0,0 +1,402 @@ +package datadog.trace.instrumentation.openaiclient; + +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan; +import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan; + +import com.openai.models.chat.completions.*; +import com.openai.models.completions.Completion; +import com.openai.models.completions.CompletionChoice; +import com.openai.models.completions.CompletionCreateParams; +import com.openai.models.embeddings.CreateEmbeddingResponse; +import com.openai.models.embeddings.Embedding; +import com.openai.models.embeddings.EmbeddingCreateParams; +import datadog.trace.api.Config; +import datadog.trace.api.llmobs.LLMObs; +import datadog.trace.api.llmobs.LLMObsSpan; +import datadog.trace.bootstrap.instrumentation.api.AgentScope; +import datadog.trace.bootstrap.instrumentation.api.AgentSpan; +import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes; +import datadog.trace.bootstrap.instrumentation.api.Tags; +import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString; +import datadog.trace.bootstrap.instrumentation.decorator.ClientDecorator; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +public class OpenAIClientDecorator extends ClientDecorator { + private static final String mlApp = Config.get().getLlmObsMlApp(); + private static final String mlProvider = "openai"; + private static final String COMPONENT_NAME = "openai"; + private static final UTF8BytesString OPENAI_REQUEST = UTF8BytesString.create("openai.request"); + + public static final OpenAIClientDecorator DECORATE = new OpenAIClientDecorator(); + + @Override + protected String[] instrumentationNames() { + return new String[] {"openai"}; + } + + @Override + protected CharSequence component() { + return UTF8BytesString.create(COMPONENT_NAME); + } + + @Override + protected CharSequence spanType() { + return InternalSpanTypes.HTTP_CLIENT; + } + + @Override + protected String service() { + return null; + } + + public AgentScope startChatCompletionSpan(ChatCompletionCreateParams params) { + AgentSpan span = startSpan(OPENAI_REQUEST); + span.setTag("openai.request.endpoint", "/chat/completions"); + span.setResourceName("chat.completions.create"); + span.setTag("openai.request.provider", "openai"); + extractChatCompletionRequestData(span, params); + afterStart(span); + return activateSpan(span); + } + + public AgentScope startLLMChatCompletionSpan(ChatCompletionCreateParams params) { + AgentSpan span = startSpan(OPENAI_REQUEST); + span.setTag("openai.request.endpoint", "/chat/completions"); + span.setResourceName("chat.completions.create"); + span.setTag("openai.request.provider", "openai"); + extractChatCompletionRequestData(span, params); + afterStart(span); + return activateSpan(span); + } + + public AgentScope startCompletionSpan(CompletionCreateParams params) { + AgentSpan span = startSpan(OPENAI_REQUEST); + span.setTag(Tags.COMPONENT, COMPONENT_NAME); + span.setTag("openai.request.endpoint", "/completions"); + span.setResourceName("completions.create"); + span.setTag("ai.provider", "openai"); + + extractCompletionRequestData(span, params); + afterStart(span); + return activateSpan(span); + } + + public AgentScope startEmbeddingSpan(EmbeddingCreateParams params) { + AgentSpan span = startSpan(OPENAI_REQUEST); + span.setTag("openai.request.endpoint", "/embeddings"); + span.setResourceName("embeddings.create"); + span.setTag("ai.provider", "openai"); + + extractEmbeddingRequestData(span, params); + afterStart(span); + return activateSpan(span); + } + + public LLMObsSpan startLLMObsChatCompletionSpan(ChatCompletionCreateParams params) { + String modelName = params.model().toString(); + + LLMObsSpan llmObsSpan = + LLMObs.startLLMSpan("chat_completion", modelName, mlProvider, mlApp, null); + + // Extract and set input data from chat completion params + // May need to be reformatted + StringBuilder inputData = new StringBuilder(); + List messages = params.messages(); + for (int i = 0; i < messages.size(); i++) { + ChatCompletionMessageParam messageParam = messages.get(i); + if (i > 0) { + inputData.append(" | "); + } + + if (messageParam.isUser()) { + inputData.append("User: ").append(messageParam.asUser().content()); + } else if (messageParam.isAssistant()) { + inputData.append("Assistant: ").append(messageParam.asAssistant().content()); + } else if (messageParam.isDeveloper()) { + inputData.append("Developer: ").append(messageParam.asDeveloper().content()); + } else if (messageParam.isSystem()) { + inputData.append("System: ").append(messageParam.asSystem().content()); + } else if (messageParam.isTool()) { + inputData.append("Tool: ").append(messageParam.asTool().content()); + } + } + + if (inputData.length() > 0) { + llmObsSpan.annotateIO(inputData.toString(), null); // No output yet, will be set in response + } + + Map metadata = new HashMap<>(); + metadata.put("endpoint", "/chat/completions"); + metadata.put("provider", "openai"); + metadata.put("model", modelName); + + params.maxTokens().ifPresent(tokens -> metadata.put("max_tokens", tokens)); + params.temperature().ifPresent(temp -> metadata.put("temperature", temp)); + + llmObsSpan.setMetadata(metadata); + + return llmObsSpan; + } + + public void finishLLMObsChatCompletionSpan( + LLMObsSpan llmObsSpan, ChatCompletion response, Throwable throwable) { + try { + if (throwable != null) { + // Set error information + Map errorMetadata = new HashMap<>(); + errorMetadata.put("error.type", throwable.getClass().getSimpleName()); + errorMetadata.put("error.message", throwable.getMessage()); + llmObsSpan.setMetadata(errorMetadata); + } else if (response != null) { + StringBuilder outputData = new StringBuilder(); + List choices = response.choices(); + + for (int i = 0; i < choices.size(); i++) { + ChatCompletion.Choice choice = choices.get(i); + ChatCompletionMessage message = choice.message(); + + if (i > 0) { + outputData.append(" | "); + } + + // Extract content + Optional content = message.content(); + content.ifPresent(s -> outputData.append("Assistant: ").append(s)); + + // Extract tool calls if present + Optional> toolCalls = message.toolCalls(); + if (toolCalls.isPresent() && !toolCalls.get().isEmpty()) { + content.ifPresent(s -> outputData.append(" | ")); + outputData.append("Tool calls: "); + for (int j = 0; j < toolCalls.get().size(); j++) { + ChatCompletionMessageToolCall call = toolCalls.get().get(j); + if (j > 0) { + outputData.append(", "); + } + outputData + .append(call.function().name()) + .append("(") + .append(call.function().arguments()) + .append(")"); + } + } + } + + if (outputData.length() > 0) { + llmObsSpan.annotateIO(null, outputData.toString()); + } + Map responseMetadata = new HashMap<>(); + responseMetadata.put("response.choices_count", choices.size()); + + llmObsSpan.setMetadata(responseMetadata); + } + } catch (Exception e) { + Map errorMetadata = new HashMap<>(); + errorMetadata.put("error.type", "ResponseProcessingError"); + errorMetadata.put("error.message", "Failed to process response: " + e.getMessage()); + llmObsSpan.setMetadata(errorMetadata); + } finally { + // Always finish the span + llmObsSpan.finish(); + } + } + + public void finishSpan(AgentScope scope, Object result, Throwable throwable) { + + AgentSpan span = scope.span(); + + try { + if (throwable != null) { + onError(span, throwable); + } else if (result != null) { + extractResponseData(span, result); + } + beforeFinish(span); + } finally { + scope.close(); + span.finish(); + } + } + + private void extractChatCompletionRequestData(AgentSpan span, ChatCompletionCreateParams params) { + + span.setTag("openai.model.name", params.model().toString()); + + // Extract messages + List messages = params.messages(); + for (int i = 0; i < messages.size(); i++) { + ChatCompletionMessageParam messageParam = messages.get(i); + extractMessageData(span, messageParam, i); + } + + // Extract request parameters + extractChatCompletionParameters(span, params); + } + + private void extractCompletionRequestData( + AgentSpan span, CompletionCreateParams completionParams) { + span.setTag("openai.model.name", completionParams.model().toString()); + + // Extract prompt + Optional prompt = completionParams.prompt(); + prompt.ifPresent( + promptValue -> { + if (promptValue.isArrayOfStrings()) { + List promptArray = promptValue.asArrayOfStrings(); + int promptIndex = 0; + for (String currentPrompt : promptArray) { + span.setTag("openai.request.prompt." + promptIndex, currentPrompt); + promptIndex++; + } + } else if (promptValue.isString()) { + // Setting the index as 0 since there is only one prompt string + span.setTag("openai.request.prompt.0", promptValue.asString()); + + } else if (promptValue.isArrayOfTokenArrays() || promptValue.isArrayOfTokens()) { + // Setting the token array as the value of the prompt tag + span.setTag("openai.request.prompt.0", promptValue.asArrayOfTokens()); + } + }); + // Extract request parameters + extractCompletionParameters(span, completionParams); + } + + private void extractEmbeddingRequestData(AgentSpan span, EmbeddingCreateParams embeddingParams) { + + span.setTag("openai.model.name", embeddingParams.model().toString()); + // Extract input + EmbeddingCreateParams.Input input = embeddingParams.input(); + int inputIndex = 0; + List inputStrings = input.asArrayOfStrings(); + for (String inputItem : inputStrings) { + span.setTag("openai.request.input." + inputIndex, inputItem); + inputIndex++; + } + } + + private void extractMessageData( + AgentSpan span, ChatCompletionMessageParam messageParam, int index) { + // Handle different message parameter types + if (messageParam.isUser()) { + span.setTag("openai.request.messages." + index + ".role", "user"); + span.setTag( + "openai.request.messages." + index + ".content", + messageParam.asUser().content().toString()); + } else if (messageParam.isAssistant()) { + span.setTag("openai.request.messages." + index + ".role", "assistant"); + span.setTag( + "openai.request.messages." + index + ".content", + messageParam.asAssistant().content().toString()); + } else if (messageParam.isDeveloper()) { + span.setTag("openai.request.messages." + index + ".role", "developer"); + span.setTag( + "openai.request.messages." + index + ".content", + messageParam.asDeveloper().content().toString()); + } else if (messageParam.isSystem()) { + span.setTag("openai.request.messages." + index + ".role", "system"); + span.setTag( + "openai.request.messages." + index + ".content", + messageParam.asSystem().content().toString()); + } else if (messageParam.isTool()) { + span.setTag("openai.request.messages." + index + ".role", "tool"); + span.setTag( + "openai.request.messages." + index + ".content", + messageParam.asTool().content().toString()); + } + } + + private void extractChatCompletionParameters(AgentSpan span, ChatCompletionCreateParams params) { + // Extract max_tokens + params + .maxCompletionTokens() + .ifPresent(tokens -> span.setTag("openai.request.max_tokens", tokens)); + + // Extract temperature + params.temperature().ifPresent(temp -> span.setTag("openai.request.temperature", temp)); + } + + private void extractCompletionParameters(AgentSpan span, CompletionCreateParams params) { + // Extract max_tokens + params.maxTokens().ifPresent(tokens -> span.setTag("openai.request.max_tokens", tokens)); + // Extract temperature + params.temperature().ifPresent(temp -> span.setTag("openai.request.temperature", temp)); + } + + private void extractResponseData(AgentSpan span, Object result) { + if (result instanceof ChatCompletion) { + extractChatCompletionResponseData(span, (ChatCompletion) result); + } else if (result instanceof Completion) { + extractCompletionResponseData(span, (Completion) result); + } else if (result instanceof CreateEmbeddingResponse) { + extractEmbeddingResponseData(span, (CreateEmbeddingResponse) result); + } + } + + private void extractChatCompletionResponseData(AgentSpan span, ChatCompletion response) { + // Extract choices + List choices = response.choices(); + int choiceIndex = 0; + if (!choices.isEmpty()) { + for (ChatCompletion.Choice curChoice : choices) { + ChatCompletionMessage curMessage = curChoice.message(); + // Extract content + Optional content = curMessage.content(); + content.ifPresent( + s -> span.setTag("openai.response.choices." + choiceIndex + ".message.content", s)); + span.setTag( + "openai.response.choices." + choiceIndex + ".message.role", + curMessage._role().toString()); + Optional> toolCalls = curMessage.toolCalls(); + if (toolCalls.isPresent() && !toolCalls.get().isEmpty()) { + // Extract tool calls if present + int callIndex = 0; + for (ChatCompletionMessageToolCall call : toolCalls.get()) { + span.setTag( + "openai.response.choices." + + choiceIndex + + ".message.tool_calls." + + callIndex + + ".name", + call.id()); + span.setTag( + "openai.response.choices." + + choiceIndex + + ".message.tool_calls." + + callIndex + + ".arguments", + call.function().arguments()); + + callIndex++; + } + } + } + } + } + + private void extractCompletionResponseData(AgentSpan span, Completion response) { + // Extract choices + List choices = response.choices(); + for (CompletionChoice choice : choices) { + span.setTag("openai.response.choices." + choice.index() + ".text", choice.text()); + } + } + + private void extractEmbeddingResponseData(AgentSpan span, CreateEmbeddingResponse response) { + // Extract data (embeddings) + List data = response.data(); + span.setTag("openai.response.embeddings_count", data.size()); + int embeddingIndex = 0; + if (!data.isEmpty()) { + for (Embedding curEmbedding : data) { + // Extract embedding array + List embedding = curEmbedding.embedding(); + span.setTag( + "openai.response.embedding." + embeddingIndex + ".embedding_length", embedding.size()); + embeddingIndex++; + } + } + } +} diff --git a/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/OpenAIClientInfo.java b/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/OpenAIClientInfo.java new file mode 100644 index 00000000000..6a2854d8aaf --- /dev/null +++ b/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/OpenAIClientInfo.java @@ -0,0 +1,83 @@ +package datadog.trace.instrumentation.openaiclient; + +import com.openai.azure.credential.AzureApiKeyCredential; +import com.openai.core.ClientOptions; +import com.openai.credential.BearerTokenCredential; + +public class OpenAIClientInfo { + private String baseURL; + private String organizationID; + private String projectID; + private String apiKey; + + // This information will be used later to populate future LLMObs spans + public static OpenAIClientInfo fromClientOptions(ClientOptions options) { + if (options == null) { + return null; + } + + OpenAIClientInfo info = new OpenAIClientInfo(); + + try { + info.setBaseURL(options.baseUrl()); + + if (options.organization().isPresent()) { + info.setOrganizationID(options.organization().get()); + } + + if (options.project().isPresent()) { + info.setProjectID(options.project().get()); + } + + if (options.credential() instanceof BearerTokenCredential) { + info.setApiKey(((BearerTokenCredential) options.credential()).token()); + } else if (options.credential() instanceof AzureApiKeyCredential) { + info.setApiKey(((AzureApiKeyCredential) options.credential()).apiKey()); + } else { + info.setApiKey(null); + } + + } catch (Exception e) { + return info; + } + + return info; + } + + private OpenAIClientInfo() { + // Used by fromClientOptions method + } + + public String getBaseURL() { + return baseURL; + } + + public String getOrganizationID() { + return organizationID; + } + + public String getProjectID() { + return projectID; + } + + public String getApiKey() { + return apiKey; + } + + // Private setter methods (write properties) + private void setBaseURL(String baseURL) { + this.baseURL = baseURL; + } + + private void setOrganizationID(String organizationID) { + this.organizationID = organizationID; + } + + private void setProjectID(String projectID) { + this.projectID = projectID; + } + + private void setApiKey(String apiKey) { + this.apiKey = apiKey; + } +} diff --git a/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/OpenAIClientInstrumentation.java b/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/OpenAIClientInstrumentation.java new file mode 100644 index 00000000000..9163373d0b5 --- /dev/null +++ b/dd-java-agent/instrumentation/openai-java-2.8/src/main/java/datadog/trace/instrumentation/openaiclient/OpenAIClientInstrumentation.java @@ -0,0 +1,83 @@ +package datadog.trace.instrumentation.openaiclient; + +import static net.bytebuddy.matcher.ElementMatchers.isConstructor; +import static net.bytebuddy.matcher.ElementMatchers.takesArgument; + +import com.openai.client.OpenAIClientImpl; +import com.openai.core.ClientOptions; +import com.openai.services.blocking.CompletionService; +import com.openai.services.blocking.EmbeddingService; +import com.openai.services.blocking.chat.ChatCompletionService; +import datadog.trace.agent.tooling.Instrumenter; +import datadog.trace.agent.tooling.InstrumenterModule; +import datadog.trace.api.InstrumenterConfig; +import datadog.trace.bootstrap.InstrumentationContext; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import net.bytebuddy.asm.Advice; + +public class OpenAIClientInstrumentation extends InstrumenterModule.Tracing + implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice { + public OpenAIClientInstrumentation() { + super("openai-java", "openai-java-2.8", "openai-client"); + } + + @Override + protected boolean defaultEnabled() { + return super.defaultEnabled() + && InstrumenterConfig.get().isIntegrationEnabled(Collections.singleton("openai"), false); + } + + @Override + public String[] helperClassNames() { + return new String[] { + packageName + ".OpenAIClientInfo", + }; + } + + @Override + public String instrumentedType() { + return "com.openai.client.OpenAIClientImpl"; + } + + @Override + public Map contextStore() { + Map contextStores = new HashMap<>(3); + contextStores.put( + "com.openai.services.blocking.chat.ChatCompletionService", + OpenAIClientInfo.class.getName()); + contextStores.put( + "com.openai.services.blocking.CompletionService", OpenAIClientInfo.class.getName()); + contextStores.put( + "com.openai.services.blocking.EmbeddingService", OpenAIClientInfo.class.getName()); + return contextStores; + } + + @Override + public void methodAdvice(MethodTransformer transformer) { + transformer.applyAdvice( + isConstructor().and(takesArgument(0, ClientOptions.class)), + getClass().getName() + "$OpenAIClientAdvice"); + } + + public static class OpenAIClientAdvice { + + @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class) + public static void methodExit( + @Advice.This final OpenAIClientImpl client, + @Advice.Argument(0) final ClientOptions options, + @Advice.Thrown final Throwable throwable) { + // This information will be used later to populate future LLMObs spans + OpenAIClientInfo info = OpenAIClientInfo.fromClientOptions(options); + if (info != null) { + InstrumentationContext.get(CompletionService.class, OpenAIClientInfo.class) + .put(client.completions(), info); + InstrumentationContext.get(EmbeddingService.class, OpenAIClientInfo.class) + .put(client.embeddings(), info); + InstrumentationContext.get(ChatCompletionService.class, OpenAIClientInfo.class) + .put(client.chat().completions(), info); + } + } + } +} diff --git a/settings.gradle.kts b/settings.gradle.kts index 70bbfb5679f..be00818e1a7 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -459,6 +459,7 @@ include( ":dd-java-agent:instrumentation:okhttp-2", ":dd-java-agent:instrumentation:okhttp-3", ":dd-java-agent:instrumentation:ognl-appsec", + ":dd-java-agent:instrumentation:openai-java-2.8", ":dd-java-agent:instrumentation:opensearch", ":dd-java-agent:instrumentation:opensearch:rest", ":dd-java-agent:instrumentation:opensearch:transport",