Skip to content

赵月辉 / fastgpt-migrated

Go to a project
Commit 75b6947c by Archer Committed by GitHub
Options

perf: vector db log;perf: s3 mock (#6165) 

* stop design doc

* remove invalid doc

* remove log

* perf: s3 mock

* perf: vector db log

* update lock
parent 3ff4fda3
Showing with 218 additions and 157 deletions
packages/service/common/file/read/utils.ts
......@@ -172,14 +172,6 @@ export const readS3FileContentByBuffer = async ({
};
// Custom read file service
const pdfParseFn = async (): Promise<ReadFileResponse> => {
console.log(
'global.systemEnv.customPdfParse?.textinAppId',
global.systemEnv.customPdfParse?.textinAppId
);
console.log(
'global.systemEnv.customPdfParse?.doc2xKey',
global.systemEnv.customPdfParse?.doc2xKey
);
if (!customPdfParse) return systemParse();
if (global.systemEnv.customPdfParse?.url) return parsePdfFromCustomService();
if (global.systemEnv.customPdfParse?.textinAppId) return parsePdfFromTextin();
......
packages/service/common/vectorDB/controller.d.ts deleted 100644 → 0
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 & {
vectors: 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[];
};
packages/service/common/vectorDB/controller.ts
......@@ -2,8 +2,7 @@
import { PgVectorCtrl } from './pg';
import { ObVectorCtrl } from './oceanbase';
import { getVectorsByText } from '../../core/ai/embedding';
import type { EmbeddingRecallCtrlProps } from './controller.d';
import { type DelDatasetVectorCtrlProps, type InsertVectorProps } from './controller.d';
import type { VectorControllerType, InsertVectorControllerPropsType } from './type';
import { type EmbeddingModelItemType } from '@fastgpt/global/core/ai/model.d';
import { MILVUS_ADDRESS, PG_ADDRESS, OCEANBASE_ADDRESS } from './constants';
import { MilvusCtrl } from './milvus';
......@@ -60,33 +59,14 @@ const teamVectorCache = {
const Vector = getVectorObj();
export const initVectorStore = Vector.init;
export const recallFromVectorStore = (props: EmbeddingRecallCtrlProps) =>
export const recallFromVectorStore: VectorControllerType['embRecall'] = (props) =>
retryFn(() => Vector.embRecall(props));
export const getVectorDataByTime = Vector.getVectorDataByTime;
// Count vector
export const getVectorCountByTeamId = async (teamId: string) => {
const cacheCount = await teamVectorCache.get(teamId);
if (cacheCount !== undefined) {
return cacheCount;
}
const count = await Vector.getVectorCount({ teamId });
teamVectorCache.set({
teamId,
count
});
return count;
};
export const getVectorCount = Vector.getVectorCount;
export const insertDatasetDataVector = async ({
model,
inputs,
...props
}: InsertVectorProps & {
}: InsertVectorControllerPropsType & {
inputs: string[];
model: EmbeddingModelItemType;
}) => {
......@@ -110,8 +90,28 @@ export const insertDatasetDataVector = async ({
};
};
export const deleteDatasetDataVector = async (props: DelDatasetVectorCtrlProps) => {
export const deleteDatasetDataVector: VectorControllerType['delete'] = async (props) => {
const result = await retryFn(() => Vector.delete(props));
teamVectorCache.delete(props.teamId);
return result;
};
export const getVectorDataByTime = Vector.getVectorDataByTime;
// Count vector
export const getVectorCountByTeamId = async (teamId: string) => {
const cacheCount = await teamVectorCache.get(teamId);
if (cacheCount !== undefined) {
return cacheCount;
}
const count = await Vector.getVectorCount({ teamId });
teamVectorCache.set({
teamId,
count
});
return count;
};
export const getVectorCount = Vector.getVectorCount;
packages/service/common/vectorDB/milvus/index.ts
......@@ -11,11 +11,12 @@ import type {
EmbeddingRecallResponse,
InsertVectorControllerProps
} from '../controller.d';
import type { VectorControllerType } from '../type';
import { retryFn } from '@fastgpt/global/common/system/utils';
import { addLog } from '../../system/log';
import { customNanoid } from '@fastgpt/global/common/string/tools';
export class MilvusCtrl {
export class MilvusCtrl implements VectorControllerType {
constructor() {}
getClient = async () => {
if (!MILVUS_ADDRESS) {
......
packages/service/common/vectorDB/oceanbase/index.ts
......@@ -8,10 +8,11 @@ import {
type EmbeddingRecallResponse,
type InsertVectorControllerProps
} from '../controller.d';
import type { VectorControllerType } from '../type';
import dayjs from 'dayjs';
import { addLog } from '../../system/log';
export class ObVectorCtrl {
export class ObVectorCtrl implements VectorControllerType {
constructor() {}
init = async () => {
try {
......
packages/service/common/vectorDB/pg/controller.ts
......@@ -9,39 +9,47 @@ export const connectPg = async (): Promise<Pool> => {
return global.pgClient;
}
global.pgClient = new Pool({
const pool = new Pool({
connectionString: PG_ADDRESS,
max: Number(process.env.DB_MAX_LINK || 20),
min: 10,
// 连接池配置
max: Number(process.env.DB_MAX_LINK || 30), // 增加到 30，支持更高并发
min: 15, // 调整为 max 的 50%
keepAlive: true,
idleTimeoutMillis: 600000,
connectionTimeoutMillis: 20000,
query_timeout: 30000,
statement_timeout: 40000,
idle_in_transaction_session_timeout: 60000
// 超时配置
idleTimeoutMillis: 1800000, // 30分钟，减少频繁重连
connectionTimeoutMillis: 30000, // 30秒，给予充足的连接获取时间
query_timeout: 60000, // 60秒，向量检索可能需要更长时间
statement_timeout: 90000, // 90秒，比 query_timeout 长
idle_in_transaction_session_timeout: 60000, // 保持 60秒
// 额外推荐配置
allowExitOnIdle: false, // 防止连接池过早关闭
application_name: 'fastgpt-vector-db' // 便于数据库监控识别
});
global.pgClient = pool;
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();
addLog.error(`[PG] error`, err);
});
global.pgClient.on('connect', async () => {
addLog.info(`[PG] connect`);
});
global.pgClient.on('remove', async (client) => {
addLog.warn('[PG] Connection removed from pool');
});
try {
await global.pgClient.connect();
console.log('pg connected');
return global.pgClient;
} catch (error) {
addLog.error(`pg connect error`, error);
addLog.error(`[PG] connect error`, error);
global.pgClient?.removeAllListeners();
global.pgClient?.end();
global.pgClient = null;
await delay(1000);
addLog.info(`Retry connect pg`);
addLog.warn(`[PG] retry connect`);
return connectPg();
}
......@@ -109,6 +117,25 @@ class PgClass {
)
.join(',');
}
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 > 1000) {
const safeSql = sql.replace(/'\[[^\]]*?\]'/g, "'[x]'");
addLog.warn(`[PG slow 2] time: ${time}ms, sql: ${safeSql}`);
} else if (time > 300) {
const safeSql = sql.replace(/'\[[^\]]*?\]'/g, "'[x]'");
addLog.warn(`[PG slow 1] time: ${time}ms, sql: ${safeSql}`);
}
return res;
});
}
async select<T extends QueryResultRow = any>(table: string, props: GetProps) {
const sql = `SELECT ${
!props.fields || props.fields?.length === 0 ? '*' : props.fields?.join(',')
......@@ -123,8 +150,7 @@ class PgClass {
LIMIT ${props.limit || 10} OFFSET ${props.offset || 0}
`;
const pg = await connectPg();
return pg.query<T>(sql);
return this.query<T>(sql);
}
async count(table: string, props: GetProps) {
const sql = `SELECT COUNT(${props?.fields?.[0] || '*'})
......@@ -132,13 +158,11 @@ class PgClass {
${this.getWhereStr(props.where)}
`;
const pg = await connectPg();
return pg.query(sql).then((res) => Number(res.rows[0]?.count || 0));
return this.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);
return this.query(sql);
}
async update(table: string, props: UpdateProps) {
if (props.values.length === 0) {
......@@ -150,8 +174,7 @@ class PgClass {
const sql = `UPDATE ${table} SET ${this.getUpdateValStr(props.values)} ${this.getWhereStr(
props.where
)}`;
const pg = await connectPg();
return pg.query(sql);
return this.query(sql);
}
async insert(table: string, props: InsertProps) {
if (props.values.length === 0) {
......@@ -166,22 +189,7 @@ class PgClass {
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) {
const safeSql = sql.replace(/'\[[^\]]*?\]'/g, "'[x]'");
addLog.warn(`pg query time: ${time}ms, sql: ${safeSql}`);
}
return res;
});
return this.query<{ id: string }>(sql);
}
}
......
packages/service/common/vectorDB/pg/index.ts
/* pg vector crud */
import { DatasetVectorTableName } from '../constants';
import { delay, retryFn } from '@fastgpt/global/common/system/utils';
import { PgClient, connectPg } from './controller';
import { type PgSearchRawType } from '@fastgpt/global/core/dataset/api';
import type {
DelDatasetVectorCtrlProps,
EmbeddingRecallCtrlProps,
EmbeddingRecallResponse,
InsertVectorControllerProps
} from '../controller.d';
import type { VectorControllerType } from '../type';
import dayjs from 'dayjs';
import { addLog } from '../../system/log';
export class PgVectorCtrl {
export class PgVectorCtrl implements VectorControllerType {
constructor() {}
init = async () => {
try {
......@@ -65,7 +59,7 @@ export class PgVectorCtrl {
addLog.error('init pg error', error);
}
};
insert = async (props: InsertVectorControllerProps): Promise<{ insertIds: string[] }> => {
insert: VectorControllerType['insert'] = async (props) => {
const { teamId, datasetId, collectionId, vectors } = props;
const values = vectors.map((vector) => [
......@@ -87,7 +81,7 @@ export class PgVectorCtrl {
insertIds: rows.map((row) => row.id)
};
};
delete = async (props: DelDatasetVectorCtrlProps): Promise<any> => {
delete: VectorControllerType['delete'] = async (props) => {
const { teamId } = props;
const teamIdWhere = `team_id='${String(teamId)}' AND`;
......@@ -122,7 +116,7 @@ export class PgVectorCtrl {
where: [where]
});
};
embRecall = async (props: EmbeddingRecallCtrlProps): Promise<EmbeddingRecallResponse> => {
embRecall: VectorControllerType['embRecall'] = async (props) => {
const { teamId, datasetIds, vector, limit, forbidCollectionIdList, filterCollectionIdList } =
props;
......@@ -186,7 +180,8 @@ export class PgVectorCtrl {
}))
};
};
getVectorDataByTime = async (start: Date, end: Date) => {
getVectorDataByTime: VectorControllerType['getVectorDataByTime'] = async (start, end) => {
const { rows } = await PgClient.query<{
id: string;
team_id: string;
......@@ -204,12 +199,7 @@ export class PgVectorCtrl {
datasetId: item.dataset_id
}));
};
getVectorCount = async (props: {
teamId?: string;
datasetId?: string;
collectionId?: string;
}) => {
getVectorCount: VectorControllerType['getVectorCount'] = async (props) => {
const { teamId, datasetId, collectionId } = props;
// Build where conditions dynamically
......
packages/service/common/vectorDB/type.d.ts deleted 100644 → 0
import type { Pool } from 'pg';
import type { Pool as MysqlPool } from 'mysql2/promise';
import type { 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;
};
packages/service/common/vectorDB/type.ts 0 → 100644
import type { Pool } from 'pg';
import type { Pool as MysqlPool } from 'mysql2/promise';
import type { MilvusClient } from '@zilliz/milvus2-sdk-node';
import { z } from 'zod';
// Embedding recall item schema
export const EmbeddingRecallItemSchema = z.object({
id: z.string(),
collectionId: z.string(),
score: z.number()
});
export type EmbeddingRecallItemType = z.infer<typeof EmbeddingRecallItemSchema>;
// Insert vector props schema
export const InsertVectorControllerPropsSchema = z.object({
teamId: z.string(),
datasetId: z.string(),
collectionId: z.string(),
vectors: z.array(z.array(z.number()))
});
export type InsertVectorControllerPropsType = z.infer<typeof InsertVectorControllerPropsSchema>;
// Insert vector response schema
export const InsertVectorResponseSchema = z.object({
insertIds: z.array(z.string())
});
export type InsertVectorResponseType = z.infer<typeof InsertVectorResponseSchema>;
// Delete vector props schema (union type for different delete scenarios)
export const DelDatasetVectorCtrlPropsSchema = z.union([
z.object({
teamId: z.string(),
id: z.string(),
retry: z.number().optional()
}),
z.object({
teamId: z.string(),
datasetIds: z.array(z.string()),
collectionIds: z.array(z.string()).optional(),
retry: z.number().optional()
}),
z.object({
teamId: z.string(),
idList: z.array(z.string()),
retry: z.number().optional()
})
]);
export type DelDatasetVectorCtrlPropsType = z.infer<typeof DelDatasetVectorCtrlPropsSchema>;
// Embedding recall props schema
export const EmbeddingRecallCtrlPropsSchema = z.object({
teamId: z.string(),
datasetIds: z.array(z.string()),
vector: z.array(z.number()),
limit: z.number(),
forbidCollectionIdList: z.array(z.string()),
filterCollectionIdList: z.array(z.string()).optional(),
retry: z.number().optional()
});
export type EmbeddingRecallCtrlPropsType = z.infer<typeof EmbeddingRecallCtrlPropsSchema>;
// Embedding recall response schema
export const EmbeddingRecallResponseSchema = z.object({
results: z.array(EmbeddingRecallItemSchema)
});
export type EmbeddingRecallResponseType = z.infer<typeof EmbeddingRecallResponseSchema>;
// Get vector data by time response schema
export const GetVectorDataByTimeResponseSchema = z.array(
z.object({
id: z.string(),
teamId: z.string(),
datasetId: z.string()
})
);
export type GetVectorDataByTimeResponseType = z.infer<typeof GetVectorDataByTimeResponseSchema>;
// Get vector count props schema
export const GetVectorCountPropsSchema = z.object({
teamId: z.string().optional(),
datasetId: z.string().optional(),
collectionId: z.string().optional()
});
export type GetVectorCountPropsType = z.infer<typeof GetVectorCountPropsSchema>;
// ==================== Vector Controller Interface ====================
export interface VectorControllerType {
/**
* Initialize vector database (create tables, indexes, etc.)
*/
init(): Promise<void>;
/**
* Insert vectors into the database
*/
insert(props: InsertVectorControllerPropsType): Promise<InsertVectorResponseType>;
/**
* Delete vectors from the database
*/
delete(props: DelDatasetVectorCtrlPropsType): Promise<void>;
/**
* Embedding recall/search vectors
*/
embRecall(props: EmbeddingRecallCtrlPropsType): Promise<EmbeddingRecallResponseType>;
/**
* Get vector data by time range
*/
getVectorDataByTime(start: Date, end: Date): Promise<GetVectorDataByTimeResponseType>;
/**
* Get vector count by filters
*/
getVectorCount(props: GetVectorCountPropsType): Promise<number>;
}
declare global {
var pgClient: Pool | null;
var obClient: MysqlPool | null;
var milvusClient: MilvusClient | null;
}
pnpm-lock.yaml
......@@ -727,7 +727,7 @@ importers:
specifier: ^5.1.3
version: 5.8.2
vitest:
specifier: ^3.0.2
specifier: ^3.0.9
version: 3.1.1(@types/debug@4.1.12)(@types/node@20.17.24)(jiti@2.6.0)(lightningcss@1.30.1)(sass@1.85.1)(terser@5.39.0)(tsx@4.20.6)(yaml@2.8.1)
projects/marketplace:
......@@ -15677,14 +15677,6 @@ snapshots:
chai: 5.2.0
tinyrainbow: 2.0.0
'@vitest/mocker@3.1.1(vite@6.2.2(@types/node@20.17.24)(jiti@2.6.0)(lightningcss@1.30.1)(sass@1.85.1)(terser@5.39.0)(tsx@4.20.6)(yaml@2.8.1))':
dependencies:
'@vitest/spy': 3.1.1
estree-walker: 3.0.3
magic-string: 0.30.17
optionalDependencies:
vite: 6.2.2(@types/node@20.17.24)(jiti@2.6.0)(lightningcss@1.30.1)(sass@1.85.1)(terser@5.39.0)(tsx@4.20.6)(yaml@2.8.1)
'@vitest/mocker@3.1.1(vite@6.2.2(@types/node@24.0.13)(jiti@2.6.0)(lightningcss@1.30.1)(sass@1.85.1)(terser@5.39.0)(tsx@4.20.6)(yaml@2.8.1))':
dependencies:
'@vitest/spy': 3.1.1
......@@ -24115,7 +24107,7 @@ snapshots:
vitest@3.1.1(@types/debug@4.1.12)(@types/node@20.17.24)(jiti@2.6.0)(lightningcss@1.30.1)(sass@1.85.1)(terser@5.39.0)(tsx@4.20.6)(yaml@2.8.1):
dependencies:
'@vitest/expect': 3.1.1
'@vitest/mocker': 3.1.1(vite@6.2.2(@types/node@20.17.24)(jiti@2.6.0)(lightningcss@1.30.1)(sass@1.85.1)(terser@5.39.0)(tsx@4.20.6)(yaml@2.8.1))
'@vitest/mocker': 3.1.1(vite@6.2.2(@types/node@24.0.13)(jiti@2.6.0)(lightningcss@1.30.1)(sass@1.85.1)(terser@5.39.0)(tsx@4.20.6)(yaml@2.8.1))
'@vitest/pretty-format': 3.1.1
'@vitest/runner': 3.1.1
'@vitest/snapshot': 3.1.1
......
projects/app/package.json
......@@ -92,6 +92,6 @@
"eslint-config-next": "14.2.26",
"tsx": "^4.20.6",
"typescript": "^5.1.3",
"vitest": "^3.0.2"
"vitest": "^3.0.9"
}
}
test/mocks/common/s3.ts
......@@ -174,3 +174,9 @@ vi.mock('@fastgpt/service/common/s3', () => ({
}),
initS3MQWorker: vi.fn().mockResolvedValue(undefined)
}));
// Mock S3 MQ (Message Queue) operations
vi.mock('@fastgpt/service/common/s3/mq', () => ({
prefixDel: vi.fn().mockResolvedValue(undefined),
addDeleteJob: vi.fn().mockResolvedValue(undefined)
}));
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment