V4.9.7 feature (#4669)

* update doc

* feat: Add coupon redemption feature for team subscriptions (#4595)

* feat: Add coupon redemption feature for team subscriptions

- Introduced `TeamCouponSub` and `TeamCouponSchema` types
- Added `redeemCoupon` API endpoint
- Updated UI to include a modal for coupon redemption
- Added new icon and translations for "Redeem coupon"

* perf: remove field teamId

* perf: use dynamic import

* refactor: move to page component

* perf: coupon code

* perf: mcp server

* perf: test

* auto layout (#4634)

* fix 4.9.6 (#4631)

* fix debug quote list

* delete next text node match

* fix extract default boolean value

* export latest 100 chat items

* fix quote item ui

* doc

* fix doc

* feat: auto layout

* perf: auto layout

* fix: auto layout null

* add start node

---------

Co-authored-by: heheer <heheer@sealos.io>

* fix: share link (#4644)

* Add workflow run duration;Get audio duration (#4645)

* add duration

* get audio duration

* Custom config path (#4649)

* feat: 通过环境变量DATA_PATH获取配置文件目录 (#4622)

通过环境变量DATA_PATH获取配置文件目录,以应对不同的部署方式的多样化需求

* feat: custom configjson path

* doc

---------

Co-authored-by: John Chen <sss1991@163.com>

* 程序api调用场景下,如果大量调用带有图片或视频,产生的聊天记录会导致后台mongo数据库异常。这个修改给api客户端一个禁止生成聊天记录的选项,避免这个后果。 (#3964)

* update special chatId

* perf: vector db rename

* update operationLog (#4647)

* update operationLog

* combine operationLogMap

* solve operationI18nLogMap bug

* remoce log

* feat: Rerank usage (#4654)

* refresh concat when update (#4655)

* fix: refresh code

* perf: timer lock

* Fix operationLog (#4657)

* perf: http streamable mcp

* add alipay (#4630)

* perf: subplan ui

* perf: pay code

* hiden bank tip

* Fix: pay error (#4665)

* fix quote number (#4666)

* remove log

---------

Co-authored-by: a.e. <49438478+I-Info@users.noreply.github.com>
Co-authored-by: heheer <heheer@sealos.io>
Co-authored-by: John Chen <sss1991@163.com>
Co-authored-by: gaord <bengao168@msn.com>
Co-authored-by: gggaaallleee <91131304+gggaaallleee@users.noreply.github.com>
This commit is contained in:
Archer
2025-04-26 16:17:21 +08:00
committed by GitHub
parent a669a60fe6
commit 0720bbe4da
143 changed files with 2067 additions and 1093 deletions

View File

@@ -0,0 +1,7 @@
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;

View File

@@ -0,0 +1,38 @@
import type { EmbeddingRecallItemType } from './type';
export type DeleteDatasetVectorProps = (
| { id: string }
| { datasetIds: string[]; collectionIds?: string[] }
| { idList: string[] }
) & {
teamId: string;
};
export type DelDatasetVectorCtrlProps = DeleteDatasetVectorProps & {
retry?: number;
};
export type InsertVectorProps = {
teamId: string;
datasetId: string;
collectionId: string;
};
export type InsertVectorControllerProps = InsertVectorProps & {
vector: number[];
retry?: number;
};
export type EmbeddingRecallProps = {
teamId: string;
datasetIds: string[];
forbidCollectionIdList: string[];
filterCollectionIdList?: string[];
};
export type EmbeddingRecallCtrlProps = EmbeddingRecallProps & {
vector: number[];
limit: number;
retry?: number;
};
export type EmbeddingRecallResponse = {
results: EmbeddingRecallItemType[];
};

View File

@@ -0,0 +1,83 @@
/* vector crud */
import { PgVectorCtrl } from './pg';
import { ObVectorCtrl } from './oceanbase';
import { getVectorsByText } from '../../core/ai/embedding';
import { DelDatasetVectorCtrlProps, InsertVectorProps } from './controller.d';
import { EmbeddingModelItemType } from '@fastgpt/global/core/ai/model.d';
import { MILVUS_ADDRESS, PG_ADDRESS, OCEANBASE_ADDRESS } from './constants';
import { MilvusCtrl } from './milvus';
import { setRedisCache, getRedisCache, delRedisCache, CacheKeyEnum } from '../redis/cache';
import { throttle } from 'lodash';
import { retryFn } from '@fastgpt/global/common/system/utils';
const getVectorObj = () => {
if (PG_ADDRESS) return new PgVectorCtrl();
if (OCEANBASE_ADDRESS) return new ObVectorCtrl();
if (MILVUS_ADDRESS) return new MilvusCtrl();
return new PgVectorCtrl();
};
const getChcheKey = (teamId: string) => `${CacheKeyEnum.team_vector_count}:${teamId}`;
const onDelCache = throttle((teamId: string) => delRedisCache(getChcheKey(teamId)), 30000, {
leading: true,
trailing: true
});
const Vector = getVectorObj();
export const initVectorStore = Vector.init;
export const recallFromVectorStore = Vector.embRecall;
export const getVectorDataByTime = Vector.getVectorDataByTime;
export const getVectorCountByTeamId = async (teamId: string) => {
const key = getChcheKey(teamId);
const countStr = await getRedisCache(key);
if (countStr) {
return Number(countStr);
}
const count = await Vector.getVectorCountByTeamId(teamId);
await setRedisCache(key, count, 30 * 60);
return count;
};
export const getVectorCountByDatasetId = Vector.getVectorCountByDatasetId;
export const getVectorCountByCollectionId = Vector.getVectorCountByCollectionId;
export const insertDatasetDataVector = async ({
model,
query,
...props
}: InsertVectorProps & {
query: string;
model: EmbeddingModelItemType;
}) => {
return retryFn(async () => {
const { vectors, tokens } = await getVectorsByText({
model,
input: query,
type: 'db'
});
const { insertId } = await Vector.insert({
...props,
vector: vectors[0]
});
onDelCache(props.teamId);
return {
tokens,
insertId
};
});
};
export const deleteDatasetDataVector = async (props: DelDatasetVectorCtrlProps) => {
const result = await Vector.delete(props);
onDelCache(props.teamId);
return result;
};

View File

@@ -0,0 +1,365 @@
import { DataType, LoadState, MilvusClient } from '@zilliz/milvus2-sdk-node';
import {
DatasetVectorDbName,
DatasetVectorTableName,
MILVUS_ADDRESS,
MILVUS_TOKEN
} from '../constants';
import type {
DelDatasetVectorCtrlProps,
EmbeddingRecallCtrlProps,
EmbeddingRecallResponse,
InsertVectorControllerProps
} from '../controller.d';
import { delay } from '@fastgpt/global/common/system/utils';
import { addLog } from '../../system/log';
import { customNanoid } from '@fastgpt/global/common/string/tools';
export class MilvusCtrl {
constructor() {}
getClient = async () => {
if (!MILVUS_ADDRESS) {
return Promise.reject('MILVUS_ADDRESS is not set');
}
if (global.milvusClient) return global.milvusClient;
global.milvusClient = new MilvusClient({
address: MILVUS_ADDRESS,
token: MILVUS_TOKEN
});
addLog.info(`Milvus connected`);
return global.milvusClient;
};
init = async () => {
const client = await this.getClient();
// init db(zilliz cloud will error)
try {
const { db_names } = await client.listDatabases();
if (!db_names.includes(DatasetVectorDbName)) {
await client.createDatabase({
db_name: DatasetVectorDbName
});
}
await client.useDatabase({
db_name: DatasetVectorDbName
});
} catch (error) {}
// init collection and index
const { value: hasCollection } = await client.hasCollection({
collection_name: DatasetVectorTableName
});
if (!hasCollection) {
const result = await client.createCollection({
collection_name: DatasetVectorTableName,
description: 'Store dataset vector',
enableDynamicField: true,
fields: [
{
name: 'id',
data_type: DataType.Int64,
is_primary_key: true,
autoID: false // disable auto id, and we need to set id in insert
},
{
name: 'vector',
data_type: DataType.FloatVector,
dim: 1536
},
{ name: 'teamId', data_type: DataType.VarChar, max_length: 64 },
{ name: 'datasetId', data_type: DataType.VarChar, max_length: 64 },
{ name: 'collectionId', data_type: DataType.VarChar, max_length: 64 },
{
name: 'createTime',
data_type: DataType.Int64
}
],
index_params: [
{
field_name: 'vector',
index_name: 'vector_HNSW',
index_type: 'HNSW',
metric_type: 'IP',
params: { efConstruction: 32, M: 64 }
},
{
field_name: 'teamId',
index_type: 'Trie'
},
{
field_name: 'datasetId',
index_type: 'Trie'
},
{
field_name: 'collectionId',
index_type: 'Trie'
},
{
field_name: 'createTime',
index_type: 'STL_SORT'
}
]
});
addLog.info(`Create milvus collection: `, result);
}
const { state: colLoadState } = await client.getLoadState({
collection_name: DatasetVectorTableName
});
if (
colLoadState === LoadState.LoadStateNotExist ||
colLoadState === LoadState.LoadStateNotLoad
) {
await client.loadCollectionSync({
collection_name: DatasetVectorTableName
});
addLog.info(`Milvus collection load success`);
}
};
insert = async (props: InsertVectorControllerProps): Promise<{ insertId: string }> => {
const client = await this.getClient();
const { teamId, datasetId, collectionId, vector, retry = 3 } = props;
const generateId = () => {
// in js, the max safe integer is 2^53 - 1: 9007199254740991
// so we can generate a random number between 1-8 as the first digit
// and the rest 15 digits can be random
const firstDigit = customNanoid('12345678', 1);
const restDigits = customNanoid('1234567890', 15);
return Number(`${firstDigit}${restDigits}`);
};
const id = generateId();
try {
const result = await client.insert({
collection_name: DatasetVectorTableName,
data: [
{
id,
vector,
teamId: String(teamId),
datasetId: String(datasetId),
collectionId: String(collectionId),
createTime: Date.now()
}
]
});
const insertId = (() => {
if ('int_id' in result.IDs) {
return `${result.IDs.int_id.data?.[0]}`;
}
return `${result.IDs.str_id.data?.[0]}`;
})();
return {
insertId: insertId
};
} 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 client = await this.getClient();
const teamIdWhere = `(teamId=="${String(teamId)}")`;
const where = await (() => {
if ('id' in props && props.id) return `(id==${props.id})`;
if ('datasetIds' in props && props.datasetIds) {
const datasetIdWhere = `(datasetId in [${props.datasetIds
.map((id) => `"${String(id)}"`)
.join(',')}])`;
if ('collectionIds' in props && props.collectionIds) {
return `${datasetIdWhere} and (collectionId in [${props.collectionIds
.map((id) => `"${String(id)}"`)
.join(',')}])`;
}
return `${datasetIdWhere}`;
}
if ('idList' in props && Array.isArray(props.idList)) {
if (props.idList.length === 0) return;
return `(id in [${props.idList.map((id) => String(id)).join(',')}])`;
}
return Promise.reject('deleteDatasetData: no where');
})();
if (!where) return;
const concatWhere = `${teamIdWhere} and ${where}`;
try {
await client.delete({
collection_name: DatasetVectorTableName,
filter: concatWhere
});
} 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 client = await this.getClient();
const {
teamId,
datasetIds,
vector,
limit,
forbidCollectionIdList,
filterCollectionIdList,
retry = 2
} = props;
// Forbid collection
const formatForbidCollectionIdList = (() => {
if (!filterCollectionIdList) return forbidCollectionIdList;
const list = forbidCollectionIdList
.map((id) => String(id))
.filter((id) => !filterCollectionIdList.includes(id));
return list;
})();
const forbidColQuery =
formatForbidCollectionIdList.length > 0
? `and (collectionId not in [${formatForbidCollectionIdList.map((id) => `"${id}"`).join(',')}])`
: '';
// filter collection id
const formatFilterCollectionId = (() => {
if (!filterCollectionIdList) return;
return filterCollectionIdList
.map((id) => String(id))
.filter((id) => !forbidCollectionIdList.includes(id));
})();
const collectionIdQuery = formatFilterCollectionId
? `and (collectionId in [${formatFilterCollectionId.map((id) => `"${id}"`)}])`
: ``;
// Empty data
if (formatFilterCollectionId && formatFilterCollectionId.length === 0) {
return { results: [] };
}
try {
const { results } = await client.search({
collection_name: DatasetVectorTableName,
data: vector,
limit,
filter: `(teamId == "${teamId}") and (datasetId in [${datasetIds.map((id) => `"${id}"`).join(',')}]) ${collectionIdQuery} ${forbidColQuery}`,
output_fields: ['collectionId']
});
const rows = results as {
score: number;
id: string;
collectionId: string;
}[];
return {
results: rows.map((item) => ({
id: String(item.id),
collectionId: item.collectionId,
score: item.score
}))
};
} catch (error) {
if (retry <= 0) {
return Promise.reject(error);
}
return this.embRecall({
...props,
retry: retry - 1
});
}
};
getVectorCountByTeamId = async (teamId: string) => {
const client = await this.getClient();
const result = await client.query({
collection_name: DatasetVectorTableName,
output_fields: ['count(*)'],
filter: `teamId == "${String(teamId)}"`
});
const total = result.data?.[0]?.['count(*)'] as number;
return total;
};
getVectorCountByDatasetId = async (teamId: string, datasetId: string) => {
const client = await this.getClient();
const result = await client.query({
collection_name: DatasetVectorTableName,
output_fields: ['count(*)'],
filter: `(teamId == "${String(teamId)}") and (dataset == "${String(datasetId)}")`
});
const total = result.data?.[0]?.['count(*)'] as number;
return total;
};
getVectorCountByCollectionId = async (
teamId: string,
datasetId: string,
collectionId: string
) => {
const client = await this.getClient();
const result = await client.query({
collection_name: DatasetVectorTableName,
output_fields: ['count(*)'],
filter: `(teamId == "${String(teamId)}") and (datasetId == "${String(datasetId)}") and (collectionId == "${String(collectionId)}")`
});
const total = result.data?.[0]?.['count(*)'] as number;
return total;
};
getVectorDataByTime = async (start: Date, end: Date) => {
const client = await this.getClient();
const startTimestamp = new Date(start).getTime();
const endTimestamp = new Date(end).getTime();
const result = await client.query({
collection_name: DatasetVectorTableName,
output_fields: ['id', 'teamId', 'datasetId'],
filter: `(createTime >= ${startTimestamp}) and (createTime <= ${endTimestamp})`
});
const rows = result.data as {
id: string;
teamId: string;
datasetId: string;
}[];
return rows.map((item) => ({
id: String(item.id),
teamId: item.teamId,
datasetId: item.datasetId
}));
};
}

View File

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

View File

@@ -0,0 +1,254 @@
/* oceanbase vector crud */
import { DatasetVectorTableName } from '../constants';
import { delay } from '@fastgpt/global/common/system/utils';
import { ObClient } from './controller';
import { RowDataPacket } 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;
};
}

View File

@@ -0,0 +1,188 @@
import { delay } from '@fastgpt/global/common/system/utils';
import { addLog } from '../../system/log';
import { Pool } from 'pg';
import type { QueryResultRow } from 'pg';
import { PG_ADDRESS } from '../constants';
export const connectPg = async (): Promise<Pool> => {
if (global.pgClient) {
return global.pgClient;
}
global.pgClient = new Pool({
connectionString: PG_ADDRESS,
max: Number(process.env.DB_MAX_LINK || 20),
min: 10,
keepAlive: true,
idleTimeoutMillis: 600000,
connectionTimeoutMillis: 20000,
query_timeout: 30000,
statement_timeout: 40000,
idle_in_transaction_session_timeout: 60000
});
global.pgClient.on('error', async (err) => {
addLog.error(`pg error`, err);
global.pgClient?.end();
global.pgClient = null;
await delay(1000);
addLog.info(`Retry connect pg`);
connectPg();
});
try {
await global.pgClient.connect();
console.log('pg connected');
return global.pgClient;
} catch (error) {
addLog.error(`pg connect error`, error);
global.pgClient?.end();
global.pgClient = null;
await delay(1000);
addLog.info(`Retry connect pg`);
return connectPg();
}
};
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 PgClass {
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 QueryResultRow = 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 pg = await connectPg();
return pg.query<T>(sql);
}
async count(table: string, props: GetProps) {
const sql = `SELECT COUNT(${props?.fields?.[0] || '*'})
FROM ${table}
${this.getWhereStr(props.where)}
`;
const pg = await connectPg();
return pg.query(sql).then((res) => Number(res.rows[0]?.count || 0));
}
async delete(table: string, props: DeleteProps) {
const sql = `DELETE FROM ${table} ${this.getWhereStr(props.where)}`;
const pg = await connectPg();
return pg.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 pg = await connectPg();
return pg.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
)} RETURNING id`;
const pg = await connectPg();
return pg.query<{ id: string }>(sql);
}
async query<T extends QueryResultRow = any>(sql: string) {
const pg = await connectPg();
const start = Date.now();
return pg.query<T>(sql).then((res) => {
const time = Date.now() - start;
if (time > 300) {
addLog.warn(`pg query time: ${time}ms, sql: ${sql}`);
}
return res;
});
}
}
export const PgClient = new PgClass();
export const Pg = global.pgClient;

View File

@@ -0,0 +1,276 @@
/* pg vector crud */
import { DatasetVectorTableName } from '../constants';
import { delay } from '@fastgpt/global/common/system/utils';
import { PgClient, connectPg } from './controller';
import { PgSearchRawType } from '@fastgpt/global/core/dataset/api';
import type {
DelDatasetVectorCtrlProps,
EmbeddingRecallCtrlProps,
EmbeddingRecallResponse,
InsertVectorControllerProps
} from '../controller.d';
import dayjs from 'dayjs';
import { addLog } from '../../system/log';
export class PgVectorCtrl {
constructor() {}
init = async () => {
try {
await connectPg();
await PgClient.query(`
CREATE EXTENSION IF NOT EXISTS vector;
CREATE TABLE IF NOT EXISTS ${DatasetVectorTableName} (
id BIGSERIAL 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 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);`
);
// 10w rows
// await PgClient.query(`
// ALTER TABLE modeldata SET (
// autovacuum_vacuum_scale_factor = 0.1,
// autovacuum_analyze_scale_factor = 0.05,
// autovacuum_vacuum_threshold = 50,
// autovacuum_analyze_threshold = 50,
// autovacuum_vacuum_cost_delay = 20,
// autovacuum_vacuum_cost_limit = 200
// );`);
// 100w rows
// await PgClient.query(`
// ALTER TABLE modeldata SET (
// autovacuum_vacuum_scale_factor = 0.01,
// autovacuum_analyze_scale_factor = 0.02,
// autovacuum_vacuum_threshold = 1000,
// autovacuum_analyze_threshold = 1000,
// autovacuum_vacuum_cost_delay = 10,
// autovacuum_vacuum_cost_limit = 2000
// );`)
addLog.info('init pg successful');
} catch (error) {
addLog.error('init pg error', error);
}
};
insert = async (props: InsertVectorControllerProps): Promise<{ insertId: string }> => {
const { teamId, datasetId, collectionId, vector, retry = 3 } = props;
try {
const { rowCount, rows } = await PgClient.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 PgClient.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 results: any = await PgClient.query(
`BEGIN;
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
from ${DatasetVectorTableName}
where dataset_id IN (${datasetIds.map((id) => `'${String(id)}'`).join(',')})
${filterCollectionIdSql}
${forbidCollectionSql}
order by score limit ${limit}
) SELECT id, collection_id, score FROM relaxed_results ORDER BY score;
COMMIT;`
);
const rows = results?.[3]?.rows as PgSearchRawType[];
if (!Array.isArray(rows)) {
return {
results: []
};
}
return {
results: rows.map((item) => ({
id: String(item.id),
collectionId: item.collection_id,
score: item.score * -1
}))
};
} 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 PgClient.query<{
id: string;
team_id: string;
dataset_id: string;
}>(`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')}';
`);
return rows.map((item) => ({
id: String(item.id),
teamId: item.team_id,
datasetId: item.dataset_id
}));
};
getVectorCountByTeamId = async (teamId: string) => {
const total = await PgClient.count(DatasetVectorTableName, {
where: [['team_id', String(teamId)]]
});
return total;
};
getVectorCountByDatasetId = async (teamId: string, datasetId: string) => {
const total = await PgClient.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 PgClient.count(DatasetVectorTableName, {
where: [
['team_id', String(teamId)],
'and',
['dataset_id', String(datasetId)],
'and',
['collection_id', String(collectionId)]
]
});
return total;
};
}

View File

@@ -0,0 +1,15 @@
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;
}
export type EmbeddingRecallItemType = {
id: string;
collectionId: string;
score: number;
};