feat(ts-sdk): add AWS Bedrock embedding provider (#6185)
This commit is contained in:
Kartik
@@ -0,0 +1,517 @@
|
||||
import type { InvokeModelCommand } from "@aws-sdk/client-bedrock-runtime";
|
||||
import { AWSBedrockEmbedder } from "../src/embeddings/aws_bedrock";
|
||||
import type { Embedder } from "../src/embeddings/base";
|
||||
import { EmbedderFactory } from "../src/utils/factory";
|
||||
|
||||
/**
|
||||
* Only the network boundary is faked: `BedrockRuntimeClient.send` never leaves
|
||||
* the process. `InvokeModelCommand` stays the real class from the AWS SDK, so
|
||||
* every assertion below runs against the exact payload Bedrock would receive.
|
||||
*/
|
||||
const mockSend = jest.fn();
|
||||
const mockClientConfigs: any[] = [];
|
||||
const mockClientConstructor = jest
|
||||
.fn()
|
||||
.mockImplementation((config: unknown) => {
|
||||
mockClientConfigs.push(config);
|
||||
return { send: mockSend };
|
||||
});
|
||||
|
||||
jest.mock("@aws-sdk/client-bedrock-runtime", () => {
|
||||
const actual = jest.requireActual("@aws-sdk/client-bedrock-runtime");
|
||||
return {
|
||||
...actual,
|
||||
BedrockRuntimeClient: mockClientConstructor,
|
||||
};
|
||||
});
|
||||
|
||||
const encode = (payload: unknown) => ({
|
||||
body: new TextEncoder().encode(JSON.stringify(payload)),
|
||||
});
|
||||
|
||||
const titanReply = (embedding: number[]) =>
|
||||
encode({ embedding, inputTextTokenCount: embedding.length });
|
||||
|
||||
const cohereReply = (embeddings: number[][]) =>
|
||||
encode({ embeddings, id: "req-1", response_type: "embeddings_floats" });
|
||||
|
||||
const commandAt = (index: number): InvokeModelCommand =>
|
||||
mockSend.mock.calls[index][0];
|
||||
|
||||
const requestBodyAt = (index: number) =>
|
||||
JSON.parse(new TextDecoder().decode(commandAt(index).input.body));
|
||||
|
||||
describe("AWSBedrockEmbedder", () => {
|
||||
const savedEnv = { ...process.env };
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
mockClientConfigs.length = 0;
|
||||
delete process.env.AWS_REGION;
|
||||
});
|
||||
|
||||
afterAll(() => {
|
||||
process.env = savedEnv;
|
||||
});
|
||||
|
||||
describe("Titan models", () => {
|
||||
it("sends inputText and returns the embedding vector", async () => {
|
||||
mockSend.mockResolvedValueOnce(titanReply([0.1, 0.2, 0.3]));
|
||||
const embedder = new AWSBedrockEmbedder({});
|
||||
|
||||
const embedding = await embedder.embed("hello world");
|
||||
|
||||
expect(embedding).toEqual([0.1, 0.2, 0.3]);
|
||||
expect(mockSend).toHaveBeenCalledTimes(1);
|
||||
expect(commandAt(0).input.modelId).toBe("amazon.titan-embed-text-v1");
|
||||
expect(commandAt(0).input.contentType).toBe("application/json");
|
||||
expect(commandAt(0).input.accept).toBe("application/json");
|
||||
expect(requestBodyAt(0)).toEqual({ inputText: "hello world" });
|
||||
});
|
||||
|
||||
it("forwards dimensions to Titan V2 when embeddingDims is set", async () => {
|
||||
mockSend.mockResolvedValueOnce(titanReply([0.1, 0.2]));
|
||||
const embedder = new AWSBedrockEmbedder({
|
||||
model: "amazon.titan-embed-text-v2:0",
|
||||
embeddingDims: 512,
|
||||
});
|
||||
|
||||
await embedder.embed("hello");
|
||||
|
||||
expect(requestBodyAt(0)).toEqual({ inputText: "hello", dimensions: 512 });
|
||||
});
|
||||
|
||||
it("omits dimensions on Titan V1, which rejects the field", async () => {
|
||||
mockSend.mockResolvedValueOnce(titanReply([0.1, 0.2]));
|
||||
const embedder = new AWSBedrockEmbedder({
|
||||
model: "amazon.titan-embed-text-v1",
|
||||
embeddingDims: 512,
|
||||
});
|
||||
|
||||
await embedder.embed("hello");
|
||||
|
||||
expect(requestBodyAt(0)).toEqual({ inputText: "hello" });
|
||||
});
|
||||
|
||||
// F6: the old guard was `model.includes("v2")`, which would also match
|
||||
// any future/other Titan model whose id merely contains "v2" somewhere
|
||||
// (e.g. an image model), wrongly sending `dimensions` to a model that may
|
||||
// reject it. Only Titan Text Embeddings V2 should get the field.
|
||||
it("does not forward dimensions to a non-Titan-V2 model whose name merely contains v2", async () => {
|
||||
mockSend.mockResolvedValueOnce(titanReply([0.1, 0.2]));
|
||||
const embedder = new AWSBedrockEmbedder({
|
||||
model: "amazon.titan-embed-image-v2:0",
|
||||
embeddingDims: 512,
|
||||
});
|
||||
|
||||
await embedder.embed("hello");
|
||||
|
||||
expect(requestBodyAt(0)).toEqual({ inputText: "hello" });
|
||||
});
|
||||
|
||||
it("embedBatch issues one request per text and preserves order", async () => {
|
||||
mockSend
|
||||
.mockResolvedValueOnce(titanReply([1, 1]))
|
||||
.mockResolvedValueOnce(titanReply([2, 2]));
|
||||
const embedder = new AWSBedrockEmbedder({});
|
||||
|
||||
const embeddings = await embedder.embedBatch(["first", "second"]);
|
||||
|
||||
expect(embeddings).toEqual([
|
||||
[1, 1],
|
||||
[2, 2],
|
||||
]);
|
||||
expect(mockSend).toHaveBeenCalledTimes(2);
|
||||
expect(requestBodyAt(0)).toEqual({ inputText: "first" });
|
||||
expect(requestBodyAt(1)).toEqual({ inputText: "second" });
|
||||
});
|
||||
});
|
||||
|
||||
describe("Cohere models", () => {
|
||||
it("sends texts with a search_document input type", async () => {
|
||||
mockSend.mockResolvedValueOnce(cohereReply([[0.4, 0.5]]));
|
||||
const embedder = new AWSBedrockEmbedder({
|
||||
model: "cohere.embed-english-v3",
|
||||
});
|
||||
|
||||
const embedding = await embedder.embed("hello");
|
||||
|
||||
expect(embedding).toEqual([0.4, 0.5]);
|
||||
expect(requestBodyAt(0)).toEqual({
|
||||
texts: ["hello"],
|
||||
input_type: "search_document",
|
||||
});
|
||||
});
|
||||
|
||||
it("embedBatch sends every text in a single request", async () => {
|
||||
mockSend.mockResolvedValueOnce(cohereReply([[1], [2], [3]]));
|
||||
const embedder = new AWSBedrockEmbedder({
|
||||
model: "cohere.embed-multilingual-v3",
|
||||
});
|
||||
|
||||
const embeddings = await embedder.embedBatch(["a", "b", "c"]);
|
||||
|
||||
expect(embeddings).toEqual([[1], [2], [3]]);
|
||||
expect(mockSend).toHaveBeenCalledTimes(1);
|
||||
expect(requestBodyAt(0).texts).toEqual(["a", "b", "c"]);
|
||||
});
|
||||
|
||||
it("embedBatch splits requests at Cohere's 96 text limit", async () => {
|
||||
const texts = Array.from({ length: 100 }, (_, i) => `text-${i}`);
|
||||
mockSend
|
||||
.mockResolvedValueOnce(
|
||||
cohereReply(texts.slice(0, 96).map((_, i) => [i])),
|
||||
)
|
||||
.mockResolvedValueOnce(cohereReply(texts.slice(96).map((_, i) => [i])));
|
||||
const embedder = new AWSBedrockEmbedder({
|
||||
model: "cohere.embed-english-v3",
|
||||
});
|
||||
|
||||
const embeddings = await embedder.embedBatch(texts);
|
||||
|
||||
expect(embeddings).toHaveLength(100);
|
||||
expect(mockSend).toHaveBeenCalledTimes(2);
|
||||
expect(requestBodyAt(0).texts).toHaveLength(96);
|
||||
expect(requestBodyAt(1).texts).toEqual([
|
||||
"text-96",
|
||||
"text-97",
|
||||
"text-98",
|
||||
"text-99",
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("Cohere Embed v4", () => {
|
||||
// F5: only Embed v4 understands embedding_types / output_dimension, and
|
||||
// (when embedding_types is requested) replies with a nested
|
||||
// `{ embeddings: { float: [...] } }` shape instead of v3's flat array.
|
||||
it("requests embedding_types and output_dimension for v4 models", async () => {
|
||||
mockSend.mockResolvedValueOnce(
|
||||
encode({ embeddings: { float: [[0.1, 0.2, 0.3]] } }),
|
||||
);
|
||||
const embedder = new AWSBedrockEmbedder({
|
||||
model: "cohere.embed-v4:0",
|
||||
embeddingDims: 512,
|
||||
});
|
||||
|
||||
const embedding = await embedder.embed("hello");
|
||||
|
||||
expect(embedding).toEqual([0.1, 0.2, 0.3]);
|
||||
expect(requestBodyAt(0)).toEqual({
|
||||
texts: ["hello"],
|
||||
input_type: "search_document",
|
||||
embedding_types: ["float"],
|
||||
output_dimension: 512,
|
||||
});
|
||||
});
|
||||
|
||||
it("omits output_dimension for v4 when embeddingDims is unset", async () => {
|
||||
mockSend.mockResolvedValueOnce(
|
||||
encode({ embeddings: { float: [[0.1, 0.2]] } }),
|
||||
);
|
||||
const embedder = new AWSBedrockEmbedder({ model: "cohere.embed-v4:0" });
|
||||
|
||||
await embedder.embed("hello");
|
||||
|
||||
expect(requestBodyAt(0)).toEqual({
|
||||
texts: ["hello"],
|
||||
input_type: "search_document",
|
||||
embedding_types: ["float"],
|
||||
});
|
||||
});
|
||||
|
||||
it("parses the nested embeddings.float response shape", async () => {
|
||||
mockSend.mockResolvedValueOnce(
|
||||
encode({
|
||||
embeddings: {
|
||||
float: [
|
||||
[1, 2],
|
||||
[3, 4],
|
||||
],
|
||||
},
|
||||
}),
|
||||
);
|
||||
const embedder = new AWSBedrockEmbedder({ model: "cohere.embed-v4:0" });
|
||||
|
||||
const embeddings = await embedder.embedBatch(["a", "b"]);
|
||||
|
||||
expect(embeddings).toEqual([
|
||||
[1, 2],
|
||||
[3, 4],
|
||||
]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("client configuration", () => {
|
||||
it("defaults to the us-west-2 region", async () => {
|
||||
mockSend.mockResolvedValueOnce(titanReply([1]));
|
||||
await new AWSBedrockEmbedder({}).embed("hello");
|
||||
|
||||
expect(mockClientConfigs[0].region).toBe("us-west-2");
|
||||
});
|
||||
|
||||
it("prefers awsRegion over the AWS_REGION environment variable", async () => {
|
||||
process.env.AWS_REGION = "eu-central-1";
|
||||
mockSend.mockResolvedValueOnce(titanReply([1]));
|
||||
await new AWSBedrockEmbedder({ awsRegion: "ap-south-1" }).embed("hello");
|
||||
|
||||
expect(mockClientConfigs[0].region).toBe("ap-south-1");
|
||||
});
|
||||
|
||||
it("falls back to the AWS_REGION environment variable", async () => {
|
||||
process.env.AWS_REGION = "eu-central-1";
|
||||
mockSend.mockResolvedValueOnce(titanReply([1]));
|
||||
await new AWSBedrockEmbedder({}).embed("hello");
|
||||
|
||||
expect(mockClientConfigs[0].region).toBe("eu-central-1");
|
||||
});
|
||||
|
||||
it("passes explicitly configured credentials to the client", async () => {
|
||||
mockSend.mockResolvedValueOnce(titanReply([1]));
|
||||
await new AWSBedrockEmbedder({
|
||||
awsAccessKeyId: "AKIA_TEST",
|
||||
awsSecretAccessKey: "secret",
|
||||
awsSessionToken: "token",
|
||||
}).embed("hello");
|
||||
|
||||
expect(mockClientConfigs[0].credentials).toEqual({
|
||||
accessKeyId: "AKIA_TEST",
|
||||
secretAccessKey: "secret",
|
||||
sessionToken: "token",
|
||||
});
|
||||
});
|
||||
|
||||
it("leaves credentials unset so the AWS default credential chain applies", async () => {
|
||||
mockSend.mockResolvedValueOnce(titanReply([1]));
|
||||
await new AWSBedrockEmbedder({}).embed("hello");
|
||||
|
||||
expect(mockClientConfigs[0].credentials).toBeUndefined();
|
||||
});
|
||||
|
||||
it("rejects a half-configured credential pair", () => {
|
||||
expect(
|
||||
() => new AWSBedrockEmbedder({ awsAccessKeyId: "AKIA_TEST" }),
|
||||
).toThrow(/awsAccessKeyId and awsSecretAccessKey/);
|
||||
expect(
|
||||
() => new AWSBedrockEmbedder({ awsSecretAccessKey: "secret" }),
|
||||
).toThrow(/awsAccessKeyId and awsSecretAccessKey/);
|
||||
});
|
||||
|
||||
// Silently ignoring a lone session token would fall back to the ambient
|
||||
// credential chain, embedding under an identity the caller never chose.
|
||||
it("rejects a session token supplied without the key pair", () => {
|
||||
expect(
|
||||
() => new AWSBedrockEmbedder({ awsSessionToken: "token" }),
|
||||
).toThrow(/awsAccessKeyId and awsSecretAccessKey/);
|
||||
});
|
||||
});
|
||||
|
||||
describe("provider registration", () => {
|
||||
it("is constructed by EmbedderFactory for the aws_bedrock provider", () => {
|
||||
const embedder = EmbedderFactory.create("aws_bedrock", {
|
||||
model: "amazon.titan-embed-text-v2:0",
|
||||
});
|
||||
|
||||
expect(embedder).toBeInstanceOf(AWSBedrockEmbedder);
|
||||
});
|
||||
});
|
||||
|
||||
describe("error handling", () => {
|
||||
it("wraps Bedrock failures with the model id", async () => {
|
||||
mockSend.mockRejectedValueOnce(new Error("AccessDeniedException"));
|
||||
const embedder = new AWSBedrockEmbedder({});
|
||||
|
||||
await expect(embedder.embed("hello")).rejects.toThrow(
|
||||
"Error getting embedding from AWS Bedrock model amazon.titan-embed-text-v1: AccessDeniedException",
|
||||
);
|
||||
});
|
||||
|
||||
it("fails when the response carries no embedding", async () => {
|
||||
mockSend.mockResolvedValueOnce(encode({ inputTextTokenCount: 3 }));
|
||||
const embedder = new AWSBedrockEmbedder({});
|
||||
|
||||
await expect(embedder.embed("hello")).rejects.toThrow(
|
||||
/returned no embedding/,
|
||||
);
|
||||
});
|
||||
|
||||
// F7: `[]` is truthy, so `payload.embedding && [payload.embedding]` used to
|
||||
// turn `{"embedding": []}` into `[[]]` -- length 1, which satisfied the
|
||||
// length check for a single-input call and handed the caller an empty vector.
|
||||
it("rejects a zero-length Titan embedding as no embedding", async () => {
|
||||
mockSend.mockResolvedValueOnce(encode({ embedding: [] }));
|
||||
const embedder = new AWSBedrockEmbedder({});
|
||||
|
||||
await expect(embedder.embed("hello")).rejects.toThrow(
|
||||
/returned no embedding/,
|
||||
);
|
||||
});
|
||||
|
||||
it("returns an empty array for an empty batch without calling Bedrock", async () => {
|
||||
const embedder = new AWSBedrockEmbedder({});
|
||||
|
||||
await expect(embedder.embedBatch([])).resolves.toEqual([]);
|
||||
expect(mockSend).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe("dynamic SDK import", () => {
|
||||
// These tests replace the module registered for
|
||||
// @aws-sdk/client-bedrock-runtime for a single resolution. Restore the
|
||||
// working mock afterward so every other test in this file keeps getting
|
||||
// the mocked client instead of hitting module resolution for real.
|
||||
afterEach(() => {
|
||||
jest.resetModules();
|
||||
jest.doMock("@aws-sdk/client-bedrock-runtime", () => {
|
||||
const actual = jest.requireActual("@aws-sdk/client-bedrock-runtime");
|
||||
return { ...actual, BedrockRuntimeClient: mockClientConstructor };
|
||||
});
|
||||
});
|
||||
|
||||
// F1: loadSdk()'s catch used to rewrite *every* import failure into the
|
||||
// "package is required" hint, even when the package is installed but
|
||||
// failed to load for an unrelated reason. That discarded the real error.
|
||||
it("propagates a non-resolution import error unchanged", async () => {
|
||||
jest.resetModules();
|
||||
jest.doMock("@aws-sdk/client-bedrock-runtime", () => {
|
||||
const err: any = new Error("boom: unrelated crash while loading");
|
||||
err.code = "ERR_SOMETHING_ELSE";
|
||||
throw err;
|
||||
});
|
||||
const embedder = new AWSBedrockEmbedder({});
|
||||
|
||||
await expect(embedder.embed("hello")).rejects.toThrow(
|
||||
"boom: unrelated crash while loading",
|
||||
);
|
||||
});
|
||||
|
||||
// F1: a genuine resolution failure should still get the friendly install
|
||||
// hint, with the original error preserved as `cause` for debugging.
|
||||
it("gives an install hint for a genuine module-not-found error, preserving the cause", async () => {
|
||||
jest.resetModules();
|
||||
jest.doMock("@aws-sdk/client-bedrock-runtime", () => {
|
||||
const err: any = new Error(
|
||||
"Cannot find module '@aws-sdk/client-bedrock-runtime'",
|
||||
);
|
||||
err.code = "MODULE_NOT_FOUND";
|
||||
throw err;
|
||||
});
|
||||
const embedder = new AWSBedrockEmbedder({});
|
||||
|
||||
await expect(embedder.embed("hello")).rejects.toThrow(
|
||||
/npm install @aws-sdk\/client-bedrock-runtime/,
|
||||
);
|
||||
await expect(embedder.embed("hello")).rejects.toMatchObject({
|
||||
cause: expect.objectContaining({ code: "MODULE_NOT_FOUND" }),
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("client promise retry", () => {
|
||||
// F2: getClient() used to memoize the client promise before it resolved,
|
||||
// so a rejected construction (e.g. a transient credentials failure) was
|
||||
// cached forever -- every later embed() call on that instance would
|
||||
// reject immediately without ever retrying.
|
||||
it("retries client construction after a failure instead of caching the rejection", async () => {
|
||||
mockClientConstructor.mockImplementationOnce(() => {
|
||||
throw new Error("credentials not ready");
|
||||
});
|
||||
mockSend.mockResolvedValueOnce(titanReply([1, 2, 3]));
|
||||
const embedder = new AWSBedrockEmbedder({});
|
||||
|
||||
await expect(embedder.embed("hello")).rejects.toThrow(
|
||||
"credentials not ready",
|
||||
);
|
||||
await expect(embedder.embed("hello")).resolves.toEqual([1, 2, 3]);
|
||||
});
|
||||
});
|
||||
|
||||
describe("memoryAction -> Cohere input_type", () => {
|
||||
// F3: buildRequestBody() used to hardcode `input_type: "search_document"`
|
||||
// regardless of the caller's action, so `Memory.search()` (which calls
|
||||
// `embed(query, "search")`) embedded the query in document mode.
|
||||
//
|
||||
// Typed as `Embedder` (not `AWSBedrockEmbedder`) because that is how
|
||||
// memory/index.ts actually calls it: the interface already declares an
|
||||
// optional `memoryAction` second parameter, so a narrower concrete
|
||||
// `embed(text: string)` satisfies it structurally and tsc stays silent --
|
||||
// the bug is a silent behavioral one, not a compile error.
|
||||
it("sends search_query for a search action", async () => {
|
||||
mockSend.mockResolvedValueOnce(cohereReply([[0.1]]));
|
||||
const embedder: Embedder = new AWSBedrockEmbedder({
|
||||
model: "cohere.embed-english-v3",
|
||||
});
|
||||
|
||||
await embedder.embed("query text", "search");
|
||||
|
||||
expect(requestBodyAt(0)).toEqual({
|
||||
texts: ["query text"],
|
||||
input_type: "search_query",
|
||||
});
|
||||
});
|
||||
|
||||
it("sends search_document for add and update actions", async () => {
|
||||
mockSend
|
||||
.mockResolvedValueOnce(cohereReply([[0.1]]))
|
||||
.mockResolvedValueOnce(cohereReply([[0.2]]));
|
||||
const embedder: Embedder = new AWSBedrockEmbedder({
|
||||
model: "cohere.embed-english-v3",
|
||||
});
|
||||
|
||||
await embedder.embed("doc one", "add");
|
||||
await embedder.embed("doc two", "update");
|
||||
|
||||
expect(requestBodyAt(0).input_type).toBe("search_document");
|
||||
expect(requestBodyAt(1).input_type).toBe("search_document");
|
||||
});
|
||||
|
||||
// Titan has no input_type concept; buildRequestBody() must not add one
|
||||
// even when a memoryAction is explicitly passed through.
|
||||
it("Titan ignores memoryAction and never sends input_type", async () => {
|
||||
mockSend.mockResolvedValueOnce(titanReply([1, 2]));
|
||||
const embedder: Embedder = new AWSBedrockEmbedder({});
|
||||
|
||||
await embedder.embed("hello", "search");
|
||||
|
||||
expect(requestBodyAt(0)).toEqual({ inputText: "hello" });
|
||||
});
|
||||
});
|
||||
|
||||
describe("Titan embedBatch concurrency", () => {
|
||||
afterEach(() => {
|
||||
// This test sets a persistent mockImplementation (not a *Once), so
|
||||
// clear it explicitly -- jest.clearAllMocks() in the top beforeEach
|
||||
// clears call data but not implementations.
|
||||
mockSend.mockReset();
|
||||
});
|
||||
|
||||
// F4: embedBatch() used to Promise.all-fan-out one InvokeModel call per
|
||||
// text with no cap, so a large batch could open hundreds of concurrent
|
||||
// requests at once. TITAN_MAX_CONCURRENCY bounds this to a small pool
|
||||
// while still preserving output order.
|
||||
it("never runs more than TITAN_MAX_CONCURRENCY Titan requests at once, and preserves order", async () => {
|
||||
let active = 0;
|
||||
let peak = 0;
|
||||
mockSend.mockImplementation(async (command: InvokeModelCommand) => {
|
||||
active++;
|
||||
peak = Math.max(peak, active);
|
||||
const body = JSON.parse(
|
||||
new TextDecoder().decode(command.input.body as Uint8Array),
|
||||
);
|
||||
await new Promise((resolve) => setTimeout(resolve, 10));
|
||||
active--;
|
||||
const index = Number(body.inputText.split("-")[1]);
|
||||
return titanReply([index]);
|
||||
});
|
||||
const embedder = new AWSBedrockEmbedder({});
|
||||
const texts = Array.from({ length: 10 }, (_, i) => `text-${i}`);
|
||||
|
||||
const embeddings = await embedder.embedBatch(texts);
|
||||
|
||||
expect(peak).toBeGreaterThan(1);
|
||||
expect(peak).toBeLessThanOrEqual(4);
|
||||
expect(embeddings).toEqual(texts.map((_, i) => [i]));
|
||||
expect(mockSend).toHaveBeenCalledTimes(10);
|
||||
});
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user