fix(llms/aws_bedrock): resolve provider for application inference profile ARNs (#6899)

mem0-ts/src/oss/src/llms/aws_bedrock.ts

@@ -29,7 +29,18 @@ const PROVIDERS = [
  * Extract the model-family provider from a Bedrock model id
  * (e.g. `anthropic.claude-3-sonnet-...` -> `anthropic`).
  */
-export function extractProvider(model: string): string {
+export function extractProvider(
+  model: string,
+  providerOverride?: string,
+): string {
+  if (providerOverride) {
+    if (!PROVIDERS.includes(providerOverride)) {
+      throw new Error(
+        `Unknown providerOverride '${providerOverride}'. Valid providers: ${PROVIDERS.join(", ")}`,
+      );
+    }
+    return providerOverride;
+  }
   for (const provider of PROVIDERS) {
     const re = new RegExp(
       `\\b${provider.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}\\b`,
@@ -79,7 +90,7 @@ export class AWSBedrockLLM implements LLM {
     this.model =
       (typeof config.model === "string" && config.model) ||
       "anthropic.claude-3-5-sonnet-20240620-v1:0";
-    this.provider = extractProvider(this.model);
+    this.provider = extractProvider(this.model, config.providerOverride);
     this.temperature = config.temperature ?? 0.1;
     this.maxTokens = config.maxTokens ?? 2000;
     this.topP = config.topP;

mem0-ts/src/oss/src/tests/aws_bedrock.test.ts

@@ -68,6 +68,29 @@ describe("extractProvider", () => {
       /Unknown provider/,
     );
   });
+
+  const ARN =
+    "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz";
+
+  it("throws for an application inference profile ARN without providerOverride", () => {
+    expect(() => extractProvider(ARN)).toThrow(/Unknown provider in model/);
+  });
+
+  it("resolves an application inference profile ARN via providerOverride", () => {
+    expect(extractProvider(ARN, "anthropic")).toBe("anthropic");
+  });
+
+  it("lets providerOverride take precedence over regex detection", () => {
+    expect(
+      extractProvider("anthropic.claude-3-5-sonnet-20240620-v1:0", "amazon"),
+    ).toBe("amazon");
+  });
+
+  it("throws on a misspelled providerOverride", () => {
+    expect(() => extractProvider(ARN, "anthorpic")).toThrow(
+      /Unknown providerOverride 'anthorpic'/,
+    );
+  });
 });
 
 describe("AWSBedrockLLM", () => {
@@ -176,6 +199,28 @@ describe("AWSBedrockLLM", () => {
     expect(res).toEqual({ content: "hello from bedrock", role: "assistant" });
   });
 
+  it("uses providerOverride to resolve inference profile ARNs (topP omitted like anthropic)", async () => {
+    const arn =
+      "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz";
+    const client = new FakeBedrockClient(textResponse);
+    const llm = makeLLM(client, {
+      model: arn,
+      providerOverride: "anthropic",
+      topP: 0.9,
+    });
+    await llm.generateResponse([{ role: "user", content: "hi" }]);
+    expect(client.lastInput.modelId).toBe(arn);
+    expect(client.lastInput.inferenceConfig.topP).toBeUndefined();
+  });
+
+  it("throws when constructed with an inference profile ARN and no providerOverride", () => {
+    const arn =
+      "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123xyz";
+    expect(() => new AWSBedrockLLM({ model: arn })).toThrow(
+      /Unknown provider in model/,
+    );
+  });
+
   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.

mem0-ts/src/oss/src/types/index.ts

@@ -78,6 +78,7 @@ export interface LLMConfig {
   awsAccessKeyId?: string;
   awsSecretAccessKey?: string;
   awsSessionToken?: string;
+  providerOverride?: string;
   // Optional pre-constructed client (e.g. BedrockRuntimeClient) for DI/testing.
   client?: any;
 }