Skip to content

Commit d992714

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: Add ChatCompletionsResponse to LlmResponse conversion
This is part of a larger chain of commits for adding chat completion API support to the Apigee model. PiperOrigin-RevId: 890876819
1 parent 3091156 commit d992714

File tree

3 files changed

+611
-3
lines changed

3 files changed

+611
-3
lines changed

core/src/main/java/com/google/adk/models/chat/ChatCompletionsCommon.java

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,13 @@
1919
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
2020
import com.fasterxml.jackson.annotation.JsonInclude;
2121
import com.fasterxml.jackson.annotation.JsonProperty;
22+
import com.fasterxml.jackson.core.type.TypeReference;
23+
import com.fasterxml.jackson.databind.ObjectMapper;
24+
import com.google.genai.types.FunctionCall;
25+
import com.google.genai.types.Part;
26+
import java.util.Base64;
2227
import java.util.Map;
28+
import org.jspecify.annotations.Nullable;
2329

2430
/** Shared models for Chat Completions Request and Response. */
2531
@JsonIgnoreProperties(ignoreUnknown = true)
@@ -28,6 +34,17 @@ final class ChatCompletionsCommon {
2834

2935
private ChatCompletionsCommon() {}
3036

37+
private static final ObjectMapper objectMapper = new ObjectMapper();
38+
39+
public static final String ROLE_ASSISTANT = "assistant";
40+
public static final String ROLE_MODEL = "model";
41+
42+
public static final String METADATA_KEY_ID = "id";
43+
public static final String METADATA_KEY_CREATED = "created";
44+
public static final String METADATA_KEY_OBJECT = "object";
45+
public static final String METADATA_KEY_SYSTEM_FINGERPRINT = "system_fingerprint";
46+
public static final String METADATA_KEY_SERVICE_TIER = "service_tier";
47+
3148
/**
3249
* See
3350
* https://developers.openai.com/api/reference/resources/chat#(resource)%20chat.completions%20%3E%20(model)%20chat_completion_message_tool_call%20%3E%20(schema)
@@ -56,6 +73,43 @@ static class ToolCall {
5673
*/
5774
@JsonProperty("extra_content")
5875
public Map<String, Object> extraContent;
76+
77+
/**
78+
* Converts the tool call to a {@link Part}.
79+
*
80+
* @return a {@link Part} containing the function call, or {@code null} if this tool call does
81+
* not contain a function call.
82+
*/
83+
public @Nullable Part toPart() {
84+
if (function != null) {
85+
FunctionCall fc = function.toFunctionCall(id);
86+
Part part = Part.builder().functionCall(fc).build();
87+
return applyThoughtSignature(part);
88+
}
89+
return null;
90+
}
91+
92+
/**
93+
* Applies the thought signature from {@code extraContent} to the given {@link Part} if present.
94+
* This is used to support the Google Gemini/Vertex AI implementation of the chat/completions
95+
* API.
96+
*
97+
* @param part the {@link Part} to modify.
98+
* @return a new {@link Part} with the thought signature applied, or the original {@link Part}
99+
* if no thought signature is found.
100+
*/
101+
public Part applyThoughtSignature(Part part) {
102+
if (extraContent != null && extraContent.containsKey("google")) {
103+
Object googleObj = extraContent.get("google");
104+
if (googleObj instanceof Map<?, ?> googleMap) {
105+
Object sigObj = googleMap.get("thought_signature");
106+
if (sigObj instanceof String sig) {
107+
return part.toBuilder().thoughtSignature(Base64.getDecoder().decode(sig)).build();
108+
}
109+
}
110+
}
111+
return part;
112+
}
59113
}
60114

61115
/**
@@ -70,6 +124,33 @@ static class Function {
70124

71125
/** See class definition for more details. */
72126
public String arguments; // JSON string
127+
128+
/**
129+
* Converts this function to a {@link FunctionCall}.
130+
*
131+
* @param toolCallId the ID of the tool call, or {@code null} if not applicable.
132+
* @return the {@link FunctionCall} object.
133+
*/
134+
public FunctionCall toFunctionCall(@Nullable String toolCallId) {
135+
FunctionCall.Builder fcBuilder = FunctionCall.builder();
136+
if (name != null) {
137+
fcBuilder.name(name);
138+
}
139+
if (arguments != null) {
140+
try {
141+
Map<String, Object> args =
142+
objectMapper.readValue(arguments, new TypeReference<Map<String, Object>>() {});
143+
fcBuilder.args(args);
144+
} catch (Exception e) {
145+
throw new IllegalArgumentException(
146+
"Failed to parse function arguments JSON: " + arguments, e);
147+
}
148+
}
149+
if (toolCallId != null) {
150+
fcBuilder.id(toolCallId);
151+
}
152+
return fcBuilder.build();
153+
}
73154
}
74155

75156
/**

core/src/main/java/com/google/adk/models/chat/ChatCompletionsResponse.java

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,16 @@
1919
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
2020
import com.fasterxml.jackson.annotation.JsonInclude;
2121
import com.fasterxml.jackson.annotation.JsonProperty;
22+
import com.google.adk.models.LlmResponse;
23+
import com.google.genai.types.Content;
24+
import com.google.genai.types.CustomMetadata;
25+
import com.google.genai.types.FinishReason;
26+
import com.google.genai.types.FinishReason.Known;
27+
import com.google.genai.types.GenerateContentResponseUsageMetadata;
28+
import com.google.genai.types.Part;
29+
import java.util.ArrayList;
2230
import java.util.List;
31+
import org.jspecify.annotations.Nullable;
2332

2433
/**
2534
* Data Transfer Objects for Chat Completion and Chat Completion Chunk API responses.
@@ -62,6 +71,162 @@ static class ChatCompletion {
6271

6372
/** See class definition for more details. */
6473
public Usage usage;
74+
75+
/**
76+
* Converts this chat completion to a {@link LlmResponse}.
77+
*
78+
* @return the {@link LlmResponse} object.
79+
*/
80+
public LlmResponse toLlmResponse() {
81+
Choice choice = (choices != null && !choices.isEmpty()) ? choices.get(0) : null;
82+
Content content = mapChoiceToContent(choice);
83+
84+
LlmResponse.Builder builder = LlmResponse.builder().content(content);
85+
86+
if (choice != null) {
87+
builder.finishReason(mapFinishReason(choice.finishReason));
88+
}
89+
90+
if (model != null) {
91+
builder.modelVersion(model);
92+
}
93+
94+
if (usage != null) {
95+
builder.usageMetadata(mapUsage(usage));
96+
}
97+
98+
List<CustomMetadata> customMetadataList = buildCustomMetadata();
99+
return builder.customMetadata(customMetadataList).build();
100+
}
101+
102+
/**
103+
* Maps the finish reason string to a {@link FinishReason}.
104+
*
105+
* @param reason the finish reason string.
106+
* @return the {@link FinishReason}, or {@code null} if the input reason is null.
107+
*/
108+
private @Nullable FinishReason mapFinishReason(String reason) {
109+
if (reason == null) {
110+
return null;
111+
}
112+
return switch (reason) {
113+
case "stop", "tool_calls" -> new FinishReason(Known.STOP.toString());
114+
case "length" -> new FinishReason(Known.MAX_TOKENS.toString());
115+
case "content_filter" -> new FinishReason(Known.SAFETY.toString());
116+
default -> new FinishReason(Known.OTHER.toString());
117+
};
118+
}
119+
120+
private GenerateContentResponseUsageMetadata mapUsage(Usage usage) {
121+
GenerateContentResponseUsageMetadata.Builder builder =
122+
GenerateContentResponseUsageMetadata.builder();
123+
if (usage.promptTokens != null) {
124+
builder.promptTokenCount(usage.promptTokens);
125+
}
126+
if (usage.completionTokens != null) {
127+
builder.candidatesTokenCount(usage.completionTokens);
128+
}
129+
if (usage.totalTokens != null) {
130+
builder.totalTokenCount(usage.totalTokens);
131+
}
132+
if (usage.thoughtsTokenCount != null) {
133+
builder.thoughtsTokenCount(usage.thoughtsTokenCount);
134+
} else if (usage.completionTokensDetails != null
135+
&& usage.completionTokensDetails.reasoningTokens != null) {
136+
builder.thoughtsTokenCount(usage.completionTokensDetails.reasoningTokens);
137+
}
138+
return builder.build();
139+
}
140+
141+
/**
142+
* Maps the chosen completion to a {@link Content} object.
143+
*
144+
* @param choice the completion choice to map, or {@code null}.
145+
* @return the {@link Content} object, which will be empty if the choice or its message is null.
146+
*/
147+
private Content mapChoiceToContent(@Nullable Choice choice) {
148+
Content.Builder contentBuilder = Content.builder();
149+
if (choice != null && choice.message != null) {
150+
contentBuilder.role(mapRole(choice.message.role)).parts(mapMessageToParts(choice.message));
151+
}
152+
return contentBuilder.build();
153+
}
154+
155+
private String mapRole(@Nullable String role) {
156+
return (role != null && role.equals(ChatCompletionsCommon.ROLE_ASSISTANT))
157+
? ChatCompletionsCommon.ROLE_MODEL
158+
: role;
159+
}
160+
161+
private List<Part> mapMessageToParts(Message message) {
162+
List<Part> parts = new ArrayList<>();
163+
if (message.content != null) {
164+
parts.add(Part.fromText(message.content));
165+
}
166+
if (message.refusal != null) {
167+
parts.add(Part.fromText(message.refusal));
168+
}
169+
if (message.toolCalls != null) {
170+
parts.addAll(mapToolCallsToParts(message.toolCalls));
171+
}
172+
return parts;
173+
}
174+
175+
private List<Part> mapToolCallsToParts(List<ChatCompletionsCommon.ToolCall> toolCalls) {
176+
List<Part> parts = new ArrayList<>();
177+
for (ChatCompletionsCommon.ToolCall toolCall : toolCalls) {
178+
Part part = toolCall.toPart();
179+
if (part != null) {
180+
parts.add(part);
181+
}
182+
}
183+
return parts;
184+
}
185+
186+
/**
187+
* Builds the list of custom metadata from the chat completion fields.
188+
*
189+
* @return a list of {@link CustomMetadata}, which will be empty if no relevant fields are set.
190+
*/
191+
private List<CustomMetadata> buildCustomMetadata() {
192+
List<CustomMetadata> customMetadataList = new ArrayList<>();
193+
if (id != null) {
194+
customMetadataList.add(
195+
CustomMetadata.builder()
196+
.key(ChatCompletionsCommon.METADATA_KEY_ID)
197+
.stringValue(id)
198+
.build());
199+
}
200+
if (created != null) {
201+
customMetadataList.add(
202+
CustomMetadata.builder()
203+
.key(ChatCompletionsCommon.METADATA_KEY_CREATED)
204+
.stringValue(created.toString())
205+
.build());
206+
}
207+
if (object != null) {
208+
customMetadataList.add(
209+
CustomMetadata.builder()
210+
.key(ChatCompletionsCommon.METADATA_KEY_OBJECT)
211+
.stringValue(object)
212+
.build());
213+
}
214+
if (systemFingerprint != null) {
215+
customMetadataList.add(
216+
CustomMetadata.builder()
217+
.key(ChatCompletionsCommon.METADATA_KEY_SYSTEM_FINGERPRINT)
218+
.stringValue(systemFingerprint)
219+
.build());
220+
}
221+
if (serviceTier != null) {
222+
customMetadataList.add(
223+
CustomMetadata.builder()
224+
.key(ChatCompletionsCommon.METADATA_KEY_SERVICE_TIER)
225+
.stringValue(serviceTier)
226+
.build());
227+
}
228+
return customMetadataList;
229+
}
65230
}
66231

67232
/**

0 commit comments

Comments
 (0)