@lancedb/lancedb 0.4.20 → 0.5.1

package/dist/connection.js

@@ -15,9 +15,9 @@
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.Connection = exports.connect = void 0;
 const arrow_1 = require("./arrow");
+const registry_1 = require("./embedding/registry");
 const native_1 = require("./native");
 const table_1 = require("./table");
-const apache_arrow_1 = require("apache-arrow");
 /**
  * Connect to a LanceDB instance at the given URI.
  *
@@ -109,13 +109,13 @@ class Connection {
             mode = "exist_ok";
         }
         let table;
-        if (data instanceof apache_arrow_1.Table) {
+        if ((0, arrow_1.isArrowTable)(data)) {
             table = data;
         }
         else {
-            table = (0, arrow_1.makeArrowTable)(data);
+            table = (0, arrow_1.makeArrowTable)(data, options);
         }
-        const buf = await (0, arrow_1.fromTableToBuffer)(table);
+        const buf = await (0, arrow_1.fromTableToBuffer)(table, options?.embeddingFunction, options?.schema);
         const innerTable = await this.inner.createTable(name, buf, mode, cleanseStorageOptions(options?.storageOptions));
         return new table_1.Table(innerTable);
     }
@@ -130,7 +130,13 @@ class Connection {
         if (mode === "create" && existOk) {
             mode = "exist_ok";
         }
-        const table = (0, arrow_1.makeEmptyTable)(schema);
+        let metadata = undefined;
+        if (options?.embeddingFunction !== undefined) {
+            const embeddingFunction = options.embeddingFunction;
+            const registry = (0, registry_1.getRegistry)();
+            metadata = registry.getTableMetadata([embeddingFunction]);
+        }
+        const table = (0, arrow_1.makeEmptyTable)(schema, metadata);
         const buf = await (0, arrow_1.fromTableToBuffer)(table);
         const innerTable = await this.inner.createEmptyTable(name, buf, mode, cleanseStorageOptions(options?.storageOptions));
         return new table_1.Table(innerTable);

package/dist/embedding/embedding_function.d.ts

@@ -1,45 +1,71 @@
-import { type Float } from "apache-arrow";
+import "reflect-metadata";
+import { DataType, Float } from "../arrow";
+/**
+ * Options for a given embedding function
+ */
+export interface FunctionOptions {
+    [key: string]: any;
+}
 /**
  * An embedding function that automatically creates vector representation for a given column.
  */
-export interface EmbeddingFunction<T> {
+export declare abstract class EmbeddingFunction<T = any, M extends FunctionOptions = FunctionOptions> {
     /**
-     * The name of the column that will be used as input for the Embedding Function.
-     */
-    sourceColumn: string;
-    /**
-     * The data type of the embedding
+     * Convert the embedding function to a JSON object
+     * It is used to serialize the embedding function to the schema
+     * It's important that any object returned by this method contains all the necessary
+     * information to recreate the embedding function
      *
-     * The embedding function should return `number`.  This will be converted into
-     * an Arrow float array.  By default this will be Float32 but this property can
-     * be used to control the conversion.
-     */
-    embeddingDataType?: Float;
-    /**
-     * The dimension of the embedding
+     * It should return the same object that was passed to the constructor
+     * If it does not, the embedding function will not be able to be recreated, or could be recreated incorrectly
      *
-     * This is optional, normally this can be determined by looking at the results of
-     * `embed`.  If this is not specified, and there is an attempt to apply the embedding
-     * to an empty table, then that process will fail.
+     * @example
+     * ```ts
+     * class MyEmbeddingFunction extends EmbeddingFunction {
+     *   constructor(options: {model: string, timeout: number}) {
+     *     super();
+     *     this.model = options.model;
+     *     this.timeout = options.timeout;
+     *   }
+     *   toJSON() {
+     *     return {
+     *       model: this.model,
+     *       timeout: this.timeout,
+     *     };
+     * }
+     * ```
      */
-    embeddingDimension?: number;
+    abstract toJSON(): Partial<M>;
     /**
-     * The name of the column that will contain the embedding
+     * sourceField is used in combination with `LanceSchema` to provide a declarative data model
+     *
+     * @param optionsOrDatatype - The options for the field or the datatype
      *
-     * By default this is "vector"
+     * @see {@link lancedb.LanceSchema}
      */
-    destColumn?: string;
+    sourceField(optionsOrDatatype: Partial<FieldOptions> | DataType): [DataType, Map<string, EmbeddingFunction>];
     /**
-     * Should the source column be excluded from the resulting table
+     * vectorField is used in combination with `LanceSchema` to provide a declarative data model
+     *
+     * @param options - The options for the field
      *
-     * By default the source column is included.  Set this to true and
-     * only the embedding will be stored.
+     * @see {@link lancedb.LanceSchema}
      */
-    excludeSource?: boolean;
+    vectorField(options?: Partial<FieldOptions>): [DataType, Map<string, EmbeddingFunction>];
+    /** The number of dimensions of the embeddings */
+    ndims(): number | undefined;
+    /** The datatype of the embeddings */
+    abstract embeddingDataType(): Float;
     /**
      * Creates a vector representation for the given values.
      */
-    embed: (data: T[]) => Promise<number[][]>;
+    abstract computeSourceEmbeddings(data: T[]): Promise<number[][] | Float32Array[] | Float64Array[]>;
+    /**
+    Compute the embeddings for a single query
+   */
+    computeQueryEmbeddings(data: T): Promise<number[] | Float32Array | Float64Array>;
+}
+export interface FieldOptions<T extends DataType = DataType> {
+    datatype: T;
+    dims?: number;
 }
-/** Test if the input seems to be an embedding function */
-export declare function isEmbeddingFunction<T>(value: unknown): value is EmbeddingFunction<T>;

package/dist/embedding/embedding_function.js

@@ -1,5 +1,5 @@
 "use strict";
-// Copyright 2023 Lance Developers.
+// 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.
@@ -13,15 +13,76 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.isEmbeddingFunction = void 0;
-/** Test if the input seems to be an embedding function */
-function isEmbeddingFunction(value) {
-    if (typeof value !== "object" || value === null) {
-        return false;
+exports.EmbeddingFunction = void 0;
+require("reflect-metadata");
+const arrow_1 = require("../arrow");
+const sanitize_1 = require("../sanitize");
+/**
+ * An embedding function that automatically creates vector representation for a given column.
+ */
+class EmbeddingFunction {
+    /**
+     * sourceField is used in combination with `LanceSchema` to provide a declarative data model
+     *
+     * @param optionsOrDatatype - The options for the field or the datatype
+     *
+     * @see {@link lancedb.LanceSchema}
+     */
+    sourceField(optionsOrDatatype) {
+        let datatype = (0, arrow_1.isDataType)(optionsOrDatatype)
+            ? optionsOrDatatype
+            : optionsOrDatatype?.datatype;
+        if (!datatype) {
+            throw new Error("Datatype is required");
+        }
+        datatype = (0, sanitize_1.sanitizeType)(datatype);
+        const metadata = new Map();
+        metadata.set("source_column_for", this);
+        return [datatype, metadata];
     }
-    if (!("sourceColumn" in value) || !("embed" in value)) {
-        return false;
+    /**
+     * vectorField is used in combination with `LanceSchema` to provide a declarative data model
+     *
+     * @param options - The options for the field
+     *
+     * @see {@link lancedb.LanceSchema}
+     */
+    vectorField(options) {
+        let dtype;
+        const dims = this.ndims() ?? options?.dims;
+        if (!options?.datatype) {
+            if (dims === undefined) {
+                throw new Error("ndims is required for vector field");
+            }
+            dtype = new arrow_1.FixedSizeList(dims, new arrow_1.Field("item", new arrow_1.Float32(), true));
+        }
+        else {
+            if ((0, arrow_1.isFixedSizeList)(options.datatype)) {
+                dtype = options.datatype;
+            }
+            else if ((0, arrow_1.isFloat)(options.datatype)) {
+                if (dims === undefined) {
+                    throw new Error("ndims is required for vector field");
+                }
+                dtype = (0, arrow_1.newVectorType)(dims, options.datatype);
+            }
+            else {
+                throw new Error("Expected FixedSizeList or Float as datatype for vector field");
+            }
+        }
+        const metadata = new Map();
+        metadata.set("vector_column_for", this);
+        return [dtype, metadata];
+    }
+    /** The number of dimensions of the embeddings */
+    ndims() {
+        return undefined;
+    }
+    /**
+    Compute the embeddings for a single query
+   */
+    async computeQueryEmbeddings(data) {
+        return this.computeSourceEmbeddings([data]).then((embeddings) => embeddings[0]);
     }
-    return (typeof value.sourceColumn === "string" && typeof value.embed === "function");
 }
-exports.isEmbeddingFunction = isEmbeddingFunction;
+exports.EmbeddingFunction = EmbeddingFunction;

package/dist/embedding/index.d.ts

@@ -1,2 +1,28 @@
-export { EmbeddingFunction, isEmbeddingFunction } from "./embedding_function";
-export { OpenAIEmbeddingFunction } from "./openai";
+import { Schema } from "../arrow";
+import { EmbeddingFunction } from "./embedding_function";
+export { EmbeddingFunction } from "./embedding_function";
+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 declare function LanceSchema(fields: Record<string, [object, Map<string, EmbeddingFunction>] | object>): Schema;

package/dist/embedding/index.js

@@ -1,7 +1,114 @@
 "use strict";
+// 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.
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __exportStar = (this && this.__exportStar) || function(m, exports) {
+    for (var p in m) if (p !== "default" && !Object.prototype.hasOwnProperty.call(exports, p)) __createBinding(exports, m, p);
+};
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.OpenAIEmbeddingFunction = exports.isEmbeddingFunction = void 0;
+exports.LanceSchema = exports.EmbeddingFunction = void 0;
+const arrow_1 = require("../arrow");
+const arrow_2 = require("../arrow");
+const sanitize_1 = require("../sanitize");
+const registry_1 = require("./registry");
 var embedding_function_1 = require("./embedding_function");
-Object.defineProperty(exports, "isEmbeddingFunction", { enumerable: true, get: function () { return embedding_function_1.isEmbeddingFunction; } });
-var openai_1 = require("./openai");
-Object.defineProperty(exports, "OpenAIEmbeddingFunction", { enumerable: true, get: function () { return openai_1.OpenAIEmbeddingFunction; } });
+Object.defineProperty(exports, "EmbeddingFunction", { enumerable: true, get: function () { return embedding_function_1.EmbeddingFunction; } });
+// We need to explicitly export '*' so that the `register` decorator actually registers the class.
+__exportStar(require("./openai"), exports);
+__exportStar(require("./registry"), exports);
+/**
+ * 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 });
+ * ```
+ */
+function LanceSchema(fields) {
+    const arrowFields = [];
+    const embeddingFunctions = new Map();
+    Object.entries(fields).forEach(([key, value]) => {
+        if ((0, arrow_2.isDataType)(value)) {
+            arrowFields.push(new arrow_1.Field(key, (0, sanitize_1.sanitizeType)(value), true));
+        }
+        else {
+            const [dtype, metadata] = value;
+            arrowFields.push(new arrow_1.Field(key, (0, sanitize_1.sanitizeType)(dtype), true));
+            parseEmbeddingFunctions(embeddingFunctions, key, metadata);
+        }
+    });
+    const registry = (0, registry_1.getRegistry)();
+    const metadata = registry.getTableMetadata(Array.from(embeddingFunctions.values()));
+    const schema = new arrow_1.Schema(arrowFields, metadata);
+    return schema;
+}
+exports.LanceSchema = LanceSchema;
+function parseEmbeddingFunctions(embeddingFunctions, key, metadata) {
+    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/dist/embedding/openai.d.ts

@@ -1,8 +1,17 @@
-import { type EmbeddingFunction } from "./embedding_function";
-export declare class OpenAIEmbeddingFunction implements EmbeddingFunction<string> {
-    private readonly _openai;
-    private readonly _modelName;
-    constructor(sourceColumn: string, openAIKey: string, modelName?: string);
-    embed(data: string[]): Promise<number[][]>;
-    sourceColumn: string;
+import { Float } from "../arrow";
+import { EmbeddingFunction } from "./embedding_function";
+export type OpenAIOptions = {
+    apiKey?: string;
+    model?: string;
+};
+export declare class OpenAIEmbeddingFunction extends EmbeddingFunction<string, OpenAIOptions> {
+    #private;
+    constructor(options?: OpenAIOptions);
+    toJSON(): {
+        model: string;
+    };
+    ndims(): number;
+    embeddingDataType(): Float;
+    computeSourceEmbeddings(data: string[]): Promise<number[][]>;
+    computeQueryEmbeddings(data: string): Promise<number[]>;
 }

package/dist/embedding/openai.js

@@ -12,12 +12,30 @@
 // 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.
+var __decorate = (this && this.__decorate) || function (decorators, target, key, desc) {
+    var c = arguments.length, r = c < 3 ? target : desc === null ? desc = Object.getOwnPropertyDescriptor(target, key) : desc, d;
+    if (typeof Reflect === "object" && typeof Reflect.decorate === "function") r = Reflect.decorate(decorators, target, key, desc);
+    else for (var i = decorators.length - 1; i >= 0; i--) if (d = decorators[i]) r = (c < 3 ? d(r) : c > 3 ? d(target, key, r) : d(target, key)) || r;
+    return c > 3 && r && Object.defineProperty(target, key, r), r;
+};
+var __metadata = (this && this.__metadata) || function (k, v) {
+    if (typeof Reflect === "object" && typeof Reflect.metadata === "function") return Reflect.metadata(k, v);
+};
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.OpenAIEmbeddingFunction = void 0;
-class OpenAIEmbeddingFunction {
-    _openai;
-    _modelName;
-    constructor(sourceColumn, openAIKey, modelName = "text-embedding-ada-002") {
+const arrow_1 = require("../arrow");
+const embedding_function_1 = require("./embedding_function");
+const registry_1 = require("./registry");
+let OpenAIEmbeddingFunction = class OpenAIEmbeddingFunction extends embedding_function_1.EmbeddingFunction {
+    #openai;
+    #modelName;
+    constructor(options = { 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";
         /**
          * @type {import("openai").default}
          */
@@ -30,16 +48,35 @@ class OpenAIEmbeddingFunction {
         catch {
             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() {
+        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;
+        }
     }
-    async embed(data) {
-        const response = await this._openai.embeddings.create({
-            model: this._modelName,
+    embeddingDataType() {
+        return new arrow_1.Float32();
+    }
+    async computeSourceEmbeddings(data) {
+        const response = await this.#openai.embeddings.create({
+            model: this.#modelName,
             input: data,
         });
         const embeddings = [];
@@ -48,6 +85,19 @@ class OpenAIEmbeddingFunction {
         }
         return embeddings;
     }
-    sourceColumn;
-}
+    async computeQueryEmbeddings(data) {
+        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;
+    }
+};
 exports.OpenAIEmbeddingFunction = OpenAIEmbeddingFunction;
+exports.OpenAIEmbeddingFunction = OpenAIEmbeddingFunction = __decorate([
+    (0, registry_1.register)("openai"),
+    __metadata("design:paramtypes", [Object])
+], OpenAIEmbeddingFunction);

package/dist/embedding/registry.d.ts

@@ -0,0 +1,54 @@
+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 declare class EmbeddingFunctionRegistry {
+    #private;
+    /**
+     * 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): (ctor: T) => any;
+    /**
+     * Fetch an embedding function by name
+     * @param name The name of the function
+     */
+    get<T extends EmbeddingFunction<unknown> = EmbeddingFunction>(name: string): EmbeddingFunctionCreate<T> | undefined;
+    /**
+     * reset the registry to the initial state
+     */
+    reset(this: EmbeddingFunctionRegistry): void;
+    parseFunctions(this: EmbeddingFunctionRegistry, metadata: Map<string, string>): Map<string, EmbeddingFunctionConfig>;
+    functionToMetadata(conf: EmbeddingFunctionConfig): Record<string, any>;
+    getTableMetadata(functions: EmbeddingFunctionConfig[]): Map<string, string>;
+}
+export declare function register(name?: string): (ctor: EmbeddingFunctionFactory<EmbeddingFunction<any, import("./embedding_function").FunctionOptions>>) => any;
+/**
+ * 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 declare function getRegistry(): EmbeddingFunctionRegistry;
+export interface EmbeddingFunctionConfig {
+    sourceColumn: string;
+    vectorColumn?: string;
+    function: EmbeddingFunction;
+}
+export {};