Skip to content

Commit f5a9bea

Browse files
ygreeygree
committed
TestOpenAiLlmObs::test_responses_create_tool_call
1 parent c783318 commit f5a9bea

File tree

5 files changed

+125
-22
lines changed

5 files changed

+125
-22
lines changed

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

Lines changed: 81 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,12 @@
2020
import com.openai.models.embeddings.EmbeddingCreateParams;
2121
import com.openai.models.responses.Response;
2222
import com.openai.models.responses.ResponseCreateParams;
23+
import com.openai.models.responses.ResponseFunctionToolCall;
24+
import com.openai.models.responses.ResponseOutputItem;
25+
import com.openai.models.responses.ResponseOutputMessage;
26+
import com.openai.models.responses.ResponseOutputText;
2327
import com.openai.models.responses.ResponseStreamEvent;
28+
import com.openai.models.responses.ResponseUsage;
2429
import datadog.trace.api.DDSpanId;
2530
import datadog.trace.api.llmobs.LLMObs;
2631
import datadog.trace.api.llmobs.LLMObsContext;
@@ -348,7 +353,8 @@ public void withCreateEmbeddingResponse(AgentSpan span, CreateEmbeddingResponse
348353
}
349354
}
350355

351-
public void withResponseCreateParams(AgentSpan span, ResponseCreateParams params) {
356+
public void withResponseCreateParams(
357+
AgentSpan span, ResponseCreateParams params, boolean stream) {
352358
span.setTag("_ml_obs_tag.span.kind", Tags.LLMOBS_LLM_SPAN_KIND);
353359
span.setResourceName(RESPONSES_CREATE);
354360
span.setTag("openai.request.endpoint", "v1/responses");
@@ -360,12 +366,77 @@ public void withResponseCreateParams(AgentSpan span, ResponseCreateParams params
360366
// ResponsesModel to Optional<ResponsesModel> in
361367
// https://github.com/openai/openai-java/commit/87dd64658da6cec7564f3b571e15ec0e2db0660b
362368
span.setTag(REQUEST_MODEL, extractResponseModel(params._model()));
369+
370+
Optional<String> textOpt = params._input().asString();
371+
if (textOpt.isPresent()) {
372+
LLMObs.LLMMessage msg = LLMObs.LLMMessage.from("user", textOpt.get());
373+
span.setTag("_ml_obs_tag.input", Collections.singletonList(msg));
374+
}
375+
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);
381+
}
382+
span.setTag("_ml_obs_tag.metadata", metadata);
363383
}
364384

365385
public void withResponse(AgentSpan span, Response response) {
366-
span.setTag(RESPONSE_MODEL, extractResponseModel(response._model()));
386+
String modelName = extractResponseModel(response._model());
387+
span.setTag(RESPONSE_MODEL, modelName);
388+
span.setTag("_ml_obs_tag.model_name", modelName);
389+
span.setTag("_ml_obs_tag.model_provider", "openai");
367390

368-
// TODO set LLMObs tags
391+
List<LLMObs.LLMMessage> outputMessages = extractResponseOutputMessages(response.output());
392+
if (!outputMessages.isEmpty()) {
393+
span.setTag("_ml_obs_tag.output", outputMessages);
394+
}
395+
396+
response.usage().ifPresent(usage -> withResponseUsage(span, usage));
397+
}
398+
399+
private List<LLMObs.LLMMessage> extractResponseOutputMessages(List<ResponseOutputItem> output) {
400+
List<LLMObs.LLMMessage> messages = new ArrayList<>();
401+
List<LLMObs.ToolCall> toolCalls = new ArrayList<>();
402+
String textContent = null;
403+
404+
for (ResponseOutputItem item : output) {
405+
if (item.isFunctionCall()) {
406+
ResponseFunctionToolCall functionCall = item.asFunctionCall();
407+
LLMObs.ToolCall toolCall = ToolCallExtractor.getToolCall(functionCall);
408+
if (toolCall != null) {
409+
toolCalls.add(toolCall);
410+
}
411+
} else if (item.isMessage()) {
412+
ResponseOutputMessage message = item.asMessage();
413+
textContent = extractMessageContent(message);
414+
}
415+
}
416+
417+
messages.add(LLMObs.LLMMessage.from("assistant", textContent, toolCalls));
418+
419+
return messages;
420+
}
421+
422+
private String extractMessageContent(ResponseOutputMessage message) {
423+
StringBuilder contentBuilder = new StringBuilder();
424+
for (ResponseOutputMessage.Content content : message.content()) {
425+
if (content.isOutputText()) {
426+
ResponseOutputText outputText = content.asOutputText();
427+
contentBuilder.append(outputText.text());
428+
}
429+
}
430+
String result = contentBuilder.toString();
431+
return result.isEmpty() ? null : result;
432+
}
433+
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());
369440
}
370441

371442
private String extractResponseModel(JsonField<ResponsesModel> model) {
@@ -396,11 +467,13 @@ private String extractResponseModel(JsonField<ResponsesModel> model) {
396467
}
397468

398469
public void withResponseStreamEvent(AgentSpan span, List<ResponseStreamEvent> events) {
399-
if (!events.isEmpty()) {
400-
// ResponseStreamEvent responseStreamEvent = events.get(0);
401-
// span.setTag(RESPONSE_MODEL, responseStreamEvent.res()); // TODO there is no model
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+
}
402477
}
403-
404-
// TODO set LLMObs tags
405478
}
406479
}

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

Lines changed: 2 additions & 2 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);
55+
DECORATE.withResponseCreateParams(span, params, false);
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);
89+
DECORATE.withResponseCreateParams(span, params, true);
9090
return activateSpan(span);
9191
}
9292

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

Lines changed: 2 additions & 2 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);
57+
DECORATE.withResponseCreateParams(span, params, false);
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);
92+
DECORATE.withResponseCreateParams(span, params, true);
9393
return activateSpan(span);
9494
}
9595

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

Lines changed: 31 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import com.fasterxml.jackson.databind.ObjectMapper;
55
import com.openai.models.chat.completions.ChatCompletionMessageFunctionToolCall;
66
import com.openai.models.chat.completions.ChatCompletionMessageToolCall;
7+
import com.openai.models.responses.ResponseFunctionToolCall;
78
import datadog.trace.api.llmobs.LLMObs;
89
import java.util.Collections;
910
import java.util.Map;
@@ -29,24 +30,46 @@ public static LLMObs.ToolCall getToolCall(ChatCompletionMessageToolCall toolCall
2930
String name = function.name();
3031
String argumentsJson = function.arguments();
3132

32-
Map<String, Object> arguments;
33-
try {
34-
arguments = MAPPER.readValue(argumentsJson, MAP_TYPE_REF);
35-
} catch (Exception e) {
36-
log.debug("Failed to parse tool call arguments as JSON: {}", argumentsJson, e);
37-
arguments = Collections.singletonMap("value", argumentsJson);
38-
}
39-
4033
String type = "function";
4134
Optional<String> typeOpt = functionToolCall._type().asString();
4235
if (typeOpt.isPresent()) {
4336
type = typeOpt.get();
4437
}
4538

39+
Map<String, Object> arguments = parseArguments(argumentsJson);
4640
return LLMObs.ToolCall.from(name, type, toolId, arguments);
4741
} catch (Exception e) {
4842
log.debug("Failed to extract tool call information", e);
4943
}
5044
return null;
5145
}
46+
47+
public static LLMObs.ToolCall getToolCall(ResponseFunctionToolCall functionCall) {
48+
try {
49+
String name = functionCall.name();
50+
String callId = functionCall.callId();
51+
String argumentsJson = functionCall.arguments();
52+
53+
String type = "function_call";
54+
Optional<String> typeOpt = functionCall._type().asString();
55+
if (typeOpt.isPresent()) {
56+
type = typeOpt.get();
57+
}
58+
59+
Map<String, Object> arguments = parseArguments(argumentsJson);
60+
return LLMObs.ToolCall.from(name, type, callId, arguments);
61+
} catch (Exception e) {
62+
log.debug("Failed to extract tool call information", e);
63+
}
64+
return null;
65+
}
66+
67+
private static Map<String, Object> parseArguments(String argumentsJson) {
68+
try {
69+
return MAPPER.readValue(argumentsJson, MAP_TYPE_REF);
70+
} catch (Exception e) {
71+
log.debug("Failed to parse tool call arguments as JSON: {}", argumentsJson, e);
72+
return Collections.singletonMap("value", argumentsJson);
73+
}
74+
}
5275
}

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

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -132,18 +132,25 @@ class ResponseServiceTest extends OpenAiTest {
132132
spanType DDSpanTypes.LLMOBS
133133
tags {
134134
"_ml_obs_tag.span.kind" "llm"
135+
"_ml_obs_tag.model_provider" "openai"
136+
"_ml_obs_tag.model_name" String
137+
"_ml_obs_tag.metadata" Map
138+
"_ml_obs_tag.output" List // TODO capture to validate tool calls
139+
"_ml_obs_metric.input_tokens" Long
140+
"_ml_obs_metric.output_tokens" Long
141+
"_ml_obs_metric.total_tokens" Long
142+
"_ml_obs_metric.cache_read_input_tokens" Long
135143
"_ml_obs_tag.parent_id" "undefined"
136144
"openai.request.method" "POST"
137145
"openai.request.endpoint" "v1/responses"
138146
"openai.api_base" openAiBaseApi
147+
"$OpenAiDecorator.RESPONSE_MODEL" "gpt-3.5-turbo-0125"
139148
if (!isStreaming) {
140149
// TODO no limit headers when streaming
141150
"openai.organization.ratelimit.requests.limit" 10000
142151
"openai.organization.ratelimit.requests.remaining" Integer
143152
"openai.organization.ratelimit.tokens.limit" 50000000
144153
"openai.organization.ratelimit.tokens.remaining" Integer
145-
// TODO no response model
146-
"$OpenAiDecorator.RESPONSE_MODEL" "gpt-3.5-turbo-0125"
147154
}
148155
"$OpenAiDecorator.OPENAI_ORGANIZATION_NAME" "datadog-staging"
149156
"$OpenAiDecorator.REQUEST_MODEL" "gpt-3.5-turbo"

0 commit comments

Comments
 (0)