V4.6.6-2 (#673)

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

@@ -1,5 +0,0 @@
-import type { Pool } from 'pg';
-
-declare global {
-  var pgClient: Pool | null;
-}

packages/service/common/system/log.ts

@@ -4,11 +4,13 @@ import dayjs from 'dayjs';
 export const addLog = {
   log(level: 'info' | 'warn' | 'error', msg: string, obj: Record<string, any> = {}) {
     console.log(
-      `[${level.toLocaleUpperCase()}] ${dayjs().format(
-        'YYYY-MM-DD HH:mm:ss'
-      )} ${msg}: ${JSON.stringify(obj)}`
+      `[${level.toLocaleUpperCase()}] ${dayjs().format('YYYY-MM-DD HH:mm:ss')} ${msg} ${
+        level !== 'error' ? JSON.stringify(obj) : ''
+      }`
     );
 
+    level === 'error' && console.error(obj);
+
     const lokiUrl = process.env.LOKI_LOG_URL as string;
     if (!lokiUrl) return;
 

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

@@ -0,0 +1,19 @@
+export type DeleteDatasetVectorProps = {
+  id?: string;
+  datasetIds?: string[];
+  collectionIds?: string[];
+  dataIds?: string[];
+};
+
+export type InsertVectorProps = {
+  teamId: string;
+  tmbId: string;
+  datasetId: string;
+  collectionId: string;
+  dataId: string;
+};
+
+export type EmbeddingRecallProps = {
+  similarity?: number;
+  datasetIds: string[];
+};

packages/service/common/vectorStore/controller.ts

@@ -0,0 +1,62 @@
+/* vector crud */
+import { PgVector } from './pg/class';
+import { getVectorsByText } from '../../core/ai/embedding';
+import { InsertVectorProps } from './controller.d';
+
+const getVectorObj = () => {
+  return new PgVector();
+};
+
+export const initVectorStore = getVectorObj().init;
+export const deleteDatasetDataVector = getVectorObj().delete;
+export const recallFromVectorStore = getVectorObj().recall;
+export const getVectorDataByTime = getVectorObj().getVectorDataByTime;
+export const getVectorCountByTeamId = getVectorObj().getVectorCountByTeamId;
+
+export const insertDatasetDataVector = async ({
+  model,
+  query,
+  ...props
+}: InsertVectorProps & {
+  query: string;
+  model: string;
+}) => {
+  const { vectors, tokens } = await getVectorsByText({
+    model,
+    input: query
+  });
+  const { insertId } = await getVectorObj().insert({
+    ...props,
+    vectors
+  });
+
+  return {
+    tokens,
+    insertId
+  };
+};
+
+export const updateDatasetDataVector = async ({
+  id,
+  query,
+  model
+}: {
+  id: string;
+  query: string;
+  model: string;
+}) => {
+  // get vector
+  const { vectors, tokens } = await getVectorsByText({
+    model,
+    input: [query]
+  });
+
+  await getVectorObj().update({
+    id,
+    vectors
+  });
+
+  return {
+    tokens
+  };
+};

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

@@ -0,0 +1,20 @@
+import {
+  initPg,
+  insertDatasetDataVector,
+  updateDatasetDataVector,
+  deleteDatasetDataVector,
+  embeddingRecall,
+  getVectorDataByTime,
+  getVectorCountByTeamId
+} from './controller';
+
+export class PgVector {
+  constructor() {}
+  init = initPg;
+  insert = insertDatasetDataVector;
+  update = updateDatasetDataVector;
+  delete = deleteDatasetDataVector;
+  recall = embeddingRecall;
+  getVectorCountByTeamId = getVectorCountByTeamId;
+  getVectorDataByTime = getVectorDataByTime;
+}

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

@@ -0,0 +1,199 @@
+/* pg vector crud */
+import { PgDatasetTableName } from '@fastgpt/global/common/vectorStore/constants';
+import { delay } from '@fastgpt/global/common/system/utils';
+import { PgClient, connectPg } from './index';
+import { PgSearchRawType } from '@fastgpt/global/core/dataset/api';
+import { EmbeddingRecallItemType } from '../type';
+import { DeleteDatasetVectorProps, EmbeddingRecallProps } from '../controller.d';
+import dayjs from 'dayjs';
+
+export async function initPg() {
+  try {
+    await connectPg();
+    await PgClient.query(`
+      CREATE EXTENSION IF NOT EXISTS vector;
+      CREATE TABLE IF NOT EXISTS ${PgDatasetTableName} (
+          id BIGSERIAL PRIMARY KEY,
+          vector VECTOR(1536) NOT NULL,
+          team_id VARCHAR(50) NOT NULL,
+          tmb_id VARCHAR(50) NOT NULL,
+          dataset_id VARCHAR(50) NOT NULL,
+          collection_id VARCHAR(50) NOT NULL,
+          data_id VARCHAR(50) NOT NULL,
+          createTime TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+      );
+      CREATE INDEX IF NOT EXISTS vector_index ON ${PgDatasetTableName} USING hnsw (vector vector_ip_ops) WITH (m = 32, ef_construction = 64);
+    `);
+
+    console.log('init pg successful');
+  } catch (error) {
+    console.log('init pg error', error);
+  }
+}
+
+export const insertDatasetDataVector = async (props: {
+  teamId: string;
+  tmbId: string;
+  datasetId: string;
+  collectionId: string;
+  dataId: string;
+  vectors: number[][];
+  retry?: number;
+}): Promise<{ insertId: string }> => {
+  const { dataId, teamId, tmbId, datasetId, collectionId, vectors, retry = 3 } = props;
+  try {
+    const { rows } = await PgClient.insert(PgDatasetTableName, {
+      values: [
+        [
+          { key: 'vector', value: `[${vectors[0]}]` },
+          { key: 'team_id', value: String(teamId) },
+          { key: 'tmb_id', value: String(tmbId) },
+          { key: 'dataset_id', value: datasetId },
+          { key: 'collection_id', value: collectionId },
+          { key: 'data_id', value: String(dataId) }
+        ]
+      ]
+    });
+    return {
+      insertId: rows[0].id
+    };
+  } catch (error) {
+    if (retry <= 0) {
+      return Promise.reject(error);
+    }
+    await delay(500);
+    return insertDatasetDataVector({
+      ...props,
+      retry: retry - 1
+    });
+  }
+};
+
+export const updateDatasetDataVector = async (props: {
+  id: string;
+  vectors: number[][];
+  retry?: number;
+}): Promise<void> => {
+  const { id, vectors, retry = 2 } = props;
+  try {
+    // update pg
+    await PgClient.update(PgDatasetTableName, {
+      where: [['id', id]],
+      values: [{ key: 'vector', value: `[${vectors[0]}]` }]
+    });
+  } catch (error) {
+    if (retry <= 0) {
+      return Promise.reject(error);
+    }
+    await delay(500);
+    return updateDatasetDataVector({
+      ...props,
+      retry: retry - 1
+    });
+  }
+};
+
+export const deleteDatasetDataVector = async (
+  props: DeleteDatasetVectorProps & {
+    retry?: number;
+  }
+): Promise<any> => {
+  const { id, datasetIds, collectionIds, dataIds, retry = 2 } = props;
+
+  const where = await (() => {
+    if (id) return `id=${id}`;
+    if (datasetIds) return `dataset_id IN (${datasetIds.map((id) => `'${String(id)}'`).join(',')})`;
+    if (collectionIds)
+      return `collection_id IN (${collectionIds.map((id) => `'${String(id)}'`).join(',')})`;
+    if (dataIds) return `data_id IN (${dataIds.map((id) => `'${String(id)}'`).join(',')})`;
+    return Promise.reject('deleteDatasetData: no where');
+  })();
+
+  try {
+    await PgClient.delete(PgDatasetTableName, {
+      where: [where]
+    });
+  } catch (error) {
+    if (retry <= 0) {
+      return Promise.reject(error);
+    }
+    await delay(500);
+    return deleteDatasetDataVector({
+      ...props,
+      retry: retry - 1
+    });
+  }
+};
+
+export const embeddingRecall = async (
+  props: EmbeddingRecallProps & {
+    vectors: number[][];
+    limit: number;
+    retry?: number;
+  }
+): Promise<{
+  results: EmbeddingRecallItemType[];
+}> => {
+  const { vectors, limit, similarity = 0, datasetIds, retry = 2 } = props;
+
+  try {
+    const results: any = await PgClient.query(
+      `BEGIN;
+        SET LOCAL hnsw.ef_search = ${global.systemEnv.pgHNSWEfSearch || 100};
+        select id, collection_id, data_id, (vector <#> '[${vectors[0]}]') * -1 AS score 
+          from ${PgDatasetTableName} 
+          where dataset_id IN (${datasetIds.map((id) => `'${String(id)}'`).join(',')})
+              AND vector <#> '[${vectors[0]}]' < -${similarity}
+          order by score desc limit ${limit};
+        COMMIT;`
+    );
+
+    const rows = results?.[2]?.rows as PgSearchRawType[];
+
+    // concat same data_id
+    const filterRows: PgSearchRawType[] = [];
+    let set = new Set<string>();
+    for (const row of rows) {
+      if (!set.has(row.data_id)) {
+        filterRows.push(row);
+        set.add(row.data_id);
+      }
+    }
+
+    return {
+      results: filterRows.map((item) => ({
+        id: item.id,
+        collectionId: item.collection_id,
+        dataId: item.data_id,
+        score: item.score
+      }))
+    };
+  } catch (error) {
+    if (retry <= 0) {
+      return Promise.reject(error);
+    }
+    return embeddingRecall(props);
+  }
+};
+
+// bill
+export const getVectorCountByTeamId = async (teamId: string) => {
+  const total = await PgClient.count(PgDatasetTableName, {
+    where: [['team_id', String(teamId)]]
+  });
+
+  return total;
+};
+export const getVectorDataByTime = async (start: Date, end: Date) => {
+  const { rows } = await PgClient.query<{ id: string; data_id: string }>(`SELECT id, data_id
+  FROM ${PgDatasetTableName}
+  WHERE createTime BETWEEN '${dayjs(start).format('YYYY-MM-DD')}' AND '${dayjs(end).format(
+    'YYYY-MM-DD 23:59:59'
+  )}';
+  `);
+
+  return rows.map((item) => ({
+    id: item.id,
+    dataId: item.data_id
+  }));
+};

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

@@ -1,6 +1,5 @@
 import { Pool } from 'pg';
 import type { QueryResultRow } from 'pg';
-import { PgDatasetTableName } from '@fastgpt/global/core/dataset/constant';
 
 export const connectPg = async (): Promise<Pool> => {
   if (global.pgClient) {
@@ -117,6 +116,7 @@ class PgClass {
       FROM ${table}
       ${this.getWhereStr(props.where)}
     `;
+
     const pg = await connectPg();
     return pg.query(sql).then((res) => Number(res.rows[0]?.count || 0));
   }
@@ -160,29 +160,5 @@ class PgClass {
   }
 }
 
-export async function initPg() {
-  try {
-    await connectPg();
-    await PgClient.query(`
-      CREATE EXTENSION IF NOT EXISTS vector;
-      CREATE TABLE IF NOT EXISTS ${PgDatasetTableName} (
-          id BIGSERIAL PRIMARY KEY,
-          vector VECTOR(1536) NOT NULL,
-          team_id VARCHAR(50) NOT NULL,
-          tmb_id VARCHAR(50) NOT NULL,
-          dataset_id VARCHAR(50) NOT NULL,
-          collection_id VARCHAR(50) NOT NULL,
-          data_id VARCHAR(50) NOT NULL,
-          createTime TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-      );
-      CREATE INDEX IF NOT EXISTS vector_index ON ${PgDatasetTableName} USING hnsw (vector vector_ip_ops) WITH (m = 32, ef_construction = 64);
-    `);
-
-    console.log('init pg successful');
-  } catch (error) {
-    console.log('init pg error', error);
-  }
-}
-
 export const PgClient = new PgClass();
 export const Pg = global.pgClient;

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

@@ -0,0 +1,12 @@
+import type { Pool } from 'pg';
+
+declare global {
+  var pgClient: Pool | null;
+}
+
+export type EmbeddingRecallItemType = {
+  id: string;
+  collectionId: string;
+  dataId: string;
+  score: number;
+};

packages/service/core/ai/embedding/index.ts

@@ -0,0 +1,73 @@
+import { getAIApi } from '../config';
+
+export type GetVectorProps = {
+  model: string;
+  input: string | string[];
+};
+
+// text to vector
+export async function getVectorsByText({
+  model = 'text-embedding-ada-002',
+  input
+}: GetVectorProps) {
+  if (typeof input === 'string' && !input) {
+    return Promise.reject({
+      code: 500,
+      message: 'input is empty'
+    });
+  } else if (Array.isArray(input)) {
+    for (let i = 0; i < input.length; i++) {
+      if (!input[i]) {
+        return Promise.reject({
+          code: 500,
+          message: 'input array is empty'
+        });
+      }
+    }
+  }
+
+  try {
+    // 获取 chatAPI
+    const ai = getAIApi();
+
+    // 把输入的内容转成向量
+    const result = await ai.embeddings
+      .create({
+        model,
+        input
+      })
+      .then(async (res) => {
+        if (!res.data) {
+          return Promise.reject('Embedding API 404');
+        }
+        if (!res?.data?.[0]?.embedding) {
+          console.log(res?.data);
+          // @ts-ignore
+          return Promise.reject(res.data?.err?.message || 'Embedding API Error');
+        }
+        return {
+          tokens: res.usage.total_tokens || 0,
+          vectors: await Promise.all(res.data.map((item) => unityDimensional(item.embedding)))
+        };
+      });
+
+    return result;
+  } catch (error) {
+    console.log(`Embedding Error`, error);
+
+    return Promise.reject(error);
+  }
+}
+
+function unityDimensional(vector: number[]) {
+  if (vector.length > 1536) {
+    console.log(`当前向量维度为: ${vector.length}, 向量维度不能超过 1536, 已自动截取前 1536 维度`);
+    return vector.slice(0, 1536);
+  }
+  let resultVector = vector;
+  const vectorLen = vector.length;
+
+  const zeroVector = new Array(1536 - vectorLen).fill(0);
+
+  return resultVector.concat(zeroVector);
+}

packages/service/core/ai/functions/createQuestionGuide.ts

@@ -26,7 +26,8 @@ export async function createQuestionGuide({
   });
 
   const answer = data.choices?.[0]?.message?.content || '';
-  const totalTokens = data.usage?.total_tokens || 0;
+  const inputTokens = data.usage?.prompt_tokens || 0;
+  const outputTokens = data.usage?.completion_tokens || 0;
 
   const start = answer.indexOf('[');
   const end = answer.lastIndexOf(']');
@@ -34,7 +35,8 @@ export async function createQuestionGuide({
   if (start === -1 || end === -1) {
     return {
       result: [],
-      tokens: totalTokens
+      inputTokens,
+      outputTokens
     };
   }
 
@@ -46,12 +48,14 @@ export async function createQuestionGuide({
   try {
     return {
       result: JSON.parse(jsonStr),
-      tokens: totalTokens
+      inputTokens,
+      outputTokens
     };
   } catch (error) {
     return {
       result: [],
-      tokens: totalTokens
+      inputTokens,
+      outputTokens
     };
   }
 }

packages/service/core/dataset/data/controller.ts

@@ -1,11 +1,11 @@
 import { MongoDatasetData } from './schema';
-import { deletePgDataById } from './pg';
 import { MongoDatasetTraining } from '../training/schema';
 import { delFileByFileIdList, delFileByMetadata } from '../../../common/file/gridfs/controller';
 import { BucketNameEnum } from '@fastgpt/global/common/file/constants';
 import { MongoDatasetCollection } from '../collection/schema';
 import { delay } from '@fastgpt/global/common/system/utils';
 import { delImgByFileIdList } from '../../../common/file/image/controller';
+import { deleteDatasetDataVector } from '../../../common/vectorStore/controller';
 
 /* delete all data by datasetIds */
 export async function delDatasetRelevantData({ datasetIds }: { datasetIds: string[] }) {
@@ -21,7 +21,7 @@ export async function delDatasetRelevantData({ datasetIds }: { datasetIds: strin
   // delete dataset.datas
   await MongoDatasetData.deleteMany({ datasetId: { $in: datasetIds } });
   // delete pg data
-  await deletePgDataById(`dataset_id IN ('${datasetIds.join("','")}')`);
+  await deleteDatasetDataVector({ datasetIds });
 
   // delete collections
   await MongoDatasetCollection.deleteMany({
@@ -56,7 +56,7 @@ export async function delCollectionRelevantData({
   // delete dataset.datas
   await MongoDatasetData.deleteMany({ collectionId: { $in: collectionIds } });
   // delete pg data
-  await deletePgDataById(`collection_id IN ('${collectionIds.join("','")}')`);
+  await deleteDatasetDataVector({ collectionIds });
 
   // delete collections
   await MongoDatasetCollection.deleteMany({
@@ -76,6 +76,6 @@ export async function delCollectionRelevantData({
  * delete one data by mongoDataId
  */
 export async function delDatasetDataByDataId(mongoDataId: string) {
-  await deletePgDataById(['data_id', mongoDataId]);
+  await deleteDatasetDataVector({ dataIds: [mongoDataId] });
   await MongoDatasetData.findByIdAndDelete(mongoDataId);
 }

packages/service/core/dataset/data/pg.ts

@@ -1,28 +0,0 @@
-import { PgDatasetTableName } from '@fastgpt/global/core/dataset/constant';
-import { delay } from '@fastgpt/global/common/system/utils';
-import { PgClient } from '../../../common/pg';
-
-export async function deletePgDataById(
-  where: ['id' | 'dataset_id' | 'collection_id' | 'data_id', string] | string
-) {
-  let retry = 2;
-  async function deleteData(): Promise<any> {
-    try {
-      await PgClient.delete(PgDatasetTableName, {
-        where: [where]
-      });
-    } catch (error) {
-      if (--retry < 0) {
-        return Promise.reject(error);
-      }
-      await delay(500);
-      return deleteData();
-    }
-  }
-
-  await deleteData();
-
-  return {
-    tokenLen: 0
-  };
-}

packages/service/core/dataset/data/schema.ts

@@ -85,7 +85,6 @@ const DatasetDataSchema = new Schema({
 });
 
 try {
-  DatasetDataSchema.index({ teamId: 1 });
   DatasetDataSchema.index({ datasetId: 1 });
   DatasetDataSchema.index({ collectionId: 1 });
   DatasetDataSchema.index({ updateTime: -1 });

packages/service/support/openapi/schema.ts

@@ -2,7 +2,7 @@ import { connectionMongo, type Model } from '../../common/mongo';
 const { Schema, model, models } = connectionMongo;
 import type { OpenApiSchema } from '@fastgpt/global/support/openapi/type';
 import { PRICE_SCALE } from '@fastgpt/global/support/wallet/bill/constants';
-import { formatPrice } from '@fastgpt/global/support/wallet/bill/tools';
+import { formatStorePrice2Read } from '@fastgpt/global/support/wallet/bill/tools';
 import {
   TeamCollectionName,
   TeamMemberCollectionName
@@ -48,7 +48,7 @@ const OpenApiSchema = new Schema(
       // total usage. value from bill total
       type: Number,
       default: 0,
-      get: (val: number) => formatPrice(val)
+      get: (val: number) => formatStorePrice2Read(val)
     },
     limit: {
       expiredTime: {
@@ -59,7 +59,7 @@ const OpenApiSchema = new Schema(
         type: Number,
         default: -1,
         set: (val: number) => val * PRICE_SCALE,
-        get: (val: number) => formatPrice(val)
+        get: (val: number) => formatStorePrice2Read(val)
       }
     }
   },

packages/service/support/user/team/teamMemberSchema.ts

@@ -43,6 +43,7 @@ const TeamMemberSchema = new Schema({
 });
 
 try {
+  TeamMemberSchema.index({ teamId: 1 });
 } catch (error) {
   console.log(error);
 }

packages/service/support/user/team/teamSchema.ts

@@ -29,10 +29,14 @@ const TeamSchema = new Schema({
   maxSize: {
     type: Number,
     default: 5
+  },
+  lastDatasetBillTime: {
+    type: Date
   }
 });
 
 try {
+  TeamSchema.index({ lastDatasetBillTime: -1 });
 } catch (error) {
   console.log(error);
 }

packages/service/support/wallet/bill/controller.ts

@@ -25,14 +25,16 @@ export const createTrainingBill = async ({
       {
         moduleName: 'wallet.moduleName.index',
         model: vectorModel,
-        amount: 0,
-        tokenLen: 0
+        inputTokens: 0,
+        outputTokens: 0,
+        amount: 0
       },
       {
         moduleName: 'wallet.moduleName.qa',
         model: agentModel,
-        amount: 0,
-        tokenLen: 0
+        inputTokens: 0,
+        outputTokens: 0,
+        amount: 0
       }
     ],
     total: 0

packages/service/support/wallet/bill/schema.ts

@@ -52,7 +52,8 @@ const BillSchema = new Schema({
 });
 
 try {
-  BillSchema.index({ userId: 1 });
+  BillSchema.index({ teamId: 1 });
+  BillSchema.index({ tmbId: 1 });
   BillSchema.index({ time: 1 }, { expireAfterSeconds: 90 * 24 * 60 * 60 });
 } catch (error) {
   console.log(error);