support changing chat model after chat session has begun (#4189)

agent/bindings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/protocol_generated/SerializedChatMessage.kt

@@ -9,6 +9,7 @@ data class SerializedChatMessage(
   val editorState: Any? = null,
   val speaker: SpeakerEnum, // Oneof: human, assistant, system
   val text: String? = null,
+  val model: String? = null,
 ) {
 
   enum class SpeakerEnum {

agent/src/index.test.ts

@@ -190,6 +190,7 @@ describe('Agent', () => {
             expect(lastMessage).toMatchInlineSnapshot(
                 `
               {
+                "model": "anthropic/claude-3-sonnet-20240229",
                 "speaker": "assistant",
                 "text": "Hello there! I'm Claude, an AI assistant created by Anthropic. It's nice to meet you. How can I help you today?",
               }
@@ -339,6 +340,7 @@ describe('Agent', () => {
   "interactions": [
     {
       "assistantMessage": {
+        "model": "anthropic/claude-2.0",
         "speaker": "assistant",
         "text": " I'm Claude, an AI assistant created by Anthropic.",
       },

lib/shared/src/chat/transcript/index.ts

@@ -8,7 +8,16 @@ export interface SerializedChatTranscript {
     /** A unique and opaque identifier for this transcript. */
     id: string
 
+    /**
+     * The chat model used for the chat session.
+     *
+     * @deprecated The chat model is now a per-message property instead of applying to the entire
+     * chat session, to allow users to change the model partway through the chat. Use the last
+     * {@link SerializedChatMessage.model} instead. Can remove in v1.22 (after 1 major release,
+     * v1.20, where both fields are written).
+     */
     chatModel?: string
+
     chatTitle?: string
     interactions: SerializedChatInteraction[]
     lastInteractionTimestamp: string

lib/shared/src/chat/transcript/messages.ts

@@ -15,6 +15,15 @@ export interface ChatMessage extends Message {
      * back to editing the text for invalid values.
      */
     editorState?: unknown
+
+    /**
+     * The model used to generate this chat message response. Not set on human messages.
+     *
+     * NOTE: The chat model used to be a property of the entire chat session and was not changeable
+     * after the chat session had begun. Now that field, {@link SerializedChatTranscript.chatModel},
+     * is deprecated, and this field should be used instead.
+     */
+    model?: string
 }
 
 // An unsafe version of the {@link ChatMessage} that has the PromptString
@@ -28,6 +37,7 @@ export interface SerializedChatMessage {
     editorState?: unknown
     speaker: 'human' | 'assistant' | 'system'
     text?: string // Changed from PromptString
+    model?: string
 }
 
 export interface ChatError {

vscode/CHANGELOG.md

@@ -13,6 +13,7 @@ This is a log of all notable changes to Cody for VS Code. [Unreleased] changes a
 ### Changed
 
 - Chat: Pressing <kbd>Space</kbd> no longer accepts an @-mention item. Press <kbd>Tab</kbd> or <kbd>Enter</kbd> instead. [pull/4154](https://github.com/sourcegraph/cody/pull/4154)
+- Chat: You can now change the model after you send a chat message. Subsequent messages will be sent using your selected model. [pull/4189](https://github.com/sourcegraph/cody/pull/4189)
 
 ## [1.18.0]
 

vscode/src/chat/chat-view/SimpleChatModel.ts

@@ -67,6 +67,7 @@ export class SimpleChatModel {
             error = this.messages.pop()?.error
         }
         this.messages.push({
+            model: this.modelID,
             ...message,
             speaker: 'assistant',
             error,
@@ -80,6 +81,7 @@ export class SimpleChatModel {
             lastMessage?.speaker === 'assistant' ? this.messages.pop() : undefined
         // Then add a new assistant message with error added
         this.messages.push({
+            model: this.modelID,
             ...(lastAssistantMessage ?? {}),
             speaker: 'assistant',
             error: errorToChatError(error),
@@ -162,10 +164,12 @@ export class SimpleChatModel {
         }
         const result: SerializedChatTranscript = {
             id: this.sessionID,
-            chatModel: this.modelID,
             chatTitle: this.getCustomChatTitle(),
             lastInteractionTimestamp: this.sessionID,
             interactions,
+
+            // DEPRECATED: Write this for backcompat.
+            chatModel: this.modelID,
         }
         if (this.selectedRepos) {
             result.enhancedContext = {

vscode/src/completions/inline-completion-item-provider.test.ts

@@ -251,6 +251,7 @@ describe('InlineCompletionItemProvider', () => {
             get: (key: string) => localStorageData[key],
             update: (key: string, value: unknown) => {
                 localStorageData[key] = value
+                return Promise.resolve()
             },
         } as any as vscode.Memento)
 

vscode/src/services/LocalStorageProvider.test.ts

@@ -1,7 +1,7 @@
 import { beforeEach, describe, expect, it } from 'vitest'
 import type * as vscode from 'vscode'
 
-import type { AuthStatus } from '@sourcegraph/cody-shared'
+import type { AuthStatus, UserLocalHistory } from '@sourcegraph/cody-shared'
 
 import { localStorage } from './LocalStorageProvider'
 
@@ -22,12 +22,47 @@ describe('LocalStorageProvider', () => {
 
     it('sets and gets chat history', async () => {
         await localStorage.setChatHistory(DUMMY_AUTH_STATUS, {
-            chat: { a: null as any },
+            chat: { a: { id: 'a', lastInteractionTimestamp: '123', interactions: [] } },
         })
 
         const loadedHistory = localStorage.getChatHistory(DUMMY_AUTH_STATUS)
-        expect(loadedHistory).toEqual({
-            chat: { a: null as any },
+        expect(loadedHistory).toEqual<UserLocalHistory>({
+            chat: { a: { id: 'a', lastInteractionTimestamp: '123', interactions: [] } },
+        })
+    })
+
+    it('converts chat history to use ChatMessage.model not just SerializedChatTranscript.chatModel', async () => {
+        const chatHistory: UserLocalHistory = {
+            chat: {
+                a: {
+                    id: 'a',
+                    lastInteractionTimestamp: '123',
+                    chatModel: 'my-model',
+                    interactions: [
+                        {
+                            humanMessage: { speaker: 'human', text: 'hello' },
+                            assistantMessage: { speaker: 'assistant', text: 'hi' },
+                        },
+                    ],
+                },
+            },
+        }
+        await localStorage.setChatHistory(DUMMY_AUTH_STATUS, chatHistory)
+        const loadedHistory = localStorage.getChatHistory(DUMMY_AUTH_STATUS)
+        expect(loadedHistory).toEqual<UserLocalHistory>({
+            chat: {
+                a: {
+                    id: 'a',
+                    lastInteractionTimestamp: '123',
+                    chatModel: 'my-model',
+                    interactions: [
+                        {
+                            humanMessage: { speaker: 'human', text: 'hello' },
+                            assistantMessage: { speaker: 'assistant', text: 'hi', model: 'my-model' },
+                        },
+                    ],
+                },
+            },
         })
     })
 })

vscode/src/services/LocalStorageProvider.ts

@@ -1,7 +1,14 @@
 import * as uuid from 'uuid'
 import type { Memento } from 'vscode'
 
-import type { AuthStatus, ChatHistory, UserLocalHistory } from '@sourcegraph/cody-shared'
+import {
+    type AuthStatus,
+    type ChatHistory,
+    type SerializedChatInteraction,
+    type SerializedChatMessage,
+    type UserLocalHistory,
+    logDebug,
+} from '@sourcegraph/cody-shared'
 
 import { isSourcegraphToken } from '../chat/protocol'
 
@@ -93,7 +100,16 @@ class LocalStorage {
     public getChatHistory(authStatus: AuthStatus): UserLocalHistory {
         const history = this.storage.get<AccountKeyedChatHistory | null>(this.KEY_LOCAL_HISTORY, null)
         const accountKey = getKeyForAuthStatus(authStatus)
-        return history?.[accountKey] ?? { chat: {} }
+
+        // Migrate chat history to set the `ChatMessage.model` property on each assistant message
+        // instead of `chatModel` on the overall transcript. Can remove when
+        // `SerializedChatTranscript.chatModel` property is removed in v1.22.
+        const migratedHistory = migrateHistoryForChatModelProperty(history)
+        if (history !== migratedHistory) {
+            this.storage.update(this.KEY_LOCAL_HISTORY, migratedHistory).then(() => {}, console.error)
+        }
+
+        return migratedHistory?.[accountKey] ?? { chat: {} }
     }
 
     public async setChatHistory(authStatus: AuthStatus, history: UserLocalHistory): Promise<void> {
@@ -203,3 +219,68 @@ export const localStorage = new LocalStorage()
 function getKeyForAuthStatus(authStatus: AuthStatus): ChatHistoryKey {
     return `${authStatus.endpoint}-${authStatus.username}`
 }
+
+/**
+ * Migrate chat history to set the {@link ChatMessage.model} property on each assistant message
+ * instead of {@link SerializedChatTranscript.chatModel} on the overall transcript. Can remove when
+ * {@link SerializedChatTranscript.chatModel} property is removed in v1.22.
+ */
+function migrateHistoryForChatModelProperty(
+    history: AccountKeyedChatHistory | null
+): AccountKeyedChatHistory | null {
+    if (!history) {
+        return null
+    }
+
+    let neededMigration = 0
+    function migrateAssistantMessage(
+        assistantMessage: SerializedChatMessage,
+        model: string | undefined
+    ): SerializedChatMessage {
+        if (assistantMessage.model) {
+            return assistantMessage
+        }
+        neededMigration++
+        return {
+            ...assistantMessage,
+            model: model ?? 'unknown',
+        }
+    }
+
+    const migratedHistory = Object.fromEntries(
+        Object.entries(history).map(([accountKey, userLocalHistory]) => [
+            accountKey,
+            {
+                chat: userLocalHistory.chat
+                    ? Object.fromEntries(
+                          Object.entries(userLocalHistory.chat).map(([id, transcript]) => [
+                              id,
+                              transcript
+                                  ? {
+                                        ...transcript,
+                                        interactions: transcript.interactions.map(
+                                            interaction =>
+                                                ({
+                                                    ...interaction,
+                                                    assistantMessage: interaction.assistantMessage
+                                                        ? migrateAssistantMessage(
+                                                              interaction.assistantMessage,
+                                                              transcript.chatModel
+                                                          )
+                                                        : null,
+                                                }) satisfies SerializedChatInteraction
+                                        ),
+                                    }
+                                  : transcript,
+                          ])
+                      )
+                    : {},
+            },
+        ])
+    )
+    if (neededMigration) {
+        logDebug('migrateHistoryForChatModelProperty', `${neededMigration} chat messages migrated`)
+        return migratedHistory
+    }
+    return history
+}

vscode/test/e2e/chat-input.test.ts

@@ -187,19 +187,30 @@ test.extend<DotcomUrlOverride>({ dotcomUrl: mockServer.SERVER_URL }).extend<Expe
         'cody.chatResponse:noCode',
     ],
 })('chat model selector', async ({ page, sidebar }) => {
+    await fetch(`${mockServer.SERVER_URL}/.test/currentUser/codyProEnabled`, { method: 'POST' })
+
     await sidebarSignin(page, sidebar)
 
     const [chatFrame, chatInput] = await createEmptyChatPanel(page)
 
     const modelSelect = chatFrame.getByRole('combobox', { name: 'Choose a model' })
 
-    // Model selector is initially enabled.
     await expect(modelSelect).toBeEnabled()
+    await expect(modelSelect).toHaveText(/^Claude 3 Sonnet/)
+
+    await chatInput.fill('to model1')
+    await chatInput.press('Enter')
+    await expect(chatFrame.getByRole('row').getByTitle('Claude 3 Sonnet by Anthropic')).toBeVisible()
 
-    // Immediately after submitting the first message, the model selector is disabled.
-    await chatInput.fill('Hello')
+    // Change model and send another message.
+    await expect(modelSelect).toBeEnabled()
+    await modelSelect.click()
+    const modelChoices = chatFrame.getByRole('listbox', { name: 'Suggestions' })
+    await modelChoices.getByRole('option', { name: 'GPT-4o' }).click()
+    await expect(modelSelect).toHaveText(/^GPT-4o/)
+    await chatInput.fill('to model2')
     await chatInput.press('Enter')
-    await expect(modelSelect).toBeDisabled()
+    await expect(chatFrame.getByRole('row').getByTitle('GPT-4o by OpenAI')).toBeVisible()
 })
 
 test.extend<ExpectedEvents>({