AIX: Gemini: collapse empty text in between reasoning

Author: enricoros

src/modules/aix/client/ContentReassembler.ts

@@ -317,13 +317,13 @@ export class ContentReassembler {
 
   }
 
-  private onAppendReasoningText({ _t /*, weak*/ }: Extract<AixWire_Particles.PartParticleOp, { p: 'tr_' }>): void {
+  private onAppendReasoningText({ _t, restart }: Extract<AixWire_Particles.PartParticleOp, { p: 'tr_' }>): void {
     // Break text accumulation
     this.currentTextFragmentIndex = null;
 
     // append to existing ModelAuxVoidFragment if possible
     const currentFragment = this.accumulator.fragments[this.accumulator.fragments.length - 1];
-    if (currentFragment && isVoidFragment(currentFragment) && isModelAuxPart(currentFragment.part)) {
+    if (!restart && currentFragment && isVoidFragment(currentFragment) && isModelAuxPart(currentFragment.part)) {
       const appendedPart = { ...currentFragment.part, aText: (currentFragment.part.aText || '') + _t } satisfies DVoidModelAuxPart;
       this.accumulator.fragments[this.accumulator.fragments.length - 1] = { ...currentFragment, part: appendedPart };
       return;

src/modules/aix/server/api/aix.wiretypes.ts

@@ -671,7 +671,7 @@ export namespace AixWire_Particles {
 
   export type PartParticleOp =
     | { p: '❤' } // heart beat
-    | { p: 'tr_', _t: string, weak?: 'tag' } // reasoning text, incremental; could be a 'weak' detection, e.g. heuristic from '<think>' rather than API-provided
+    | { p: 'tr_', _t: string, weak?: 'tag', restart?: boolean } // reasoning text, incremental; could be a 'weak' detection, e.g. heuristic from '<think>' rather than API-provided
     | { p: 'trs', signature: string } // reasoning signature
     | { p: 'trr_', _data: string } // reasoning raw (or redacted) data
     // | { p: 'ii', mimeType: string, i_b64?: string /* never undefined */ }

src/modules/aix/server/dispatch/chatGenerate/ChatGenerateTransmitter.ts

@@ -265,15 +265,16 @@ export class ChatGenerateTransmitter implements IParticleTransmitter {
   }
 
   /** Appends reasoning text, which is its own kind of content */
-  appendReasoningText(textChunk: string, weak?: Extract<AixWire_Particles.PartParticleOp, { p: 'tr_' }>['weak']) {
+  appendReasoningText(textChunk: string, options?: { weak?: 'tag', restart?: boolean }) {
     // NOTE: don't skip on empty chunks, as we want to transition states
     // if there was another Part in the making, queue it
     if (this.currentPart)
       this.endMessagePart();
     this.currentPart = {
       p: 'tr_',
       _t: textChunk,
-      ...(weak ? { weak } : {}),
+      ...(options?.weak ? { weak: options.weak } : {}),
+      ...(options?.restart ? { restart: true } : {}),
     };
     // [throttle] send it immediately for now
     this._queueParticleS();
@@ -326,12 +327,12 @@ export class ChatGenerateTransmitter implements IParticleTransmitter {
         const closingIdx = remaining.indexOf('</think>');
         if (closingIdx >= 0) {
           const reasoningText = remaining.substring(0, closingIdx);
-          this.appendReasoningText(reasoningText, 'tag');
+          this.appendReasoningText(reasoningText, { weak: 'tag' });
           this.isThinkingText = false;
           remaining = remaining.substring(closingIdx + '</think>'.length);
           // this is the only branch that can still loop
         } else {
-          this.appendReasoningText(remaining, 'tag');
+          this.appendReasoningText(remaining, { weak: 'tag' });
           return;
         }
       } else {

src/modules/aix/server/dispatch/chatGenerate/parsers/IParticleTransmitter.ts

@@ -23,7 +23,7 @@ export interface IParticleTransmitter {
   appendText(textChunk: string): void;
 
   /** Appends reasoning text, creating a part if missing [throttled] */
-  appendReasoningText(textChunk: string, weak?: 'tag'): void;
+  appendReasoningText(textChunk: string, options?: { weak?: 'tag', restart?: boolean }): void;
 
   /** Sets a reasoning signature, associated with the current reasoning text */
   setReasoningSignature(signature: string): void;

src/modules/aix/server/dispatch/chatGenerate/parsers/gemini.parser.ts

@@ -9,6 +9,7 @@ import { geminiConvertPCM2WAV } from './gemini.audioutils';
 
 
 // configuration
+const COLLAPSE_EMPTY_TEXT_PARTS = true;
 const ENABLE_RECITATIONS_AS_CITATIONS = false;
 
 
@@ -34,6 +35,7 @@ export function createGeminiGenerateContentResponseParser(requestedModelName: st
   let sentRequestedModelName = false;
   let sentActualModelName = false;
   let timeToFirstEvent: number;
+  let collapsedTextPartForReasoning = false;
   let skipComputingTotalsOnce = isStreaming;
   let groundingIndexNumber = 0;
 
@@ -118,14 +120,18 @@ export function createGeminiGenerateContentResponseParser(requestedModelName: st
           // <- TextPart
           case 'text' in mPart:
             // [Gemini, 2025-01-23] CoT support
-            if (mPart.thought)
-              pt.appendReasoningText(mPart.text || '');
-            else {
+            if (mPart.thought) {
+              pt.appendReasoningText(mPart.text || '', collapsedTextPartForReasoning ? { restart: true } : undefined);
+              collapsedTextPartForReasoning = false;
+            } else {
               // NOTE: considering the below, but not yet
               // don't send an empty text part, which may happen in between reasoning parts
               // and this way we can merge them
               // if (mPart.text?.length)
-              pt.appendText(mPart.text || '');
+              if (!COLLAPSE_EMPTY_TEXT_PARTS || mPart.text)
+                pt.appendText(mPart.text || '');
+              else
+                collapsedTextPartForReasoning = true;
             }
             break;