update embedding tests

Author: raindrop93

.gitignore

@@ -1,4 +1,4 @@
-node_modules/
+**/node_modules/
 dist/
 *.log
 .DS_Store
@@ -12,5 +12,6 @@ coverage/
 .pnpm-store
 .cursor
 *.db
+examples/seekdb-prisma/generated/**
 spec/
 .vscode

examples/seekdb-prisma/package.json

@@ -5,7 +5,7 @@
   "scripts": {
     "start": "tsx index.ts",
     "start:embedded": "tsx index-embedded.ts",
-    "db:generate": "prisma generate",
+    "postinstall": "prisma generate",
     "db:push": "prisma db push"
   },
   "dependencies": {

examples/seekdb-prisma/schema.prisma

@@ -1,6 +1,7 @@
 // Prisma schema: relational tables only. Vector tables are managed by seekdb Collection.
 generator client {
   provider = "prisma-client-js"
+  output   = "./generated/client"
 }
 
 datasource db {

examples/seekdb-prisma/tsconfig.json

@@ -0,0 +1,11 @@
+{
+  "compilerOptions": {
+    "target": "ES2022",
+    "module": "NodeNext",
+    "moduleResolution": "NodeNext",
+    "strict": true,
+    "noEmit": true,
+    "skipLibCheck": true
+  },
+  "include": ["*.ts", "generated/**/*.ts"]
+}

package.json

@@ -6,7 +6,7 @@
     "build": "pnpm --filter seekdb run build && pnpm --filter '@seekdb/*' run build",
     "build:seekdb": "pnpm --filter seekdb run build",
     "build:embeddings": "pnpm --filter '@seekdb/*' run build",
-    "test": "pnpm --filter seekdb run test -- run && pnpm --filter @seekdb/prisma-adapter run test",
+    "test": "pnpm --filter seekdb run test && pnpm --filter @seekdb/prisma-adapter run test",
     "lint": "pnpm -r run lint",
     "type-check": "pnpm -r run type-check",
     "prettier": "prettier --write .",

packages/embeddings/amazon-bedrock/index.test.ts

@@ -1,83 +1,153 @@
-import { describe, it, expect, beforeEach, vi } from "vitest";
+import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
 import { AmazonBedrockEmbeddingFunction } from "./index";
 
-// Mock AWS Bedrock Runtime client
-vi.mock("@aws-sdk/client-bedrock-runtime", () => {
-  const mockSend = vi.fn().mockResolvedValue({
-    body: new TextEncoder().encode(
-      JSON.stringify({
-        embedding: Array(1536)
-          .fill(0)
-          .map((_, i) => i / 1000),
-      })
-    ),
-  });
-
-  return {
-    BedrockRuntimeClient: vi.fn().mockImplementation(() => ({
-      send: mockSend,
-    })),
-    InvokeModelCommand: vi.fn().mockImplementation((params) => params),
-  };
-});
+const mockFetch = vi.hoisted(() => vi.fn());
 
 describe("AmazonBedrockEmbeddingFunction", () => {
   beforeEach(() => {
+    vi.stubGlobal("fetch", mockFetch);
+    let call = 0;
+    mockFetch.mockImplementation(async () => ({
+      ok: true,
+      status: 200,
+      statusText: "OK",
+      text: async () => "",
+      json: async () => ({
+        embedding: Array(1024)
+          .fill(0)
+          .map((_, i) => (i + call++) / 1000),
+      }),
+    }));
     vi.clearAllMocks();
+    process.env.AMAZON_BEDROCK_API_KEY = "test-api-key";
+  });
+
+  afterEach(() => {
+    vi.unstubAllGlobals();
+    delete process.env.AMAZON_BEDROCK_API_KEY;
   });
 
   it("should initialize with default parameters", () => {
-    const embedder = new AmazonBedrockEmbeddingFunction({});
-    expect(embedder.name).toBe("bedrock");
+    const embedder = new AmazonBedrockEmbeddingFunction({
+      region: "us-east-1",
+    });
+    expect(embedder.name).toBe("amazon_bedrock");
 
     const config = embedder.getConfig();
-    expect(config.model_id).toBe("amazon.titan-embed-text-v1");
+    expect(config.model_name).toBe("amazon.titan-embed-text-v2");
+    expect(config.region).toBe("us-east-1");
+    expect(config.api_key_env).toBe("AMAZON_BEDROCK_API_KEY");
   });
 
   it("should initialize with custom parameters", () => {
     const embedder = new AmazonBedrockEmbeddingFunction({
+      apiKey: "direct-key",
       region: "us-west-2",
-      modelId: "custom-model-id",
+      modelName: "amazon.titan-embed-text-v1",
+      apiKeyEnv: "CUSTOM_BEDROCK_KEY",
     });
 
     const config = embedder.getConfig();
     expect(config.region).toBe("us-west-2");
-    expect(config.model_id).toBe("custom-model-id");
+    expect(config.model_name).toBe("amazon.titan-embed-text-v1");
+    expect(config.api_key_env).toBe("CUSTOM_BEDROCK_KEY");
+  });
+
+  it("should throw when API key is missing", () => {
+    delete process.env.AMAZON_BEDROCK_API_KEY;
+    expect(() => {
+      new AmazonBedrockEmbeddingFunction({ region: "us-east-1" });
+    }).toThrow(/apiKey is required/);
+  });
+
+  it("should throw when region is missing", () => {
+    expect(() => {
+      new AmazonBedrockEmbeddingFunction({ apiKey: "k" });
+    }).toThrow(/region is required/);
   });
 
   it("should generate embeddings", async () => {
-    const embedder = new AmazonBedrockEmbeddingFunction({});
+    const embedder = new AmazonBedrockEmbeddingFunction({
+      region: "us-east-1",
+    });
     const texts = ["Hello world", "Test text"];
     const embeddings = await embedder.generate(texts);
 
     expect(embeddings.length).toBe(texts.length);
     embeddings.forEach((embedding) => {
-      expect(embedding.length).toBe(1536);
+      expect(embedding.length).toBe(1024);
     });
+    expect(embeddings[0]).not.toEqual(embeddings[1]);
+  });
+
+  it("should return empty array for empty input", async () => {
+    const embedder = new AmazonBedrockEmbeddingFunction({
+      region: "us-east-1",
+    });
+    expect(await embedder.generate([])).toEqual([]);
+    expect(mockFetch).not.toHaveBeenCalled();
   });
 
   it("should build from config", () => {
     const snakeCaseConfig = {
       region: "eu-west-1",
-      model_id: "amazon.titan-embed-text-v2:0",
+      model_name: "amazon.titan-embed-text-v2:0" as const,
+      api_key_env: "AMAZON_BEDROCK_API_KEY",
     };
 
     const embedder =
       AmazonBedrockEmbeddingFunction.buildFromConfig(snakeCaseConfig);
 
     expect(embedder).toBeInstanceOf(AmazonBedrockEmbeddingFunction);
-    expect(embedder.name).toBe("bedrock");
+    expect(embedder.name).toBe("amazon_bedrock");
 
     const config = embedder.getConfig();
     expect(config.region).toBe("eu-west-1");
-    expect(config.model_id).toBe("amazon.titan-embed-text-v2:0");
+    expect(config.model_name).toBe("amazon.titan-embed-text-v2:0");
   });
 
-  it("should handle empty config", () => {
-    const embedder = AmazonBedrockEmbeddingFunction.buildFromConfig({});
+  it("should build instance with defaults from partial config", () => {
+    const embedder = AmazonBedrockEmbeddingFunction.buildFromConfig({
+      region: "us-east-1",
+    });
 
     expect(embedder).toBeInstanceOf(AmazonBedrockEmbeddingFunction);
     const config = embedder.getConfig();
-    expect(config.model_id).toBe("amazon.titan-embed-text-v1");
+    expect(config.model_name).toBe("amazon.titan-embed-text-v2");
+  });
+
+  it("should call fetch with correct URL and headers", async () => {
+    const embedder = new AmazonBedrockEmbeddingFunction({
+      region: "us-east-1",
+      modelName: "amazon.titan-embed-text-v2",
+    });
+
+    await embedder.generate(["hello"]);
+
+    expect(mockFetch).toHaveBeenCalledWith(
+      "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-embed-text-v2/invoke",
+      expect.objectContaining({
+        method: "POST",
+        headers: expect.objectContaining({
+          Authorization: "Bearer test-api-key",
+          "Content-Type": "application/json",
+          Accept: "application/json",
+        }),
+        body: JSON.stringify({ inputText: "hello" }),
+      })
+    );
+  });
+
+  it("should expose dimension and getModelDimensions", () => {
+    const embedder = new AmazonBedrockEmbeddingFunction({
+      region: "us-east-1",
+      modelName: "amazon.titan-embed-text-v1",
+    });
+    expect(embedder.dimension).toBe(1536);
+    expect(
+      AmazonBedrockEmbeddingFunction.getModelDimensions()[
+        "amazon.titan-embed-text-v2"
+      ]
+    ).toBe(1024);
   });
 });

packages/embeddings/amazon-bedrock/index.ts

@@ -139,12 +139,6 @@ export class AmazonBedrockEmbeddingFunction implements EmbeddingFunction {
   static buildFromConfig(
     config: EmbeddingConfig
   ): AmazonBedrockEmbeddingFunction {
-    if (!config.api_key_env) {
-      throw new Error(
-        "Building Amazon bedrock embedding function from config: api_key_env is required in config."
-      );
-    }
-
     return new AmazonBedrockEmbeddingFunction({
       apiKeyEnv: config.api_key_env,
       region: config.region,