Skip to content

Commit 23c1a09

Browse files
ygreeygree
committed
TestOpenAiLlmObs::test_responses_create
1 parent f5a9bea commit 23c1a09

File tree

6 files changed

+266
-41
lines changed

6 files changed

+266
-41
lines changed

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/OpenAiDecorator.java‎

Lines changed: 106 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import com.openai.models.responses.ResponseOutputMessage;
2626
import com.openai.models.responses.ResponseOutputText;
2727
import com.openai.models.responses.ResponseStreamEvent;
28-
import com.openai.models.responses.ResponseUsage;
2928
import datadog.trace.api.DDSpanId;
3029
import datadog.trace.api.llmobs.LLMObs;
3130
import datadog.trace.api.llmobs.LLMObsContext;
@@ -353,8 +352,7 @@ public void withCreateEmbeddingResponse(AgentSpan span, CreateEmbeddingResponse
353352
}
354353
}
355354

356-
public void withResponseCreateParams(
357-
AgentSpan span, ResponseCreateParams params, boolean stream) {
355+
public void withResponseCreateParams(AgentSpan span, ResponseCreateParams params) {
358356
span.setTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND);
359357
span.setResourceName(RESPONSES_CREATE);
360358
span.setTag("openai.request.endpoint", "v1/responses");
@@ -365,24 +363,54 @@ public void withResponseCreateParams(
365363
// Use ResponseCreateParams._model() b/o ResponseCreateParams.model() changed type from
366364
// ResponsesModel to Optional<ResponsesModel> in
367365
// https://github.com/openai/openai-java/commit/87dd64658da6cec7564f3b571e15ec0e2db0660b
368-
span.setTag(REQUEST_MODEL, extractResponseModel(params._model()));
366+
String modelName = extractResponseModel(params._model());
367+
span.setTag(REQUEST_MODEL, modelName);
369368

369+
// Set model_name and model_provider as fallback (will be overridden by withResponse if called)
370+
// span.setTag("_ml_obs_tag.model_name", modelName);
371+
// span.setTag("_ml_obs_tag.model_provider", "openai");
372+
373+
List<LLMObs.LLMMessage> inputMessages = new ArrayList<>();
374+
375+
// Add instructions as system message first (if present)
376+
params
377+
.instructions()
378+
.ifPresent(
379+
instructions -> {
380+
inputMessages.add(LLMObs.LLMMessage.from("system", instructions));
381+
});
382+
383+
// Add user input message
370384
Optional<String> textOpt = params._input().asString();
371385
if (textOpt.isPresent()) {
372-
LLMObs.LLMMessage msg = LLMObs.LLMMessage.from("user", textOpt.get());
373-
span.setTag("_ml_obs_tag.input", Collections.singletonList(msg));
386+
inputMessages.add(LLMObs.LLMMessage.from("user", textOpt.get()));
374387
}
375388

376-
Map<String, Object> metadata = new HashMap<>();
377-
params.maxOutputTokens().ifPresent(v -> metadata.put("max_tokens", v));
378-
params.temperature().ifPresent(v -> metadata.put("temperature", v));
379-
if (stream) {
380-
metadata.put("stream", true);
389+
if (!inputMessages.isEmpty()) {
390+
span.setTag("_ml_obs_tag.input", inputMessages);
381391
}
382-
span.setTag("_ml_obs_tag.metadata", metadata);
383392
}
384393

385394
public void withResponse(AgentSpan span, Response response) {
395+
withResponse(span, response, false);
396+
}
397+
398+
public void withResponseStreamEvents(AgentSpan span, List<ResponseStreamEvent> events) {
399+
for (ResponseStreamEvent event : events) {
400+
if (event.isCompleted()) {
401+
Response response = event.asCompleted().response();
402+
withResponse(span, response, true);
403+
return;
404+
}
405+
if (event.isIncomplete()) {
406+
Response response = event.asIncomplete().response();
407+
withResponse(span, response, true);
408+
return;
409+
}
410+
}
411+
}
412+
413+
private void withResponse(AgentSpan span, Response response, boolean stream) {
386414
String modelName = extractResponseModel(response._model());
387415
span.setTag(RESPONSE_MODEL, modelName);
388416
span.setTag("_ml_obs_tag.model_name", modelName);
@@ -393,7 +421,72 @@ public void withResponse(AgentSpan span, Response response) {
393421
span.setTag("_ml_obs_tag.output", outputMessages);
394422
}
395423

396-
response.usage().ifPresent(usage -> withResponseUsage(span, usage));
424+
Map<String, Object> metadata = new HashMap<>();
425+
426+
response.maxOutputTokens().ifPresent(v -> metadata.put("max_output_tokens", v));
427+
response.temperature().ifPresent(v -> metadata.put("temperature", v));
428+
response.topP().ifPresent(v -> metadata.put("top_p", v));
429+
430+
// Extract tool_choice as string
431+
Response.ToolChoice toolChoice = response.toolChoice();
432+
if (toolChoice.isOptions()) {
433+
metadata.put("tool_choice", toolChoice.asOptions()._value().asString().orElse(null));
434+
} else if (toolChoice.isTypes()) {
435+
metadata.put("tool_choice", toolChoice.asTypes().type().toString().toLowerCase());
436+
} else if (toolChoice.isFunction()) {
437+
metadata.put("tool_choice", "function");
438+
}
439+
440+
// Extract truncation as string
441+
response
442+
.truncation()
443+
.ifPresent(
444+
(Response.Truncation t) ->
445+
metadata.put("truncation", t._value().asString().orElse(null)));
446+
447+
// Extract text format
448+
response
449+
.text()
450+
.ifPresent(
451+
textConfig -> {
452+
textConfig
453+
.format()
454+
.ifPresent(
455+
format -> {
456+
Map<String, Object> textMap = new HashMap<>();
457+
Map<String, String> formatMap = new HashMap<>();
458+
if (format.isText()) {
459+
formatMap.put("type", "text");
460+
// metadata.put("text.format.type", "text");
461+
// } else if (format.isJsonSchema()) {
462+
// formatMap.put("type", "json_schema");
463+
// } else if (format.isJsonObject()) {
464+
// formatMap.put("type", "json_object");
465+
}
466+
textMap.put("format", formatMap);
467+
metadata.put("text", textMap);
468+
});
469+
});
470+
471+
if (stream) {
472+
metadata.put("stream", true);
473+
}
474+
475+
response
476+
.usage()
477+
.ifPresent(
478+
usage -> {
479+
span.setTag("_ml_obs_metric.input_tokens", usage.inputTokens());
480+
span.setTag("_ml_obs_metric.output_tokens", usage.outputTokens());
481+
span.setTag("_ml_obs_metric.total_tokens", usage.totalTokens());
482+
span.setTag(
483+
"_ml_obs_metric.cache_read_input_tokens",
484+
usage.inputTokensDetails().cachedTokens());
485+
long reasoningTokens = usage.outputTokensDetails().reasoningTokens();
486+
metadata.put("reasoning_tokens", reasoningTokens);
487+
});
488+
489+
span.setTag("_ml_obs_tag.metadata", metadata);
397490
}
398491

399492
private List<LLMObs.LLMMessage> extractResponseOutputMessages(List<ResponseOutputItem> output) {
@@ -431,14 +524,6 @@ private String extractMessageContent(ResponseOutputMessage message) {
431524
return result.isEmpty() ? null : result;
432525
}
433526

434-
private static void withResponseUsage(AgentSpan span, ResponseUsage usage) {
435-
span.setTag("_ml_obs_metric.input_tokens", usage.inputTokens());
436-
span.setTag("_ml_obs_metric.output_tokens", usage.outputTokens());
437-
span.setTag("_ml_obs_metric.total_tokens", usage.totalTokens());
438-
span.setTag(
439-
"_ml_obs_metric.cache_read_input_tokens", usage.inputTokensDetails().cachedTokens());
440-
}
441-
442527
private String extractResponseModel(JsonField<ResponsesModel> model) {
443528
Optional<String> str = model.asString();
444529
if (str.isPresent()) {
@@ -465,15 +550,4 @@ private String extractResponseModel(JsonField<ResponsesModel> model) {
465550
}
466551
return null;
467552
}
468-
469-
public void withResponseStreamEvent(AgentSpan span, List<ResponseStreamEvent> events) {
470-
// Find the completed event which contains the full response
471-
for (ResponseStreamEvent event : events) {
472-
if (event.isCompleted()) {
473-
Response response = event.asCompleted().response();
474-
withResponse(span, response);
475-
return;
476-
}
477-
}
478-
}
479553
}

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/ResponseServiceAsyncInstrumentation.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ public static AgentScope enter(
5252
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
5353
DECORATE.afterStart(span);
5454
DECORATE.withClientOptions(span, clientOptions);
55-
DECORATE.withResponseCreateParams(span, params, false);
55+
DECORATE.withResponseCreateParams(span, params);
5656
return activateSpan(span);
5757
}
5858

@@ -86,7 +86,7 @@ public static AgentScope enter(
8686
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
8787
DECORATE.afterStart(span);
8888
DECORATE.withClientOptions(span, clientOptions);
89-
DECORATE.withResponseCreateParams(span, params, true);
89+
DECORATE.withResponseCreateParams(span, params);
9090
return activateSpan(span);
9191
}
9292

@@ -104,7 +104,7 @@ public static void exit(
104104
if (future != null) {
105105
future =
106106
ResponseWrappers.wrapFutureStreamResponse(
107-
future, span, DECORATE::withResponseStreamEvent);
107+
future, span, DECORATE::withResponseStreamEvents);
108108
} else {
109109
span.finish();
110110
}

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/main/java/datadog/trace/instrumentation/openai_java/ResponseServiceInstrumentation.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ public static AgentScope enter(
5454
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
5555
DECORATE.afterStart(span);
5656
DECORATE.withClientOptions(span, clientOptions);
57-
DECORATE.withResponseCreateParams(span, params, false);
57+
DECORATE.withResponseCreateParams(span, params);
5858
return activateSpan(span);
5959
}
6060

@@ -89,7 +89,7 @@ public static AgentScope enter(
8989
AgentSpan span = startSpan(OpenAiDecorator.INSTRUMENTATION_NAME, OpenAiDecorator.SPAN_NAME);
9090
DECORATE.afterStart(span);
9191
DECORATE.withClientOptions(span, clientOptions);
92-
DECORATE.withResponseCreateParams(span, params, true);
92+
DECORATE.withResponseCreateParams(span, params);
9393
return activateSpan(span);
9494
}
9595

@@ -107,7 +107,7 @@ public static void exit(
107107
if (response != null) {
108108
response =
109109
ResponseWrappers.wrapStreamResponse(
110-
response, span, DECORATE::withResponseStreamEvent);
110+
response, span, DECORATE::withResponseStreamEvents);
111111
} else {
112112
span.finish();
113113
}

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/OpenAiTest.groovy‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,14 @@ abstract class OpenAiTest extends LlmObsSpecification {
138138
.build()
139139
}
140140

141+
ResponseCreateParams responseCreateParamsWithMaxOutputTokens() {
142+
ResponseCreateParams.builder()
143+
.model("gpt-3.5-turbo")
144+
.input("Do not continue the Evan Li slander!")
145+
.maxOutputTokens(30)
146+
.build()
147+
}
148+
141149
ChatCompletionCreateParams chatCompletionCreateParamsWithTools() {
142150
ChatCompletionCreateParams.builder()
143151
.model(ChatModel.GPT_4O_MINI)

‎dd-java-agent/instrumentation/openai-java/openai-java-3.0/src/test/groovy/ResponseServiceTest.groovy‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,9 +37,9 @@ class ResponseServiceTest extends OpenAiTest {
3737
assertResponseTrace(false)
3838
}
3939

40-
def "create streaming response test"() {
40+
def "create streaming response test (#scenario)"() {
4141
runnableUnderTrace("parent") {
42-
StreamResponse<ResponseStreamEvent> streamResponse = openAiClient.responses().createStreaming(responseCreateParams())
42+
StreamResponse<ResponseStreamEvent> streamResponse = openAiClient.responses().createStreaming(params)
4343
try (Stream stream = streamResponse.stream()) {
4444
stream.forEach {
4545
// consume the stream
@@ -49,6 +49,11 @@ class ResponseServiceTest extends OpenAiTest {
4949
5050
expect:
5151
assertResponseTrace(true)
52+
53+
where:
54+
scenario | params
55+
"complete" | responseCreateParams()
56+
"incomplete" | responseCreateParamsWithMaxOutputTokens()
5257
}
5358
5459
def "create streaming response test withRawResponse"() {
@@ -146,7 +151,6 @@ class ResponseServiceTest extends OpenAiTest {
146151
"openai.api_base" openAiBaseApi
147152
"$OpenAiDecorator.RESPONSE_MODEL" "gpt-3.5-turbo-0125"
148153
if (!isStreaming) {
149-
// TODO no limit headers when streaming
150154
"openai.organization.ratelimit.requests.limit" 10000
151155
"openai.organization.ratelimit.requests.remaining" Integer
152156
"openai.organization.ratelimit.tokens.limit" 50000000

0 commit comments

Comments
 (0)