Fix missing thoughts in GenerativeModelSession.streamResponse (#16075)

Co-authored-by: Andrew Heard <andrewheard@google.com>

Author: peterfriese

FirebaseAI/CHANGELOG.md

@@ -3,6 +3,7 @@
 - [fixed] Fix unfound file warnings from `swift build`. (#16012)
 - [fixed] Fixed a `no member 'autoFunctionDeclaration'` compilation error on
   unofficially supported Xcode versions older than 26.2. (#16037)
+- [fixed] Fixed missing thought summary output in `GenerativeModelSession.streamResponse`. (#16075)
 
 # 12.12.0
 - [feature] Added support for automatic function calling in

FirebaseAI/Sources/GenerativeModelSession.swift

@@ -361,7 +361,8 @@
               }
 
               // 2. If we have pending data, we now know it wasn't the last chunk.
-              if let pending = pendingChunkData, !pending.text.isEmpty {
+              if let pending = pendingChunkData,
+                 !pending.text.isEmpty || pending.response.thoughtSummary != nil {
                 let rawContent = try Self.makeRawContent(
                   from: pending.text,
                   generationID: pending.id,
@@ -413,7 +414,8 @@
 
               if !functionResponses.isEmpty {
                 // Yield any pending text if it's not empty, but mark it as NOT complete yet.
-                if let pending = pendingChunkData, !pending.text.isEmpty {
+                if let pending = pendingChunkData,
+                   !pending.text.isEmpty || pending.response.thoughtSummary != nil {
                   let rawContent = try Self.makeRawContent(
                     from: pending.text,
                     generationID: pending.id,
@@ -523,6 +525,15 @@
                                        hasSchema: Bool, isComplete: Bool) throws
       -> FirebaseAI.GeneratedContent {
       if hasSchema {
+        if text.isEmpty && !isComplete {
+          return FirebaseAI.GeneratedContent(
+            // TODO: Set `kind:` to `.array(...)`, `.bool()`, `.number()` based on schema type.
+            kind: .structure(properties: [:], orderedKeys: []),
+            id: generationID,
+            isComplete: isComplete
+          )
+        }
+
         return try FirebaseAI.GeneratedContent(json: text, id: generationID, isComplete: isComplete)
       }
 

FirebaseAI/Tests/TestApp/Tests/Integration/GenerativeModelSessionTests.swift

@@ -92,8 +92,15 @@
         let firebaseAI = FirebaseAI.componentInstance(config)
         let session = firebaseAI.generativeModelSession(model: ModelNames.gemini2_5_FlashLite)
         let prompt = "Generate a Ragdoll kitten"
+        let config = GenerationConfig(
+          thinkingConfig: ThinkingConfig(thinkingBudget: -1, includeThoughts: true)
+        )
 
-        let response = try await session.respond(to: prompt, generating: CatProfile.self)
+        let response = try await session.respond(
+          to: prompt,
+          generating: CatProfile.self,
+          options: config
+        )
 
         let catProfile = response.content
         #expect(!catProfile.name.isEmpty)
@@ -102,6 +109,10 @@
         #expect(!catProfile.profile.isEmpty)
         #expect(response.rawContent.isComplete)
         #expect(response.rawContent.generationID != nil)
+        let thoughtSummary = try #require(
+          response.rawResponse.thoughtSummary, "No thought summary was generated."
+        )
+        #expect(!thoughtSummary.isEmpty)
       }
 
       @Generable
@@ -438,15 +449,20 @@
         let firebaseAI = FirebaseAI.componentInstance(config)
         let session = firebaseAI.generativeModelSession(model: ModelNames.gemini2_5_FlashLite)
         let prompt = "Generate a Ragdoll kitten"
+        let config = GenerationConfig(
+          thinkingConfig: ThinkingConfig(thinkingBudget: -1, includeThoughts: true)
+        )
 
         let stream = session.streamResponse(
           to: prompt,
-          generating: CatProfile.self
+          generating: CatProfile.self,
+          options: config
         )
 
         var generationID: FirebaseAI.GenerationID?
         var id: FoundationModels.GenerationID?
         var isComplete = false
+        var thoughtSummary = ""
         for try await snapshot in stream {
           #expect(!isComplete, "Stream yielded more elements after a snapshot was marked complete.")
           let partial = snapshot.content
@@ -477,11 +493,17 @@
           } else {
             id = snapshot.content.id
           }
+          if let partialThoughtSummary = snapshot.rawResponse.thoughtSummary {
+            thoughtSummary += partialThoughtSummary
+          }
           isComplete = snapshot.rawContent.isComplete
         }
         #expect(
           isComplete, "The stream finished, but the final snapshot was not marked as complete."
         )
+        #expect(
+          !thoughtSummary.isEmpty, "The stream finished, but no thought summary was generated."
+        )
 
         let response = try await stream.collect()
         #expect(response.rawContent.isComplete, "The final response was not marked as complete.")