Skip to content

Commit b30fc96

Browse files
Smartsheet-JB-Brownggoranov-smar
Smartsheet-JB-Brown
authored and
ggoranov-smar
committed
Improvements to AWS Bedrock embeddings support
- Enhanced bedrock.ts embedder implementation - Added comprehensive test coverage in bedrock.spec.ts - Updated config-manager.ts for better Bedrock configuration handling - Improved service-factory.ts integration - Updated embeddingModels.ts with Bedrock models - Enhanced CodeIndexPopover.tsx UI for Bedrock options - Added auto-populate test for CodeIndexPopover - Updated pnpm-lock.yaml dependencies
1 parent 840d95a commit b30fc96

File tree

8 files changed

+557
-28
lines changed

8 files changed

+557
-28
lines changed

‎pnpm-lock.yaml‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎src/services/code-index/config-manager.ts‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -143,10 +143,10 @@ export class CodeIndexConfigManager {
143143
this.geminiOptions = geminiApiKey ? { apiKey: geminiApiKey } : undefined
144144
this.mistralOptions = mistralApiKey ? { apiKey: mistralApiKey } : undefined
145145
this.vercelAiGatewayOptions = vercelAiGatewayApiKey ? { apiKey: vercelAiGatewayApiKey } : undefined
146-
// Set bedrockOptions only if both region and profile are provided
147-
this.bedrockOptions =
148-
bedrockRegion && bedrockProfile ? { region: bedrockRegion, profile: bedrockProfile } : undefined
149-
this.openRouterOptions = openRouterApiKey ? { apiKey: openRouterApiKey } : undefined
146+
// Set bedrockOptions if region is provided (profile is optional)
147+
this.bedrockOptions = bedrockRegion
148+
? { region: bedrockRegion, profile: bedrockProfile || undefined }
149+
: undefined
150150
}
151151

152152
/**
@@ -260,11 +260,10 @@ export class CodeIndexConfigManager {
260260
const isConfigured = !!(apiKey && qdrantUrl)
261261
return isConfigured
262262
} else if (this.embedderProvider === "bedrock") {
263-
// Both region and profile are required for Bedrock
263+
// Only region is required for Bedrock (profile is optional)
264264
const region = this.bedrockOptions?.region
265-
const profile = this.bedrockOptions?.profile
266265
const qdrantUrl = this.qdrantUrl
267-
const isConfigured = !!(region && profile && qdrantUrl)
266+
const isConfigured = !!(region && qdrantUrl)
268267
return isConfigured
269268
} else if (this.embedderProvider === "openrouter") {
270269
const apiKey = this.openRouterOptions?.apiKey

‎src/services/code-index/embedders/__tests__/bedrock.spec.ts‎

Lines changed: 136 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -88,12 +88,9 @@ describe("BedrockEmbedder", () => {
8888
expect(embedder.embedderInfo.name).toBe("bedrock")
8989
})
9090

91-
it("should require both region and profile", () => {
91+
it("should require region", () => {
9292
expect(() => new BedrockEmbedder("", "profile", "model")).toThrow(
93-
"Both region and profile are required for Amazon Bedrock embedder",
94-
)
95-
expect(() => new BedrockEmbedder("us-east-1", "", "model")).toThrow(
96-
"Both region and profile are required for Amazon Bedrock embedder",
93+
"Region is required for AWS Bedrock embedder",
9794
)
9895
})
9996

@@ -202,6 +199,140 @@ describe("BedrockEmbedder", () => {
202199
})
203200
})
204201

202+
it("should create embeddings with Nova multimodal model", async () => {
203+
const novaMultimodalEmbedder = new BedrockEmbedder(
204+
"us-east-1",
205+
"test-profile",
206+
"amazon.nova-2-multimodal-embeddings-v1:0",
207+
)
208+
const testTexts = ["Hello world"]
209+
const mockResponse = {
210+
body: new TextEncoder().encode(
211+
JSON.stringify({
212+
embeddings: [
213+
{
214+
embedding: [0.1, 0.2, 0.3],
215+
},
216+
],
217+
inputTextTokenCount: 2,
218+
}),
219+
),
220+
}
221+
mockSend.mockResolvedValue(mockResponse)
222+
223+
const result = await novaMultimodalEmbedder.createEmbeddings(testTexts)
224+
225+
expect(mockSend).toHaveBeenCalled()
226+
const command = mockSend.mock.calls[0][0] as any
227+
expect(command.input.modelId).toBe("amazon.nova-2-multimodal-embeddings-v1:0")
228+
const bodyStr =
229+
typeof command.input.body === "string"
230+
? command.input.body
231+
: new TextDecoder().decode(command.input.body as Uint8Array)
232+
// Nova multimodal embeddings use a task-based format with nested text object
233+
expect(JSON.parse(bodyStr || "{}")).toEqual({
234+
taskType: "SINGLE_EMBEDDING",
235+
singleEmbeddingParams: {
236+
embeddingPurpose: "GENERIC_INDEX",
237+
embeddingDimension: 1024,
238+
text: {
239+
truncationMode: "END",
240+
value: "Hello world",
241+
},
242+
},
243+
})
244+
245+
expect(result).toEqual({
246+
embeddings: [[0.1, 0.2, 0.3]],
247+
usage: { promptTokens: 2, totalTokens: 2 },
248+
})
249+
})
250+
251+
it("should handle Nova multimodal model with multiple texts", async () => {
252+
const novaMultimodalEmbedder = new BedrockEmbedder(
253+
"us-east-1",
254+
"test-profile",
255+
"amazon.nova-2-multimodal-embeddings-v1:0",
256+
)
257+
const testTexts = ["Hello world", "Another text"]
258+
const mockResponses = [
259+
{
260+
body: new TextEncoder().encode(
261+
JSON.stringify({
262+
embeddings: [
263+
{
264+
embedding: [0.1, 0.2, 0.3],
265+
},
266+
],
267+
inputTextTokenCount: 2,
268+
}),
269+
),
270+
},
271+
{
272+
body: new TextEncoder().encode(
273+
JSON.stringify({
274+
embeddings: [
275+
{
276+
embedding: [0.4, 0.5, 0.6],
277+
},
278+
],
279+
inputTextTokenCount: 3,
280+
}),
281+
),
282+
},
283+
]
284+
285+
mockSend.mockResolvedValueOnce(mockResponses[0]).mockResolvedValueOnce(mockResponses[1])
286+
287+
const result = await novaMultimodalEmbedder.createEmbeddings(testTexts)
288+
289+
expect(mockSend).toHaveBeenCalledTimes(2)
290+
291+
// Verify the request format for both texts
292+
const firstCommand = mockSend.mock.calls[0][0] as any
293+
const firstBodyStr =
294+
typeof firstCommand.input.body === "string"
295+
? firstCommand.input.body
296+
: new TextDecoder().decode(firstCommand.input.body as Uint8Array)
297+
// Nova multimodal embeddings use a task-based format with nested text object
298+
expect(JSON.parse(firstBodyStr || "{}")).toEqual({
299+
taskType: "SINGLE_EMBEDDING",
300+
singleEmbeddingParams: {
301+
embeddingPurpose: "GENERIC_INDEX",
302+
embeddingDimension: 1024,
303+
text: {
304+
truncationMode: "END",
305+
value: "Hello world",
306+
},
307+
},
308+
})
309+
310+
const secondCommand = mockSend.mock.calls[1][0] as any
311+
const secondBodyStr =
312+
typeof secondCommand.input.body === "string"
313+
? secondCommand.input.body
314+
: new TextDecoder().decode(secondCommand.input.body as Uint8Array)
315+
expect(JSON.parse(secondBodyStr || "{}")).toEqual({
316+
taskType: "SINGLE_EMBEDDING",
317+
singleEmbeddingParams: {
318+
embeddingPurpose: "GENERIC_INDEX",
319+
embeddingDimension: 1024,
320+
text: {
321+
truncationMode: "END",
322+
value: "Another text",
323+
},
324+
},
325+
})
326+
327+
expect(result).toEqual({
328+
embeddings: [
329+
[0.1, 0.2, 0.3],
330+
[0.4, 0.5, 0.6],
331+
],
332+
usage: { promptTokens: 5, totalTokens: 5 },
333+
})
334+
})
335+
205336
it("should use custom model when provided", async () => {
206337
const testTexts = ["Hello world"]
207338
const customModel = "amazon.titan-embed-text-v1"

‎src/services/code-index/embedders/bedrock.ts‎

Lines changed: 46 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -23,20 +23,21 @@ export class BedrockEmbedder implements IEmbedder {
2323
/**
2424
* Creates a new Amazon Bedrock embedder
2525
* @param region AWS region for Bedrock service (required)
26-
* @param profile AWS profile name for credentials (required)
26+
* @param profile AWS profile name for credentials (optional - uses default credential chain if not provided)
2727
* @param modelId Optional model ID override
2828
*/
2929
constructor(
3030
private readonly region: string,
31-
private readonly profile: string,
31+
private readonly profile?: string,
3232
modelId?: string,
3333
) {
34-
if (!region || !profile) {
35-
throw new Error("Both region and profile are required for Amazon Bedrock embedder")
34+
if (!region) {
35+
throw new Error("Region is required for AWS Bedrock embedder")
3636
}
3737

38-
// Initialize the Bedrock client with credentials from the specified profile
39-
const credentials = fromIni({ profile: this.profile })
38+
// Initialize the Bedrock client with credentials
39+
// If profile is specified, use it; otherwise use default credential chain
40+
const credentials = this.profile ? fromIni({ profile: this.profile }) : fromEnv()
4041

4142
this.bedrockClient = new BedrockRuntimeClient({
4243
region: this.region,
@@ -188,7 +189,22 @@ export class BedrockEmbedder implements IEmbedder {
188189
let modelId = model
189190

190191
// Prepare the request body based on the model
191-
if (model.startsWith("amazon.titan-embed")) {
192+
if (model.startsWith("amazon.nova-2-multimodal")) {
193+
// Nova multimodal embeddings use a task-based format with embeddingParams
194+
// Reference: https://docs.aws.amazon.com/bedrock/latest/userguide/embeddings-nova.html
195+
requestBody = {
196+
taskType: "SINGLE_EMBEDDING",
197+
singleEmbeddingParams: {
198+
embeddingPurpose: "GENERIC_INDEX",
199+
embeddingDimension: 1024, // Nova supports 1024 or 3072
200+
text: {
201+
truncationMode: "END",
202+
value: text,
203+
},
204+
},
205+
}
206+
console.log(`[BedrockEmbedder] Nova multimodal request for model ${model}:`, JSON.stringify(requestBody))
207+
} else if (model.startsWith("amazon.titan-embed")) {
192208
requestBody = {
193209
inputText: text,
194210
}
@@ -211,14 +227,35 @@ export class BedrockEmbedder implements IEmbedder {
211227
accept: "application/json",
212228
}
213229

230+
console.log(`[BedrockEmbedder] Sending request to model ${modelId}`)
231+
console.log(`[BedrockEmbedder] Request body:`, requestBody)
232+
214233
const command = new InvokeModelCommand(params)
215-
const response = await this.bedrockClient.send(command)
234+
235+
let response
236+
try {
237+
response = await this.bedrockClient.send(command)
238+
} catch (error: any) {
239+
console.error(`[BedrockEmbedder] API error for model ${modelId}:`, error)
240+
console.error(`[BedrockEmbedder] Error name:`, error.name)
241+
console.error(`[BedrockEmbedder] Error message:`, error.message)
242+
console.error(`[BedrockEmbedder] Error details:`, JSON.stringify(error, null, 2))
243+
throw error
244+
}
216245

217246
// Parse the response
218247
const responseBody = JSON.parse(new TextDecoder().decode(response.body))
248+
console.log(`[BedrockEmbedder] Response for model ${modelId}:`, responseBody)
219249

220250
// Extract embedding based on model type
221-
if (model.startsWith("amazon.titan-embed")) {
251+
if (model.startsWith("amazon.nova-2-multimodal")) {
252+
// Nova multimodal returns { embeddings: [{ embedding: [...] }] }
253+
// Reference: AWS Bedrock documentation
254+
return {
255+
embedding: responseBody.embeddings?.[0]?.embedding || responseBody.embedding,
256+
inputTextTokenCount: responseBody.inputTextTokenCount,
257+
}
258+
} else if (model.startsWith("amazon.titan-embed")) {
222259
return {
223260
embedding: responseBody.embedding,
224261
inputTextTokenCount: responseBody.inputTextTokenCount,

‎src/services/code-index/service-factory.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -82,8 +82,8 @@ export class CodeIndexServiceFactory {
8282
}
8383
return new VercelAiGatewayEmbedder(config.vercelAiGatewayOptions.apiKey, config.modelId)
8484
} else if (provider === "bedrock") {
85-
// Both region and profile are required for Bedrock
86-
if (!config.bedrockOptions?.region || !config.bedrockOptions?.profile) {
85+
// Only region is required for Bedrock (profile is optional)
86+
if (!config.bedrockOptions?.region) {
8787
throw new Error(t("embeddings:serviceFactory.bedrockConfigMissing"))
8888
}
8989
return new BedrockEmbedder(config.bedrockOptions.region, config.bedrockOptions.profile, config.modelId)

‎src/shared/embeddingModels.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,8 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = {
8383
"amazon.titan-embed-text-v1": { dimension: 1536, scoreThreshold: 0.4 },
8484
"amazon.titan-embed-text-v2:0": { dimension: 1024, scoreThreshold: 0.4 },
8585
"amazon.titan-embed-image-v1": { dimension: 1024, scoreThreshold: 0.4 },
86+
// Amazon Nova Embed models
87+
"amazon.nova-2-multimodal-embeddings-v1:0": { dimension: 1024, scoreThreshold: 0.4 },
8688
// Cohere models available through Bedrock
8789
"cohere.embed-english-v3": { dimension: 1024, scoreThreshold: 0.4 },
8890
"cohere.embed-multilingual-v3": { dimension: 1024, scoreThreshold: 0.4 },

‎webview-ui/src/components/chat/CodeIndexPopover.tsx‎

Lines changed: 32 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -158,9 +158,7 @@ const createValidationSchema = (provider: EmbedderProvider, t: any) => {
158158
case "bedrock":
159159
return baseSchema.extend({
160160
codebaseIndexBedrockRegion: z.string().min(1, t("settings:codeIndex.validation.bedrockRegionRequired")),
161-
codebaseIndexBedrockProfile: z
162-
.string()
163-
.min(1, t("settings:codeIndex.validation.bedrockProfileRequired")),
161+
codebaseIndexBedrockProfile: z.string().optional(),
164162
codebaseIndexEmbedderModelId: z
165163
.string()
166164
.min(1, t("settings:codeIndex.validation.modelSelectionRequired")),
@@ -187,7 +185,7 @@ export const CodeIndexPopover: React.FC<CodeIndexPopoverProps> = ({
187185
}) => {
188186
const SECRET_PLACEHOLDER = "••••••••••••••••"
189187
const { t } = useAppTranslation()
190-
const { codebaseIndexConfig, codebaseIndexModels, cwd } = useExtensionState()
188+
const { codebaseIndexConfig, codebaseIndexModels, cwd, apiConfiguration } = useExtensionState()
191189
const [open, setOpen] = useState(false)
192190
const [isAdvancedSettingsOpen, setIsAdvancedSettingsOpen] = useState(false)
193191
const [isSetupSettingsOpen, setIsSetupSettingsOpen] = useState(false)
@@ -689,6 +687,33 @@ export const CodeIndexPopover: React.FC<CodeIndexPopoverProps> = ({
689687
updateSetting("codebaseIndexEmbedderProvider", value)
690688
// Clear model selection when switching providers
691689
updateSetting("codebaseIndexEmbedderModelId", "")
690+
691+
// Auto-populate Region and Profile when switching to Bedrock
692+
// if the main API provider is also configured for Bedrock
693+
if (
694+
value === "bedrock" &&
695+
apiConfiguration?.apiProvider === "bedrock"
696+
) {
697+
// Only populate if currently empty
698+
if (
699+
!currentSettings.codebaseIndexBedrockRegion &&
700+
apiConfiguration.awsRegion
701+
) {
702+
updateSetting(
703+
"codebaseIndexBedrockRegion",
704+
apiConfiguration.awsRegion,
705+
)
706+
}
707+
if (
708+
!currentSettings.codebaseIndexBedrockProfile &&
709+
apiConfiguration.awsProfile
710+
) {
711+
updateSetting(
712+
"codebaseIndexBedrockProfile",
713+
apiConfiguration.awsProfile,
714+
)
715+
}
716+
}
692717
}}>
693718
<SelectTrigger className="w-full">
694719
<SelectValue />
@@ -1206,6 +1231,9 @@ export const CodeIndexPopover: React.FC<CodeIndexPopoverProps> = ({
12061231
<div className="space-y-2">
12071232
<label className="text-sm font-medium">
12081233
{t("settings:codeIndex.bedrockProfileLabel")}
1234+
<span className="text-xs text-vscode-descriptionForeground ml-1">
1235+
({t("settings:codeIndex.optional")})
1236+
</span>
12091237
</label>
12101238
<VSCodeTextField
12111239
value={currentSettings.codebaseIndexBedrockProfile || ""}

0 commit comments

Comments
 (0)