|
23 | 23 | import com.google.errorprone.annotations.CanIgnoreReturnValue; |
24 | 24 | import com.google.genai.Client; |
25 | 25 | import com.google.genai.ResponseStream; |
26 | | -import com.google.genai.types.Candidate; |
27 | 26 | import com.google.genai.types.Content; |
28 | 27 | import com.google.genai.types.FinishReason; |
29 | 28 | import com.google.genai.types.GenerateContentConfig; |
@@ -239,114 +238,131 @@ public Flowable<LlmResponse> generateContent(LlmRequest llmRequest, boolean stre |
239 | 238 | } |
240 | 239 |
|
241 | 240 | static Flowable<LlmResponse> processRawResponses(Flowable<GenerateContentResponse> rawResponses) { |
242 | | - final StringBuilder accumulatedText = new StringBuilder(); |
243 | | - final StringBuilder accumulatedThoughtText = new StringBuilder(); |
244 | | - // Array to bypass final local variable reassignment in lambda. |
245 | | - final GenerateContentResponse[] lastRawResponseHolder = {null}; |
246 | | - return rawResponses |
247 | | - .concatMap( |
248 | | - rawResponse -> { |
249 | | - lastRawResponseHolder[0] = rawResponse; |
250 | | - logger.trace("Raw streaming response: {}", rawResponse); |
251 | | - |
252 | | - List<LlmResponse> responsesToEmit = new ArrayList<>(); |
253 | | - LlmResponse currentProcessedLlmResponse = LlmResponse.create(rawResponse); |
254 | | - Optional<Part> part = GeminiUtil.getPart0FromLlmResponse(currentProcessedLlmResponse); |
255 | | - String currentTextChunk = part.flatMap(Part::text).orElse(""); |
256 | | - |
257 | | - if (!currentTextChunk.isBlank()) { |
258 | | - if (part.get().thought().orElse(false)) { |
259 | | - accumulatedThoughtText.append(currentTextChunk); |
260 | | - responsesToEmit.add( |
261 | | - thinkingResponseFromText(currentTextChunk).toBuilder() |
262 | | - .usageMetadata(currentProcessedLlmResponse.usageMetadata().orElse(null)) |
263 | | - .partial(true) |
264 | | - .build()); |
265 | | - } else { |
266 | | - accumulatedText.append(currentTextChunk); |
267 | | - responsesToEmit.add( |
268 | | - responseFromText(currentTextChunk).toBuilder() |
269 | | - .usageMetadata(currentProcessedLlmResponse.usageMetadata().orElse(null)) |
270 | | - .partial(true) |
271 | | - .build()); |
272 | | - } |
273 | | - } else { |
274 | | - if (accumulatedThoughtText.length() > 0 |
275 | | - && GeminiUtil.shouldEmitAccumulatedText(currentProcessedLlmResponse)) { |
276 | | - LlmResponse aggregatedThoughtResponse = |
277 | | - thinkingResponseFromText(accumulatedThoughtText.toString()); |
278 | | - responsesToEmit.add(aggregatedThoughtResponse); |
279 | | - accumulatedThoughtText.setLength(0); |
280 | | - } |
281 | | - if (accumulatedText.length() > 0 |
282 | | - && GeminiUtil.shouldEmitAccumulatedText(currentProcessedLlmResponse)) { |
283 | | - LlmResponse aggregatedTextResponse = responseFromText(accumulatedText.toString()); |
284 | | - responsesToEmit.add(aggregatedTextResponse); |
285 | | - accumulatedText.setLength(0); |
286 | | - } |
287 | | - if (isEmptyTextOnlyResponse(currentProcessedLlmResponse)) { |
288 | | - // Strip the empty-text content while preserving any carried metadata |
289 | | - // (`usageMetadata`, `finishReason`, `modelVersion`, etc.) by emitting a |
290 | | - // content-less response marked as `partial`. This handles the trailing |
291 | | - // `{parts:[{text:""}], finishReason:STOP}` chunk emitted by some Gemini |
292 | | - // preview models (e.g. 3.1-flash-lite) after a function call: keeping |
293 | | - // the chunk as-is would propagate it as a non-partial event whose |
294 | | - // Event#finalResponse() returns true and prematurely terminate |
295 | | - // BaseLlmFlow#run before the function response is sent back to the |
296 | | - // model; dropping it entirely would lose the carried metadata. If the |
297 | | - // chunk carries no useful metadata at all, suppress it outright. |
298 | | - LlmResponse metadataOnly = |
299 | | - currentProcessedLlmResponse.toBuilder() |
300 | | - .content((Content) null) |
301 | | - .partial(true) |
302 | | - .build(); |
303 | | - if (hasUsefulMetadata(metadataOnly)) { |
304 | | - responsesToEmit.add(metadataOnly); |
305 | | - } |
306 | | - } else { |
307 | | - responsesToEmit.add(currentProcessedLlmResponse); |
308 | | - } |
309 | | - } |
310 | | - logger.debug("Responses to emit: {}", responsesToEmit); |
311 | | - return Flowable.fromIterable(responsesToEmit); |
312 | | - }) |
313 | | - .concatWith( |
314 | | - Flowable.defer( |
315 | | - () -> { |
316 | | - GenerateContentResponse finalRawResp = lastRawResponseHolder[0]; |
317 | | - if (finalRawResp == null) { |
318 | | - return Flowable.empty(); |
319 | | - } |
320 | | - boolean isStop = |
321 | | - finalRawResp |
322 | | - .candidates() |
323 | | - .flatMap(candidates -> candidates.stream().findFirst()) |
324 | | - .flatMap(Candidate::finishReason) |
325 | | - .map(finishReason -> finishReason.knownEnum() == FinishReason.Known.STOP) |
326 | | - .orElse(false); |
327 | | - |
328 | | - if (isStop) { |
329 | | - List<LlmResponse> finalResponses = new ArrayList<>(); |
330 | | - if (accumulatedThoughtText.length() > 0) { |
331 | | - finalResponses.add( |
332 | | - thinkingResponseFromText(accumulatedThoughtText.toString()).toBuilder() |
333 | | - .usageMetadata( |
334 | | - accumulatedText.length() > 0 |
335 | | - ? null |
336 | | - : finalRawResp.usageMetadata().orElse(null)) |
337 | | - .build()); |
338 | | - } |
339 | | - if (accumulatedText.length() > 0) { |
340 | | - finalResponses.add( |
341 | | - responseFromText(accumulatedText.toString()).toBuilder() |
342 | | - .usageMetadata(finalRawResp.usageMetadata().orElse(null)) |
343 | | - .build()); |
344 | | - } |
345 | | - |
346 | | - return Flowable.fromIterable(finalResponses); |
347 | | - } |
348 | | - return Flowable.empty(); |
349 | | - })); |
| 241 | + return Flowable.defer(() -> new StreamingResponseAggregator().process(rawResponses)); |
| 242 | + } |
| 243 | + |
| 244 | + private static final class StreamingResponseAggregator { |
| 245 | + private final StringBuilder accumulatedText = new StringBuilder(); |
| 246 | + private final StringBuilder accumulatedThoughtText = new StringBuilder(); |
| 247 | + private final List<Part> accumulatedFunctionCalls = new ArrayList<>(); |
| 248 | + private GenerateContentResponse lastRawResponse = null; |
| 249 | + |
| 250 | + private Flowable<LlmResponse> process(Flowable<GenerateContentResponse> rawResponses) { |
| 251 | + return rawResponses |
| 252 | + .concatMap(this::processRawResponse) |
| 253 | + .concatWith(Flowable.defer(this::processFinalResponse)); |
| 254 | + } |
| 255 | + |
| 256 | + private Flowable<LlmResponse> processRawResponse(GenerateContentResponse rawResponse) { |
| 257 | + lastRawResponse = rawResponse; |
| 258 | + logger.trace("Raw streaming response: {}", rawResponse); |
| 259 | + |
| 260 | + List<LlmResponse> responsesToEmit = new ArrayList<>(); |
| 261 | + LlmResponse currentProcessedLlmResponse = LlmResponse.create(rawResponse); |
| 262 | + Optional<Part> partOpt = GeminiUtil.getPart0FromLlmResponse(currentProcessedLlmResponse); |
| 263 | + String currentTextChunk = partOpt.flatMap(Part::text).orElse(""); |
| 264 | + |
| 265 | + if (!currentTextChunk.isBlank()) { |
| 266 | + if (partOpt.get().thought().orElse(false)) { |
| 267 | + accumulatedThoughtText.append(currentTextChunk); |
| 268 | + } else { |
| 269 | + accumulatedText.append(currentTextChunk); |
| 270 | + } |
| 271 | + responsesToEmit.add(currentProcessedLlmResponse.toBuilder().partial(true).build()); |
| 272 | + } else { |
| 273 | + if (accumulatedThoughtText.length() > 0 |
| 274 | + && GeminiUtil.shouldEmitAccumulatedText(currentProcessedLlmResponse)) { |
| 275 | + responsesToEmit.add(thinkingResponseFromText(accumulatedThoughtText.toString())); |
| 276 | + accumulatedThoughtText.setLength(0); |
| 277 | + } |
| 278 | + if (accumulatedText.length() > 0 |
| 279 | + && GeminiUtil.shouldEmitAccumulatedText(currentProcessedLlmResponse)) { |
| 280 | + responsesToEmit.add(responseFromText(accumulatedText.toString())); |
| 281 | + accumulatedText.setLength(0); |
| 282 | + } |
| 283 | + |
| 284 | + if (partOpt.isPresent() && partOpt.get().functionCall().isPresent()) { |
| 285 | + accumulatedFunctionCalls.add(partOpt.get()); |
| 286 | + responsesToEmit.add(currentProcessedLlmResponse.toBuilder().partial(true).build()); |
| 287 | + } else if (!responsesToEmit.isEmpty()) { |
| 288 | + LlmResponse lastResponse = responsesToEmit.get(responsesToEmit.size() - 1); |
| 289 | + responsesToEmit.set( |
| 290 | + responsesToEmit.size() - 1, |
| 291 | + mergeMetadata(lastResponse, currentProcessedLlmResponse, partOpt.orElse(null))); |
| 292 | + } else if (!accumulatedFunctionCalls.isEmpty()) { |
| 293 | + // Suppress the empty STOP chunk because processFinalResponse() will immediately emit |
| 294 | + // the final aggregated function call response carrying the final metadata. |
| 295 | + } else { |
| 296 | + responsesToEmit.add(currentProcessedLlmResponse); |
| 297 | + } |
| 298 | + } |
| 299 | + logger.info("Responses to emit: {}", responsesToEmit); |
| 300 | + return Flowable.fromIterable(responsesToEmit); |
| 301 | + } |
| 302 | + |
| 303 | + private Flowable<LlmResponse> processFinalResponse() { |
| 304 | + if (lastRawResponse == null) { |
| 305 | + return Flowable.empty(); |
| 306 | + } |
| 307 | + LlmResponse currentResponse = LlmResponse.create(lastRawResponse); |
| 308 | + boolean isStop = |
| 309 | + currentResponse |
| 310 | + .finishReason() |
| 311 | + .map(reason -> reason.knownEnum() == FinishReason.Known.STOP) |
| 312 | + .orElse(false); |
| 313 | + |
| 314 | + if (!isStop) { |
| 315 | + return Flowable.empty(); |
| 316 | + } |
| 317 | + |
| 318 | + List<LlmResponse> finalResponses = new ArrayList<>(); |
| 319 | + if (accumulatedThoughtText.length() > 0) { |
| 320 | + finalResponses.add(thinkingResponseFromText(accumulatedThoughtText.toString())); |
| 321 | + } |
| 322 | + if (accumulatedText.length() > 0) { |
| 323 | + finalResponses.add(responseFromText(accumulatedText.toString())); |
| 324 | + } |
| 325 | + if (!accumulatedFunctionCalls.isEmpty()) { |
| 326 | + finalResponses.add( |
| 327 | + LlmResponse.builder() |
| 328 | + .content(Content.builder().role("model").parts(accumulatedFunctionCalls).build()) |
| 329 | + .partial(false) |
| 330 | + .build()); |
| 331 | + } |
| 332 | + |
| 333 | + if (!finalResponses.isEmpty()) { |
| 334 | + LlmResponse lastResponse = finalResponses.get(finalResponses.size() - 1); |
| 335 | + finalResponses.set( |
| 336 | + finalResponses.size() - 1, mergeMetadata(lastResponse, currentResponse, null)); |
| 337 | + } |
| 338 | + return Flowable.fromIterable(finalResponses); |
| 339 | + } |
| 340 | + |
| 341 | + private static LlmResponse mergeMetadata( |
| 342 | + LlmResponse lastResponse, LlmResponse currentResponse, Part currentPart) { |
| 343 | + LlmResponse.Builder mergedBuilder = |
| 344 | + lastResponse.toBuilder() |
| 345 | + .usageMetadata(currentResponse.usageMetadata().orElse(null)) |
| 346 | + .finishReason(currentResponse.finishReason().orElse(null)) |
| 347 | + .modelVersion(currentResponse.modelVersion().orElse(null)) |
| 348 | + .errorCode(currentResponse.errorCode().orElse(null)) |
| 349 | + .groundingMetadata(currentResponse.groundingMetadata().orElse(null)) |
| 350 | + .inputTranscription(currentResponse.inputTranscription().orElse(null)) |
| 351 | + .outputTranscription(currentResponse.outputTranscription().orElse(null)); |
| 352 | + |
| 353 | + if (currentPart != null && currentPart.thoughtSignature().isPresent()) { |
| 354 | + Content lastContent = lastResponse.content().orElse(null); |
| 355 | + if (lastContent != null |
| 356 | + && lastContent.parts().isPresent() |
| 357 | + && !lastContent.parts().get().isEmpty()) { |
| 358 | + Part lastPart = lastContent.parts().get().get(0); |
| 359 | + Part mergedPart = |
| 360 | + lastPart.toBuilder().thoughtSignature(currentPart.thoughtSignature().get()).build(); |
| 361 | + mergedBuilder.content(lastContent.toBuilder().parts(mergedPart).build()); |
| 362 | + } |
| 363 | + } |
| 364 | + return mergedBuilder.build(); |
| 365 | + } |
350 | 366 | } |
351 | 367 |
|
352 | 368 | private static LlmResponse responseFromText(String accumulatedText) { |
|
0 commit comments