feat: vector store support oceanbase (#4356)

* feat: vector store support oceanbase

* chore(config): Rename pgHNSWEfSearch to hnswEfSearch to work for pg and oceanbase both

packages/global/common/system/types/index.d.ts

@@ -112,7 +112,7 @@ export type SystemEnvType = {
   vectorMaxProcess: number;
   qaMaxProcess: number;
   vlmMaxProcess: number;
-  pgHNSWEfSearch: number;
+  hnswEfSearch: number;
   tokenWorkers: number; // token count max worker
 
   oneapiUrl?: string;

packages/service/common/vectorStore/constants.ts

@@ -2,5 +2,6 @@ export const DatasetVectorDbName = 'fastgpt';
 export const DatasetVectorTableName = 'modeldata';
 
 export const PG_ADDRESS = process.env.PG_URL;
+export const OCEANBASE_ADDRESS = process.env.OCEANBASE_URL;
 export const MILVUS_ADDRESS = process.env.MILVUS_ADDRESS;
 export const MILVUS_TOKEN = process.env.MILVUS_TOKEN;

packages/service/common/vectorStore/controller.ts

@@ -1,13 +1,15 @@
 /* vector crud */
 import { PgVectorCtrl } from './pg/class';
+import { ObVectorCtrl } from './oceanbase/class';
 import { getVectorsByText } from '../../core/ai/embedding';
 import { InsertVectorProps } from './controller.d';
 import { EmbeddingModelItemType } from '@fastgpt/global/core/ai/model.d';
-import { MILVUS_ADDRESS, PG_ADDRESS } from './constants';
+import { MILVUS_ADDRESS, PG_ADDRESS, OCEANBASE_ADDRESS } from './constants';
 import { MilvusCtrl } from './milvus/class';
 
 const getVectorObj = () => {
   if (PG_ADDRESS) return new PgVectorCtrl();
+  if (OCEANBASE_ADDRESS) return new ObVectorCtrl();
   if (MILVUS_ADDRESS) return new MilvusCtrl();
 
   return new PgVectorCtrl();

packages/service/common/vectorStore/oceanbase/class.ts

@@ -0,0 +1,254 @@
+/* oceanbase vector crud */
+import { DatasetVectorTableName } from '../constants';
+import { delay } from '@fastgpt/global/common/system/utils';
+import { ObClient } from './index';
+import { RowDataPacket, ResultSetHeader } from 'mysql2/promise';
+import {
+  DelDatasetVectorCtrlProps,
+  EmbeddingRecallCtrlProps,
+  EmbeddingRecallResponse,
+  InsertVectorControllerProps
+} from '../controller.d';
+import dayjs from 'dayjs';
+import { addLog } from '../../system/log';
+
+export class ObVectorCtrl {
+  constructor() {}
+  init = async () => {
+    try {
+      await ObClient.query(`
+        CREATE TABLE IF NOT EXISTS ${DatasetVectorTableName} (
+            id BIGINT AUTO_INCREMENT PRIMARY KEY,
+            vector VECTOR(1536) NOT NULL,
+            team_id VARCHAR(50) NOT NULL,
+            dataset_id VARCHAR(50) NOT NULL,
+            collection_id VARCHAR(50) NOT NULL,
+            createtime TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+        );
+      `);
+
+      await ObClient.query(
+        `CREATE VECTOR INDEX IF NOT EXISTS vector_index ON ${DatasetVectorTableName}(vector) WITH (distance=inner_product, type=hnsw, m=32, ef_construction=128);`
+      );
+      await ObClient.query(
+        `CREATE INDEX IF NOT EXISTS team_dataset_collection_index ON ${DatasetVectorTableName}(team_id, dataset_id, collection_id);`
+      );
+      await ObClient.query(
+        `CREATE INDEX IF NOT EXISTS create_time_index ON ${DatasetVectorTableName}(createtime);`
+      );
+
+      addLog.info('init oceanbase successful');
+    } catch (error) {
+      addLog.error('init oceanbase error', error);
+    }
+  };
+  insert = async (props: InsertVectorControllerProps): Promise<{ insertId: string }> => {
+    const { teamId, datasetId, collectionId, vector, retry = 3 } = props;
+
+    try {
+      const { rowCount, rows } = await ObClient.insert(DatasetVectorTableName, {
+        values: [
+          [
+            { key: 'vector', value: `[${vector}]` },
+            { key: 'team_id', value: String(teamId) },
+            { key: 'dataset_id', value: String(datasetId) },
+            { key: 'collection_id', value: String(collectionId) }
+          ]
+        ]
+      });
+
+      if (rowCount === 0) {
+        return Promise.reject('insertDatasetData: no insert');
+      }
+
+      return {
+        insertId: rows[0].id
+      };
+    } catch (error) {
+      if (retry <= 0) {
+        return Promise.reject(error);
+      }
+      await delay(500);
+      return this.insert({
+        ...props,
+        retry: retry - 1
+      });
+    }
+  };
+  delete = async (props: DelDatasetVectorCtrlProps): Promise<any> => {
+    const { teamId, retry = 2 } = props;
+
+    const teamIdWhere = `team_id='${String(teamId)}' AND`;
+
+    const where = await (() => {
+      if ('id' in props && props.id) return `${teamIdWhere} id=${props.id}`;
+
+      if ('datasetIds' in props && props.datasetIds) {
+        const datasetIdWhere = `dataset_id IN (${props.datasetIds
+          .map((id) => `'${String(id)}'`)
+          .join(',')})`;
+
+        if ('collectionIds' in props && props.collectionIds) {
+          return `${teamIdWhere} ${datasetIdWhere} AND collection_id IN (${props.collectionIds
+            .map((id) => `'${String(id)}'`)
+            .join(',')})`;
+        }
+
+        return `${teamIdWhere} ${datasetIdWhere}`;
+      }
+
+      if ('idList' in props && Array.isArray(props.idList)) {
+        if (props.idList.length === 0) return;
+        return `${teamIdWhere} id IN (${props.idList.map((id) => String(id)).join(',')})`;
+      }
+      return Promise.reject('deleteDatasetData: no where');
+    })();
+
+    if (!where) return;
+
+    try {
+      await ObClient.delete(DatasetVectorTableName, {
+        where: [where]
+      });
+    } catch (error) {
+      if (retry <= 0) {
+        return Promise.reject(error);
+      }
+      await delay(500);
+      return this.delete({
+        ...props,
+        retry: retry - 1
+      });
+    }
+  };
+  embRecall = async (props: EmbeddingRecallCtrlProps): Promise<EmbeddingRecallResponse> => {
+    const {
+      teamId,
+      datasetIds,
+      vector,
+      limit,
+      forbidCollectionIdList,
+      filterCollectionIdList,
+      retry = 2
+    } = props;
+
+    // Get forbid collection
+    const formatForbidCollectionIdList = (() => {
+      if (!filterCollectionIdList) return forbidCollectionIdList;
+      const list = forbidCollectionIdList
+        .map((id) => String(id))
+        .filter((id) => !filterCollectionIdList.includes(id));
+      return list;
+    })();
+    const forbidCollectionSql =
+      formatForbidCollectionIdList.length > 0
+        ? `AND collection_id NOT IN (${formatForbidCollectionIdList.map((id) => `'${id}'`).join(',')})`
+        : '';
+
+    // Filter by collectionId
+    const formatFilterCollectionId = (() => {
+      if (!filterCollectionIdList) return;
+
+      return filterCollectionIdList
+        .map((id) => String(id))
+        .filter((id) => !forbidCollectionIdList.includes(id));
+    })();
+    const filterCollectionIdSql = formatFilterCollectionId
+      ? `AND collection_id IN (${formatFilterCollectionId.map((id) => `'${id}'`).join(',')})`
+      : '';
+    // Empty data
+    if (formatFilterCollectionId && formatFilterCollectionId.length === 0) {
+      return { results: [] };
+    }
+
+    try {
+      const rows = await ObClient.query<
+        ({
+          id: string;
+          collection_id: string;
+          score: number;
+        } & RowDataPacket)[][]
+      >(
+        `BEGIN;
+          SET ob_hnsw_ef_search = ${global.systemEnv?.hnswEfSearch || 100};
+          SELECT id, collection_id, inner_product(vector, [${vector}]) AS score
+            FROM ${DatasetVectorTableName}
+            WHERE team_id='${teamId}'
+              AND dataset_id IN (${datasetIds.map((id) => `'${String(id)}'`).join(',')})
+              ${filterCollectionIdSql}
+              ${forbidCollectionSql}
+            ORDER BY score desc APPROXIMATE LIMIT ${limit};
+        COMMIT;`
+      ).then(([rows]) => rows[2]);
+
+      return {
+        results: rows.map((item) => ({
+          id: String(item.id),
+          collectionId: item.collection_id,
+          score: item.score
+        }))
+      };
+    } catch (error) {
+      if (retry <= 0) {
+        return Promise.reject(error);
+      }
+      return this.embRecall({
+        ...props,
+        retry: retry - 1
+      });
+    }
+  };
+  getVectorDataByTime = async (start: Date, end: Date) => {
+    const rows = await ObClient.query<
+      ({
+        id: string;
+        team_id: string;
+        dataset_id: string;
+      } & RowDataPacket)[]
+    >(
+      `SELECT id, team_id, dataset_id
+    FROM ${DatasetVectorTableName}
+    WHERE createtime BETWEEN '${dayjs(start).format('YYYY-MM-DD HH:mm:ss')}' AND '${dayjs(
+      end
+    ).format('YYYY-MM-DD HH:mm:ss')}';
+    `
+    ).then(([rows]) => rows);
+
+    return rows.map((item) => ({
+      id: String(item.id),
+      teamId: item.team_id,
+      datasetId: item.dataset_id
+    }));
+  };
+  getVectorCountByTeamId = async (teamId: string) => {
+    const total = await ObClient.count(DatasetVectorTableName, {
+      where: [['team_id', String(teamId)]]
+    });
+
+    return total;
+  };
+  getVectorCountByDatasetId = async (teamId: string, datasetId: string) => {
+    const total = await ObClient.count(DatasetVectorTableName, {
+      where: [['team_id', String(teamId)], 'and', ['dataset_id', String(datasetId)]]
+    });
+
+    return total;
+  };
+  getVectorCountByCollectionId = async (
+    teamId: string,
+    datasetId: string,
+    collectionId: string
+  ) => {
+    const total = await ObClient.count(DatasetVectorTableName, {
+      where: [
+        ['team_id', String(teamId)],
+        'and',
+        ['dataset_id', String(datasetId)],
+        'and',
+        ['collection_id', String(collectionId)]
+      ]
+    });
+
+    return total;
+  };
+}

packages/service/common/vectorStore/oceanbase/index.ts

@@ -0,0 +1,173 @@
+import mysql, { Pool, QueryResult, RowDataPacket, ResultSetHeader } from 'mysql2/promise';
+import { addLog } from '../../system/log';
+import { OCEANBASE_ADDRESS } from '../constants';
+
+export const getClient = async (): Promise<Pool> => {
+  if (!OCEANBASE_ADDRESS) {
+    return Promise.reject('OCEANBASE_ADDRESS is not set');
+  }
+
+  if (global.obClient) {
+    return global.obClient;
+  }
+
+  global.obClient = mysql.createPool({
+    uri: OCEANBASE_ADDRESS,
+    waitForConnections: true,
+    connectionLimit: Number(process.env.DB_MAX_LINK || 20),
+    connectTimeout: 20000,
+    idleTimeout: 60000,
+    queueLimit: 0,
+    enableKeepAlive: true,
+    keepAliveInitialDelay: 0
+  });
+
+  addLog.info(`oceanbase connected`);
+
+  return global.obClient;
+};
+
+type WhereProps = (string | [string, string | number])[];
+type GetProps = {
+  fields?: string[];
+  where?: WhereProps;
+  order?: { field: string; mode: 'DESC' | 'ASC' | string }[];
+  limit?: number;
+  offset?: number;
+};
+
+type DeleteProps = {
+  where: WhereProps;
+};
+
+type ValuesProps = { key: string; value?: string | number }[];
+type UpdateProps = {
+  values: ValuesProps;
+  where: WhereProps;
+};
+type InsertProps = {
+  values: ValuesProps[];
+};
+
+class ObClass {
+  private getWhereStr(where?: WhereProps) {
+    return where
+      ? `WHERE ${where
+          .map((item) => {
+            if (typeof item === 'string') {
+              return item;
+            }
+            const val = typeof item[1] === 'number' ? item[1] : `'${String(item[1])}'`;
+            return `${item[0]}=${val}`;
+          })
+          .join(' ')}`
+      : '';
+  }
+  private getUpdateValStr(values: ValuesProps) {
+    return values
+      .map((item) => {
+        const val =
+          typeof item.value === 'number'
+            ? item.value
+            : `'${String(item.value).replace(/\'/g, '"')}'`;
+
+        return `${item.key}=${val}`;
+      })
+      .join(',');
+  }
+  private getInsertValStr(values: ValuesProps[]) {
+    return values
+      .map(
+        (items) =>
+          `(${items
+            .map((item) =>
+              typeof item.value === 'number'
+                ? item.value
+                : `'${String(item.value).replace(/\'/g, '"')}'`
+            )
+            .join(',')})`
+      )
+      .join(',');
+  }
+  async select<T extends QueryResult = any>(table: string, props: GetProps) {
+    const sql = `SELECT ${
+      !props.fields || props.fields?.length === 0 ? '*' : props.fields?.join(',')
+    }
+      FROM ${table}
+      ${this.getWhereStr(props.where)}
+      ${
+        props.order
+          ? `ORDER BY ${props.order.map((item) => `${item.field} ${item.mode}`).join(',')}`
+          : ''
+      }
+      LIMIT ${props.limit || 10} OFFSET ${props.offset || 0}
+    `;
+
+    const client = await getClient();
+    return client.query<T>(sql);
+  }
+  async count(table: string, props: GetProps) {
+    const sql = `SELECT COUNT(${props?.fields?.[0] || '*'})
+      FROM ${table}
+      ${this.getWhereStr(props.where)}
+    `;
+
+    const client = await getClient();
+    return client
+      .query<({ count: number } & RowDataPacket)[]>(sql)
+      .then(([rows]) => Number(rows[0]?.count || 0));
+  }
+  async delete(table: string, props: DeleteProps) {
+    const sql = `DELETE FROM ${table} ${this.getWhereStr(props.where)}`;
+    const client = await getClient();
+    return client.query(sql);
+  }
+  async update(table: string, props: UpdateProps) {
+    if (props.values.length === 0) {
+      return {
+        rowCount: 0
+      };
+    }
+
+    const sql = `UPDATE ${table} SET ${this.getUpdateValStr(props.values)} ${this.getWhereStr(
+      props.where
+    )}`;
+    const client = await getClient();
+    return client.query(sql);
+  }
+  async insert(table: string, props: InsertProps) {
+    if (props.values.length === 0) {
+      return {
+        rowCount: 0,
+        rows: []
+      };
+    }
+
+    const fields = props.values[0].map((item) => item.key).join(',');
+    const sql = `INSERT INTO ${table} (${fields}) VALUES ${this.getInsertValStr(props.values)}`;
+
+    const client = await getClient();
+    return client.query<ResultSetHeader>(sql).then(([result]) => {
+      return {
+        rowCount: result.affectedRows,
+        rows: [{ id: String(result.insertId) }]
+      };
+    });
+  }
+  async query<T extends QueryResult = any>(sql: string) {
+    const client = await getClient();
+    const start = Date.now();
+    return client.query<T>(sql).then((res) => {
+      const time = Date.now() - start;
+
+      if (time > 300) {
+        addLog.warn(`oceanbase query time: ${time}ms, sql: ${sql}`);
+      }
+
+      return res;
+    });
+  }
+}
+
+export const ObClient = new ObClass();
+export const Oceanbase = global.obClient;

packages/service/common/vectorStore/pg/class.ts

@@ -187,7 +187,7 @@ export class PgVectorCtrl {
     try {
       const results: any = await PgClient.query(
         `BEGIN;
-          SET LOCAL hnsw.ef_search = ${global.systemEnv?.pgHNSWEfSearch || 100};
+          SET LOCAL hnsw.ef_search = ${global.systemEnv?.hnswEfSearch || 100};
           SET LOCAL hnsw.iterative_scan = relaxed_order;
           WITH relaxed_results AS MATERIALIZED (
             select id, collection_id, vector <#> '[${vector}]' AS score

packages/service/common/vectorStore/type.d.ts

@@ -1,8 +1,10 @@
 import type { Pool } from 'pg';
+import { Pool as MysqlPool } from 'mysql2/promise';
 import { MilvusClient } from '@zilliz/milvus2-sdk-node';
 
 declare global {
   var pgClient: Pool | null;
+  var obClient: MysqlPool | null;
   var milvusClient: MilvusClient | null;
 }
 

packages/service/package.json

@@ -26,6 +26,7 @@
     "mammoth": "^1.6.0",
     "mongoose": "^8.10.1",
     "multer": "1.4.5-lts.1",
+    "mysql2": "^3.11.3",
     "next": "14.2.25",
     "nextjs-cors": "^2.2.0",
     "node-cron": "^3.0.3",