Chat: waitlist for OpenAI-o1 & OpenAI-o1 mini (#5508)

CLOSE https://linear.app/sourcegraph/issue/CODY-3681/improve-support-for-previewexperimental-models Related to sourcegraph/sourcegraph#323 - Added `OpenAI o1` and `OpenAI o1 mini` to dropdown list where users can join the waitlist - Clicking on `Join Waitlist` will open the blog post, fire a telemetry event called `cody.joinLlmWaitlist` and update the `Join Waitlist` label to `On Waitlist` - Added a new `ModelTag.StreamDisabled` tag to indicate models that do not support streaming - Updated the `AssistantMessageCell` and `HumanMessageEditor` components to handle models without streaming support - Displayed a message to the user when a non-streaming model is used - Filtered out the initial codebase context when using a non-streaming model to avoid longer processing times - Display `Try again with a different model` instead of `Try again with different context` below each assistant response - Updated the `ModelSelectField` component to display an "Early Access" badge for models with the `ModelTag.Preview` tag - Expanded the `ModelRef` interface and `ModelsService` to better handle model ID matching - Adds the ability to duplicate the current chat session in the Cody chat interface. - When the "Try again with a different model" button is clicked, a new chat session is created that is a copy of the current session, allowing the user to resubmit their request with a different model. - The changes include: - Adding a new `chatSession` command to the webview message protocol, with `duplicate` and `new` actions - Implementing the `duplicateSession` method in the `ChatController` to create a new session with a unique ID based on the current session - Updating the `ContextFocusActions` component to use the new `chatSession` command when the "Try again with a different model" button is clicked - Models with the `ModelTag.StreamDisabled` tag will no longer display the initial codebase context to avoid longer processing times. - Updated the current `OpenAI o1` and `OpenAI o1 mini` models to use the stable version instead of `-latest` versions - Fix the model provider logo not matching issue ## Test plan  ![image](https://github.com/user-attachments/assets/172810a8-2a44-4062-a808-81b7476460e5) Build from this branch and verify the following in your debug modes: 1. `OpenAI o1` & `OpenAI o1 mini` are added to the model dropdown list 2. These models have the `Join Waitlist` label 3. Clicking on one of the model will open a webpage in browser 4. The `Join Waitlist` label now turned to `On Waitlist` ![image](https://github.com/user-attachments/assets/2f67f24f-c68c-4595-9696-7f65d78eb566) When using these models: 1. When submitting a question using one of these models, you will see `Model without streaming support takes longer to response.` when waiting for the LLM response 1. You should not find `Try again with different context` underneath the LLM response when using these models. 1. You should see `Try again with a different model` instead 1. Verify that when the "Try again with a different model" button is clicked, a new chat session is created that is a copy of the current session, and the user is able to resubmit their request with a different model. ## Changelog  feat(chat): add support for preview models `Gemini 1.5 Pro Latest` & `Gemini 1.5 Flash Latest` feat(chat): Added ability to duplicate chat sessions
sourcegraph · Sep 12, 2024 · 162162d · 162162d
1 parent 33ea5a7
commit 162162d
Show file tree

Hide file tree

Showing 22 changed files with 358 additions and 109 deletions.
diff --git a/...ngs/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Constants.kt b/...ngs/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/Constants.kt
@@ -28,6 +28,7 @@ object Constants {
   const val `delete-file` = "delete-file"
   const val deprecated = "deprecated"
   const val dev = "dev"
+  const val `early-access` = "early-access"
   const val edit = "edit"
   const val `edit-file` = "edit-file"
   const val editor = "editor"
@@ -57,6 +58,7 @@ object Constants {
   const val notification = "notification"
   const val `object-encoded` = "object-encoded"
   const val ollama = "ollama"
+  const val `on-waitlist` = "on-waitlist"
   const val openctx = "openctx"
   const val power = "power"
   const val pro = "pro"
@@ -72,6 +74,7 @@ object Constants {
   const val single = "single"
   const val speed = "speed"
   const val stateless = "stateless"
+  const val `stream-disabled` = "stream-disabled"
   const val streaming = "streaming"
   const val `string-encoded` = "string-encoded"
   const val suggestion = "suggestion"
@@ -83,6 +86,7 @@ object Constants {
   const val unified = "unified"
   const val use = "use"
   const val user = "user"
+  const val waitlist = "waitlist"
   const val warning = "warning"
   const val workspace = "workspace"
 }
diff --git a/...ings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/ModelTag.kt b/...ings/kotlin/lib/src/main/kotlin/com/sourcegraph/cody/agent/protocol_generated/ModelTag.kt
@@ -1,5 +1,5 @@
 @file:Suppress("FunctionName", "ClassName", "unused", "EnumEntryName", "UnusedImport")
 package com.sourcegraph.cody.agent.protocol_generated;
 
-typealias ModelTag = String // One of: power, speed, balanced, recommended, deprecated, experimental, pro, free, enterprise, gateway, byok, local, ollama, dev
+typealias ModelTag = String // One of: power, speed, balanced, recommended, deprecated, experimental, waitlist, on-waitlist, early-access, pro, free, enterprise, gateway, byok, local, ollama, dev, stream-disabled
 
diff --git a/agent/src/TestClient.ts b/agent/src/TestClient.ts
@@ -728,7 +728,7 @@ export class TestClient extends MessageHandler {
     public async reset(id: string): Promise<void> {
         await this.request('webview/receiveMessage', {
             id,
-            message: { command: 'reset' },
+            message: { command: 'chatSession', action: 'new' },
         })
     }
 

diff --git a/lib/shared/src/experimentation/FeatureFlagProvider.ts b/lib/shared/src/experimentation/FeatureFlagProvider.ts
@@ -74,6 +74,9 @@ export enum FeatureFlag {
 
     /** Enable experimental One Box feature in Cody */
     CodyExperimentalOneBox = 'cody-experimental-one-box',
+
+    /** Whether user has access to early-acess models. */
+    CodyEarlyAccess = 'cody-early-access',
 }
 
 const ONE_HOUR = 60 * 60 * 1000

diff --git a/lib/shared/src/models/dotcom.ts b/lib/shared/src/models/dotcom.ts
@@ -63,19 +63,51 @@ export const DEFAULT_DOT_COM_MODELS = [
     },
     {
         title: 'Gemini 1.5 Pro',
-        id: 'google/gemini-1.5-pro-latest',
+        id: 'google/gemini-1.5-pro',
         provider: 'Google',
         usage: [ModelUsage.Chat, ModelUsage.Edit],
         contextWindow: expandedContextWindow,
         tags: [ModelTag.Gateway, ModelTag.Power],
     },
 
+    // --------------------------------
+    // Preview / Early Access
+    // --------------------------------
+    {
+        title: 'OpenAI o1',
+        id: 'openai/cody-chat-preview-001',
+        provider: 'OpenAI',
+        usage: [ModelUsage.Chat, ModelUsage.Edit],
+        contextWindow: basicContextWindow,
+        tags: [
+            ModelTag.Gateway,
+            ModelTag.Power,
+            ModelTag.Pro,
+            ModelTag.Waitlist,
+            ModelTag.StreamDisabled,
+        ],
+    },
+    {
+        title: 'OpenAI o1-mini',
+        id: 'openai/cody-chat-preview-002',
+        provider: 'OpenAI',
+        usage: [ModelUsage.Chat, ModelUsage.Edit],
+        contextWindow: expandedContextWindow,
+        tags: [
+            ModelTag.Gateway,
+            ModelTag.Power,
+            ModelTag.Pro,
+            ModelTag.Waitlist,
+            ModelTag.StreamDisabled,
+        ],
+    },
+
     // --------------------------------
     // Faster models
     // --------------------------------
     {
         title: 'Gemini 1.5 Flash',
-        id: 'google/gemini-1.5-flash-latest',
+        id: 'google/gemini-1.5-flash',
         provider: 'Google',
         usage: [ModelUsage.Chat, ModelUsage.Edit],
         contextWindow: expandedContextWindow,

diff --git a/lib/shared/src/models/index.ts b/lib/shared/src/models/index.ts
@@ -21,7 +21,13 @@ interface ModelRef {
 }
 
 export type ModelCategory = ModelTag.Power | ModelTag.Balanced | ModelTag.Speed
-type ModelStatus = ModelTag.Experimental | ModelTag.Experimental | 'stable' | ModelTag.Deprecated
+type ModelStatus =
+    | ModelTag.Experimental
+    | ModelTag.EarlyAccess
+    | ModelTag.OnWaitlist
+    | ModelTag.Waitlist
+    | 'stable'
+    | ModelTag.Deprecated
 export type ModelTier = ModelTag.Free | ModelTag.Pro | ModelTag.Enterprise
 type ModelCapability = 'chat' | 'autocomplete'
 
@@ -536,7 +542,11 @@ export class ModelsService {
         // (But in reality, Sourcegraph.com wouldn't serve any Enterprise-only models to
         // Cody Pro users anyways.)
         if (isCodyProUser(status)) {
-            return tier !== 'enterprise'
+            return (
+                tier !== 'enterprise' &&
+                !resolved.tags.includes(ModelTag.Waitlist) &&
+                !resolved.tags.includes(ModelTag.OnWaitlist)
+            )
         }
 
         return tier === 'free'
@@ -553,7 +563,10 @@ export class ModelsService {
             return modelID
         }
 
-        return this.models.find(m => modelID.includes(m.id))
+        return (
+            this.models.find(m => modelID.endsWith(m.id)) ??
+            this.models.find(m => modelID.includes(m.id))
+        )
     }
 
     /**
@@ -582,6 +595,11 @@ export class ModelsService {
         const modelsList = this.models.map(m => m.id).join(', ')
         throw new Error(`${errorMessage} Available models: ${modelsList}`)
     }
+
+    public isStreamDisabled(modelID: string): boolean {
+        const model = this.getModelByID(modelID)
+        return model?.tags.includes(ModelTag.StreamDisabled) ?? false
+    }
 }
 
 export const modelsService = singletonNotYetSet<ModelsService>()

diff --git a/lib/shared/src/models/tags.ts b/lib/shared/src/models/tags.ts
@@ -13,6 +13,9 @@ export enum ModelTag {
     Recommended = 'recommended',
     Deprecated = 'deprecated',
     Experimental = 'experimental',
+    Waitlist = 'waitlist', // join waitlist
+    OnWaitlist = 'on-waitlist', // on waitlist
+    EarlyAccess = 'early-access',
 
     // Tiers - the level of access to the model
     Pro = 'pro',
@@ -25,4 +28,7 @@ export enum ModelTag {
     Local = 'local',
     Ollama = 'ollama',
     Dev = 'dev',
+
+    // Additional Info about the model
+    StreamDisabled = 'stream-disabled', // Model does not support streaming
 }
diff --git a/lib/shared/src/models/utils.ts b/lib/shared/src/models/utils.ts
@@ -34,6 +34,10 @@ export function isCodyProModel(model: Model): boolean {
     return modelHasTag(model, ModelTag.Pro)
 }
 
+export function isWaitlistModel(model: Model): boolean {
+    return modelHasTag(model, ModelTag.Waitlist) || modelHasTag(model, ModelTag.OnWaitlist)
+}
+
 export function isCustomModel(model: Model): boolean {
     return (
         modelHasTag(model, ModelTag.Local) ||

diff --git a/lib/shared/src/sourcegraph-api/completions/types.ts b/lib/shared/src/sourcegraph-api/completions/types.ts
@@ -36,6 +36,7 @@ export interface CompletionParameters {
     topK?: number
     topP?: number
     model?: string
+    stream?: boolean
 }
 
 export interface SerializedCompletionParameters extends Omit<CompletionParameters, 'messages'> {

diff --git a/vscode/CHANGELOG.md b/vscode/CHANGELOG.md
@@ -6,6 +6,9 @@ This is a log of all notable changes to Cody for VS Code. [Unreleased] changes a
 
 ### Added
 
+- The [new OpenAI models (OpenAI O1 & OpenAI O1-mini)](https://sourcegraph.com/blog/openai-o1-for-cody) are now available to selected Cody Pro users for early access. [pull/5508](https://github.com/sourcegraph/cody/pull/5508)
+- Cody Pro users can join the waitlist for the new models by clicking the "Join Waitlist" button. [pull/5508](https://github.com/sourcegraph/cody/pull/5508)
+
 ### Fixed
 
 Chat: Fixed feedback buttons not working in chat. [pull/5509](https://github.com/sourcegraph/cody/pull/5509)

diff --git a/vscode/src/chat/chat-view/ChatController.ts b/vscode/src/chat/chat-view/ChatController.ts
@@ -11,6 +11,7 @@ import {
     type ChatMessage,
     ClientConfigSingleton,
     CodyIDE,
+    type CompletionParameters,
     type ContextItem,
     type ContextItemOpenCtx,
     ContextItemSource,
@@ -87,6 +88,7 @@ import type { LocalEmbeddingsController } from '../../local-context/local-embedd
 import type { SymfRunner } from '../../local-context/symf'
 import { logDebug } from '../../log'
 import { migrateAndNotifyForOutdatedModels } from '../../models/modelMigrator'
+import { joinModelWaitlist } from '../../models/sync'
 import { mergedPromptsAndLegacyCommands } from '../../prompts/prompts'
 import { workspaceReposMonitor } from '../../repository/repo-metadata-from-git-api'
 import { authProvider } from '../../services/AuthProvider'
@@ -107,13 +109,14 @@ import {
 } from '../clientStateBroadcaster'
 import { getChatContextItemsForMention, getMentionMenuData } from '../context/chatContext'
 import type { ContextAPIClient } from '../context/contextAPIClient'
-import type {
-    ChatSubmitType,
-    ConfigurationSubsetForWebview,
-    ExtensionMessage,
-    LocalEnv,
-    SmartApplyResult,
-    WebviewMessage,
+import {
+    CODY_BLOG_URL_o1_WAITLIST,
+    type ChatSubmitType,
+    type ConfigurationSubsetForWebview,
+    type ExtensionMessage,
+    type LocalEnv,
+    type SmartApplyResult,
+    type WebviewMessage,
 } from '../protocol'
 import { countGeneratedCode } from '../utils'
 import { chatHistory } from './ChatHistoryManager'
@@ -355,9 +358,18 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
             case 'openURI':
                 vscode.commands.executeCommand('vscode.open', message.uri)
                 break
-            case 'links':
-                void openExternalLinks(message.value)
+            case 'links': {
+                let link = message.value
+                if (message.value === 'waitlist') {
+                    const authStatus = currentAuthStatusAuthed()
+                    const waitlistURI = CODY_BLOG_URL_o1_WAITLIST
+                    waitlistURI.searchParams.append('userId', authStatus?.username)
+                    link = waitlistURI.toString()
+                    void joinModelWaitlist(authStatus)
+                }
+                void openExternalLinks(link)
                 break
+            }
             case 'openFileLink':
                 vscode.commands.executeCommand('vscode.open', message.uri, {
                     selection: message.range,
@@ -395,8 +407,15 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
                 await this.restoreSession(message.chatID)
                 this.setWebviewToChat()
                 break
-            case 'reset':
-                await this.clearAndRestartSession()
+            case 'chatSession':
+                switch (message.action) {
+                    case 'new':
+                        await this.clearAndRestartSession()
+                        break
+                    case 'duplicate':
+                        await this.duplicateSession(message.sessionID ?? this.chatModel.sessionID)
+                        break
+                }
                 break
             case 'command':
                 vscode.commands.executeCommand(message.id, message.arg)
@@ -1000,6 +1019,7 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
             this.submitOrEditOperation.abort()
             this.submitOrEditOperation = undefined
         }
+        this.saveSession()
     }
 
     /**
@@ -1104,6 +1124,18 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
         }
     }
 
+    public async handleResubmitLastUserInput(): Promise<void> {
+        const lastHumanMessage = this.chatModel.getLastHumanMessage()
+        const getLastHumanMessageText = lastHumanMessage?.text?.toString()
+        if (getLastHumanMessageText) {
+            await this.clearAndRestartSession()
+            void this.postMessage({
+                type: 'clientAction',
+                appendTextToLastPromptEditor: getLastHumanMessageText,
+            })
+        }
+    }
+
     public async handleSmartApplyResult(result: SmartApplyResult): Promise<void> {
         void this.postMessage({
             type: 'clientAction',
@@ -1356,15 +1388,15 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
         })
 
         try {
-            const stream = this.chatClient.chat(
-                prompt,
-                {
-                    model: this.chatModel.modelID,
-                    maxTokensToSample: this.chatModel.contextWindow.output,
-                },
-                abortSignal
-            )
-
+            const params = {
+                model: this.chatModel.modelID,
+                maxTokensToSample: this.chatModel.contextWindow.output,
+            } as CompletionParameters
+            // Set stream param only when the model is disabled for streaming.
+            if (modelsService.instance!.isStreamDisabled(this.chatModel.modelID)) {
+                params.stream = false
+            }
+            const stream = this.chatClient.chat(prompt, params, abortSignal)
             for await (const message of stream) {
                 switch (message.type) {
                     case 'change': {
@@ -1470,6 +1502,28 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
         }
     }
 
+    private async duplicateSession(sessionID: string): Promise<void> {
+        this.cancelSubmitOrEditOperation()
+        const transcript = chatHistory.getChat(currentAuthStatusAuthed(), sessionID)
+        if (!transcript) {
+            return
+        }
+        // Assign a new session ID to the duplicated session
+        this.chatModel = newChatModelFromSerializedChatTranscript(
+            transcript,
+            this.chatModel.modelID,
+            new Date(Date.now()).toUTCString()
+        )
+        this.postViewTranscript()
+        await this.saveSession()
+        // Move the new session to the editor
+        await vscode.commands.executeCommand('cody.chat.moveToEditor')
+        // Restore the old session in the current window
+        await this.restoreSession(sessionID)
+
+        telemetryRecorder.recordEvent('cody.duplicateSession', 'clicked')
+    }
+
     public async clearAndRestartSession(): Promise<void> {
         this.cancelSubmitOrEditOperation()
         await this.saveSession()
@@ -1757,11 +1811,12 @@ export class ChatController implements vscode.Disposable, vscode.WebviewViewProv
 
 function newChatModelFromSerializedChatTranscript(
     json: SerializedChatTranscript,
-    modelID: string
+    modelID: string,
+    newSessionID?: string
 ): ChatModel {
     return new ChatModel(
         migrateAndNotifyForOutdatedModels(modelID)!,
-        json.id,
+        newSessionID ?? json.id,
         json.interactions.flatMap((interaction: SerializedChatInteraction): ChatMessage[] =>
             [
                 PromptString.unsafe_deserializeChatMessage(interaction.humanMessage),