Skip to content

Commit 56c872d

Browse files
committed
feat: add custom model support for Vertex AI provider
- Add custom model input field to Vertex provider settings UI - Update provider settings schema to include vertexCustomModelId field - Modify useSelectedModel hook to prioritize custom models for Vertex - Update AnthropicVertexHandler to handle custom model names with thinking suffix support - Add comprehensive tests for custom model functionality - Add i18n translations for custom model UI elements Fixes #5751: Users can now input custom Vertex AI model IDs in the correct @ format (e.g., claude-sonnet-4@20250514) to avoid 404 errors
1 parent d0452d0 commit 56c872d

File tree

8 files changed

+354
-445
lines changed

8 files changed

+354
-445
lines changed

packages/types/src/provider-settings.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,7 @@ const vertexSchema = apiModelIdProviderModelSchema.extend({
116116
vertexJsonCredentials: z.string().optional(),
117117
vertexProjectId: z.string().optional(),
118118
vertexRegion: z.string().optional(),
119+
vertexCustomModelId: z.string().optional(),
119120
})
120121

121122
const openAiSchema = baseProviderSettingsSchema.extend({

src/api/providers/__tests__/anthropic-vertex.spec.ts

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -809,4 +809,124 @@ describe("VertexHandler", () => {
809809
)
810810
})
811811
})
812+
813+
describe("custom model handling", () => {
814+
it("should use custom model when vertexCustomModelId is provided", () => {
815+
handler = new AnthropicVertexHandler({
816+
apiModelId: "claude-3-5-sonnet-v2@20241022",
817+
vertexCustomModelId: "claude-sonnet-4@20250514",
818+
vertexProjectId: "test-project",
819+
vertexRegion: "us-central1",
820+
})
821+
822+
const modelInfo = handler.getModel()
823+
expect(modelInfo.id).toBe("claude-sonnet-4@20250514")
824+
// Should use default model info as fallback
825+
expect(modelInfo.info).toBeDefined()
826+
expect(modelInfo.info.maxTokens).toBe(8192)
827+
expect(modelInfo.info.contextWindow).toBe(200_000)
828+
})
829+
830+
it("should trim whitespace from custom model ID", () => {
831+
handler = new AnthropicVertexHandler({
832+
apiModelId: "claude-3-5-sonnet-v2@20241022",
833+
vertexCustomModelId: " claude-sonnet-4@20250514 ",
834+
vertexProjectId: "test-project",
835+
vertexRegion: "us-central1",
836+
})
837+
838+
const modelInfo = handler.getModel()
839+
expect(modelInfo.id).toBe("claude-sonnet-4@20250514")
840+
})
841+
842+
it("should handle custom model with thinking suffix", () => {
843+
handler = new AnthropicVertexHandler({
844+
apiModelId: "claude-3-5-sonnet-v2@20241022",
845+
vertexCustomModelId: "claude-sonnet-4@20250514:thinking",
846+
vertexProjectId: "test-project",
847+
vertexRegion: "us-central1",
848+
modelMaxTokens: 16384,
849+
modelMaxThinkingTokens: 4096,
850+
})
851+
852+
const modelInfo = handler.getModel()
853+
expect(modelInfo.id).toBe("claude-sonnet-4@20250514")
854+
// For custom models with thinking suffix, reasoning parameters should be set
855+
expect(modelInfo.reasoningBudget).toBe(4096)
856+
expect(modelInfo.temperature).toBe(1.0)
857+
})
858+
859+
it("should fall back to predefined model when custom model is empty", () => {
860+
handler = new AnthropicVertexHandler({
861+
apiModelId: "claude-3-5-sonnet-v2@20241022",
862+
vertexCustomModelId: "",
863+
vertexProjectId: "test-project",
864+
vertexRegion: "us-central1",
865+
})
866+
867+
const modelInfo = handler.getModel()
868+
expect(modelInfo.id).toBe("claude-3-5-sonnet-v2@20241022")
869+
})
870+
871+
it("should fall back to predefined model when custom model is only whitespace", () => {
872+
handler = new AnthropicVertexHandler({
873+
apiModelId: "claude-3-5-sonnet-v2@20241022",
874+
vertexCustomModelId: " ",
875+
vertexProjectId: "test-project",
876+
vertexRegion: "us-central1",
877+
})
878+
879+
const modelInfo = handler.getModel()
880+
expect(modelInfo.id).toBe("claude-3-5-sonnet-v2@20241022")
881+
})
882+
883+
it("should fall back to default model when custom model is not provided", () => {
884+
handler = new AnthropicVertexHandler({
885+
vertexProjectId: "test-project",
886+
vertexRegion: "us-central1",
887+
})
888+
889+
const modelInfo = handler.getModel()
890+
expect(modelInfo.id).toBe("claude-sonnet-4@20250514") // default model
891+
})
892+
893+
it("should use custom model in API calls", async () => {
894+
handler = new AnthropicVertexHandler({
895+
apiModelId: "claude-3-5-sonnet-v2@20241022",
896+
vertexCustomModelId: "claude-sonnet-4@20250514",
897+
vertexProjectId: "test-project",
898+
vertexRegion: "us-central1",
899+
})
900+
901+
const mockCreate = vitest.fn().mockImplementation(async (options) => {
902+
return {
903+
async *[Symbol.asyncIterator]() {
904+
yield {
905+
type: "message_start",
906+
message: {
907+
usage: {
908+
input_tokens: 10,
909+
output_tokens: 5,
910+
},
911+
},
912+
}
913+
},
914+
}
915+
})
916+
;(handler["client"].messages as any).create = mockCreate
917+
918+
const stream = handler.createMessage("You are a helpful assistant", [{ role: "user", content: "Hello" }])
919+
920+
// Consume the stream
921+
for await (const _chunk of stream) {
922+
// Just consume the stream
923+
}
924+
925+
expect(mockCreate).toHaveBeenCalledWith(
926+
expect.objectContaining({
927+
model: "claude-sonnet-4@20250514",
928+
}),
929+
)
930+
})
931+
})
812932
})

src/api/providers/anthropic-vertex.ts

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,45 @@ export class AnthropicVertexHandler extends BaseProvider implements SingleComple
163163
}
164164

165165
getModel() {
166+
// Check if a custom model is specified
167+
const customModelId = this.options.vertexCustomModelId
168+
if (customModelId && customModelId.trim()) {
169+
// For custom models, use default model info as fallback
170+
const defaultInfo: ModelInfo = vertexModels[vertexDefaultModelId]
171+
const trimmedId = customModelId.trim()
172+
173+
// Check if custom model has thinking suffix
174+
const hasThinkingSuffix = trimmedId.endsWith(":thinking")
175+
const actualModelId = hasThinkingSuffix ? trimmedId.replace(":thinking", "") : trimmedId
176+
177+
// For thinking models, create a model info that supports reasoning
178+
let modelInfo: ModelInfo = defaultInfo
179+
if (hasThinkingSuffix) {
180+
modelInfo = {
181+
...defaultInfo,
182+
supportsReasoningBudget: true,
183+
requiredReasoningBudget: true,
184+
maxThinkingTokens: defaultInfo.maxThinkingTokens || 8192,
185+
}
186+
}
187+
188+
// Use the full model ID (with :thinking suffix) for getModelParams to get proper reasoning parameters
189+
const modelIdForParams = hasThinkingSuffix ? trimmedId : actualModelId
190+
const params = getModelParams({
191+
format: "anthropic",
192+
modelId: modelIdForParams,
193+
model: modelInfo,
194+
settings: this.options,
195+
})
196+
197+
return {
198+
id: actualModelId,
199+
info: modelInfo,
200+
...params,
201+
}
202+
}
203+
204+
// Use predefined models
166205
const modelId = this.options.apiModelId
167206
let id = modelId && modelId in vertexModels ? (modelId as VertexModelId) : vertexDefaultModelId
168207
const info: ModelInfo = vertexModels[id]

webview-ui/src/components/settings/providers/Vertex.tsx

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,16 @@ export const Vertex = ({ apiConfiguration, setApiConfigurationField }: VertexPro
9191
</SelectContent>
9292
</Select>
9393
</div>
94+
<VSCodeTextField
95+
value={apiConfiguration?.vertexCustomModelId || ""}
96+
onInput={handleInputChange("vertexCustomModelId")}
97+
placeholder="claude-sonnet-4@20250514"
98+
className="w-full">
99+
<label className="block font-medium mb-1">{t("settings:providers.customModel")}</label>
100+
</VSCodeTextField>
101+
<div className="text-sm text-vscode-descriptionForeground">
102+
{t("settings:providers.vertex.customModelDescription")}
103+
</div>
94104
</>
95105
)
96106
}

0 commit comments

Comments
 (0)