This commit is contained in:
Archer
2023-12-31 14:12:51 +08:00
committed by GitHub
parent ccca0468da
commit 9ccfda47b7
270 changed files with 8182 additions and 1295 deletions

View File

@@ -1,5 +0,0 @@
import type { Pool } from 'pg';
declare global {
var pgClient: Pool | null;
}

View File

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

View 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[];
};

View 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
};
};

View 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;
}

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

View File

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

View 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;
};