feat: halfvec v1

Author: I-Info

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

@@ -21,23 +21,24 @@ export class PgVectorCtrl {
         CREATE EXTENSION IF NOT EXISTS vector;
         CREATE TABLE IF NOT EXISTS ${DatasetVectorTableName} (
             id BIGSERIAL PRIMARY KEY,
-            vector VECTOR(1536) NOT NULL,
+            halfvector HALFVEC(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 PgClient.query(
-        `CREATE INDEX CONCURRENTLY IF NOT EXISTS vector_index ON ${DatasetVectorTableName} USING hnsw (vector vector_ip_ops) WITH (m = 32, ef_construction = 128);`
-      );
       await PgClient.query(
         `CREATE INDEX CONCURRENTLY IF NOT EXISTS team_dataset_collection_index ON ${DatasetVectorTableName} USING btree(team_id, dataset_id, collection_id);`
       );
       await PgClient.query(
         `CREATE INDEX CONCURRENTLY IF NOT EXISTS create_time_index ON ${DatasetVectorTableName} USING btree(createtime);`
       );
+      // TODO: enable halfvector index
+      // await PgClient.query(
+      //   `CREATE INDEX CONCURRENTLY IF NOT EXISTS halfvector_index ON ${DatasetVectorTableName} USING hnsw (halfvector halfvec_ip_ops) WITH (m = 32, ef_construction = 128);`
+      // );
 
       addLog.info('init pg successful');
     } catch (error) {
@@ -48,10 +49,12 @@ export class PgVectorCtrl {
     const { teamId, datasetId, collectionId, vector, retry = 3 } = props;
 
     try {
+      // TODO: remove vector
       const { rowCount, rows } = await PgClient.insert(DatasetVectorTableName, {
         values: [
           [
             { key: 'vector', value: `[${vector}]` },
+            { key: 'halfvector', value: `[${vector}]` },
             { key: 'team_id', value: String(teamId) },
             { key: 'dataset_id', value: String(datasetId) },
             { key: 'collection_id', value: String(collectionId) }
@@ -177,12 +180,27 @@ export class PgVectorCtrl {
       // );
       // console.log(explan[2].rows);
 
+      // TODO: use halfvector
+      // const results: any = await PgClient.query(
+      //   `
+      //         BEGIN;
+      //           SET LOCAL hnsw.ef_search = ${global.systemEnv?.pgHNSWEfSearch || 100};
+      //           select id, collection_id, halfvector <#> '[${vector}]' AS score
+      //             from ${DatasetVectorTableName}
+      //             where team_id='${teamId}'
+      //               AND dataset_id IN (${datasetIds.map((id) => `'${String(id)}'`).join(',')})
+      //               ${filterCollectionIdSql}
+      //               ${forbidCollectionSql}
+      //             order by score limit ${limit};
+      //         COMMIT;`
+      // );
+
       const results: any = await PgClient.query(
         `
         BEGIN;
           SET LOCAL hnsw.ef_search = ${global.systemEnv?.pgHNSWEfSearch || 100};
-          select id, collection_id, vector <#> '[${vector}]' AS score 
-            from ${DatasetVectorTableName} 
+          select id, collection_id, vector <#> '[${vector}]' AS score
+            from ${DatasetVectorTableName}
             where team_id='${teamId}'
               AND dataset_id IN (${datasetIds.map((id) => `'${String(id)}'`).join(',')})
               ${filterCollectionIdSql}

projects/app/src/pages/api/admin/inithalfvec1.ts

@@ -0,0 +1,63 @@
+import type { NextApiRequest, NextApiResponse } from 'next';
+import { jsonRes } from '@fastgpt/service/common/response';
+import { connectToDatabase } from '@/service/mongo';
+import { authCert } from '@fastgpt/service/support/permission/auth/common';
+import { PgClient } from '@fastgpt/service/common/vectorStore/pg';
+import { DatasetVectorTableName } from '@fastgpt/service/common/vectorStore/constants';
+
+export default async function handler(req: NextApiRequest, res: NextApiResponse) {
+  try {
+    await connectToDatabase();
+    await authCert({ req, authRoot: true });
+
+    // pg 新建字段：halfvector
+    const columnExists = await PgClient.query(`
+      SELECT column_name
+      FROM information_schema.columns
+      WHERE table_name='${DatasetVectorTableName}' AND column_name='halfvector';
+    `);
+
+    if (columnExists.rows.length === 0) {
+      await PgClient.query(`
+        BEGIN;
+          ALTER TABLE ${DatasetVectorTableName} ADD COLUMN halfvector halfvec(1536);
+        COMMIT;
+      `);
+      console.log('halfvector column added');
+    }
+
+    let rowsUpdated;
+    do {
+      rowsUpdated = await PgClient.query(`
+        WITH updated AS (
+          UPDATE ${DatasetVectorTableName}
+          SET halfvector = vector::halfvec(1536)
+          WHERE id IN (
+              SELECT id
+              FROM ${DatasetVectorTableName}
+              WHERE halfvector IS NULL
+              LIMIT 1000
+          )
+          RETURNING 1
+        )
+        SELECT count(*) FROM updated;
+      `);
+      console.log('rowsUpdated:', rowsUpdated.rows[0].count);
+      // 每批次更新后休眠一段时间
+      if (rowsUpdated.rows[0].count > 0) {
+        await new Promise((resolve) => setTimeout(resolve, 100)); // 休眠100毫秒
+      }
+    } while (rowsUpdated.rows[0].count > 0);
+
+    jsonRes(res, {
+      message: 'success'
+    });
+  } catch (error) {
+    console.log(error);
+
+    jsonRes(res, {
+      code: 500,
+      error
+    });
+  }
+}

projects/app/src/pages/api/admin/inithalfvec2.ts

@@ -0,0 +1,34 @@
+import type { NextApiRequest, NextApiResponse } from 'next';
+import { jsonRes } from '@fastgpt/service/common/response';
+import { connectToDatabase } from '@/service/mongo';
+import { authCert } from '@fastgpt/service/support/permission/auth/common';
+import { PgClient } from '@fastgpt/service/common/vectorStore/pg';
+import { DatasetVectorTableName } from '@fastgpt/service/common/vectorStore/constants';
+
+export default async function handler(req: NextApiRequest, res: NextApiResponse) {
+  try {
+    await connectToDatabase();
+    await authCert({ req, authRoot: true });
+
+    // 设置halfvector字段为非空
+    await PgClient.query(
+      `BEGIN;
+      ALTER TABLE ${DatasetVectorTableName} ALTER COLUMN halfvector SET NOT NULL;
+      DROP INDEX IF EXISTS vector_index;
+      ALTER TABLE ${DatasetVectorTableName} DROP COLUMN IF EXISTS vector;
+      COMMIT;
+      `
+    );
+
+    jsonRes(res, {
+      message: 'success'
+    });
+  } catch (error) {
+    console.log(error);
+
+    jsonRes(res, {
+      code: 500,
+      error
+    });
+  }
+}