feat(ts-sdk): add AWS Bedrock LLM provider (closes #5765) (#5890)

Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>
This commit is contained in:
Bartok
2026-07-10 19:14:15 +05:30
committed by GitHub
co-authored by kartik-mem0
parent df9d5cc4b1
commit ba9054e0c8
13 changed files with 1087 additions and 193 deletions
+19
View File
@@ -137,6 +137,11 @@ export class ConfigManager {
userConf?.maxTokens ?? (llmRaw?.max_tokens as number | undefined);
return {
// Spread user-provided config first so any additional fields
// (e.g. future aws_bedrock options) pass through without a
// manager.ts edit, matching the vectorStore.config pattern above
// and making the schema's .passthrough() on llm.config meaningful.
...userConf,
baseURL: llmBaseURL,
url: userConf?.url,
apiKey:
@@ -151,6 +156,20 @@ export class ConfigManager {
temperature,
topP,
maxTokens,
// Pass through AWS Bedrock fields so the aws_bedrock provider works
// through the standard Memory config path (snake_case tolerated).
awsRegion:
userConf?.awsRegion ?? (llmRaw?.aws_region as string | undefined),
awsAccessKeyId:
userConf?.awsAccessKeyId ??
(llmRaw?.aws_access_key_id as string | undefined),
awsSecretAccessKey:
userConf?.awsSecretAccessKey ??
(llmRaw?.aws_secret_access_key as string | undefined),
awsSessionToken:
userConf?.awsSessionToken ??
(llmRaw?.aws_session_token as string | undefined),
client: userConf?.client,
};
})(),
},
+1
View File
@@ -24,6 +24,7 @@ export * from "./llms/mistral";
export * from "./llms/langchain";
export * from "./llms/litellm";
export * from "./llms/vllm";
export * from "./llms/aws_bedrock";
export * from "./vector_stores/base";
export * from "./vector_stores/memory";
export * from "./vector_stores/baidu";
+294
View File
@@ -0,0 +1,294 @@
import { LLM, LLMResponse } from "./base";
import { LLMConfig, Message } from "../types";
/**
* Providers recognised in Bedrock model identifiers, mirroring the Python
* provider's `PROVIDERS` list (`mem0/llms/aws_bedrock.py`).
*/
const PROVIDERS = [
"ai21",
"amazon",
"anthropic",
"cohere",
"meta",
"mistral",
"stability",
"writer",
"deepseek",
"gpt-oss",
"perplexity",
"snowflake",
"titan",
"command",
"j2",
"llama",
"minimax",
];
/**
* Extract the model-family provider from a Bedrock model id
* (e.g. `anthropic.claude-3-sonnet-...` -> `anthropic`).
*/
export function extractProvider(model: string): string {
for (const provider of PROVIDERS) {
const re = new RegExp(
`\\b${provider.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}\\b`,
);
if (re.test(model)) return provider;
}
throw new Error(`Unknown provider in model: ${model}`);
}
/**
* AWS Bedrock fields (awsRegion / awsAccessKeyId / awsSecretAccessKey /
* awsSessionToken / client) now live on the shared `LLMConfig`, so the
* provider is configurable through the standard typed `Memory` config path.
*/
type AWSBedrockConfig = LLMConfig;
/**
* AWS Bedrock LLM provider for the TypeScript OSS SDK.
*
* Mirrors `mem0/llms/aws_bedrock.py`. Uses the Bedrock **Converse API**
* (`ConverseCommand`), which provides a uniform message/tool interface across
* the Anthropic / Amazon (Nova) / Meta / Mistral / Cohere model families, so a
* single code path serves them all (the Python provider keeps per-family
* `invoke_model` branches for legacy reasons; Converse supersedes them).
*
* The `@aws-sdk/client-bedrock-runtime` dependency is loaded on first use via
* dynamic `import()` so the package stays optional. Credentials resolve via the
* standard AWS chain unless provided explicitly in config.
*/
interface BedrockSDK {
BedrockRuntimeClient: new (config: Record<string, any>) => any;
ConverseCommand: new (input: Record<string, any>) => any;
}
export class AWSBedrockLLM implements LLM {
private model: string;
private provider: string;
private temperature: number;
private maxTokens: number;
private topP?: number;
private clientConfig: Record<string, any>;
private clientOverride?: any;
private sdkPromise?: Promise<BedrockSDK>;
private clientPromise?: Promise<any>;
constructor(config: AWSBedrockConfig = {}) {
this.model =
(typeof config.model === "string" && config.model) ||
"anthropic.claude-3-5-sonnet-20240620-v1:0";
this.provider = extractProvider(this.model);
this.temperature = config.temperature ?? 0.1;
this.maxTokens = config.maxTokens ?? 2000;
this.topP = config.topP;
const region =
config.awsRegion ||
process.env.AWS_REGION ||
process.env.AWS_DEFAULT_REGION;
const clientConfig: Record<string, any> = {};
if (region) clientConfig.region = region;
if (config.awsAccessKeyId && config.awsSecretAccessKey) {
clientConfig.credentials = {
accessKeyId: config.awsAccessKeyId,
secretAccessKey: config.awsSecretAccessKey,
...(config.awsSessionToken && {
sessionToken: config.awsSessionToken,
}),
};
}
this.clientConfig = clientConfig;
this.clientOverride = config.client;
}
/**
* Load the optional AWS SDK on first use.
*
* This MUST be a dynamic `import()`, never `require()`: tsup/esbuild rewrite
* `require()` in the published ESM bundle (`dist/oss/index.mjs`) into a
* `__require` shim that throws `Dynamic require of "..." is not supported`,
* so every ESM consumer would hit a dead provider even with the SDK installed.
*/
private async getSDK(): Promise<BedrockSDK> {
if (!this.sdkPromise) {
this.sdkPromise = import("@aws-sdk/client-bedrock-runtime").then(
(sdk) => sdk as unknown as BedrockSDK,
(err) => {
// Let a later call retry rather than caching the rejection forever.
this.sdkPromise = undefined;
const detail = err instanceof Error ? err.message : String(err);
throw new Error(
"The '@aws-sdk/client-bedrock-runtime' package is required to use the AWS Bedrock LLM provider. " +
`Install it with: npm install @aws-sdk/client-bedrock-runtime (original error: ${detail})`,
);
},
);
}
return this.sdkPromise;
}
/** Memoized Bedrock client; an injected `config.client` short-circuits the SDK. */
private async getClient(): Promise<any> {
if (this.clientOverride) return this.clientOverride;
if (!this.clientPromise) {
this.clientPromise = this.getSDK().then(
({ BedrockRuntimeClient }) =>
new BedrockRuntimeClient(this.clientConfig),
);
}
return this.clientPromise;
}
/**
* Split messages into a top-level `system` block (Converse passes system
* prompts separately) and role-tagged content blocks for everything else.
*/
private formatMessages(messages: Message[]): {
system?: { text: string }[];
converseMessages: { role: string; content: { text: string }[] }[];
} {
const systemParts: string[] = [];
const converseMessages: { role: string; content: { text: string }[] }[] =
[];
for (const msg of messages) {
const role = msg.role;
const content =
typeof msg.content === "string"
? msg.content
: JSON.stringify(msg.content);
if (role === "system") {
systemParts.push(content);
} else {
converseMessages.push({
role: role === "assistant" ? "assistant" : "user",
content: [{ text: content }],
});
}
}
if (converseMessages.length === 0) {
converseMessages.push({ role: "user", content: [{ text: "" }] });
}
return {
system: systemParts.length
? [{ text: systemParts.join("\n") }]
: undefined,
converseMessages,
};
}
/**
* Build the Converse `inferenceConfig`. Anthropic and MiniMax reasoning
* models reject requests carrying both `temperature` and `topP`, so `topP`
* is omitted for those families (mirrors the Python `_build_inference_config`).
*/
private buildInferenceConfig(): Record<string, any> {
const inferenceConfig: Record<string, any> = {
maxTokens: this.maxTokens,
temperature: this.temperature,
};
if (
this.topP != null &&
!["anthropic", "minimax"].includes(this.provider)
) {
inferenceConfig.topP = this.topP;
}
return inferenceConfig;
}
/** Convert OpenAI-style tools to the Converse `toolConfig` shape. */
private convertToolsToConverse(tools: any[]): any | undefined {
if (!tools || tools.length === 0) return undefined;
const converseTools = tools
.filter((t) => t?.type === "function" && t.function)
.map((t) => ({
toolSpec: {
name: t.function.name,
description: t.function.description || "",
inputSchema: { json: t.function.parameters || {} },
},
}));
return converseTools.length ? { tools: converseTools } : undefined;
}
private async converse(messages: Message[], tools?: any[]): Promise<any> {
const { system, converseMessages } = this.formatMessages(messages);
const input: Record<string, any> = {
modelId: this.model,
messages: converseMessages,
inferenceConfig: this.buildInferenceConfig(),
};
if (system) input.system = system;
const toolConfig = tools ? this.convertToolsToConverse(tools) : undefined;
if (toolConfig) input.toolConfig = toolConfig;
const [{ ConverseCommand }, client] = await Promise.all([
this.getSDK(),
this.getClient(),
]);
return client.send(new ConverseCommand(input));
}
/** Pull the first text block out of a Converse response. */
private parseText(response: any): string {
const content = response?.output?.message?.content || [];
for (const block of content) {
if (block && typeof block.text === "string") return block.text;
}
return "";
}
/** Collect any toolUse blocks out of a Converse response. */
private parseToolCalls(response: any): { name: string; arguments: string }[] {
const content = response?.output?.message?.content || [];
const calls: { name: string; arguments: string }[] = [];
for (const block of content) {
if (block?.toolUse) {
calls.push({
name: block.toolUse.name,
arguments: JSON.stringify(block.toolUse.input ?? {}),
});
}
}
return calls;
}
async generateResponse(
messages: Message[],
_responseFormat?: { type: string },
tools?: any[],
): Promise<string | LLMResponse> {
try {
const response = await this.converse(messages, tools);
if (tools && tools.length) {
const toolCalls = this.parseToolCalls(response);
if (toolCalls.length) {
return {
content: this.parseText(response),
role: "assistant",
toolCalls,
};
}
}
return this.parseText(response);
} catch (err) {
const message = err instanceof Error ? err.message : String(err);
throw new Error(`AWS Bedrock LLM failed: ${message}`);
}
}
async generateChat(messages: Message[]): Promise<LLMResponse> {
try {
const response = await this.converse(messages);
return { content: this.parseText(response), role: "assistant" };
} catch (err) {
const message = err instanceof Error ? err.message : String(err);
throw new Error(`AWS Bedrock LLM failed: ${message}`);
}
}
}
@@ -0,0 +1,214 @@
import { readFileSync } from "fs";
import { join } from "path";
import { AWSBedrockLLM, extractProvider } from "../llms/aws_bedrock";
/**
* Fake BedrockRuntimeClient capturing the Converse command input and returning
* a scripted response — exercises request shaping + response parsing without
* the AWS SDK or live credentials.
*/
class FakeBedrockClient {
public lastInput: any = null;
public response: any;
constructor(response: any) {
this.response = response;
}
async send(command: any) {
// ConverseCommand stores its input on `.input` (mirrored by our fake command).
this.lastInput = command?.input ?? command;
return this.response;
}
}
// Counts real client constructions so we can prove the SDK stays untouched
// until the first request. (jest only lets mock factories close over `mock*`.)
let mockClientConstructions = 0;
// The provider loads the AWS SDK on first use via dynamic import(). Mock the
// module so it resolves and ConverseCommand simply wraps its input.
jest.mock(
"@aws-sdk/client-bedrock-runtime",
() => ({
BedrockRuntimeClient: class {
constructor(public config: any) {
mockClientConstructions++;
}
},
ConverseCommand: class {
input: any;
constructor(input: any) {
this.input = input;
}
},
}),
{ virtual: true },
);
function makeLLM(client: FakeBedrockClient, overrides: any = {}) {
return new AWSBedrockLLM({
model: "anthropic.claude-3-5-sonnet-20240620-v1:0",
client,
...overrides,
});
}
describe("extractProvider", () => {
it("detects the model family from the Bedrock model id", () => {
expect(extractProvider("anthropic.claude-3-5-sonnet-20240620-v1:0")).toBe(
"anthropic",
);
expect(extractProvider("amazon.nova-pro-v1:0")).toBe("amazon");
expect(extractProvider("meta.llama3-70b-instruct-v1:0")).toBe("meta");
expect(extractProvider("mistral.mistral-large-2407-v1:0")).toBe("mistral");
});
it("throws on an unknown provider", () => {
expect(() => extractProvider("totally-unknown-model")).toThrow(
/Unknown provider/,
);
});
});
describe("AWSBedrockLLM", () => {
const textResponse = {
output: { message: { content: [{ text: "hello from bedrock" }] } },
};
it("returns assistant text from a Converse response", async () => {
const client = new FakeBedrockClient(textResponse);
const llm = makeLLM(client);
const out = await llm.generateResponse([{ role: "user", content: "hi" }]);
expect(out).toBe("hello from bedrock");
});
it("lifts system messages into the top-level system block", async () => {
const client = new FakeBedrockClient(textResponse);
const llm = makeLLM(client);
await llm.generateResponse([
{ role: "system", content: "be terse" },
{ role: "user", content: "hi" },
]);
expect(client.lastInput.system).toEqual([{ text: "be terse" }]);
expect(client.lastInput.messages).toEqual([
{ role: "user", content: [{ text: "hi" }] },
]);
});
it("omits topP for anthropic models even when configured", async () => {
const client = new FakeBedrockClient(textResponse);
const llm = makeLLM(client, { topP: 0.9, temperature: 0.2 });
await llm.generateResponse([{ role: "user", content: "hi" }]);
expect(client.lastInput.inferenceConfig.topP).toBeUndefined();
expect(client.lastInput.inferenceConfig.temperature).toBe(0.2);
});
it("includes topP for non-anthropic models", async () => {
const client = new FakeBedrockClient(textResponse);
const llm = makeLLM(client, {
model: "meta.llama3-70b-instruct-v1:0",
topP: 0.9,
});
await llm.generateResponse([{ role: "user", content: "hi" }]);
expect(client.lastInput.inferenceConfig.topP).toBe(0.9);
});
it("converts OpenAI-style tools to Converse toolConfig and parses toolUse", async () => {
const toolResponse = {
output: {
message: {
content: [{ toolUse: { name: "add_memory", input: { text: "x" } } }],
},
},
};
const client = new FakeBedrockClient(toolResponse);
const llm = makeLLM(client);
const tools = [
{
type: "function",
function: {
name: "add_memory",
description: "store a memory",
parameters: { type: "object", properties: {} },
},
},
];
const out = await llm.generateResponse(
[{ role: "user", content: "remember x" }],
undefined,
tools,
);
expect(client.lastInput.toolConfig.tools[0].toolSpec.name).toBe(
"add_memory",
);
expect(typeof out).toBe("object");
expect((out as any).toolCalls[0]).toEqual({
name: "add_memory",
arguments: JSON.stringify({ text: "x" }),
});
});
it("never sends an empty messages array", async () => {
const client = new FakeBedrockClient(textResponse);
const llm = makeLLM(client);
await llm.generateResponse([{ role: "system", content: "only system" }]);
expect(client.lastInput.messages).toEqual([
{ role: "user", content: [{ text: "" }] },
]);
});
it("wraps SDK errors with a provider-tagged message", async () => {
const client = {
send: async () => {
throw new Error("AccessDeniedException");
},
} as any;
const llm = makeLLM(client);
await expect(
llm.generateResponse([{ role: "user", content: "hi" }]),
).rejects.toThrow(/AWS Bedrock LLM failed: AccessDeniedException/);
});
it("generateChat returns text content with assistant role", async () => {
const client = new FakeBedrockClient(textResponse);
const llm = makeLLM(client);
const res = await llm.generateChat([{ role: "user", content: "hi" }]);
expect(res).toEqual({ content: "hello from bedrock", role: "assistant" });
});
it("does not construct the Bedrock client until the first request", () => {
const before = mockClientConstructions;
// No injected client: the old constructor eagerly built a real one.
new AWSBedrockLLM({
model: "anthropic.claude-3-5-sonnet-20240620-v1:0",
awsRegion: "us-west-2",
});
expect(mockClientConstructions).toBe(before);
});
});
/**
* ts-jest downlevels this suite to CommonJS, where `require()` works fine — so
* no runtime test here can observe the ESM failure. Guard the source invariant
* that causes it instead: esbuild rewrites `require()` in the published ESM
* bundle (dist/oss/index.mjs) into a `__require` shim that throws
* `Dynamic require of "..." is not supported`, stranding every ESM consumer.
*/
describe("aws_bedrock.ts ESM safety", () => {
const source = readFileSync(
join(__dirname, "..", "llms", "aws_bedrock.ts"),
"utf8",
);
// Comments legitimately mention require(); only real code should be checked.
const code = source
.replace(/\/\*[\s\S]*?\*\//g, "")
.replace(/(^|[^:])\/\/.*$/gm, "$1");
it("never calls require() — the ESM bundle turns it into a throwing shim", () => {
expect(code).not.toMatch(/\brequire\s*\(/);
});
it("loads the optional AWS SDK through a dynamic import()", () => {
expect(code).toMatch(/import\("@aws-sdk\/client-bedrock-runtime"\)/);
});
});
+29 -13
View File
@@ -67,6 +67,14 @@ export interface LLMConfig {
temperature?: number;
topP?: number;
maxTokens?: number;
// AWS Bedrock provider config (used when provider === "aws_bedrock").
// Credentials otherwise resolve via the standard AWS credential chain.
awsRegion?: string;
awsAccessKeyId?: string;
awsSecretAccessKey?: string;
awsSessionToken?: string;
// Optional pre-constructed client (e.g. BedrockRuntimeClient) for DI/testing.
client?: any;
}
export interface RerankerConfig {
@@ -203,21 +211,29 @@ export const MemoryConfigSchema = z.object({
})
.passthrough(),
}),
llm: z.object({
provider: z.string(),
config: z.object({
apiKey: z.string().optional(),
model: z.union([z.string(), z.any()]).optional(),
modelProperties: z.record(z.string(), z.any()).optional(),
baseURL: z.string().optional(),
vllmBaseURL: z.string().optional(),
vllm_base_url: z.string().optional(),
url: z.string().optional(),
timeout: z.number().optional(),
temperature: z.number().optional(),
topP: z.number().optional(),
maxTokens: z.number().optional(),
}),
config: z
.object({
apiKey: z.string().optional(),
model: z.union([z.string(), z.any()]).optional(),
modelProperties: z.record(z.string(), z.any()).optional(),
baseURL: z.string().optional(),
vllmBaseURL: z.string().optional(),
vllm_base_url: z.string().optional(),
url: z.string().optional(),
timeout: z.number().optional(),
temperature: z.number().optional(),
topP: z.number().optional(),
maxTokens: z.number().optional(),
awsRegion: z.string().optional(),
awsAccessKeyId: z.string().optional(),
awsSecretAccessKey: z.string().optional(),
awsSessionToken: z.string().optional(),
client: z.any().optional(),
})
.passthrough(),
}),
historyDbPath: z.string().optional(),
customInstructions: z.string().optional(),
+3
View File
@@ -34,6 +34,7 @@ import { LMStudioLLM } from "../llms/lmstudio";
import { DeepSeekLLM } from "../llms/deepseek";
import { XAILLM } from "../llms/xai";
import { SarvamLLM } from "../llms/sarvam";
import { AWSBedrockLLM } from "../llms/aws_bedrock";
import { LiteLLM } from "../llms/litellm";
import { MiniMaxLLM } from "../llms/minimax";
import { TogetherLLM } from "../llms/together";
@@ -128,6 +129,8 @@ export class LLMFactory {
return new XAILLM(config);
case "sarvam":
return new SarvamLLM(config);
case "aws_bedrock":
return new AWSBedrockLLM(config);
case "litellm":
return new LiteLLM(config);
case "minimax":
@@ -378,6 +378,55 @@ describe("ConfigManager", () => {
});
});
describe("mergeConfig - AWS Bedrock credential passthrough", () => {
const baseEmbedder = { provider: "openai", config: { apiKey: "k" } };
const baseVectorStore = { provider: "memory", config: {} };
it("preserves explicit AWS Bedrock credentials through MemoryConfigSchema.parse", () => {
const cfg = ConfigManager.mergeConfig({
embedder: baseEmbedder,
vectorStore: baseVectorStore,
llm: {
provider: "aws_bedrock",
config: {
model: "anthropic.claude-3-sonnet-20240229-v1:0",
awsRegion: "us-east-1",
awsAccessKeyId: "AKIAEXAMPLE",
awsSecretAccessKey: "secret-example",
awsSessionToken: "session-example",
} as any,
},
});
// Regression: zod default parse strips undeclared keys, which silently
// dropped credentials before .passthrough() + typed fields were added.
expect(cfg.llm.config.awsRegion).toBe("us-east-1");
expect(cfg.llm.config.awsAccessKeyId).toBe("AKIAEXAMPLE");
expect(cfg.llm.config.awsSecretAccessKey).toBe("secret-example");
expect(cfg.llm.config.awsSessionToken).toBe("session-example");
});
it("normalizes snake_case AWS credentials from the raw config", () => {
const cfg = ConfigManager.mergeConfig({
embedder: baseEmbedder,
vectorStore: baseVectorStore,
llm: {
provider: "aws_bedrock",
config: {
model: "anthropic.claude-3-sonnet-20240229-v1:0",
aws_region: "eu-west-1",
aws_access_key_id: "AKIASNAKE",
aws_secret_access_key: "snake-secret",
} as any,
},
});
expect(cfg.llm.config.awsRegion).toBe("eu-west-1");
expect(cfg.llm.config.awsAccessKeyId).toBe("AKIASNAKE");
expect(cfg.llm.config.awsSecretAccessKey).toBe("snake-secret");
});
});
describe("mergeConfig - full OpenClaw-style LM Studio config", () => {
it("handles the exact config from issue #4235", () => {
const cfg = ConfigManager.mergeConfig({
@@ -713,6 +762,47 @@ describe("Memory – LM Studio end-to-end flow", () => {
);
});
it("passes AWS Bedrock credentials to LLMFactory through the Memory stack", async () => {
const mem = new MemoryClass({
embedder: {
provider: "lmstudio",
config: {
model: "nomic-embed-text-v1.5",
baseURL: "http://localhost:1234/v1",
embeddingDims: 768,
},
},
vectorStore: {
provider: "memory",
config: { collectionName: "test", dimension: 768 },
},
llm: {
provider: "aws_bedrock",
config: {
model: "anthropic.claude-3-sonnet-20240229-v1:0",
awsRegion: "us-east-1",
awsAccessKeyId: "AKIAEXAMPLE",
awsSecretAccessKey: "secret-example",
},
},
disableHistory: true,
});
await mem.getAll({ filters: { user_id: "u1" } });
// Regression: credentials must survive ConfigManager -> Memory -> LLMFactory,
// otherwise Bedrock silently falls back to the ambient AWS credential chain.
expect(mockLlmFactory.create).toHaveBeenCalledWith(
"aws_bedrock",
expect.objectContaining({
model: "anthropic.claude-3-sonnet-20240229-v1:0",
awsRegion: "us-east-1",
awsAccessKeyId: "AKIAEXAMPLE",
awsSecretAccessKey: "secret-example",
}),
);
});
it("add flow works with lmstudio LLM for fact extraction", async () => {
mockLlm.generateResponse.mockResolvedValueOnce(
'{"facts":["User loves sushi"]}',