mirror of
https://github.com/labring/FastGPT.git
synced 2025-10-19 01:54:04 +00:00
V4.6.6-2 (#673)
This commit is contained in:
5
packages/service/common/pg/type.d.ts
vendored
5
packages/service/common/pg/type.d.ts
vendored
@@ -1,5 +0,0 @@
|
||||
import type { Pool } from 'pg';
|
||||
|
||||
declare global {
|
||||
var pgClient: Pool | null;
|
||||
}
|
@@ -4,11 +4,13 @@ import dayjs from 'dayjs';
|
||||
export const addLog = {
|
||||
log(level: 'info' | 'warn' | 'error', msg: string, obj: Record<string, any> = {}) {
|
||||
console.log(
|
||||
`[${level.toLocaleUpperCase()}] ${dayjs().format(
|
||||
'YYYY-MM-DD HH:mm:ss'
|
||||
)} ${msg}: ${JSON.stringify(obj)}`
|
||||
`[${level.toLocaleUpperCase()}] ${dayjs().format('YYYY-MM-DD HH:mm:ss')} ${msg} ${
|
||||
level !== 'error' ? JSON.stringify(obj) : ''
|
||||
}`
|
||||
);
|
||||
|
||||
level === 'error' && console.error(obj);
|
||||
|
||||
const lokiUrl = process.env.LOKI_LOG_URL as string;
|
||||
if (!lokiUrl) return;
|
||||
|
||||
|
19
packages/service/common/vectorStore/controller.d.ts
vendored
Normal file
19
packages/service/common/vectorStore/controller.d.ts
vendored
Normal file
@@ -0,0 +1,19 @@
|
||||
export type DeleteDatasetVectorProps = {
|
||||
id?: string;
|
||||
datasetIds?: string[];
|
||||
collectionIds?: string[];
|
||||
dataIds?: string[];
|
||||
};
|
||||
|
||||
export type InsertVectorProps = {
|
||||
teamId: string;
|
||||
tmbId: string;
|
||||
datasetId: string;
|
||||
collectionId: string;
|
||||
dataId: string;
|
||||
};
|
||||
|
||||
export type EmbeddingRecallProps = {
|
||||
similarity?: number;
|
||||
datasetIds: string[];
|
||||
};
|
62
packages/service/common/vectorStore/controller.ts
Normal file
62
packages/service/common/vectorStore/controller.ts
Normal file
@@ -0,0 +1,62 @@
|
||||
/* vector crud */
|
||||
import { PgVector } from './pg/class';
|
||||
import { getVectorsByText } from '../../core/ai/embedding';
|
||||
import { InsertVectorProps } from './controller.d';
|
||||
|
||||
const getVectorObj = () => {
|
||||
return new PgVector();
|
||||
};
|
||||
|
||||
export const initVectorStore = getVectorObj().init;
|
||||
export const deleteDatasetDataVector = getVectorObj().delete;
|
||||
export const recallFromVectorStore = getVectorObj().recall;
|
||||
export const getVectorDataByTime = getVectorObj().getVectorDataByTime;
|
||||
export const getVectorCountByTeamId = getVectorObj().getVectorCountByTeamId;
|
||||
|
||||
export const insertDatasetDataVector = async ({
|
||||
model,
|
||||
query,
|
||||
...props
|
||||
}: InsertVectorProps & {
|
||||
query: string;
|
||||
model: string;
|
||||
}) => {
|
||||
const { vectors, tokens } = await getVectorsByText({
|
||||
model,
|
||||
input: query
|
||||
});
|
||||
const { insertId } = await getVectorObj().insert({
|
||||
...props,
|
||||
vectors
|
||||
});
|
||||
|
||||
return {
|
||||
tokens,
|
||||
insertId
|
||||
};
|
||||
};
|
||||
|
||||
export const updateDatasetDataVector = async ({
|
||||
id,
|
||||
query,
|
||||
model
|
||||
}: {
|
||||
id: string;
|
||||
query: string;
|
||||
model: string;
|
||||
}) => {
|
||||
// get vector
|
||||
const { vectors, tokens } = await getVectorsByText({
|
||||
model,
|
||||
input: [query]
|
||||
});
|
||||
|
||||
await getVectorObj().update({
|
||||
id,
|
||||
vectors
|
||||
});
|
||||
|
||||
return {
|
||||
tokens
|
||||
};
|
||||
};
|
20
packages/service/common/vectorStore/pg/class.ts
Normal file
20
packages/service/common/vectorStore/pg/class.ts
Normal file
@@ -0,0 +1,20 @@
|
||||
import {
|
||||
initPg,
|
||||
insertDatasetDataVector,
|
||||
updateDatasetDataVector,
|
||||
deleteDatasetDataVector,
|
||||
embeddingRecall,
|
||||
getVectorDataByTime,
|
||||
getVectorCountByTeamId
|
||||
} from './controller';
|
||||
|
||||
export class PgVector {
|
||||
constructor() {}
|
||||
init = initPg;
|
||||
insert = insertDatasetDataVector;
|
||||
update = updateDatasetDataVector;
|
||||
delete = deleteDatasetDataVector;
|
||||
recall = embeddingRecall;
|
||||
getVectorCountByTeamId = getVectorCountByTeamId;
|
||||
getVectorDataByTime = getVectorDataByTime;
|
||||
}
|
199
packages/service/common/vectorStore/pg/controller.ts
Normal file
199
packages/service/common/vectorStore/pg/controller.ts
Normal file
@@ -0,0 +1,199 @@
|
||||
/* pg vector crud */
|
||||
import { PgDatasetTableName } from '@fastgpt/global/common/vectorStore/constants';
|
||||
import { delay } from '@fastgpt/global/common/system/utils';
|
||||
import { PgClient, connectPg } from './index';
|
||||
import { PgSearchRawType } from '@fastgpt/global/core/dataset/api';
|
||||
import { EmbeddingRecallItemType } from '../type';
|
||||
import { DeleteDatasetVectorProps, EmbeddingRecallProps } from '../controller.d';
|
||||
import dayjs from 'dayjs';
|
||||
|
||||
export async function initPg() {
|
||||
try {
|
||||
await connectPg();
|
||||
await PgClient.query(`
|
||||
CREATE EXTENSION IF NOT EXISTS vector;
|
||||
CREATE TABLE IF NOT EXISTS ${PgDatasetTableName} (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
vector VECTOR(1536) NOT NULL,
|
||||
team_id VARCHAR(50) NOT NULL,
|
||||
tmb_id VARCHAR(50) NOT NULL,
|
||||
dataset_id VARCHAR(50) NOT NULL,
|
||||
collection_id VARCHAR(50) NOT NULL,
|
||||
data_id VARCHAR(50) NOT NULL,
|
||||
createTime TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS vector_index ON ${PgDatasetTableName} USING hnsw (vector vector_ip_ops) WITH (m = 32, ef_construction = 64);
|
||||
`);
|
||||
|
||||
console.log('init pg successful');
|
||||
} catch (error) {
|
||||
console.log('init pg error', error);
|
||||
}
|
||||
}
|
||||
|
||||
export const insertDatasetDataVector = async (props: {
|
||||
teamId: string;
|
||||
tmbId: string;
|
||||
datasetId: string;
|
||||
collectionId: string;
|
||||
dataId: string;
|
||||
vectors: number[][];
|
||||
retry?: number;
|
||||
}): Promise<{ insertId: string }> => {
|
||||
const { dataId, teamId, tmbId, datasetId, collectionId, vectors, retry = 3 } = props;
|
||||
try {
|
||||
const { rows } = await PgClient.insert(PgDatasetTableName, {
|
||||
values: [
|
||||
[
|
||||
{ key: 'vector', value: `[${vectors[0]}]` },
|
||||
{ key: 'team_id', value: String(teamId) },
|
||||
{ key: 'tmb_id', value: String(tmbId) },
|
||||
{ key: 'dataset_id', value: datasetId },
|
||||
{ key: 'collection_id', value: collectionId },
|
||||
{ key: 'data_id', value: String(dataId) }
|
||||
]
|
||||
]
|
||||
});
|
||||
return {
|
||||
insertId: rows[0].id
|
||||
};
|
||||
} catch (error) {
|
||||
if (retry <= 0) {
|
||||
return Promise.reject(error);
|
||||
}
|
||||
await delay(500);
|
||||
return insertDatasetDataVector({
|
||||
...props,
|
||||
retry: retry - 1
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
export const updateDatasetDataVector = async (props: {
|
||||
id: string;
|
||||
vectors: number[][];
|
||||
retry?: number;
|
||||
}): Promise<void> => {
|
||||
const { id, vectors, retry = 2 } = props;
|
||||
try {
|
||||
// update pg
|
||||
await PgClient.update(PgDatasetTableName, {
|
||||
where: [['id', id]],
|
||||
values: [{ key: 'vector', value: `[${vectors[0]}]` }]
|
||||
});
|
||||
} catch (error) {
|
||||
if (retry <= 0) {
|
||||
return Promise.reject(error);
|
||||
}
|
||||
await delay(500);
|
||||
return updateDatasetDataVector({
|
||||
...props,
|
||||
retry: retry - 1
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
export const deleteDatasetDataVector = async (
|
||||
props: DeleteDatasetVectorProps & {
|
||||
retry?: number;
|
||||
}
|
||||
): Promise<any> => {
|
||||
const { id, datasetIds, collectionIds, dataIds, retry = 2 } = props;
|
||||
|
||||
const where = await (() => {
|
||||
if (id) return `id=${id}`;
|
||||
if (datasetIds) return `dataset_id IN (${datasetIds.map((id) => `'${String(id)}'`).join(',')})`;
|
||||
if (collectionIds)
|
||||
return `collection_id IN (${collectionIds.map((id) => `'${String(id)}'`).join(',')})`;
|
||||
if (dataIds) return `data_id IN (${dataIds.map((id) => `'${String(id)}'`).join(',')})`;
|
||||
return Promise.reject('deleteDatasetData: no where');
|
||||
})();
|
||||
|
||||
try {
|
||||
await PgClient.delete(PgDatasetTableName, {
|
||||
where: [where]
|
||||
});
|
||||
} catch (error) {
|
||||
if (retry <= 0) {
|
||||
return Promise.reject(error);
|
||||
}
|
||||
await delay(500);
|
||||
return deleteDatasetDataVector({
|
||||
...props,
|
||||
retry: retry - 1
|
||||
});
|
||||
}
|
||||
};
|
||||
|
||||
export const embeddingRecall = async (
|
||||
props: EmbeddingRecallProps & {
|
||||
vectors: number[][];
|
||||
limit: number;
|
||||
retry?: number;
|
||||
}
|
||||
): Promise<{
|
||||
results: EmbeddingRecallItemType[];
|
||||
}> => {
|
||||
const { vectors, limit, similarity = 0, datasetIds, retry = 2 } = props;
|
||||
|
||||
try {
|
||||
const results: any = await PgClient.query(
|
||||
`BEGIN;
|
||||
SET LOCAL hnsw.ef_search = ${global.systemEnv.pgHNSWEfSearch || 100};
|
||||
select id, collection_id, data_id, (vector <#> '[${vectors[0]}]') * -1 AS score
|
||||
from ${PgDatasetTableName}
|
||||
where dataset_id IN (${datasetIds.map((id) => `'${String(id)}'`).join(',')})
|
||||
AND vector <#> '[${vectors[0]}]' < -${similarity}
|
||||
order by score desc limit ${limit};
|
||||
COMMIT;`
|
||||
);
|
||||
|
||||
const rows = results?.[2]?.rows as PgSearchRawType[];
|
||||
|
||||
// concat same data_id
|
||||
const filterRows: PgSearchRawType[] = [];
|
||||
let set = new Set<string>();
|
||||
for (const row of rows) {
|
||||
if (!set.has(row.data_id)) {
|
||||
filterRows.push(row);
|
||||
set.add(row.data_id);
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
results: filterRows.map((item) => ({
|
||||
id: item.id,
|
||||
collectionId: item.collection_id,
|
||||
dataId: item.data_id,
|
||||
score: item.score
|
||||
}))
|
||||
};
|
||||
} catch (error) {
|
||||
if (retry <= 0) {
|
||||
return Promise.reject(error);
|
||||
}
|
||||
return embeddingRecall(props);
|
||||
}
|
||||
};
|
||||
|
||||
// bill
|
||||
export const getVectorCountByTeamId = async (teamId: string) => {
|
||||
const total = await PgClient.count(PgDatasetTableName, {
|
||||
where: [['team_id', String(teamId)]]
|
||||
});
|
||||
|
||||
return total;
|
||||
};
|
||||
export const getVectorDataByTime = async (start: Date, end: Date) => {
|
||||
const { rows } = await PgClient.query<{ id: string; data_id: string }>(`SELECT id, data_id
|
||||
FROM ${PgDatasetTableName}
|
||||
WHERE createTime BETWEEN '${dayjs(start).format('YYYY-MM-DD')}' AND '${dayjs(end).format(
|
||||
'YYYY-MM-DD 23:59:59'
|
||||
)}';
|
||||
`);
|
||||
|
||||
return rows.map((item) => ({
|
||||
id: item.id,
|
||||
dataId: item.data_id
|
||||
}));
|
||||
};
|
@@ -1,6 +1,5 @@
|
||||
import { Pool } from 'pg';
|
||||
import type { QueryResultRow } from 'pg';
|
||||
import { PgDatasetTableName } from '@fastgpt/global/core/dataset/constant';
|
||||
|
||||
export const connectPg = async (): Promise<Pool> => {
|
||||
if (global.pgClient) {
|
||||
@@ -117,6 +116,7 @@ class PgClass {
|
||||
FROM ${table}
|
||||
${this.getWhereStr(props.where)}
|
||||
`;
|
||||
|
||||
const pg = await connectPg();
|
||||
return pg.query(sql).then((res) => Number(res.rows[0]?.count || 0));
|
||||
}
|
||||
@@ -160,29 +160,5 @@ class PgClass {
|
||||
}
|
||||
}
|
||||
|
||||
export async function initPg() {
|
||||
try {
|
||||
await connectPg();
|
||||
await PgClient.query(`
|
||||
CREATE EXTENSION IF NOT EXISTS vector;
|
||||
CREATE TABLE IF NOT EXISTS ${PgDatasetTableName} (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
vector VECTOR(1536) NOT NULL,
|
||||
team_id VARCHAR(50) NOT NULL,
|
||||
tmb_id VARCHAR(50) NOT NULL,
|
||||
dataset_id VARCHAR(50) NOT NULL,
|
||||
collection_id VARCHAR(50) NOT NULL,
|
||||
data_id VARCHAR(50) NOT NULL,
|
||||
createTime TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS vector_index ON ${PgDatasetTableName} USING hnsw (vector vector_ip_ops) WITH (m = 32, ef_construction = 64);
|
||||
`);
|
||||
|
||||
console.log('init pg successful');
|
||||
} catch (error) {
|
||||
console.log('init pg error', error);
|
||||
}
|
||||
}
|
||||
|
||||
export const PgClient = new PgClass();
|
||||
export const Pg = global.pgClient;
|
12
packages/service/common/vectorStore/type.d.ts
vendored
Normal file
12
packages/service/common/vectorStore/type.d.ts
vendored
Normal file
@@ -0,0 +1,12 @@
|
||||
import type { Pool } from 'pg';
|
||||
|
||||
declare global {
|
||||
var pgClient: Pool | null;
|
||||
}
|
||||
|
||||
export type EmbeddingRecallItemType = {
|
||||
id: string;
|
||||
collectionId: string;
|
||||
dataId: string;
|
||||
score: number;
|
||||
};
|
Reference in New Issue
Block a user