pg table name

This commit is contained in:
archer
2023-08-14 10:36:15 +08:00
parent c5fd5706a1
commit 3a49efd46d
12 changed files with 37 additions and 22 deletions

View File

@@ -5,6 +5,7 @@ import { getVector } from '@/pages/api/openapi/plugin/vector';
import { countModelPrice } from '@/service/events/pushBill';
import type { SelectedKbType } from '@/types/plugin';
import type { QuoteItemType } from '@/types/chat';
import { PgTrainingTableName } from '@/constants/plugin';
type KBSearchProps = {
kbList: SelectedKbType;
@@ -48,7 +49,7 @@ export async function dispatchKBSearch(props: Record<string, any>): Promise<KBSe
const res: any = await PgClient.query(
`BEGIN;
SET LOCAL ivfflat.probes = ${global.systemEnv.pgIvfflatProbe || 10};
select kb_id,id,q,a,source from modelData where kb_id IN (${kbList
select kb_id,id,q,a,source from ${PgTrainingTableName} where kb_id IN (${kbList
.map((item) => `'${item.kbId}'`)
.join(',')}) AND vector <#> '[${vectors[0]}]' < -${similarity} order by vector <#> '[${
vectors[0]

View File

@@ -6,6 +6,7 @@ import { User } from './models/user';
import { PRICE_SCALE } from '@/constants/common';
import { connectPg, PgClient } from './pg';
import { createHashPassword } from '@/utils/tools';
import { PgTrainingTableName } from '@/constants/plugin';
/**
* connect MongoDB and init data
@@ -92,7 +93,7 @@ async function initPg() {
await connectPg();
await PgClient.query(`
CREATE EXTENSION IF NOT EXISTS vector;
CREATE TABLE IF NOT EXISTS modelData (
CREATE TABLE IF NOT EXISTS ${PgTrainingTableName} (
id BIGSERIAL PRIMARY KEY,
vector VECTOR(1536) NOT NULL,
user_id VARCHAR(50) NOT NULL,
@@ -101,9 +102,9 @@ async function initPg() {
q TEXT NOT NULL,
a TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS modelData_userId_index ON modelData USING HASH (user_id);
CREATE INDEX IF NOT EXISTS modelData_kbId_index ON modelData USING HASH (kb_id);
CREATE INDEX IF NOT EXISTS idx_model_data_md5_q_a_user_id_kb_id ON modelData (md5(q), md5(a), user_id, kb_id);
CREATE INDEX IF NOT EXISTS modelData_userId_index ON ${PgTrainingTableName} USING HASH (user_id);
CREATE INDEX IF NOT EXISTS modelData_kbId_index ON ${PgTrainingTableName} USING HASH (kb_id);
CREATE INDEX IF NOT EXISTS idx_model_data_md5_q_a_user_id_kb_id ON ${PgTrainingTableName} (md5(q), md5(a), user_id, kb_id);
`);
console.log('init pg successful');
} catch (error) {

View File

@@ -1,5 +1,6 @@
import { Pool } from 'pg';
import type { QueryResultRow } from 'pg';
import { PgTrainingTableName } from '@/constants/plugin';
export const connectPg = async () => {
if (global.pgClient) {
@@ -173,7 +174,7 @@ export const insertKbItem = ({
source?: string;
}[];
}) => {
return PgClient.insert('modelData', {
return PgClient.insert(PgTrainingTableName, {
values: data.map((item) => [
{ key: 'user_id', value: userId },
{ key: 'kb_id', value: kbId },