|
19 | 19 | import com.fasterxml.jackson.annotation.JsonIgnoreProperties; |
20 | 20 | import com.fasterxml.jackson.annotation.JsonInclude; |
21 | 21 | 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; |
22 | 30 | import java.util.List; |
| 31 | +import org.jspecify.annotations.Nullable; |
23 | 32 |
|
24 | 33 | /** |
25 | 34 | * Data Transfer Objects for Chat Completion and Chat Completion Chunk API responses. |
@@ -62,6 +71,162 @@ static class ChatCompletion { |
62 | 71 |
|
63 | 72 | /** See class definition for more details. */ |
64 | 73 | 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 | + } |
65 | 230 | } |
66 | 231 |
|
67 | 232 | /** |
|
0 commit comments