V4.6.6-2 (#673)

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

@@ -48,6 +48,8 @@ export type FastGPTFeConfigsType = {
   };
   scripts?: { [key: string]: string }[];
   favicon?: string;
+  customApiDomain?: string;
+  customSharePageDomain?: string;
 };
 
 export type SystemEnvType = {

packages/global/common/vectorStore/constants.ts

@@ -0,0 +1 @@
+export const PgDatasetTableName = 'modeldata';

packages/global/core/ai/model.d.ts

@@ -3,7 +3,8 @@ export type LLMModelItemType = {
   name: string;
   maxContext: number;
   maxResponse: number;
-  price: number;
+  inputPrice: number;
+  outputPrice: number;
 };
 export type ChatModelItemType = LLMModelItemType & {
   quoteMaxToken: number;
@@ -22,7 +23,8 @@ export type VectorModelItemType = {
   model: string;
   name: string;
   defaultToken: number;
-  price: number;
+  inputPrice: number;
+  outputPrice: number;
   maxToken: number;
   weight: number;
 };
@@ -30,7 +32,8 @@ export type VectorModelItemType = {
 export type ReRankModelItemType = {
   model: string;
   name: string;
-  price: number;
+  inputPrice: number;
+  outputPrice?: number;
   requestUrl?: string;
   requestAuth?: string;
 };
@@ -38,12 +41,14 @@ export type ReRankModelItemType = {
 export type AudioSpeechModelType = {
   model: string;
   name: string;
-  price: number;
+  inputPrice: number;
+  outputPrice?: number;
   voices: { label: string; value: string; bufferId: string }[];
 };
 
 export type WhisperModelType = {
   model: string;
   name: string;
-  price: number;
+  inputPrice: number;
+  outputPrice?: number;
 };

packages/global/core/ai/model.ts

@@ -6,7 +6,8 @@ export const defaultQAModels: LLMModelItemType[] = [
     name: 'GPT35-16k',
     maxContext: 16000,
     maxResponse: 16000,
-    price: 0
+    inputPrice: 0,
+    outputPrice: 0
   }
 ];
 
@@ -14,7 +15,8 @@ export const defaultVectorModels: VectorModelItemType[] = [
   {
     model: 'text-embedding-ada-002',
     name: 'Embedding-2',
-    price: 0,
+    inputPrice: 0,
+    outputPrice: 0,
     defaultToken: 500,
     maxToken: 3000,
     weight: 100

packages/global/core/app/type.d.ts

@@ -65,6 +65,7 @@ export type AppSimpleEditFormType = {
     similarity: number;
     limit: number;
     searchMode: `${DatasetSearchModeEnum}`;
+    usingReRank: boolean;
     searchEmptyText: string;
   };
   cfr: {
@@ -112,6 +113,7 @@ export type AppSimpleEditConfigTemplateType = {
       similarity?: boolean;
       limit?: boolean;
       searchMode: `${DatasetSearchModeEnum}`;
+      usingReRank: boolean;
       searchEmptyText?: boolean;
     };
     cfr?: {

packages/global/core/app/utils.ts

@@ -26,7 +26,8 @@ export const getDefaultAppForm = (templateId = 'fastgpt-universal'): AppSimpleEd
       similarity: 0.4,
       limit: 1500,
       searchEmptyText: '',
-      searchMode: DatasetSearchModeEnum.embedding
+      searchMode: DatasetSearchModeEnum.embedding,
+      usingReRank: false
     },
     userGuide: {
       welcomeText: '',
@@ -95,6 +96,10 @@ export const appModules2Form = ({
       defaultAppForm.dataset.searchMode =
         findInputValueByKey(module.inputs, ModuleInputKeyEnum.datasetSearchMode) ||
         DatasetSearchModeEnum.embedding;
+      defaultAppForm.dataset.usingReRank = !!findInputValueByKey(
+        module.inputs,
+        ModuleInputKeyEnum.datasetSearchUsingReRank
+      );
 
       // empty text
       const emptyOutputs =

packages/global/core/chat/type.d.ts

@@ -89,7 +89,8 @@ export type moduleDispatchResType = {
   moduleLogo?: string;
   price?: number;
   runningTime?: number;
-  tokens?: number;
+  inputTokens?: number;
+  outputTokens?: number;
   model?: string;
   query?: string;
   contextTotalLen?: number;
@@ -105,6 +106,7 @@ export type moduleDispatchResType = {
   similarity?: number;
   limit?: number;
   searchMode?: `${DatasetSearchModeEnum}`;
+  searchUsingReRank?: boolean;
 
   // cq
   cqList?: ClassifyQuestionAgentItemType[];
@@ -124,6 +126,9 @@ export type moduleDispatchResType = {
 
   // tf switch
   tfSwitchResult?: boolean;
+
+  // abandon
+  tokens?: number;
 };
 
 export type ChatHistoryItemResType = moduleDispatchResType & {

packages/global/core/dataset/constant.ts

@@ -1,5 +1,3 @@
-export const PgDatasetTableName = 'modeldata';
-
 /* ------------ dataset -------------- */
 export enum DatasetTypeEnum {
   folder = 'folder',
@@ -119,8 +117,8 @@ export const TrainingTypeMap = {
 /* ------------ search -------------- */
 export enum DatasetSearchModeEnum {
   embedding = 'embedding',
-  embeddingReRank = 'embeddingReRank',
-  embFullTextReRank = 'embFullTextReRank'
+  fullTextRecall = 'fullTextRecall',
+  mixedRecall = 'mixedRecall'
 }
 
 export const DatasetSearchModeMap = {
@@ -130,18 +128,25 @@ export const DatasetSearchModeMap = {
     desc: 'core.dataset.search.mode.embedding desc',
     value: DatasetSearchModeEnum.embedding
   },
-  [DatasetSearchModeEnum.embeddingReRank]: {
-    icon: 'core/dataset/modeEmbeddingRerank',
-    title: 'core.dataset.search.mode.embeddingReRank',
-    desc: 'core.dataset.search.mode.embeddingReRank desc',
-    value: DatasetSearchModeEnum.embeddingReRank
+  [DatasetSearchModeEnum.fullTextRecall]: {
+    icon: 'core/dataset/fullTextRecall',
+    title: 'core.dataset.search.mode.fullTextRecall',
+    desc: 'core.dataset.search.mode.fullTextRecall desc',
+    value: DatasetSearchModeEnum.fullTextRecall
   },
-  [DatasetSearchModeEnum.embFullTextReRank]: {
-    icon: 'core/dataset/modeEmbFTRerank',
-    title: 'core.dataset.search.mode.embFullTextReRank',
-    desc: 'core.dataset.search.mode.embFullTextReRank desc',
-    value: DatasetSearchModeEnum.embFullTextReRank
+  [DatasetSearchModeEnum.mixedRecall]: {
+    icon: 'core/dataset/mixedRecall',
+    title: 'core.dataset.search.mode.mixedRecall',
+    desc: 'core.dataset.search.mode.mixedRecall desc',
+    value: DatasetSearchModeEnum.mixedRecall
   }
 };
 
+export enum SearchScoreTypeEnum {
+  embedding = 'embedding',
+  fullText = 'fullText',
+  reRank = 'reRank',
+  rrf = 'rrf'
+}
+
 export const FolderAvatarSrc = '/imgs/files/folder.svg';

packages/global/core/dataset/type.d.ts

@@ -6,6 +6,7 @@ import {
   DatasetDataIndexTypeEnum,
   DatasetStatusEnum,
   DatasetTypeEnum,
+  SearchScoreTypeEnum,
   TrainingModeEnum
 } from './constant';
 
@@ -161,5 +162,6 @@ export type DatasetFileSchema = {
 
 /* ============= search =============== */
 export type SearchDataResponseItemType = Omit<DatasetDataItemType, 'isOwner' | 'canWrite'> & {
-  score: number;
+  score: { type: `${SearchScoreTypeEnum}`; value: number; index: number }[];
+  // score: number;
 };

packages/global/core/module/constants.ts

@@ -63,6 +63,7 @@ export enum ModuleInputKeyEnum {
   datasetSimilarity = 'similarity',
   datasetLimit = 'limit',
   datasetSearchMode = 'searchMode',
+  datasetSearchUsingReRank = 'usingReRank',
   datasetParamsModal = 'datasetParamsModal',
 
   // context extract

packages/global/core/module/template/system/datasetSearch.ts

@@ -64,12 +64,21 @@ export const DatasetSearchModule: FlowModuleTemplateType = {
     {
       key: ModuleInputKeyEnum.datasetSearchMode,
       type: FlowNodeInputTypeEnum.hidden,
-      label: 'core.dataset.search.Mode',
+      label: '',
       valueType: ModuleIOValueTypeEnum.string,
       showTargetInApp: false,
       showTargetInPlugin: false,
       value: DatasetSearchModeEnum.embedding
     },
+    {
+      key: ModuleInputKeyEnum.datasetSearchUsingReRank,
+      type: FlowNodeInputTypeEnum.hidden,
+      label: '',
+      valueType: ModuleIOValueTypeEnum.boolean,
+      showTargetInApp: false,
+      showTargetInPlugin: false,
+      value: false
+    },
     {
       key: ModuleInputKeyEnum.datasetParamsModal,
       type: FlowNodeInputTypeEnum.selectDatasetParamsModal,

packages/global/support/user/team/type.d.ts

@@ -9,6 +9,7 @@ export type TeamSchema = {
   createTime: Date;
   balance: number;
   maxSize: number;
+  lastDatasetBillTime: Date;
 };
 
 export type TeamMemberSchema = {

packages/global/support/wallet/bill/api.d.ts

@@ -1,5 +1,5 @@
 import { BillSourceEnum } from './constants';
-import { BillListItemType } from './type';
+import { BillListItemCountType, BillListItemType } from './type';
 
 export type CreateTrainingBillProps = {
   name: string;
@@ -7,13 +7,12 @@ export type CreateTrainingBillProps = {
   agentModel?: string;
 };
 
-export type ConcatBillProps = {
+export type ConcatBillProps = BillListItemCountType & {
   teamId: string;
   tmbId: string;
   billId?: string;
   total: number;
   listIndex?: number;
-  tokens?: number;
 };
 
 export type CreateBillProps = {

packages/global/support/wallet/bill/constants.ts

@@ -1,16 +1,19 @@
-// ￥1 = 100000
+// model price: xxx/1k tokens
+// ￥1 = 100000.
 export const PRICE_SCALE = 100000;
 
 export enum BillSourceEnum {
   fastgpt = 'fastgpt',
   api = 'api',
   shareLink = 'shareLink',
-  training = 'training'
+  training = 'training',
+  datasetStore = 'datasetStore'
 }
 
 export const BillSourceMap: Record<`${BillSourceEnum}`, string> = {
   [BillSourceEnum.fastgpt]: '在线使用',
   [BillSourceEnum.api]: 'Api',
   [BillSourceEnum.shareLink]: '免登录链接',
-  [BillSourceEnum.training]: '数据训练'
+  [BillSourceEnum.training]: '数据训练',
+  [BillSourceEnum.datasetStore]: '知识库存储'
 };

packages/global/support/wallet/bill/tools.ts

@@ -6,9 +6,12 @@ import { AuthUserTypeEnum } from '../../permission/constant';
 /**
  * dataset price / PRICE_SCALE = real price
  */
-export const formatPrice = (val = 0, multiple = 1) => {
+export const formatStorePrice2Read = (val = 0, multiple = 1) => {
   return Number(((val / PRICE_SCALE) * multiple).toFixed(10));
 };
+export const formatModelPrice2Read = (val = 0) => {
+  return Number((val / 1000).toFixed(10));
+};
 
 export const getBillSourceByAuthType = ({
   shareId,

packages/global/support/wallet/bill/type.d.ts

@@ -1,11 +1,20 @@
 import { CreateBillProps } from './api';
 import { BillSourceEnum } from './constants';
 
-export type BillListItemType = {
+export type BillListItemCountType = {
+  inputTokens?: number;
+  outputTokens?: number;
+  textLen?: number;
+  duration?: number;
+  dataLen?: number;
+
+  // abandon
+  tokenLen?: number;
+};
+export type BillListItemType = BillListItemCountType & {
   moduleName: string;
   amount: number;
   model?: string;
-  tokenLen?: number;
 };
 
 export type BillSchema = CreateBillProps & {

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);