diff --git a/nodejs/__test__/embedding.test.ts b/nodejs/__test__/embedding.test.ts index 0133905dc..a1bf98bea 100644 --- a/nodejs/__test__/embedding.test.ts +++ b/nodejs/__test__/embedding.test.ts @@ -184,7 +184,7 @@ describe("embedding functions", () => { const vector0 = JSON.parse(JSON.stringify(arr[0].vector)); expect(vector0).toEqual([1, 2, 3]); }); - it("should append using embedding metadata created by Python", async () => { + it("should append multiple Python embeddings with the same alias", async () => { @register("python-mock") // biome-ignore lint/correctness/noUnusedVariables: the decorator registers this class class MockEmbeddingFunction extends EmbeddingFunction { @@ -198,21 +198,29 @@ describe("embedding functions", () => { return [1, 2, 3]; } async computeSourceEmbeddings(data: string[]) { - return data.map(() => [1, 2, 3]); + return data.map((value) => + value === "hello world" ? [1, 2, 3] : [4, 5, 6], + ); } } const metadata = new Map([ [ "embedding_functions", - '[{"source_column":"text","vector_column":"vector","name":"python-mock","model":{}}]', + '[{"source_column":"text1","vector_column":"vector1","name":"python-mock","model":{}},{"source_column":"text2","vector_column":"vector2","name":"python-mock","model":{}}]', ], ]); const schema = new Schema( [ - new Field("text", new Utf8(), true), + new Field("text1", new Utf8(), true), + new Field("text2", new Utf8(), true), new Field( - "vector", + "vector1", + new FixedSizeList(3, new Field("item", new Float32(), true)), + true, + ), + new Field( + "vector2", new FixedSizeList(3, new Field("item", new Float32(), true)), true, ), @@ -222,10 +230,11 @@ describe("embedding functions", () => { const db = await connect(tmpDir.name); const table = await db.createEmptyTable("test", schema); - await table.add([{ text: "hello world" }]); + await table.add([{ text1: "hello world", text2: "goodbye world" }]); const rows = await table.query().toArray(); - expect(JSON.parse(JSON.stringify(rows[0].vector))).toEqual([1, 2, 3]); + expect(JSON.parse(JSON.stringify(rows[0].vector1))).toEqual([1, 2, 3]); + expect(JSON.parse(JSON.stringify(rows[0].vector2))).toEqual([4, 5, 6]); }); it("should error when appending to a table with an unregistered embedding function", async () => { @register("mock") diff --git a/nodejs/lancedb/embedding/registry.ts b/nodejs/lancedb/embedding/registry.ts index d17613805..46d8f675c 100644 --- a/nodejs/lancedb/embedding/registry.ts +++ b/nodejs/lancedb/embedding/registry.ts @@ -118,7 +118,7 @@ export class EmbeddingFunctionRegistry { } const func = await this.get(f.name)!.create(f.model); return [ - f.name, + f.vectorColumn ?? "vector", { sourceColumn: f.sourceColumn, vectorColumn: f.vectorColumn,