rewrite index

Signed-off-by: cbh778899 <[email protected]>

database/index.js

@@ -1,42 +1,38 @@
-import * as lancedb from "@lancedb/lancedb";
-import { get, post } from "../tools/request.js"
-import { Schema, Field, FixedSizeList, Int16, Float16, Utf8 } from "apache-arrow";
+import { connect } from "@lancedb/lancedb";
+import { 
+    Schema, Field, FixedSizeList, 
+    Float32, Utf8, 
+    // eslint-disable-next-line
+    Table 
+} from "apache-arrow";
+import { DATASET_TABLE, SYSTEM_TABLE } from "./types";
 
 const uri = "/tmp/lancedb/";
-const db = await lancedb.connect(uri);
+const db = await connect(uri);
 
-const table = await db.createEmptyTable("rag_data", new Schema([
-    new Field("id", new Int16()),
-    new Field("vector", new FixedSizeList(384, new Field("item", new Float16(), true)), false),
-    new Field("question", new Utf8()),
-    new Field("answer", new Utf8())
-]), {
-    // mode: "overwrite",
-    existOk: true
-})
-
-export async function loadDataset(dataset_link) {
-    const {rows, http_error} = await get('', {}, { URL: dataset_link })
-    if(http_error) {
-        return false;
-    }
-    await table.add(rows.map(({ row_id, row })=>{
-        const { question, answer, question_embedding } = row;
-        return { id: row_id, question, answer, vector: question_embedding }
-    }))
-    return true;
+export async function initDB(force = false) {
+    const open_options = force ? { mode: "overwrite" } : { existOk: true }
+    // create or re-open system table to store long-lasting data
+    await db.createEmptyTable(SYSTEM_TABLE, new Schema([
+        new Field("title", new Utf8()),
+        new Field("value", new Utf8())
+    ]), open_options)
+    // create or re-open dataset table
+    await db.createEmptyTable(DATASET_TABLE, new Schema([
+        new Field("vector", new FixedSizeList(384, new Field("item", new Float32(), true)), false),
+        new Field("dataset_name", new Utf8()),
+        new Field("question", new Utf8()),
+        new Field("answer", new Utf8())
+    ]), open_options)
 }
 
-export async function searchByEmbedding(vector) {
-    const record = await table.search(vector).limit(1).toArray();
-    if(!record.length) return null;
-    const { question, answer } = record[0];
-    return { question, answer };
-}
+initDB();
 
-export async function searchByMessage(msg) {
-    const { embedding } = await post('embedding', {body: {
-        content: msg
-    }}, { eng: "embedding" });
-    return await searchByEmbedding(embedding);
+/**
+ * Open a table with table name
+ * @param {String} table_name table name to be opened
+ * @returns {Promise<Table>} Promise containes the table object.
+ */
+export async function getTable(table_name) {
+    return await db.openTable(table_name)
 }