Merge pull request #21 from epsilla-cloud/auto-embedding

Support auto embedding
epsilla-cloud · Dec 17, 2023 · db55bc6 · db55bc6
2 parents 323eabf + 82ff488
commit db55bc6
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 21 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "epsillajs",
-  "version": "0.2.4",
+  "version": "0.3.0",
   "description": "A JS library to connect Epsilla vector database",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",

diff --git a/src/cloud.ts b/src/cloud.ts
@@ -1,9 +1,10 @@
 import axios, { AxiosError } from 'axios';
-import { DeleteRecordsConfig, EpsillaResponse, PreviewConfig, QueryConfig, TableField } from './models';
+import { DeleteRecordsConfig, EpsillaResponse, Index, PreviewConfig, QueryConfig, TableField } from './models';
 
 export interface CloudClientConfig {
   projectID: string;
   apiKey: string;
+  headers?: { [key: string]: string };
 }
 
 const dispatchDomain = 'https://dispatch.epsilla.com';
@@ -21,9 +22,12 @@ export class EpsillaCloud {
   projectID: string;
   headers: any;
 
-  constructor({ projectID, apiKey }: CloudClientConfig) {
+  constructor({ projectID, apiKey, headers = {} }: CloudClientConfig) {
     this.projectID = projectID;
     this.headers = { 'Content-type': 'application/json', 'X-API-KEY': apiKey };
+    if (headers) {
+      this.headers = { ...this.headers, ...headers };
+    }
   }
 }
 
@@ -58,15 +62,19 @@ export class VectorDB {
     }
   }
 
-  async createTable(tableName: string, fields: TableField[]) {
+  async createTable(tableName: string, fields: TableField[], indices?: Index[]) {
     try {
       const domain = this.host || dispatchDomain;
+      let payload: any = {
+        name: tableName,
+        fields
+      };
+      if (indices) {
+        payload['indices'] = indices;
+      }
       const response = await axios.post(
         `${domain}/api/v3/project/${this.projectID}/vectordb/${this.dbID}/table/create`,
-        {
-          name: tableName,
-          fields
-        },
+        payload,
         { headers: this.headers }
       );
       return response.data;

diff --git a/src/models.ts b/src/models.ts
@@ -13,24 +13,34 @@ export interface TableField {
   metricType?: string;
 }
 
+export interface Index {
+  name: string;
+  field: string;
+  model?: string;
+}
+
 export interface SparseVector {
   indices: number[];
   values: number[];
 }
 
 export interface QueryPayload {
   table: string;
-  queryField: string;
-  queryVector: number[] | SparseVector;
+  query?: string;
+  queryIndex?: string;
+  queryField?: string;
+  queryVector?: number[] | SparseVector;
   limit: number;
   response?: string[];
   filter?: string;
   withDistance?: boolean;
 }
 
 export interface QueryConfig {
-  queryField: string;
-  queryVector: number[] | SparseVector;
+  query?: string;
+  queryIndex?: string;
+  queryField?: string;
+  queryVector?: number[] | SparseVector;
   limit: number;
   response?: string[];
   filter?: string;

diff --git a/src/vectordb.ts b/src/vectordb.ts
@@ -1,13 +1,14 @@
 import axios, { AxiosError } from 'axios';
 import {
-  DeleteRecordsConfig, EpsillaResponse, LoadDBPayload,
+  DeleteRecordsConfig, EpsillaResponse, Index, LoadDBPayload,
   PreviewConfig, QueryConfig, QueryPayload, TableField
 } from './models';
 
 export interface ClientConfig {
   protocol?: string;
   host?: string;
   port?: number;
+  headers?: { [key: string]: string };
 }
 
 class EpsillaDB {
@@ -18,13 +19,16 @@ class EpsillaDB {
   private baseurl: string;
   private headers: any;
 
-  constructor({ protocol = 'http', host = 'localhost', port = 8888 }: ClientConfig = {}) {
+  constructor({ protocol = 'http', host = 'localhost', port = 8888, headers = {} }: ClientConfig = {}) {
     this.protocol = protocol;
     this.host = host;
     this.port = port;
     this.db = null;
     this.baseurl = `${this.protocol}://${this.host}:${this.port}`;
     this.headers = { 'Content-type': 'application/json' };
+    if (headers) {
+      this.headers = { ...this.headers, ...headers };
+    }
   }
 
   useDB(dbName: string) {
@@ -65,17 +69,21 @@ class EpsillaDB {
     }
   }
 
-  async createTable(tableName: string, fields: TableField[]): Promise<EpsillaResponse | Error> {
+  async createTable(tableName: string, fields: TableField[], indices?: Index[]): Promise<EpsillaResponse | Error> {
     if (!this.db) {
       console.error('[ERROR] Please useDB() first!');
       return new Error('[ERROR] Please useDB() first!');
     }
     try {
+      let payload: any = {
+        name: tableName,
+        fields
+      };
+      if (indices) {
+        payload['indices'] = indices;
+      }
       const response = await axios.post(`${this.baseurl}/api/${this.db}/schema/tables`,
-        {
-          name: tableName,
-          fields
-        },
+        payload,
         { headers: this.headers }
       );
       return response.data;