mirror of
https://github.com/labring/FastGPT.git
synced 2025-07-24 22:03:54 +00:00
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:
7
packages/service/common/vectorDB/constants.ts
Normal file
7
packages/service/common/vectorDB/constants.ts
Normal 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;
|
38
packages/service/common/vectorDB/controller.d.ts
vendored
Normal file
38
packages/service/common/vectorDB/controller.d.ts
vendored
Normal 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[];
|
||||
};
|
83
packages/service/common/vectorDB/controller.ts
Normal file
83
packages/service/common/vectorDB/controller.ts
Normal 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;
|
||||
};
|
365
packages/service/common/vectorDB/milvus/index.ts
Normal file
365
packages/service/common/vectorDB/milvus/index.ts
Normal 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
|
||||
}));
|
||||
};
|
||||
}
|
173
packages/service/common/vectorDB/oceanbase/controller.ts
Normal file
173
packages/service/common/vectorDB/oceanbase/controller.ts
Normal 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;
|
254
packages/service/common/vectorDB/oceanbase/index.ts
Normal file
254
packages/service/common/vectorDB/oceanbase/index.ts
Normal 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;
|
||||
};
|
||||
}
|
188
packages/service/common/vectorDB/pg/controller.ts
Normal file
188
packages/service/common/vectorDB/pg/controller.ts
Normal 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;
|
276
packages/service/common/vectorDB/pg/index.ts
Normal file
276
packages/service/common/vectorDB/pg/index.ts
Normal 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;
|
||||
};
|
||||
}
|
15
packages/service/common/vectorDB/type.d.ts
vendored
Normal file
15
packages/service/common/vectorDB/type.d.ts
vendored
Normal 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;
|
||||
};
|
Reference in New Issue
Block a user