Chat: Assign ModelUIGroup for Enterprise model (#4607)

sourcegraph · Jun 19, 2024 · 14a9d19 · 14a9d19
1 parent 910afd6
commit 14a9d19
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 2 deletions.
diff --git a/lib/shared/src/models/utils.ts b/lib/shared/src/models/utils.ts
@@ -33,4 +33,6 @@ export const ModelUIGroup: Record<string, string> = {
     Balanced: 'Balanced (Speed & Accuracy)',
     Speed: 'Optimized for Speed',
     Ollama: 'Ollama (Local)',
+    // Used for identifying the model type but currently not used for displaying in the UI.
+    Enterprise: 'Enterprise',
 }
diff --git a/vscode/src/models/sync.test.ts b/vscode/src/models/sync.test.ts
@@ -1,5 +1,6 @@
 import {
     Model,
+    ModelUIGroup,
     ModelUsage,
     ModelsService,
     RestClient,
@@ -65,7 +66,9 @@ describe('syncModels', () => {
             new Model(
                 authStatus.configOverwrites.chatModel,
                 [ModelUsage.Chat, ModelUsage.Edit],
-                getEnterpriseContextWindow(chatModel, authStatus.configOverwrites)
+                getEnterpriseContextWindow(chatModel, authStatus.configOverwrites),
+                undefined,
+                ModelUIGroup.Enterprise
             ),
         ])
     })

diff --git a/vscode/src/models/sync.ts b/vscode/src/models/sync.ts
@@ -3,6 +3,7 @@ import {
     type AuthStatus,
     CHAT_INPUT_TOKEN_BUDGET,
     Model,
+    ModelUIGroup,
     ModelUsage,
     ModelsService,
     RestClient,
@@ -65,7 +66,9 @@ export async function syncModels(authStatus: AuthStatus): Promise<void> {
                 getEnterpriseContextWindow(
                     authStatus?.configOverwrites?.chatModel,
                     authStatus?.configOverwrites
-                )
+                ),
+                undefined,
+                ModelUIGroup.Enterprise
             ),
         ])
     } else {