-
Notifications
You must be signed in to change notification settings - Fork 324
Add openai-java v3.0+ instrumentation #9959
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
ygree
wants to merge
106
commits into
master
Choose a base branch
from
ygree/openai-java
base: master
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+5,640
−62
Open
Changes from all commits
Commits
Show all changes
106 commits
Select commit
Hold shift + click to select a range
e7bd8d9
Initial APM-only openai-java instrumentation with a unit test.
ygree 6feb61e
Start llmobs system in tests and create llmobs span.
ygree d541adb
streamed request completion test
ygree c209bd0
README
ygree a82856d
Instrument sync streamed completion
ygree 3d3def9
Mock Streamed Completion in Tests
ygree 3542b9e
Add failing test for Async Completion
ygree c12cd0a
Async Single Completion. Simplify tests. Remove println. Cleanup.
ygree 2af4e5f
Extract DDHttpResponseFor to intercept when the response is parsed in…
ygree cec33af
Instrument and test Async Stream Completion
ygree b2b4141
More tests "streamed request completion test with withRawResponse" an…
ygree b1b6db2
Fix muzzle check
ygree baa687d
Add openai APM tags, assert tags, todo tags
ygree b9ad963
Wrap HttpResponseFor instead of forcing parsing. Add TODOs
ygree 5bf16f4
Set response model tag
ygree be20fa1
Set "openai.organization.name"
ygree dd230b7
Set ratelimit metrics
ygree c3f94b0
api_base
ygree ddf0e09
"openai.request.method" & "openai.request.endpoint"
ygree fb9fd04
createChatCompletion instrumentation
ygree a4358a8
Reorder tests single then stream
ygree 99dedfe
Async Single Chat Completion
ygree d66a5ee
Async Streamed Chat Completion
ygree f8af3ae
Rename assertChatCompletionTrace
ygree 846e945
Instrument Embeddings
ygree b14bc4b
ResponseService WIP
ygree e8daf08
ResponseService synch
ygree c8a9e6e
ResponseServiceAsyncInstrumentation
ygree 9754cd8
Setup httpClient for tests WIP
ygree aeaae2a
Intercept Http req/resp with TestOpenAiHttpClient
ygree 7e50d1d
Implement req/resp recorder and mock backend for tests
ygree 8e85447
Add lockfile
ygree ea2f465
Minor changes to the records format
ygree 1f3aa6c
OpenAiHttpClient for tests. Record only if the record doesn't exist. …
ygree 0c28997
Rename TestOpenAiHttpClient
ygree 6a41720
Do not dump a record if already exists
ygree 6050515
Fix format
ygree c8a1e94
Fix linter errors
ygree 20c88d7
Fix unused imports
ygree 001ce0e
Fix format
ygree 96f4353
Fix format
ygree 46c2d3a
Fix format
ygree 29ae26c
Merge branch 'master' into ygree/openai-java
ygree 8de5108
Fix format
ygree fb74ed2
Remove unexisting helper class that failed the test
ygree 9a2bc61
Extract response model.
ygree 637f8e5
Extract response model.
ygree 6bf8a77
Fix llmObsSpanName in LLMObsSpanMapper
ygree a204bab
LLMObsState -> LLMObsContext (internal-api) to be shared with auto-in…
ygree 8a6a0d8
Experimental use of LLMObsContext in the openai-java completion instr…
ygree 277bd12
Fix bug in the mapper to write input/output fields as maps
ygree 8f2f83c
Add necessary tags to pass TestOpenAiLlmObs::test_completion
ygree 135e42a
Fix assertion when expect a class but it's null
ygree 60d2f06
Fix unit tests
ygree 5cf2669
Fix format
ygree 201e61d
Implement proper extractResponseModel
ygree 6d11609
Fix format
ygree 848d46a
Add note about instrumented code change
ygree 85dbbf2
Enable tests for 0.45.0
ygree 57212c6
:TestOpenAiLlmObs::test_chat_completion[java-test-ml-app-tcp-False] P…
ygree 81295b9
TestOpenAiLlmObs::test_chat_completion[java-test-ml-app-tcp-True] PASSED
ygree ec29979
Fix format
ygree c26909b
Fix ChatCompletionServiceTest
ygree 1850e79
Reorg/rename decorator functions
ygree d6110fc
TestOpenAiLlmObs::test_embedding[java-test-ml-app-tcp] PASSED
ygree f74d8b6
chat/completion tool call for openai-java <v3.0+
ygree f2d4428
Change the naming of HTTP records to keep the scanner quiet during a …
ygree 08b0bab
Support openai-java v3.0.0+ only because of breaking changes in tool …
ygree f718f22
Rename to openai-java-3.0
ygree a384464
Merge branch 'master' into ygree/openai-java
ygree ffe043a
fix format
ygree 6e654e5
Extract and assert chatCompletion toolCalls single and streamed
ygree c783318
Use ChatCompletionAccumulator to simplify streamed chat/response deco…
ygree f5a9bea
TestOpenAiLlmObs::test_responses_create_tool_call
ygree 23c1a09
TestOpenAiLlmObs::test_responses_create
ygree ac39251
reasoning
ygree c25d3e2
Remove the unused record file and add SET_RECORD_FILE_ATTR_ON_READ to…
ygree c29ee56
Test both JSON and typed parameters because they are accessed differe…
ygree 0f0b0a1
Split to modules b/o exceeded Muzzle.create limit
ygree 55423a0
extract startSpan
ygree 4282f65
rename Response wrappers to minimize confusion with ResponseService
ygree 972b2ac
test_responses_create_tool_input WIP
ygree b726b1f
Responses create tool input support
ygree 7476caa
withHttpResponse clean up
ygree 854d7b0
Fix unused import
ygree 2d127f2
Handle possible response parse errors
ygree 8718f94
Clean up HttpResponse wrappers
ygree a4d0421
Properly close scope, finish span, and handle errors in HttpResponse …
ygree 1ca2ae0
Refactor LLMObs getting current LLMObs parentSpanID
ygree f8d8ce5
Do not decorate LLMObs specific tags when it's disabled
ygree 1154c16
Finish span in HttpResponseWrapper when closing, if the response was …
ygree 6e5cf86
Clean up the unnecessary test subsystem for now because the unit test…
ygree b19cbb9
remove leftover
ygree b341414
Implement LLMObs "span.finished" metric
ygree 9dedc4e
Merge branch 'master' into ygree/openai-java
ygree ff22fa8
update lockfile
ygree 2313a23
Merge branch 'master' into ygree/openai-java
ygree 3a93165
Fix CompletionDecorator to properly accumulate output and collect usa…
ygree aea89e3
Code review follow-up fixes
ygree 35767b7
remove lockfile
ygree ca7ac93
multi-choice response tests
ygree 37a8fd1
clean up
ygree 84065c9
minor code review follow up fixes
ygree d9a6fcf
Add common tags
ygree cb7fc0d
Avoid finishing a span too early in HttpResponse.wrapFuture
ygree c12ac35
Merge branch 'master' into ygree/openai-java
ygree File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
37 changes: 0 additions & 37 deletions
37
dd-java-agent/agent-llmobs/src/main/java/datadog/trace/llmobs/domain/LLMObsState.java
This file was deleted.
Oops, something went wrong.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
25 changes: 25 additions & 0 deletions
25
dd-java-agent/instrumentation/openai-java/openai-java-3.0/build.gradle
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,25 @@ | ||
| apply from: "$rootDir/gradle/java.gradle" | ||
| apply plugin: 'idea' | ||
|
|
||
| def minVer = '3.0.0' | ||
|
|
||
| muzzle { | ||
| pass { | ||
| group = "com.openai" | ||
| module = "openai-java" | ||
| versions = "[$minVer,)" | ||
| } | ||
| } | ||
|
|
||
| addTestSuiteForDir('latestDepTest', 'test') | ||
|
|
||
| dependencies { | ||
| compileOnly group: 'com.openai', name: 'openai-java', version: minVer | ||
| implementation project(':internal-api') | ||
|
|
||
| testImplementation group: 'com.openai', name: 'openai-java', version: minVer | ||
| latestDepTestImplementation group: 'com.openai', name: 'openai-java', version: '+' | ||
|
|
||
| testImplementation project(':dd-java-agent:instrumentation:okhttp:okhttp-3.0') | ||
| } | ||
|
|
||
157 changes: 157 additions & 0 deletions
157
...-3.0/src/main/java/datadog/trace/instrumentation/openai_java/ChatCompletionDecorator.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,157 @@ | ||
| package datadog.trace.instrumentation.openai_java; | ||
|
|
||
| import static datadog.trace.instrumentation.openai_java.OpenAiDecorator.REQUEST_MODEL; | ||
| import static datadog.trace.instrumentation.openai_java.OpenAiDecorator.RESPONSE_MODEL; | ||
|
|
||
| import com.openai.helpers.ChatCompletionAccumulator; | ||
| import com.openai.models.chat.completions.ChatCompletion; | ||
| import com.openai.models.chat.completions.ChatCompletionChunk; | ||
| import com.openai.models.chat.completions.ChatCompletionCreateParams; | ||
| import com.openai.models.chat.completions.ChatCompletionMessage; | ||
| import com.openai.models.chat.completions.ChatCompletionMessageParam; | ||
| import com.openai.models.chat.completions.ChatCompletionMessageToolCall; | ||
| import datadog.trace.api.Config; | ||
| import datadog.trace.api.llmobs.LLMObs; | ||
| import datadog.trace.bootstrap.instrumentation.api.AgentSpan; | ||
| import datadog.trace.bootstrap.instrumentation.api.Tags; | ||
| import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString; | ||
| import java.util.ArrayList; | ||
| import java.util.Collections; | ||
| import java.util.HashMap; | ||
| import java.util.List; | ||
| import java.util.Map; | ||
| import java.util.Optional; | ||
| import java.util.stream.Collectors; | ||
|
|
||
| public class ChatCompletionDecorator { | ||
| public static final ChatCompletionDecorator DECORATE = new ChatCompletionDecorator(); | ||
| private static final CharSequence CHAT_COMPLETIONS_CREATE = | ||
| UTF8BytesString.create("createChatCompletion"); | ||
|
|
||
| private final boolean llmObsEnabled = Config.get().isLlmObsEnabled(); | ||
|
|
||
| public void withChatCompletionCreateParams( | ||
| AgentSpan span, ChatCompletionCreateParams params, boolean stream) { | ||
| span.setResourceName(CHAT_COMPLETIONS_CREATE); | ||
| span.setTag("openai.request.endpoint", "v1/chat/completions"); | ||
| span.setTag("openai.request.method", "POST"); | ||
| if (!llmObsEnabled) { | ||
| return; | ||
| } | ||
|
|
||
| span.setTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND); | ||
| if (params == null) { | ||
| return; | ||
| } | ||
| params.model()._value().asString().ifPresent(str -> span.setTag(REQUEST_MODEL, str)); | ||
|
|
||
| span.setTag( | ||
| "_ml_obs_tag.input", | ||
| params.messages().stream() | ||
| .map(ChatCompletionDecorator::llmMessage) | ||
| .collect(Collectors.toList())); | ||
|
|
||
| Map<String, Object> metadata = new HashMap<>(); | ||
| // maxTokens is deprecated but integration tests missing to provide maxCompletionTokens | ||
| params.maxTokens().ifPresent(v -> metadata.put("max_tokens", v)); | ||
| params.temperature().ifPresent(v -> metadata.put("temperature", v)); | ||
| if (stream) { | ||
| metadata.put("stream", true); | ||
| } | ||
| params | ||
| .streamOptions() | ||
| .ifPresent( | ||
| v -> { | ||
| if (v.includeUsage().orElse(false)) { | ||
| metadata.put("stream_options", Collections.singletonMap("include_usage", true)); | ||
| } | ||
| }); | ||
| span.setTag("_ml_obs_tag.metadata", metadata); | ||
| } | ||
|
|
||
| private static LLMObs.LLMMessage llmMessage(ChatCompletionMessageParam m) { | ||
| String role = "unknown"; | ||
| String content = null; | ||
| if (m.isAssistant()) { | ||
| role = "assistant"; | ||
| content = m.asAssistant().content().map(v -> v.text().orElse(null)).orElse(null); | ||
| } else if (m.isDeveloper()) { | ||
| role = "developer"; | ||
| content = m.asDeveloper().content().text().orElse(null); | ||
| } else if (m.isSystem()) { | ||
| role = "system"; | ||
| content = m.asSystem().content().text().orElse(null); | ||
| } else if (m.isTool()) { | ||
| role = "tool"; | ||
| content = m.asTool().content().text().orElse(null); | ||
| } else if (m.isUser()) { | ||
| role = "user"; | ||
| content = m.asUser().content().text().orElse(null); | ||
| } | ||
| return LLMObs.LLMMessage.from(role, content); | ||
| } | ||
|
|
||
| public void withChatCompletion(AgentSpan span, ChatCompletion completion) { | ||
| if (!llmObsEnabled) { | ||
| return; | ||
| } | ||
| String modelName = completion.model(); | ||
| span.setTag(RESPONSE_MODEL, modelName); | ||
| span.setTag("_ml_obs_tag.model_name", modelName); | ||
| span.setTag("_ml_obs_tag.model_provider", "openai"); | ||
|
|
||
| List<LLMObs.LLMMessage> output = | ||
| completion.choices().stream() | ||
| .map(ChatCompletionDecorator::llmMessage) | ||
| .collect(Collectors.toList()); | ||
| span.setTag("_ml_obs_tag.output", output); | ||
|
|
||
| completion | ||
| .usage() | ||
| .ifPresent( | ||
| usage -> { | ||
| span.setTag("_ml_obs_metric.input_tokens", usage.promptTokens()); | ||
| span.setTag("_ml_obs_metric.output_tokens", usage.completionTokens()); | ||
| span.setTag("_ml_obs_metric.total_tokens", usage.totalTokens()); | ||
| }); | ||
| } | ||
|
|
||
| private static LLMObs.LLMMessage llmMessage(ChatCompletion.Choice choice) { | ||
| ChatCompletionMessage msg = choice.message(); | ||
| Optional<?> roleOpt = msg._role().asString(); | ||
| String role = "unknown"; | ||
| if (roleOpt.isPresent()) { | ||
| role = String.valueOf(roleOpt.get()); | ||
| } | ||
| String content = msg.content().orElse(null); | ||
|
|
||
| Optional<List<ChatCompletionMessageToolCall>> toolCallsOpt = msg.toolCalls(); | ||
| if (toolCallsOpt.isPresent() && !toolCallsOpt.get().isEmpty()) { | ||
| List<LLMObs.ToolCall> toolCalls = new ArrayList<>(); | ||
| for (ChatCompletionMessageToolCall toolCall : toolCallsOpt.get()) { | ||
| LLMObs.ToolCall llmObsToolCall = ToolCallExtractor.getToolCall(toolCall); | ||
| if (llmObsToolCall != null) { | ||
| toolCalls.add(llmObsToolCall); | ||
| } | ||
| } | ||
|
|
||
| if (!toolCalls.isEmpty()) { | ||
| return LLMObs.LLMMessage.from(role, content, toolCalls); | ||
| } | ||
| } | ||
|
|
||
| return LLMObs.LLMMessage.from(role, content); | ||
| } | ||
|
|
||
| public void withChatCompletionChunks(AgentSpan span, List<ChatCompletionChunk> chunks) { | ||
| if (!llmObsEnabled) { | ||
| return; | ||
| } | ||
| ChatCompletionAccumulator accumulator = ChatCompletionAccumulator.create(); | ||
| for (ChatCompletionChunk chunk : chunks) { | ||
| accumulator.accumulate(chunk); | ||
| } | ||
| ChatCompletion chatCompletion = accumulator.chatCompletion(); | ||
| withChatCompletion(span, chatCompletion); | ||
| } | ||
| } |
34 changes: 34 additions & 0 deletions
34
...ava-3.0/src/main/java/datadog/trace/instrumentation/openai_java/ChatCompletionModule.java
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,34 @@ | ||
| package datadog.trace.instrumentation.openai_java; | ||
|
|
||
| import com.google.auto.service.AutoService; | ||
| import datadog.trace.agent.tooling.Instrumenter; | ||
| import datadog.trace.agent.tooling.InstrumenterModule; | ||
| import java.util.Arrays; | ||
| import java.util.List; | ||
|
|
||
| @AutoService(InstrumenterModule.class) | ||
| public class ChatCompletionModule extends InstrumenterModule.Tracing { | ||
| public ChatCompletionModule() { | ||
| super("openai-java"); | ||
| } | ||
|
|
||
| @Override | ||
| public String[] helperClassNames() { | ||
| return new String[] { | ||
| packageName + ".ChatCompletionDecorator", | ||
| packageName + ".OpenAiDecorator", | ||
| packageName + ".HttpResponseWrapper", | ||
| packageName + ".HttpStreamResponseWrapper", | ||
| packageName + ".HttpStreamResponseStreamWrapper", | ||
| packageName + ".ToolCallExtractor", | ||
| packageName + ".ToolCallExtractor$1" | ||
| }; | ||
| } | ||
|
|
||
| @Override | ||
| public List<Instrumenter> typeInstrumentations() { | ||
| return Arrays.asList( | ||
| new ChatCompletionServiceAsyncInstrumentation(), | ||
| new ChatCompletionServiceInstrumentation()); | ||
| } | ||
| } |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.