Add logic to account for thought and thoughtSignature in chat history

Author: hsubox76

common/api-review/ai.api.md

@@ -242,6 +242,8 @@ export interface FileDataPart {
     text?: never;
     // (undocumented)
     thought?: boolean;
+    // @internal (undocumented)
+    thoughtSignature?: never;
 }
 
 // @public
@@ -298,6 +300,8 @@ export interface FunctionCallPart {
     text?: never;
     // (undocumented)
     thought?: boolean;
+    // @internal (undocumented)
+    thoughtSignature?: never;
 }
 
 // @public
@@ -332,6 +336,8 @@ export interface FunctionResponsePart {
     text?: never;
     // (undocumented)
     thought?: boolean;
+    // @internal (undocumented)
+    thoughtSignature?: never;
 }
 
 // @public
@@ -699,6 +705,8 @@ export interface InlineDataPart {
     text?: never;
     // (undocumented)
     thought?: boolean;
+    // @internal (undocumented)
+    thoughtSignature?: never;
     videoMetadata?: VideoMetadata;
 }
 
@@ -967,6 +975,8 @@ export interface TextPart {
     text: string;
     // (undocumented)
     thought?: boolean;
+    // @internal (undocumented)
+    thoughtSignature?: string;
 }
 
 // @public

packages/ai/src/methods/chat-session-helpers.test.ts

@@ -22,7 +22,11 @@ import { FirebaseError } from '@firebase/util';
 
 describe('chat-session-helpers', () => {
   describe('validateChatHistory', () => {
-    const TCS: Array<{ history: Content[]; isValid: boolean }> = [
+    const TCS: Array<{
+      history: Content[];
+      isValid: boolean;
+      errorShouldInclude?: string;
+    }> = [
       {
         history: [{ role: 'user', parts: [{ text: 'hi' }] }],
         isValid: true
@@ -99,19 +103,23 @@ describe('chat-session-helpers', () => {
       {
         //@ts-expect-error
         history: [{ role: 'user', parts: '' }],
+        errorShouldInclude: `array of Parts`,
         isValid: false
       },
       {
         //@ts-expect-error
         history: [{ role: 'user' }],
+        errorShouldInclude: `array of Parts`,
         isValid: false
       },
       {
         history: [{ role: 'user', parts: [] }],
+        errorShouldInclude: `at least one part`,
         isValid: false
       },
       {
         history: [{ role: 'model', parts: [{ text: 'hi' }] }],
+        errorShouldInclude: `model`,
         isValid: false
       },
       {
@@ -125,13 +133,15 @@ describe('chat-session-helpers', () => {
             ]
           }
         ],
+        errorShouldInclude: `function`,
         isValid: false
       },
       {
         history: [
           { role: 'user', parts: [{ text: 'hi' }] },
           { role: 'user', parts: [{ text: 'hi' }] }
         ],
+        errorShouldInclude: `can't follow 'user'`,
         isValid: false
       },
       {
@@ -140,6 +150,45 @@ describe('chat-session-helpers', () => {
           { role: 'model', parts: [{ text: 'hi' }] },
           { role: 'model', parts: [{ text: 'hi' }] }
         ],
+        errorShouldInclude: `can't follow 'model'`,
+        isValid: false
+      },
+      {
+        history: [
+          { role: 'user', parts: [{ text: 'hi' }] },
+          {
+            role: 'model',
+            parts: [
+              { text: 'hi' },
+              {
+                text: 'thought about hi',
+                thought: true,
+                thoughtSignature: 'thought signature'
+              }
+            ]
+          }
+        ],
+        isValid: true
+      },
+      {
+        history: [
+          {
+            role: 'user',
+            parts: [{ text: 'hi', thought: true, thoughtSignature: 'sig' }]
+          },
+          {
+            role: 'model',
+            parts: [
+              { text: 'hi' },
+              {
+                text: 'thought about hi',
+                thought: true,
+                thoughtSignature: 'thought signature'
+              }
+            ]
+          }
+        ],
+        errorShouldInclude: 'thought',
         isValid: false
       }
     ];
@@ -149,7 +198,14 @@ describe('chat-session-helpers', () => {
         if (tc.isValid) {
           expect(fn).to.not.throw();
         } else {
-          expect(fn).to.throw(FirebaseError);
+          try {
+            fn();
+          } catch (e) {
+            expect(e).to.be.instanceOf(FirebaseError);
+            if (e instanceof FirebaseError && tc.errorShouldInclude) {
+              expect(e.message).to.include(tc.errorShouldInclude);
+            }
+          }
         }
       });
     });

packages/ai/src/methods/chat-session-helpers.ts

@@ -24,13 +24,15 @@ const VALID_PART_FIELDS: Array<keyof Part> = [
   'text',
   'inlineData',
   'functionCall',
-  'functionResponse'
+  'functionResponse',
+  'thought',
+  'thoughtSignature'
 ];
 
 const VALID_PARTS_PER_ROLE: { [key in Role]: Array<keyof Part> } = {
   user: ['text', 'inlineData'],
   function: ['functionResponse'],
-  model: ['text', 'functionCall'],
+  model: ['text', 'functionCall', 'thought', 'thoughtSignature'],
   // System instructions shouldn't be in history anyway.
   system: ['text']
 };
@@ -65,7 +67,7 @@ export function validateChatHistory(history: Content[]): void {
     if (!Array.isArray(parts)) {
       throw new AIError(
         AIErrorCode.INVALID_CONTENT,
-        `Content should have 'parts' but property with an array of Parts`
+        `Content should have 'parts' property with an array of Parts`
       );
     }
 
@@ -81,7 +83,8 @@ export function validateChatHistory(history: Content[]): void {
       inlineData: 0,
       functionCall: 0,
       functionResponse: 0,
-      thought: 0
+      thought: 0,
+      thoughtSignature: 0
     };
 
     for (const part of parts) {

packages/ai/src/methods/chat-session.test.ts

@@ -20,7 +20,7 @@ import { match, restore, stub, useFakeTimers } from 'sinon';
 import sinonChai from 'sinon-chai';
 import chaiAsPromised from 'chai-as-promised';
 import * as generateContentMethods from './generate-content';
-import { GenerateContentStreamResult } from '../types';
+import { Content, GenerateContentStreamResult } from '../types';
 import { ChatSession } from './chat-session';
 import { ApiSettings } from '../types/internal';
 import { VertexAIBackend } from '../backend';
@@ -54,6 +54,53 @@ describe('ChatSession', () => {
         match.any
       );
     });
+    it('adds message and response to history', async () => {
+      const fakeContent: Content = {
+        role: 'model',
+        parts: [
+          { text: 'hi' },
+          {
+            text: 'thought about hi',
+            thoughtSignature: 'thought signature'
+          }
+        ]
+      };
+      const fakeResponse = {
+        candidates: [
+          {
+            index: 1,
+            content: fakeContent
+          }
+        ]
+      };
+      const generateContentStub = stub(
+        generateContentMethods,
+        'generateContent'
+      ).resolves({
+        // @ts-ignore
+        response: fakeResponse
+      });
+      const chatSession = new ChatSession(fakeApiSettings, 'a-model');
+      const result = await chatSession.sendMessage('hello');
+      // @ts-ignore
+      expect(result.response).to.equal(fakeResponse);
+      // Test: stores history correctly?
+      const history = await chatSession.getHistory();
+      expect(history[0].role).to.equal('user');
+      expect(history[0].parts[0].text).to.equal('hello');
+      expect(history[1]).to.deep.equal(fakeResponse.candidates[0].content);
+      // Test: sends history correctly?
+      await chatSession.sendMessage('hello 2');
+      expect(generateContentStub.args[1][2].contents[0].parts[0].text).to.equal(
+        'hello'
+      );
+      expect(generateContentStub.args[1][2].contents[1]).to.deep.equal(
+        fakeResponse.candidates[0].content
+      );
+      expect(generateContentStub.args[1][2].contents[2].parts[0].text).to.equal(
+        'hello 2'
+      );
+    });
   });
   describe('sendMessageStream()', () => {
     it('generateContentStream errors should be catchable', async () => {

packages/ai/src/types/content.ts

@@ -48,6 +48,10 @@ export interface TextPart {
   functionCall?: never;
   functionResponse?: never;
   thought?: boolean;
+  /**
+   * @internal
+   */
+  thoughtSignature?: string;
 }
 
 /**
@@ -64,6 +68,10 @@ export interface InlineDataPart {
    */
   videoMetadata?: VideoMetadata;
   thought?: boolean;
+  /**
+   * @internal
+   */
+  thoughtSignature?: never;
 }
 
 /**
@@ -93,6 +101,10 @@ export interface FunctionCallPart {
   functionCall: FunctionCall;
   functionResponse?: never;
   thought?: boolean;
+  /**
+   * @internal
+   */
+  thoughtSignature?: never;
 }
 
 /**
@@ -105,6 +117,10 @@ export interface FunctionResponsePart {
   functionCall?: never;
   functionResponse: FunctionResponse;
   thought?: boolean;
+  /**
+   * @internal
+   */
+  thoughtSignature?: never;
 }
 
 /**
@@ -118,6 +134,10 @@ export interface FileDataPart {
   functionResponse?: never;
   fileData: FileData;
   thought?: boolean;
+  /**
+   * @internal
+   */
+  thoughtSignature?: never;
 }
 
 /**