Skip to content

Commit b6f25c0

Browse files
committed
fix: Improve model selection fallback logic when user's preferred model is unavailable (aws#2089)
* fix: if user preferred model does not exist, fall back to default model * fix: minor test changes
1 parent 1fac0ad commit b6f25c0

File tree

3 files changed

+22
-17
lines changed

3 files changed

+22
-17
lines changed

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3234,7 +3234,7 @@ ${' '.repeat(8)}}
32343234

32353235
const result = await chatController.onListAvailableModels({ tabId: mockTabId })
32363236

3237-
assert.strictEqual(result.selectedModelId, 'claude-4-sonnet') // MODEL_RECORD[DEFAULT_MODEL_ID].label
3237+
assert.strictEqual(result.selectedModelId, 'claude-4-sonnet') // FALLBACK_MODEL_RECORD[DEFAULT_MODEL_ID].label
32383238
// Verify session modelId is updated
32393239
assert.strictEqual(session.modelId, 'claude-4-sonnet')
32403240
})

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -220,7 +220,7 @@ import {
220220
Message as DbMessage,
221221
messageToStreamingMessage,
222222
} from './tools/chatDb/util'
223-
import { FALLBACK_MODEL_OPTIONS, MODEL_RECORD } from './constants/modelSelection'
223+
import { FALLBACK_MODEL_OPTIONS, FALLBACK_MODEL_RECORD } from './constants/modelSelection'
224224
import { DEFAULT_IMAGE_VERIFICATION_OPTIONS, verifyServerImage } from '../../shared/imageVerification'
225225
import { sanitize } from '@aws/lsp-core/out/util/path'
226226
import { ActiveUserTracker } from '../../shared/activeUserTracker'
@@ -755,7 +755,7 @@ export class AgenticChatController implements ChatHandlers {
755755

756756
// Get the first fallback model option as default
757757
const defaultModelOption = FALLBACK_MODEL_OPTIONS[1]
758-
const DEFAULT_MODEL_ID = defaultModelOption?.id || Object.keys(MODEL_RECORD)[1]
758+
const DEFAULT_MODEL_ID = defaultModelOption?.id || Object.keys(FALLBACK_MODEL_RECORD)[1]
759759

760760
const sessionResult = this.#chatSessionManagementService.getSession(params.tabId)
761761
const { data: session, success } = sessionResult
@@ -773,22 +773,27 @@ export class AgenticChatController implements ChatHandlers {
773773
let selectedModelId: string
774774
let modelId = this.#chatHistoryDb.getModelId()
775775

776+
// Helper function to get model label from FALLBACK_MODEL_RECORD
777+
const getModelLabel = (modelKey: string) =>
778+
FALLBACK_MODEL_RECORD[modelKey as keyof typeof FALLBACK_MODEL_RECORD]?.label || modelKey
779+
780+
// Determine selected model ID based on priority
776781
if (modelId) {
777-
// Case 1: User has previously selected a model - use that model or its mapped version
778-
// Check if modelId is a key in MODEL_RECORD
779-
if (modelId in MODEL_RECORD) {
780-
// If it's a valid model key, use its mapped label
781-
selectedModelId = MODEL_RECORD[modelId as keyof typeof MODEL_RECORD]?.label || modelId
782-
} else {
783-
// Otherwise use as is (might be a direct model ID from the API)
782+
// Priority 1: Use modelId if it exists in available models from backend
783+
if (models.some(model => model.id === modelId)) {
784784
selectedModelId = modelId
785785
}
786-
} else if (defaultModelId) {
787-
// Case 2: No user selection, but API provided a default model - use server recommendation
788-
selectedModelId = defaultModelId
786+
// Priority 2: Use mapped version if modelId exists in FALLBACK_MODEL_RECORD
787+
else if (modelId in FALLBACK_MODEL_RECORD) {
788+
selectedModelId = getModelLabel(modelId)
789+
}
790+
// Priority 3: Fall back to default or system default
791+
else {
792+
selectedModelId = defaultModelId || getModelLabel(DEFAULT_MODEL_ID)
793+
}
789794
} else {
790-
// Case 3: Last resort - use default model's label
791-
selectedModelId = MODEL_RECORD[DEFAULT_MODEL_ID as keyof typeof MODEL_RECORD]?.label || DEFAULT_MODEL_ID
795+
// No user-selected model - use API default or system default
796+
selectedModelId = defaultModelId || getModelLabel(DEFAULT_MODEL_ID)
792797
}
793798

794799
// Store the selected model in the session

server/aws-lsp-codewhisperer/src/language-server/agenticChat/constants/modelSelection.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,12 @@ type ModelDetails = {
1212
label: string
1313
}
1414

15-
export const MODEL_RECORD: Record<BedrockModel, ModelDetails> = {
15+
export const FALLBACK_MODEL_RECORD: Record<BedrockModel, ModelDetails> = {
1616
[BedrockModel.CLAUDE_3_7_SONNET_20250219_V1_0]: { label: 'claude-3.7-sonnet' },
1717
[BedrockModel.CLAUDE_SONNET_4_20250514_V1_0]: { label: 'claude-4-sonnet' },
1818
}
1919

20-
export const FALLBACK_MODEL_OPTIONS: ListAvailableModelsResult['models'] = Object.entries(MODEL_RECORD).map(
20+
export const FALLBACK_MODEL_OPTIONS: ListAvailableModelsResult['models'] = Object.entries(FALLBACK_MODEL_RECORD).map(
2121
([value, { label }]) => ({
2222
id: value,
2323
name: label,

0 commit comments

Comments
 (0)