@lancedb/lancedb 0.5.0 → 0.5.1

package/lancedb/embedding/index.ts

@@ -1,2 +1,113 @@
-export { EmbeddingFunction, isEmbeddingFunction } from "./embedding_function";
-export { OpenAIEmbeddingFunction } from "./openai";
+// Copyright 2023 Lance Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import { DataType, Field, Schema } from "../arrow";
+import { isDataType } from "../arrow";
+import { sanitizeType } from "../sanitize";
+import { EmbeddingFunction } from "./embedding_function";
+import { EmbeddingFunctionConfig, getRegistry } from "./registry";
+
+export { EmbeddingFunction } from "./embedding_function";
+
+// We need to explicitly export '*' so that the `register` decorator actually registers the class.
+export * from "./openai";
+export * from "./registry";
+
+/**
+ * Create a schema with embedding functions.
+ *
+ * @param fields
+ * @returns Schema
+ * @example
+ * ```ts
+ * class MyEmbeddingFunction extends EmbeddingFunction {
+ * // ...
+ * }
+ * const func = new MyEmbeddingFunction();
+ * const schema = LanceSchema({
+ *   id: new Int32(),
+ *   text: func.sourceField(new Utf8()),
+ *   vector: func.vectorField(),
+ *   // optional: specify the datatype and/or dimensions
+ *   vector2: func.vectorField({ datatype: new Float32(), dims: 3}),
+ * });
+ *
+ * const table = await db.createTable("my_table", data, { schema });
+ * ```
+ */
+export function LanceSchema(
+  fields: Record<string, [object, Map<string, EmbeddingFunction>] | object>,
+): Schema {
+  const arrowFields: Field[] = [];
+
+  const embeddingFunctions = new Map<
+    EmbeddingFunction,
+    Partial<EmbeddingFunctionConfig>
+  >();
+  Object.entries(fields).forEach(([key, value]) => {
+    if (isDataType(value)) {
+      arrowFields.push(new Field(key, sanitizeType(value), true));
+    } else {
+      const [dtype, metadata] = value as [
+        object,
+        Map<string, EmbeddingFunction>,
+      ];
+      arrowFields.push(new Field(key, sanitizeType(dtype), true));
+      parseEmbeddingFunctions(embeddingFunctions, key, metadata);
+    }
+  });
+  const registry = getRegistry();
+  const metadata = registry.getTableMetadata(
+    Array.from(embeddingFunctions.values()) as EmbeddingFunctionConfig[],
+  );
+  const schema = new Schema(arrowFields, metadata);
+  return schema;
+}
+
+function parseEmbeddingFunctions(
+  embeddingFunctions: Map<EmbeddingFunction, Partial<EmbeddingFunctionConfig>>,
+  key: string,
+  metadata: Map<string, EmbeddingFunction>,
+): void {
+  if (metadata.has("source_column_for")) {
+    const embedFunction = metadata.get("source_column_for")!;
+    const current = embeddingFunctions.get(embedFunction);
+    if (current !== undefined) {
+      embeddingFunctions.set(embedFunction, {
+        ...current,
+        sourceColumn: key,
+      });
+    } else {
+      embeddingFunctions.set(embedFunction, {
+        sourceColumn: key,
+        function: embedFunction,
+      });
+    }
+  } else if (metadata.has("vector_column_for")) {
+    const embedFunction = metadata.get("vector_column_for")!;
+
+    const current = embeddingFunctions.get(embedFunction);
+    if (current !== undefined) {
+      embeddingFunctions.set(embedFunction, {
+        ...current,
+        vectorColumn: key,
+      });
+    } else {
+      embeddingFunctions.set(embedFunction, {
+        vectorColumn: key,
+        function: embedFunction,
+      });
+    }
+  }
+}

package/lancedb/embedding/openai.ts

@@ -13,17 +13,31 @@
 // limitations under the License.
 
 import type OpenAI from "openai";
-import { type EmbeddingFunction } from "./embedding_function";
+import { Float, Float32 } from "../arrow";
+import { EmbeddingFunction } from "./embedding_function";
+import { register } from "./registry";
 
-export class OpenAIEmbeddingFunction implements EmbeddingFunction<string> {
-  private readonly _openai: OpenAI;
-  private readonly _modelName: string;
+export type OpenAIOptions = {
+  apiKey?: string;
+  model?: string;
+};
+
+@register("openai")
+export class OpenAIEmbeddingFunction extends EmbeddingFunction<
+  string,
+  OpenAIOptions
+> {
+  #openai: OpenAI;
+  #modelName: string;
+
+  constructor(options: OpenAIOptions = { model: "text-embedding-ada-002" }) {
+    super();
+    const openAIKey = options?.apiKey ?? process.env.OPENAI_API_KEY;
+    if (!openAIKey) {
+      throw new Error("OpenAI API key is required");
+    }
+    const modelName = options?.model ?? "text-embedding-ada-002";
 
-  constructor(
-    sourceColumn: string,
-    openAIKey: string,
-    modelName: string = "text-embedding-ada-002",
-  ) {
     /**
      * @type {import("openai").default}
      */
@@ -36,18 +50,40 @@ export class OpenAIEmbeddingFunction implements EmbeddingFunction<string> {
       throw new Error("please install openai@^4.24.1 using npm install openai");
     }
 
-    this.sourceColumn = sourceColumn;
     const configuration = {
       apiKey: openAIKey,
     };
 
-    this._openai = new Openai(configuration);
-    this._modelName = modelName;
+    this.#openai = new Openai(configuration);
+    this.#modelName = modelName;
+  }
+
+  toJSON() {
+    return {
+      model: this.#modelName,
+    };
+  }
+
+  ndims(): number {
+    switch (this.#modelName) {
+      case "text-embedding-ada-002":
+        return 1536;
+      case "text-embedding-3-large":
+        return 3072;
+      case "text-embedding-3-small":
+        return 1536;
+      default:
+        return null as never;
+    }
+  }
+
+  embeddingDataType(): Float {
+    return new Float32();
   }
 
-  async embed(data: string[]): Promise<number[][]> {
-    const response = await this._openai.embeddings.create({
-      model: this._modelName,
+  async computeSourceEmbeddings(data: string[]): Promise<number[][]> {
+    const response = await this.#openai.embeddings.create({
+      model: this.#modelName,
       input: data,
     });
 
@@ -58,5 +94,15 @@ export class OpenAIEmbeddingFunction implements EmbeddingFunction<string> {
     return embeddings;
   }
 
-  sourceColumn: string;
+  async computeQueryEmbeddings(data: string): Promise<number[]> {
+    if (typeof data !== "string") {
+      throw new Error("Data must be a string");
+    }
+    const response = await this.#openai.embeddings.create({
+      model: this.#modelName,
+      input: data,
+    });
+
+    return response.data[0].embedding;
+  }
 }

package/lancedb/embedding/registry.ts

@@ -0,0 +1,172 @@
+// Copyright 2024 Lance Developers.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import type { EmbeddingFunction } from "./embedding_function";
+import "reflect-metadata";
+
+export interface EmbeddingFunctionOptions {
+  [key: string]: unknown;
+}
+
+export interface EmbeddingFunctionFactory<
+  T extends EmbeddingFunction = EmbeddingFunction,
+> {
+  new (modelOptions?: EmbeddingFunctionOptions): T;
+}
+
+interface EmbeddingFunctionCreate<T extends EmbeddingFunction> {
+  create(options?: EmbeddingFunctionOptions): T;
+}
+
+/**
+ * This is a singleton class used to register embedding functions
+ * and fetch them by name. It also handles serializing and deserializing.
+ * You can implement your own embedding function by subclassing EmbeddingFunction
+ * or TextEmbeddingFunction and registering it with the registry
+ */
+export class EmbeddingFunctionRegistry {
+  #functions: Map<string, EmbeddingFunctionFactory> = new Map();
+
+  /**
+   * Register an embedding function
+   * @param name The name of the function
+   * @param func The function to register
+   */
+  register<T extends EmbeddingFunctionFactory = EmbeddingFunctionFactory>(
+    this: EmbeddingFunctionRegistry,
+    alias?: string,
+    // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+  ): (ctor: T) => any {
+    const self = this;
+    return function (ctor: T) {
+      if (!alias) {
+        alias = ctor.name;
+      }
+      if (self.#functions.has(alias)) {
+        throw new Error(
+          `Embedding function with alias "${alias}" already exists`,
+        );
+      }
+      self.#functions.set(alias, ctor);
+      Reflect.defineMetadata("lancedb::embedding::name", alias, ctor);
+      return ctor;
+    };
+  }
+
+  /**
+   * Fetch an embedding function by name
+   * @param name The name of the function
+   */
+  get<T extends EmbeddingFunction<unknown> = EmbeddingFunction>(
+    name: string,
+  ): EmbeddingFunctionCreate<T> | undefined {
+    const factory = this.#functions.get(name);
+    if (!factory) {
+      return undefined;
+    }
+    return {
+      create: function (options: EmbeddingFunctionOptions) {
+        return new factory(options) as unknown as T;
+      },
+    };
+  }
+
+  /**
+   * reset the registry to the initial state
+   */
+  reset(this: EmbeddingFunctionRegistry) {
+    this.#functions.clear();
+  }
+
+  parseFunctions(
+    this: EmbeddingFunctionRegistry,
+    metadata: Map<string, string>,
+  ): Map<string, EmbeddingFunctionConfig> {
+    if (!metadata.has("embedding_functions")) {
+      return new Map();
+    } else {
+      type FunctionConfig = {
+        name: string;
+        sourceColumn: string;
+        vectorColumn: string;
+        model: EmbeddingFunctionOptions;
+      };
+      const functions = <FunctionConfig[]>(
+        JSON.parse(metadata.get("embedding_functions")!)
+      );
+      return new Map(
+        functions.map((f) => {
+          const fn = this.get(f.name);
+          if (!fn) {
+            throw new Error(`Function "${f.name}" not found in registry`);
+          }
+          return [
+            f.name,
+            {
+              sourceColumn: f.sourceColumn,
+              vectorColumn: f.vectorColumn,
+              function: this.get(f.name)!.create(f.model),
+            },
+          ];
+        }),
+      );
+    }
+  }
+  // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+  functionToMetadata(conf: EmbeddingFunctionConfig): Record<string, any> {
+    // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+    const metadata: Record<string, any> = {};
+    const name = Reflect.getMetadata(
+      "lancedb::embedding::name",
+      conf.function.constructor,
+    );
+    metadata["sourceColumn"] = conf.sourceColumn;
+    metadata["vectorColumn"] = conf.vectorColumn ?? "vector";
+    metadata["name"] = name ?? conf.function.constructor.name;
+    metadata["model"] = conf.function.toJSON();
+    return metadata;
+  }
+
+  getTableMetadata(functions: EmbeddingFunctionConfig[]): Map<string, string> {
+    const metadata = new Map<string, string>();
+    const jsonData = functions.map((conf) => this.functionToMetadata(conf));
+    metadata.set("embedding_functions", JSON.stringify(jsonData));
+
+    return metadata;
+  }
+}
+
+const _REGISTRY = new EmbeddingFunctionRegistry();
+
+export function register(name?: string) {
+  return _REGISTRY.register(name);
+}
+
+/**
+ * Utility function to get the global instance of the registry
+ * @returns `EmbeddingFunctionRegistry` The global instance of the registry
+ * @example
+ * ```ts
+ * const registry = getRegistry();
+ * const openai = registry.get("openai").create();
+ */
+export function getRegistry(): EmbeddingFunctionRegistry {
+  return _REGISTRY;
+}
+
+export interface EmbeddingFunctionConfig {
+  sourceColumn: string;
+  vectorColumn?: string;
+  function: EmbeddingFunction;
+}

package/lancedb/query.ts

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import { Table as ArrowTable, RecordBatch, tableFromIPC } from "apache-arrow";
+import { Table as ArrowTable, RecordBatch, tableFromIPC } from "./arrow";
 import { type IvfPqOptions } from "./indices";
 import {
   RecordBatchIterator as NativeBatchIterator,
@@ -170,6 +170,7 @@ export class QueryBase<
   /** Collect the results as an array of objects. */
   async toArray(): Promise<unknown[]> {
     const tbl = await this.toArrow();
+
     // eslint-disable-next-line @typescript-eslint/no-unsafe-return
     return tbl.toArray();
   }

package/lancedb/sanitize.ts

@@ -20,6 +20,7 @@
 // comes from the exact same library instance.  This is not always the case
 // and so we must sanitize the input to ensure that it is compatible.
 
+import type { IntBitWidth, TKeys, TimeBitWidth } from "apache-arrow/type";
 import {
   Binary,
   Bool,
@@ -75,10 +76,9 @@ import {
   Uint64,
   Union,
   Utf8,
-} from "apache-arrow";
-import type { IntBitWidth, TKeys, TimeBitWidth } from "apache-arrow/type";
+} from "./arrow";
 
-function sanitizeMetadata(
+export function sanitizeMetadata(
   metadataLike?: unknown,
 ): Map<string, string> | undefined {
   if (metadataLike === undefined || metadataLike === null) {
@@ -97,7 +97,7 @@ function sanitizeMetadata(
   return metadataLike as Map<string, string>;
 }
 
-function sanitizeInt(typeLike: object) {
+export function sanitizeInt(typeLike: object) {
   if (
     !("bitWidth" in typeLike) ||
     typeof typeLike.bitWidth !== "number" ||
@@ -111,14 +111,14 @@ function sanitizeInt(typeLike: object) {
   return new Int(typeLike.isSigned, typeLike.bitWidth as IntBitWidth);
 }
 
-function sanitizeFloat(typeLike: object) {
+export function sanitizeFloat(typeLike: object) {
   if (!("precision" in typeLike) || typeof typeLike.precision !== "number") {
     throw Error("Expected a Float Type to have a `precision` property");
   }
   return new Float(typeLike.precision as Precision);
 }
 
-function sanitizeDecimal(typeLike: object) {
+export function sanitizeDecimal(typeLike: object) {
   if (
     !("scale" in typeLike) ||
     typeof typeLike.scale !== "number" ||
@@ -134,14 +134,14 @@ function sanitizeDecimal(typeLike: object) {
   return new Decimal(typeLike.scale, typeLike.precision, typeLike.bitWidth);
 }
 
-function sanitizeDate(typeLike: object) {
+export function sanitizeDate(typeLike: object) {
   if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
     throw Error("Expected a Date type to have a `unit` property");
   }
   return new Date_(typeLike.unit as DateUnit);
 }
 
-function sanitizeTime(typeLike: object) {
+export function sanitizeTime(typeLike: object) {
   if (
     !("unit" in typeLike) ||
     typeof typeLike.unit !== "number" ||
@@ -155,7 +155,7 @@ function sanitizeTime(typeLike: object) {
   return new Time(typeLike.unit, typeLike.bitWidth as TimeBitWidth);
 }
 
-function sanitizeTimestamp(typeLike: object) {
+export function sanitizeTimestamp(typeLike: object) {
   if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
     throw Error("Expected a Timestamp type to have a `unit` property");
   }
@@ -166,7 +166,7 @@ function sanitizeTimestamp(typeLike: object) {
   return new Timestamp(typeLike.unit, timezone);
 }
 
-function sanitizeTypedTimestamp(
+export function sanitizeTypedTimestamp(
   typeLike: object,
   // eslint-disable-next-line @typescript-eslint/naming-convention
   Datatype:
@@ -182,14 +182,14 @@ function sanitizeTypedTimestamp(
   return new Datatype(timezone);
 }
 
-function sanitizeInterval(typeLike: object) {
+export function sanitizeInterval(typeLike: object) {
   if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
     throw Error("Expected an Interval type to have a `unit` property");
   }
   return new Interval(typeLike.unit);
 }
 
-function sanitizeList(typeLike: object) {
+export function sanitizeList(typeLike: object) {
   if (!("children" in typeLike) || !Array.isArray(typeLike.children)) {
     throw Error(
       "Expected a List type to have an array-like `children` property",
@@ -201,7 +201,7 @@ function sanitizeList(typeLike: object) {
   return new List(sanitizeField(typeLike.children[0]));
 }
 
-function sanitizeStruct(typeLike: object) {
+export function sanitizeStruct(typeLike: object) {
   if (!("children" in typeLike) || !Array.isArray(typeLike.children)) {
     throw Error(
       "Expected a Struct type to have an array-like `children` property",
@@ -210,7 +210,7 @@ function sanitizeStruct(typeLike: object) {
   return new Struct(typeLike.children.map((child) => sanitizeField(child)));
 }
 
-function sanitizeUnion(typeLike: object) {
+export function sanitizeUnion(typeLike: object) {
   if (
     !("typeIds" in typeLike) ||
     !("mode" in typeLike) ||
@@ -234,7 +234,7 @@ function sanitizeUnion(typeLike: object) {
   );
 }
 
-function sanitizeTypedUnion(
+export function sanitizeTypedUnion(
   typeLike: object,
   // eslint-disable-next-line @typescript-eslint/naming-convention
   UnionType: typeof DenseUnion | typeof SparseUnion,
@@ -256,7 +256,7 @@ function sanitizeTypedUnion(
   );
 }
 
-function sanitizeFixedSizeBinary(typeLike: object) {
+export function sanitizeFixedSizeBinary(typeLike: object) {
   if (!("byteWidth" in typeLike) || typeof typeLike.byteWidth !== "number") {
     throw Error(
       "Expected a FixedSizeBinary type to have a `byteWidth` property",
@@ -265,7 +265,7 @@ function sanitizeFixedSizeBinary(typeLike: object) {
   return new FixedSizeBinary(typeLike.byteWidth);
 }
 
-function sanitizeFixedSizeList(typeLike: object) {
+export function sanitizeFixedSizeList(typeLike: object) {
   if (!("listSize" in typeLike) || typeof typeLike.listSize !== "number") {
     throw Error("Expected a FixedSizeList type to have a `listSize` property");
   }
@@ -283,7 +283,7 @@ function sanitizeFixedSizeList(typeLike: object) {
   );
 }
 
-function sanitizeMap(typeLike: object) {
+export function sanitizeMap(typeLike: object) {
   if (!("children" in typeLike) || !Array.isArray(typeLike.children)) {
     throw Error(
       "Expected a Map type to have an array-like `children` property",
@@ -300,14 +300,14 @@ function sanitizeMap(typeLike: object) {
   );
 }
 
-function sanitizeDuration(typeLike: object) {
+export function sanitizeDuration(typeLike: object) {
   if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
     throw Error("Expected a Duration type to have a `unit` property");
   }
   return new Duration(typeLike.unit);
 }
 
-function sanitizeDictionary(typeLike: object) {
+export function sanitizeDictionary(typeLike: object) {
   if (!("id" in typeLike) || typeof typeLike.id !== "number") {
     throw Error("Expected a Dictionary type to have an `id` property");
   }
@@ -329,7 +329,7 @@ function sanitizeDictionary(typeLike: object) {
 }
 
 // biome-ignore lint/suspicious/noExplicitAny: skip
-function sanitizeType(typeLike: unknown): DataType<any> {
+export function sanitizeType(typeLike: unknown): DataType<any> {
   if (typeof typeLike !== "object" || typeLike === null) {
     throw Error("Expected a Type but object was null/undefined");
   }
@@ -449,7 +449,7 @@ function sanitizeType(typeLike: unknown): DataType<any> {
   }
 }
 
-function sanitizeField(fieldLike: unknown): Field {
+export function sanitizeField(fieldLike: unknown): Field {
   if (fieldLike instanceof Field) {
     return fieldLike;
   }

package/lancedb/table.ts

@@ -12,8 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-import { Schema, tableFromIPC } from "apache-arrow";
-import { Data, fromDataToBuffer } from "./arrow";
+import { Data, Schema, fromDataToBuffer, tableFromIPC } from "./arrow";
+
+import { getRegistry } from "./embedding/registry";
 import { IndexOptions } from "./indices";
 import {
   AddColumnsSql,
@@ -122,8 +123,14 @@ export class Table {
    */
   async add(data: Data, options?: Partial<AddDataOptions>): Promise<void> {
     const mode = options?.mode ?? "append";
+    const schema = await this.schema();
+    const registry = getRegistry();
+    const functions = registry.parseFunctions(schema.metadata);
 
-    const buffer = await fromDataToBuffer(data);
+    const buffer = await fromDataToBuffer(
+      data,
+      functions.values().next().value,
+    );
     await this.inner.add(buffer, mode);
   }