feat!: add variable store to embeddings registry (#2112)

BREAKING CHANGE: embedding function implementations in Node need to now
call `resolveVariables()` in their constructors and should **not**
implement `toJSON()`.

This tries to address the handling of secrets. In Node, they are
currently lost. In Python, they are currently leaked into the table
schema metadata.

This PR introduces an in-memory variable store on the function registry.
It also allows embedding function definitions to label certain config
values as "sensitive", and the preprocessing logic will raise an error
if users try to pass in hard-coded values.

Closes #2110
Closes #521

---------

Co-authored-by: Weston Pace <weston.pace@gmail.com>

nodejs/__test__/embedding.test.ts

@@ -17,6 +17,8 @@ import {
 import { EmbeddingFunction, LanceSchema } from "../lancedb/embedding";
 import { getRegistry, register } from "../lancedb/embedding/registry";
 
+const testOpenAIInteg = process.env.OPENAI_API_KEY == null ? test.skip : test;
+
 describe("embedding functions", () => {
   let tmpDir: tmp.DirResult;
   beforeEach(() => {
@@ -29,9 +31,6 @@ describe("embedding functions", () => {
 
   it("should be able to create a table with an embedding function", async () => {
     class MockEmbeddingFunction extends EmbeddingFunction<string> {
-      toJSON(): object {
-        return {};
-      }
       ndims() {
         return 3;
       }
@@ -75,9 +74,6 @@ describe("embedding functions", () => {
   it("should be able to append and upsert using embedding function", async () => {
     @register()
     class MockEmbeddingFunction extends EmbeddingFunction<string> {
-      toJSON(): object {
-        return {};
-      }
       ndims() {
         return 3;
       }
@@ -143,9 +139,6 @@ describe("embedding functions", () => {
   it("should be able to create an empty table with an embedding function", async () => {
     @register()
     class MockEmbeddingFunction extends EmbeddingFunction<string> {
-      toJSON(): object {
-        return {};
-      }
       ndims() {
         return 3;
       }
@@ -194,9 +187,6 @@ describe("embedding functions", () => {
   it("should error when appending to a table with an unregistered embedding function", async () => {
     @register("mock")
     class MockEmbeddingFunction extends EmbeddingFunction<string> {
-      toJSON(): object {
-        return {};
-      }
       ndims() {
         return 3;
       }
@@ -241,13 +231,35 @@ describe("embedding functions", () => {
       `Function "mock" not found in registry`,
     );
   });
+
+  testOpenAIInteg("propagates variables through all methods", async () => {
+    delete process.env.OPENAI_API_KEY;
+    const registry = getRegistry();
+    registry.setVar("openai_api_key", "sk-...");
+    const func = registry.get("openai")?.create({
+      model: "text-embedding-ada-002",
+      apiKey: "$var:openai_api_key",
+    }) as EmbeddingFunction;
+
+    const db = await connect("memory://");
+    const wordsSchema = LanceSchema({
+      text: func.sourceField(new Utf8()),
+      vector: func.vectorField(),
+    });
+    const tbl = await db.createEmptyTable("words", wordsSchema, {
+      mode: "overwrite",
+    });
+    await tbl.add([{ text: "hello world" }, { text: "goodbye world" }]);
+
+    const query = "greetings";
+    const actual = (await tbl.search(query).limit(1).toArray())[0];
+    expect(actual).toHaveProperty("text");
+  });
+
   test.each([new Float16(), new Float32(), new Float64()])(
     "should be able to provide manual embeddings with multiple float datatype",
     async (floatType) => {
       class MockEmbeddingFunction extends EmbeddingFunction<string> {
-        toJSON(): object {
-          return {};
-        }
         ndims() {
           return 3;
         }
@@ -292,10 +304,6 @@ describe("embedding functions", () => {
     async (floatType) => {
       @register("test1")
       class MockEmbeddingFunctionWithoutNDims extends EmbeddingFunction<string> {
-        toJSON(): object {
-          return {};
-        }
-
         embeddingDataType(): Float {
           return floatType;
         }
@@ -310,9 +318,6 @@ describe("embedding functions", () => {
       }
       @register("test")
       class MockEmbeddingFunction extends EmbeddingFunction<string> {
-        toJSON(): object {
-          return {};
-        }
         ndims() {
           return 3;
         }

nodejs/__test__/registry.test.ts

@@ -11,7 +11,11 @@ import * as arrow18 from "apache-arrow-18";
 import * as tmp from "tmp";
 
 import { connect } from "../lancedb";
-import { EmbeddingFunction, LanceSchema } from "../lancedb/embedding";
+import {
+  EmbeddingFunction,
+  FunctionOptions,
+  LanceSchema,
+} from "../lancedb/embedding";
 import { getRegistry, register } from "../lancedb/embedding/registry";
 
 describe.each([arrow15, arrow16, arrow17, arrow18])("LanceSchema", (arrow) => {
@@ -39,11 +43,6 @@ describe.each([arrow15, arrow16, arrow17, arrow18])("Registry", (arrow) => {
   it("should register a new item to the registry", async () => {
     @register("mock-embedding")
     class MockEmbeddingFunction extends EmbeddingFunction<string> {
-      toJSON(): object {
-        return {
-          someText: "hello",
-        };
-      }
       constructor() {
         super();
       }
@@ -89,11 +88,6 @@ describe.each([arrow15, arrow16, arrow17, arrow18])("Registry", (arrow) => {
   });
   test("should error if registering with the same name", async () => {
     class MockEmbeddingFunction extends EmbeddingFunction<string> {
-      toJSON(): object {
-        return {
-          someText: "hello",
-        };
-      }
       constructor() {
         super();
       }
@@ -114,13 +108,9 @@ describe.each([arrow15, arrow16, arrow17, arrow18])("Registry", (arrow) => {
   });
   test("schema should contain correct metadata", async () => {
     class MockEmbeddingFunction extends EmbeddingFunction<string> {
-      toJSON(): object {
-        return {
-          someText: "hello",
-        };
-      }
-      constructor() {
+      constructor(args: FunctionOptions = {}) {
         super();
+        this.resolveVariables(args);
       }
       ndims() {
         return 3;
@@ -132,7 +122,7 @@ describe.each([arrow15, arrow16, arrow17, arrow18])("Registry", (arrow) => {
         return data.map(() => [1, 2, 3]);
       }
     }
-    const func = new MockEmbeddingFunction();
+    const func = new MockEmbeddingFunction({ someText: "hello" });
 
     const schema = LanceSchema({
       id: new arrow.Int32(),
@@ -155,3 +145,79 @@ describe.each([arrow15, arrow16, arrow17, arrow18])("Registry", (arrow) => {
     expect(schema.metadata).toEqual(expectedMetadata);
   });
 });
+
+describe("Registry.setVar", () => {
+  const registry = getRegistry();
+
+  beforeEach(() => {
+    @register("mock-embedding")
+    // biome-ignore lint/correctness/noUnusedVariables :
+    class MockEmbeddingFunction extends EmbeddingFunction<string> {
+      constructor(optionsRaw: FunctionOptions = {}) {
+        super();
+        const options = this.resolveVariables(optionsRaw);
+
+        expect(optionsRaw["someKey"].startsWith("$var:someName")).toBe(true);
+        expect(options["someKey"]).toBe("someValue");
+
+        if (options["secretKey"]) {
+          expect(optionsRaw["secretKey"]).toBe("$var:secretKey");
+          expect(options["secretKey"]).toBe("mySecret");
+        }
+      }
+      async computeSourceEmbeddings(data: string[]) {
+        return data.map(() => [1, 2, 3]);
+      }
+      embeddingDataType() {
+        return new arrow18.Float32() as apiArrow.Float;
+      }
+      protected getSensitiveKeys() {
+        return ["secretKey"];
+      }
+    }
+  });
+  afterEach(() => {
+    registry.reset();
+  });
+
+  it("Should error if the variable is not set", () => {
+    console.log(registry.get("mock-embedding"));
+    expect(() =>
+      registry.get("mock-embedding")!.create({ someKey: "$var:someName" }),
+    ).toThrow('Variable "someName" not found');
+  });
+
+  it("should use default values if not set", () => {
+    registry
+      .get("mock-embedding")!
+      .create({ someKey: "$var:someName:someValue" });
+  });
+
+  it("should set a variable that the embedding function understand", () => {
+    registry.setVar("someName", "someValue");
+    registry.get("mock-embedding")!.create({ someKey: "$var:someName" });
+  });
+
+  it("should reject secrets that aren't passed as variables", () => {
+    registry.setVar("someName", "someValue");
+    expect(() =>
+      registry
+        .get("mock-embedding")!
+        .create({ secretKey: "someValue", someKey: "$var:someName" }),
+    ).toThrow(
+      'The key "secretKey" is sensitive and cannot be set directly. Please use the $var: syntax to set it.',
+    );
+  });
+
+  it("should not serialize secrets", () => {
+    registry.setVar("someName", "someValue");
+    registry.setVar("secretKey", "mySecret");
+    const func = registry
+      .get("mock-embedding")!
+      .create({ secretKey: "$var:secretKey", someKey: "$var:someName" });
+    expect(func.toJSON()).toEqual({
+      secretKey: "$var:secretKey",
+      someKey: "$var:someName",
+    });
+  });
+});

nodejs/__test__/table.test.ts

@@ -1038,9 +1038,6 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
     test("can search using a string", async () => {
       @register()
       class MockEmbeddingFunction extends EmbeddingFunction<string> {
-        toJSON(): object {
-          return {};
-        }
         ndims() {
           return 1;
         }