From 980f910f50a05b95efcbec32dbc838474b3ea0b0 Mon Sep 17 00:00:00 2001 From: Lei Xu Date: Tue, 18 Jul 2023 16:15:27 -0700 Subject: [PATCH] [Node] initial support of nodejs remote sdk (#333) --- node/package-lock.json | 79 +++++++++++------- node/package.json | 3 +- node/src/index.ts | 140 +++++--------------------------- node/src/query.ts | 130 ++++++++++++++++++++++++++++++ node/src/remote/client.ts | 69 ++++++++++++++++ node/src/remote/index.ts | 163 ++++++++++++++++++++++++++++++++++++++ node/src/test/test.ts | 5 +- 7 files changed, 438 insertions(+), 151 deletions(-) create mode 100644 node/src/query.ts create mode 100644 node/src/remote/client.ts create mode 100644 node/src/remote/index.ts diff --git a/node/package-lock.json b/node/package-lock.json index 68c20c77..4ae03b32 100644 --- a/node/package-lock.json +++ b/node/package-lock.json @@ -20,7 +20,8 @@ "dependencies": { "@apache-arrow/ts": "^12.0.0", "@neon-rs/load": "^0.0.74", - "apache-arrow": "^12.0.0" + "apache-arrow": "^12.0.0", + "axios": "^1.4.0" }, "devDependencies": { "@neon-rs/cli": "^0.0.74", @@ -842,8 +843,7 @@ "node_modules/asynckit": { "version": "0.4.0", "resolved": "https://registry.npmjs.org/asynckit/-/asynckit-0.4.0.tgz", - "integrity": "sha512-Oei9OH4tRh0YqU3GxhX79dM/mwVgvbZJaSNaRk+bshkj0S5cfHcgYakreBjrHwatXKbz+IoIdYLxrKim2MjW0Q==", - "dev": true + "integrity": "sha512-Oei9OH4tRh0YqU3GxhX79dM/mwVgvbZJaSNaRk+bshkj0S5cfHcgYakreBjrHwatXKbz+IoIdYLxrKim2MjW0Q==" }, "node_modules/available-typed-arrays": { "version": "1.0.5", @@ -858,12 +858,13 @@ } }, "node_modules/axios": { - "version": "0.26.1", - "resolved": "https://registry.npmjs.org/axios/-/axios-0.26.1.tgz", - "integrity": "sha512-fPwcX4EvnSHuInCMItEhAGnaSEXRBjtzh9fOtsE6E1G6p7vl7edEeZe11QHf18+6+9gR5PbKV/sGKNaD8YaMeA==", - "dev": true, + "version": "1.4.0", + "resolved": "https://registry.npmjs.org/axios/-/axios-1.4.0.tgz", + "integrity": "sha512-S4XCWMEmzvo64T9GfvQDOXgYRDJ/wsSZc7Jvdgx5u1sd0JwsuPLqb3SYmusag+edF6ziyMensPVqLTSc1PiSEA==", "dependencies": { - "follow-redirects": "^1.14.8" + "follow-redirects": "^1.15.0", + "form-data": "^4.0.0", + "proxy-from-env": "^1.1.0" } }, "node_modules/balanced-match": { @@ -1094,7 +1095,6 @@ "version": "1.0.8", "resolved": "https://registry.npmjs.org/combined-stream/-/combined-stream-1.0.8.tgz", "integrity": "sha512-FQN4MRfuJeHf7cBbBMJFXhKSDq+2kAArBlmRBvcvFE5BB1HZKXtSFASDhdlz9zOYwxh8lDdnvmMOe/+5cdoEdg==", - "dev": true, "dependencies": { "delayed-stream": "~1.0.0" }, @@ -1317,7 +1317,6 @@ "version": "1.0.0", "resolved": "https://registry.npmjs.org/delayed-stream/-/delayed-stream-1.0.0.tgz", "integrity": "sha512-ZySD7Nf91aLB0RxL4KGrKHBXl7Eds1DAmEdcoVawXnLD7SDhpNgtuII2aAkg7a7QS41jxPSZ17p4VdGnMHk3MQ==", - "dev": true, "engines": { "node": ">=0.4.0" } @@ -2084,7 +2083,6 @@ "version": "1.15.2", "resolved": "https://registry.npmjs.org/follow-redirects/-/follow-redirects-1.15.2.tgz", "integrity": "sha512-VQLG33o04KaQ8uYi2tVNbdrWp1QWxNNea+nmIB4EVM28v0hmP17z7aG1+wAkNzVq4KeXTq3221ye5qTJP91JwA==", - "dev": true, "funding": [ { "type": "individual", @@ -2113,7 +2111,6 @@ "version": "4.0.0", "resolved": "https://registry.npmjs.org/form-data/-/form-data-4.0.0.tgz", "integrity": "sha512-ETEklSGi5t0QMZuiXoA/Q6vcnxcLQP5vdugSpuAyi6SVGi2clPPp+xgEhuMaHC+zGgn31Kd235W35f7Hykkaww==", - "dev": true, "dependencies": { "asynckit": "^0.4.0", "combined-stream": "^1.0.8", @@ -2987,7 +2984,6 @@ "version": "1.52.0", "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.52.0.tgz", "integrity": "sha512-sPU4uV7dYlvtWJxwwxHD0PuihVNiE7TyAbQ5SWxDCB9mUYvOgroQOwYQQOKPJ8CIbE+1ETVlOoK1UC2nU3gYvg==", - "dev": true, "engines": { "node": ">= 0.6" } @@ -2996,7 +2992,6 @@ "version": "2.1.35", "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-2.1.35.tgz", "integrity": "sha512-ZDY+bPm5zTTF+YpCrAU9nK0UgICYPT0QtT1NZWFv4s++TNkcgVaT0g6+4R2uI4MjQjzysHB1zxuWL50hzaeXiw==", - "dev": true, "dependencies": { "mime-db": "1.52.0" }, @@ -3290,6 +3285,15 @@ "form-data": "^4.0.0" } }, + "node_modules/openai/node_modules/axios": { + "version": "0.26.1", + "resolved": "https://registry.npmjs.org/axios/-/axios-0.26.1.tgz", + "integrity": "sha512-fPwcX4EvnSHuInCMItEhAGnaSEXRBjtzh9fOtsE6E1G6p7vl7edEeZe11QHf18+6+9gR5PbKV/sGKNaD8YaMeA==", + "dev": true, + "dependencies": { + "follow-redirects": "^1.14.8" + } + }, "node_modules/optionator": { "version": "0.9.1", "resolved": "https://registry.npmjs.org/optionator/-/optionator-0.9.1.tgz", @@ -3441,6 +3445,11 @@ "node": ">= 0.8.0" } }, + "node_modules/proxy-from-env": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/proxy-from-env/-/proxy-from-env-1.1.0.tgz", + "integrity": "sha512-D+zkORCbA9f1tdWRK0RaCR3GPv50cMxcrz4X8k5LTSUD1Dkw47mKJEZQNunItRTkWwgtaUSo1RVFRIG9ZXiFYg==" + }, "node_modules/punycode": { "version": "2.3.0", "resolved": "https://registry.npmjs.org/punycode/-/punycode-2.3.0.tgz", @@ -5135,8 +5144,7 @@ "asynckit": { "version": "0.4.0", "resolved": "https://registry.npmjs.org/asynckit/-/asynckit-0.4.0.tgz", - "integrity": "sha512-Oei9OH4tRh0YqU3GxhX79dM/mwVgvbZJaSNaRk+bshkj0S5cfHcgYakreBjrHwatXKbz+IoIdYLxrKim2MjW0Q==", - "dev": true + "integrity": "sha512-Oei9OH4tRh0YqU3GxhX79dM/mwVgvbZJaSNaRk+bshkj0S5cfHcgYakreBjrHwatXKbz+IoIdYLxrKim2MjW0Q==" }, "available-typed-arrays": { "version": "1.0.5", @@ -5145,12 +5153,13 @@ "dev": true }, "axios": { - "version": "0.26.1", - "resolved": "https://registry.npmjs.org/axios/-/axios-0.26.1.tgz", - "integrity": "sha512-fPwcX4EvnSHuInCMItEhAGnaSEXRBjtzh9fOtsE6E1G6p7vl7edEeZe11QHf18+6+9gR5PbKV/sGKNaD8YaMeA==", - "dev": true, + "version": "1.4.0", + "resolved": "https://registry.npmjs.org/axios/-/axios-1.4.0.tgz", + "integrity": "sha512-S4XCWMEmzvo64T9GfvQDOXgYRDJ/wsSZc7Jvdgx5u1sd0JwsuPLqb3SYmusag+edF6ziyMensPVqLTSc1PiSEA==", "requires": { - "follow-redirects": "^1.14.8" + "follow-redirects": "^1.15.0", + "form-data": "^4.0.0", + "proxy-from-env": "^1.1.0" } }, "balanced-match": { @@ -5330,7 +5339,6 @@ "version": "1.0.8", "resolved": "https://registry.npmjs.org/combined-stream/-/combined-stream-1.0.8.tgz", "integrity": "sha512-FQN4MRfuJeHf7cBbBMJFXhKSDq+2kAArBlmRBvcvFE5BB1HZKXtSFASDhdlz9zOYwxh8lDdnvmMOe/+5cdoEdg==", - "dev": true, "requires": { "delayed-stream": "~1.0.0" } @@ -5497,8 +5505,7 @@ "delayed-stream": { "version": "1.0.0", "resolved": "https://registry.npmjs.org/delayed-stream/-/delayed-stream-1.0.0.tgz", - "integrity": "sha512-ZySD7Nf91aLB0RxL4KGrKHBXl7Eds1DAmEdcoVawXnLD7SDhpNgtuII2aAkg7a7QS41jxPSZ17p4VdGnMHk3MQ==", - "dev": true + "integrity": "sha512-ZySD7Nf91aLB0RxL4KGrKHBXl7Eds1DAmEdcoVawXnLD7SDhpNgtuII2aAkg7a7QS41jxPSZ17p4VdGnMHk3MQ==" }, "diff": { "version": "4.0.2", @@ -6068,8 +6075,7 @@ "follow-redirects": { "version": "1.15.2", "resolved": "https://registry.npmjs.org/follow-redirects/-/follow-redirects-1.15.2.tgz", - "integrity": "sha512-VQLG33o04KaQ8uYi2tVNbdrWp1QWxNNea+nmIB4EVM28v0hmP17z7aG1+wAkNzVq4KeXTq3221ye5qTJP91JwA==", - "dev": true + "integrity": "sha512-VQLG33o04KaQ8uYi2tVNbdrWp1QWxNNea+nmIB4EVM28v0hmP17z7aG1+wAkNzVq4KeXTq3221ye5qTJP91JwA==" }, "for-each": { "version": "0.3.3", @@ -6084,7 +6090,6 @@ "version": "4.0.0", "resolved": "https://registry.npmjs.org/form-data/-/form-data-4.0.0.tgz", "integrity": "sha512-ETEklSGi5t0QMZuiXoA/Q6vcnxcLQP5vdugSpuAyi6SVGi2clPPp+xgEhuMaHC+zGgn31Kd235W35f7Hykkaww==", - "dev": true, "requires": { "asynckit": "^0.4.0", "combined-stream": "^1.0.8", @@ -6698,14 +6703,12 @@ "mime-db": { "version": "1.52.0", "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.52.0.tgz", - "integrity": "sha512-sPU4uV7dYlvtWJxwwxHD0PuihVNiE7TyAbQ5SWxDCB9mUYvOgroQOwYQQOKPJ8CIbE+1ETVlOoK1UC2nU3gYvg==", - "dev": true + "integrity": "sha512-sPU4uV7dYlvtWJxwwxHD0PuihVNiE7TyAbQ5SWxDCB9mUYvOgroQOwYQQOKPJ8CIbE+1ETVlOoK1UC2nU3gYvg==" }, "mime-types": { "version": "2.1.35", "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-2.1.35.tgz", "integrity": "sha512-ZDY+bPm5zTTF+YpCrAU9nK0UgICYPT0QtT1NZWFv4s++TNkcgVaT0g6+4R2uI4MjQjzysHB1zxuWL50hzaeXiw==", - "dev": true, "requires": { "mime-db": "1.52.0" } @@ -6931,6 +6934,17 @@ "requires": { "axios": "^0.26.0", "form-data": "^4.0.0" + }, + "dependencies": { + "axios": { + "version": "0.26.1", + "resolved": "https://registry.npmjs.org/axios/-/axios-0.26.1.tgz", + "integrity": "sha512-fPwcX4EvnSHuInCMItEhAGnaSEXRBjtzh9fOtsE6E1G6p7vl7edEeZe11QHf18+6+9gR5PbKV/sGKNaD8YaMeA==", + "dev": true, + "requires": { + "follow-redirects": "^1.14.8" + } + } } }, "optionator": { @@ -7039,6 +7053,11 @@ "integrity": "sha512-vkcDPrRZo1QZLbn5RLGPpg/WmIQ65qoWWhcGKf/b5eplkkarX0m9z8ppCat4mlOqUsWpyNuYgO3VRyrYHSzX5g==", "dev": true }, + "proxy-from-env": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/proxy-from-env/-/proxy-from-env-1.1.0.tgz", + "integrity": "sha512-D+zkORCbA9f1tdWRK0RaCR3GPv50cMxcrz4X8k5LTSUD1Dkw47mKJEZQNunItRTkWwgtaUSo1RVFRIG9ZXiFYg==" + }, "punycode": { "version": "2.3.0", "resolved": "https://registry.npmjs.org/punycode/-/punycode-2.3.0.tgz", diff --git a/node/package.json b/node/package.json index 7487dcd2..528067d3 100644 --- a/node/package.json +++ b/node/package.json @@ -56,7 +56,8 @@ "dependencies": { "@apache-arrow/ts": "^12.0.0", "@neon-rs/load": "^0.0.74", - "apache-arrow": "^12.0.0" + "apache-arrow": "^12.0.0", + "axios": "^1.4.0" }, "os": [ "darwin", diff --git a/node/src/index.ts b/node/src/index.ts index 0636cb95..887aaa8c 100644 --- a/node/src/index.ts +++ b/node/src/index.ts @@ -14,15 +14,15 @@ import { RecordBatchFileWriter, - type Table as ArrowTable, - tableFromIPC, - Vector + type Table as ArrowTable } from 'apache-arrow' import { fromRecordsToBuffer } from './arrow' import type { EmbeddingFunction } from './embedding/embedding_function' +import { RemoteConnection } from './remote' +import { Query } from './query' // eslint-disable-next-line @typescript-eslint/no-var-requires -const { databaseNew, databaseTableNames, databaseOpenTable, databaseDropTable, tableCreate, tableSearch, tableAdd, tableCreateVectorIndex, tableCountRows, tableDelete } = require('../native.js') +const { databaseNew, databaseTableNames, databaseOpenTable, databaseDropTable, tableCreate, tableAdd, tableCreateVectorIndex, tableCountRows, tableDelete } = require('../native.js') export type { EmbeddingFunction } export { OpenAIEmbeddingFunction } from './embedding/openai' @@ -37,7 +37,13 @@ export interface AwsCredentials { export interface ConnectionOptions { uri: string + awsCredentials?: AwsCredentials + + // API key for the remote connections + apiKey?: string + // Region to connect + region?: string } /** @@ -54,9 +60,16 @@ export async function connect (arg: string | Partial): Promis // opts = { uri: arg.uri, awsCredentials = arg.awsCredentials } opts = Object.assign({ uri: '', - awsCredentials: undefined + awsCredentials: undefined, + apiKey: undefined, + region: 'us-west-2' }, arg) } + + if (opts.uri.startsWith('db://')) { + // Remote connection + return new RemoteConnection(opts) + } const db = await databaseNew(opts.uri) return new LocalConnection(db, opts) } @@ -191,8 +204,8 @@ export class LocalConnection implements Connection { } /** - * Get the names of all tables in the database. - */ + * Get the names of all tables in the database. + */ async tableNames (): Promise { return databaseTableNames.call(this._db) } @@ -203,6 +216,7 @@ export class LocalConnection implements Connection { * @param name The name of the table. */ async openTable (name: string): Promise + /** * Open a table in the database. * @@ -308,7 +322,7 @@ export class LocalTable implements Table { * @param query The query search term */ search (query: T): Query { - return new Query(this._tbl, query, this._embeddings) + return new Query(query, this._tbl, this._embeddings) } /** @@ -430,116 +444,6 @@ export interface IvfPQIndexConfig { export type VectorIndexParams = IvfPQIndexConfig -/** - * A builder for nearest neighbor queries for LanceDB. - */ -export class Query { - private readonly _tbl: any - private readonly _query: T - private _queryVector?: number[] - private _limit: number - private _refineFactor?: number - private _nprobes: number - private _select?: string[] - private _filter?: string - private _metricType?: MetricType - private readonly _embeddings?: EmbeddingFunction - - constructor (tbl: any, query: T, embeddings?: EmbeddingFunction) { - this._tbl = tbl - this._query = query - this._limit = 10 - this._nprobes = 20 - this._refineFactor = undefined - this._select = undefined - this._filter = undefined - this._metricType = undefined - this._embeddings = embeddings - } - - /*** - * Sets the number of results that will be returned - * @param value number of results - */ - limit (value: number): Query { - this._limit = value - return this - } - - /** - * Refine the results by reading extra elements and re-ranking them in memory. - * @param value refine factor to use in this query. - */ - refineFactor (value: number): Query { - this._refineFactor = value - return this - } - - /** - * The number of probes used. A higher number makes search more accurate but also slower. - * @param value The number of probes used. - */ - nprobes (value: number): Query { - this._nprobes = value - return this - } - - /** - * A filter statement to be applied to this query. - * @param value A filter in the same format used by a sql WHERE clause. - */ - filter (value: string): Query { - this._filter = value - return this - } - - where = this.filter - - /** Return only the specified columns. - * - * @param value Only select the specified columns. If not specified, all columns will be returned. - */ - select (value: string[]): Query { - this._select = value - return this - } - - /** - * The MetricType used for this Query. - * @param value The metric to the. @see MetricType for the different options - */ - metricType (value: MetricType): Query { - this._metricType = value - return this - } - - /** - * Execute the query and return the results as an Array of Objects - */ - async execute> (): Promise { - if (this._embeddings !== undefined) { - this._queryVector = (await this._embeddings.embed([this._query]))[0] - } else { - this._queryVector = this._query as number[] - } - - const buffer = await tableSearch.call(this._tbl, this) - const data = tableFromIPC(buffer) - - return data.toArray().map((entry: Record) => { - const newObject: Record = {} - Object.keys(entry).forEach((key: string) => { - if (entry[key] instanceof Vector) { - newObject[key] = (entry[key] as Vector).toArray() - } else { - newObject[key] = entry[key] - } - }) - return newObject as unknown as T - }) - } -} - /** * Write mode for writing a table. */ diff --git a/node/src/query.ts b/node/src/query.ts new file mode 100644 index 00000000..e2717b4e --- /dev/null +++ b/node/src/query.ts @@ -0,0 +1,130 @@ +// Copyright 2023 LanceDB 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 { Vector, tableFromIPC } from 'apache-arrow' +import { type EmbeddingFunction } from './embedding/embedding_function' +import { type MetricType } from '.' + +// eslint-disable-next-line @typescript-eslint/no-var-requires +const { tableSearch } = require('../native.js') + +/** + * A builder for nearest neighbor queries for LanceDB. + */ +export class Query { + private readonly _query: T + private readonly _tbl?: any + private _queryVector?: number[] + private _limit: number + private _refineFactor?: number + private _nprobes: number + private _select?: string[] + private _filter?: string + private _metricType?: MetricType + protected readonly _embeddings?: EmbeddingFunction + + constructor (query: T, tbl?: any, embeddings?: EmbeddingFunction) { + this._tbl = tbl + this._query = query + this._limit = 10 + this._nprobes = 20 + this._refineFactor = undefined + this._select = undefined + this._filter = undefined + this._metricType = undefined + this._embeddings = embeddings + } + + /*** + * Sets the number of results that will be returned + * @param value number of results + */ + limit (value: number): Query { + this._limit = value + return this + } + + /** + * Refine the results by reading extra elements and re-ranking them in memory. + * @param value refine factor to use in this query. + */ + refineFactor (value: number): Query { + this._refineFactor = value + return this + } + + /** + * The number of probes used. A higher number makes search more accurate but also slower. + * @param value The number of probes used. + */ + nprobes (value: number): Query { + this._nprobes = value + return this + } + + /** + * A filter statement to be applied to this query. + * @param value A filter in the same format used by a sql WHERE clause. + */ + filter (value: string): Query { + this._filter = value + return this + } + + where = this.filter + + /** Return only the specified columns. + * + * @param value Only select the specified columns. If not specified, all columns will be returned. + */ + select (value: string[]): Query { + this._select = value + return this + } + + /** + * The MetricType used for this Query. + * @param value The metric to the. @see MetricType for the different options + */ + metricType (value: MetricType): Query { + this._metricType = value + return this + } + + /** + * Execute the query and return the results as an Array of Objects + */ + async execute> (): Promise { + if (this._embeddings !== undefined) { + this._queryVector = (await this._embeddings.embed([this._query]))[0] + } else { + this._queryVector = this._query as number[] + } + + const buffer = await tableSearch.call(this._tbl, this) + const data = tableFromIPC(buffer) + + return data.toArray().map((entry: Record) => { + const newObject: Record = {} + Object.keys(entry).forEach((key: string) => { + if (entry[key] instanceof Vector) { + newObject[key] = (entry[key] as Vector).toArray() + } else { + newObject[key] = entry[key] + } + }) + return newObject as unknown as T + }) + } +} diff --git a/node/src/remote/client.ts b/node/src/remote/client.ts new file mode 100644 index 00000000..a2e74ead --- /dev/null +++ b/node/src/remote/client.ts @@ -0,0 +1,69 @@ +// Copyright 2023 LanceDB 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 axios from 'axios' + +import { tableFromIPC, type Table as ArrowTable } from 'apache-arrow' + +export class HttpLancedbClient { + private readonly _url: string + + public constructor (url: string, private readonly _apiKey: string) { + this._url = url + } + + get uri (): string { + return this._url + } + + public async search ( + tableName: string, + vector: number[], + k: number, + nprobes: number, + refineFactor?: number, + columns?: string[], + filter?: string + ): Promise> { + const response = await axios.post( + `${this._url}/v1/table/${tableName}`, + { + vector, + k, + nprobes, + refineFactor, + columns, + filter + }, + { + headers: { + 'Content-Type': 'application/json', + 'x-api-key': this._apiKey + }, + responseType: 'arraybuffer', + timeout: 10000 + } + ).catch((err) => { + console.error('error: ', err) + return err.response + }) + if (response.status !== 200) { + const errorData = new TextDecoder().decode(response.data) + throw new Error(`Server Error, status: ${response.status as number}, message: ${response.statusText as string}: ${errorData}`) + } + + const table = tableFromIPC(response.data) + return table + } +} diff --git a/node/src/remote/index.ts b/node/src/remote/index.ts new file mode 100644 index 00000000..92ae066a --- /dev/null +++ b/node/src/remote/index.ts @@ -0,0 +1,163 @@ +// Copyright 2023 LanceDB 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, type Table, type VectorIndexParams, type Connection, + type ConnectionOptions +} from '../index' +import { Query } from '../query' + +import { type Table as ArrowTable, Vector } from 'apache-arrow' +import { HttpLancedbClient } from './client' + +/** + * Remote connection. + */ +export class RemoteConnection implements Connection { + private readonly _client: HttpLancedbClient + private readonly _dbName: string + + constructor (opts: ConnectionOptions) { + if (!opts.uri.startsWith('db://')) { + throw new Error(`Invalid remote DB URI: ${opts.uri}`) + } + if (opts.apiKey === undefined || opts.region === undefined) { + throw new Error('API key and region are not supported for remote connections') + } + + this._dbName = opts.uri.slice('db://'.length) + const server = `https://${this._dbName}.${opts.region}.api.lancedb.com` + this._client = new HttpLancedbClient(server, opts.apiKey) + } + + get uri (): string { + // add the lancedb+ prefix back + return 'db://' + this._client.uri + } + + async tableNames (): Promise { + throw new Error('Not implemented') + } + + async openTable (name: string): Promise
+ async openTable (name: string, embeddings: EmbeddingFunction): Promise> + async openTable (name: string, embeddings?: EmbeddingFunction): Promise> { + if (embeddings !== undefined) { + return new RemoteTable(this._client, name, embeddings) + } else { + return new RemoteTable(this._client, name) + } + } + + async createTable (name: string, data: Array>): Promise
+ async createTable (name: string, data: Array>, embeddings: EmbeddingFunction): Promise> + async createTable (name: string, data: Array>, embeddings?: EmbeddingFunction): Promise> { + throw new Error('Not implemented') + } + + async createTableArrow (name: string, table: ArrowTable): Promise
{ + throw new Error('Not implemented') + } + + async dropTable (name: string): Promise { + throw new Error('Not implemented') + } +} + +export class RemoteQuery extends Query { + constructor (query: T, private readonly _client: HttpLancedbClient, + private readonly _name: string, embeddings?: EmbeddingFunction) { + super(query, undefined, embeddings) + } + + // TODO: refactor this to a base class + queryImpl pattern + async execute>(): Promise { + // TODO: remove as any hack once we refactor + const embeddings = this._embeddings + const query = (this as any)._query + let queryVector: number[] + + if (embeddings !== undefined) { + queryVector = (await embeddings.embed([query]))[0] + } else { + queryVector = query as number[] + } + + const data = await this._client.search( + this._name, + queryVector, + (this as any)._limit, + (this as any)._nprobes, + (this as any)._refineFactor, + (this as any)._select, + (this as any)._filter + ) + + return data.toArray().map((entry: Record) => { + const newObject: Record = {} + Object.keys(entry).forEach((key: string) => { + if (entry[key] instanceof Vector) { + newObject[key] = (entry[key] as Vector).toArray() + } else { + newObject[key] = entry[key] + } + }) + return newObject as unknown as T + }) + } +} + +// we are using extend until we have next next version release +// Table and Connection has both been refactored to interfaces +export class RemoteTable implements Table { + private readonly _client: HttpLancedbClient + private readonly _embeddings?: EmbeddingFunction + private readonly _name: string + + constructor (client: HttpLancedbClient, name: string) + constructor (client: HttpLancedbClient, name: string, embeddings: EmbeddingFunction) + constructor (client: HttpLancedbClient, name: string, embeddings?: EmbeddingFunction) { + this._client = client + this._name = name + this._embeddings = embeddings + } + + get name (): string { + return this._name + } + + search (query: T): Query { + return new RemoteQuery(query, this._client, this._name)//, this._embeddings_new) + } + + async add (data: Array>): Promise { + throw new Error('Not implemented') + } + + async overwrite (data: Array>): Promise { + throw new Error('Not implemented') + } + + async createIndex (indexParams: VectorIndexParams): Promise { + throw new Error('Not implemented') + } + + async countRows (): Promise { + throw new Error('Not implemented') + } + + async delete (filter: string): Promise { + throw new Error('Not implemented') + } +} diff --git a/node/src/test/test.ts b/node/src/test/test.ts index aefb6f18..cdfca32e 100644 --- a/node/src/test/test.ts +++ b/node/src/test/test.ts @@ -18,7 +18,8 @@ import * as chai from 'chai' import * as chaiAsPromised from 'chai-as-promised' import * as lancedb from '../index' -import { type AwsCredentials, type EmbeddingFunction, MetricType, Query, WriteMode } from '../index' +import { type AwsCredentials, type EmbeddingFunction, MetricType, WriteMode } from '../index' +import { Query } from '../query' const expect = chai.expect const assert = chai.assert @@ -268,7 +269,7 @@ describe('LanceDB client', function () { describe('Query object', function () { it('sets custom parameters', async function () { - const query = new Query(undefined, [0.1, 0.3]) + const query = new Query([0.1, 0.3]) .limit(1) .metricType(MetricType.Cosine) .refineFactor(100)