v4.6 -1 (#459)

packages/service/common/api/plusRequest.ts

@@ -18,7 +18,6 @@ function requestStart(config: InternalAxiosRequestConfig): InternalAxiosRequestC
   if (config.headers) {
     config.headers.rootkey = process.env.ROOT_KEY;
   }
-
   return config;
 }
 
@@ -62,7 +61,8 @@ function responseError(err: any) {
 const instance = axios.create({
   timeout: 60000, // 超时时间
   headers: {
-    'content-type': 'application/json'
+    'content-type': 'application/json',
+    'Cache-Control': 'no-cache'
   }
 });
 
@@ -73,7 +73,7 @@ instance.interceptors.response.use(responseSuccess, (err) => Promise.reject(err)
 
 export function request(url: string, data: any, config: ConfigType, method: Method): any {
   if (!global.systemEnv?.pluginBaseUrl) {
-    return Promise.reject('商业版插件加载中...');
+    return Promise.reject('该功能为商业版特有...');
   }
 
   /* 去空 */

packages/service/common/file/gridfs/controller.ts

@@ -0,0 +1,111 @@
+import { Types, connectionMongo } from '../../mongo';
+import { BucketNameEnum } from '@fastgpt/global/common/file/constants';
+import fsp from 'fs/promises';
+import fs from 'fs';
+import { DatasetFileSchema } from '@fastgpt/global/core/dataset/type';
+
+export function getGFSCollection(bucket: `${BucketNameEnum}`) {
+  return connectionMongo.connection.db.collection(`${bucket}.files`);
+}
+export function getGridBucket(bucket: `${BucketNameEnum}`) {
+  return new connectionMongo.mongo.GridFSBucket(connectionMongo.connection.db, {
+    bucketName: bucket
+  });
+}
+
+/* crud  file */
+export async function uploadFile({
+  bucketName,
+  teamId,
+  tmbId,
+  path,
+  filename,
+  metadata = {}
+}: {
+  bucketName: `${BucketNameEnum}`;
+  teamId: string;
+  tmbId: string;
+  path: string;
+  filename: string;
+  metadata?: Record<string, any>;
+}) {
+  if (!path) return Promise.reject(`filePath is empty`);
+  if (!filename) return Promise.reject(`filename is empty`);
+
+  const stats = await fsp.stat(path);
+  if (!stats.isFile()) return Promise.reject(`${path} is not a file`);
+
+  metadata.teamId = teamId;
+  metadata.tmbId = tmbId;
+
+  // create a gridfs bucket
+  const bucket = getGridBucket(bucketName);
+
+  const stream = bucket.openUploadStream(filename, {
+    metadata,
+    contentType: metadata?.contentType
+  });
+
+  // save to gridfs
+  await new Promise((resolve, reject) => {
+    fs.createReadStream(path)
+      .pipe(stream as any)
+      .on('finish', resolve)
+      .on('error', reject);
+  });
+
+  return String(stream.id);
+}
+
+export async function getFileById({
+  bucketName,
+  fileId
+}: {
+  bucketName: `${BucketNameEnum}`;
+  fileId: string;
+}) {
+  const db = getGFSCollection(bucketName);
+  const file = await db.findOne<DatasetFileSchema>({
+    _id: new Types.ObjectId(fileId)
+  });
+
+  if (!file) {
+    return Promise.reject('File not found');
+  }
+
+  return file;
+}
+
+export async function delFileById({
+  bucketName,
+  fileId
+}: {
+  bucketName: `${BucketNameEnum}`;
+  fileId: string;
+}) {
+  const bucket = getGridBucket(bucketName);
+
+  await bucket.delete(new Types.ObjectId(fileId));
+  return true;
+}
+
+export async function getDownloadBuf({
+  bucketName,
+  fileId
+}: {
+  bucketName: `${BucketNameEnum}`;
+  fileId: string;
+}) {
+  const bucket = getGridBucket(bucketName);
+
+  const stream = bucket.openDownloadStream(new Types.ObjectId(fileId));
+
+  const buf: Buffer = await new Promise((resolve, reject) => {
+    const buffers: Buffer[] = [];
+    stream.on('data', (data) => buffers.push(data));
+    stream.on('error', reject);
+    stream.on('end', () => resolve(Buffer.concat(buffers)));
+  });
+
+  return buf;
+}

packages/service/common/mongo/controller.ts

@@ -0,0 +1,26 @@
+/* add logger */
+export const addLog = {
+  info: (msg: string, obj?: Record<string, any>) => {
+    global.logger?.info(msg, { meta: obj });
+  },
+  error: (msg: string, error?: any) => {
+    global.logger?.error(msg, {
+      meta: {
+        stack: error?.stack,
+        ...(error?.config && {
+          config: {
+            headers: error.config.headers,
+            url: error.config.url,
+            data: error.config.data
+          }
+        }),
+        ...(error?.response && {
+          response: {
+            status: error.response.status,
+            statusText: error.response.statusText
+          }
+        })
+      }
+    });
+  }
+};

packages/service/common/mongo/init.ts

@@ -1,6 +1,6 @@
 import mongoose from './index';
-import 'winston-mongodb';
 import { createLogger, format, transports } from 'winston';
+import 'winston-mongodb';
 
 /**
  * connect MongoDB and init data
@@ -19,9 +19,6 @@ export async function connectMongo({
 
   beforeHook && (await beforeHook());
 
-  // logger
-  initLogger();
-
   console.log('mongo start connect');
   try {
     mongoose.set('strictQuery', true);
@@ -35,9 +32,11 @@ export async function connectMongo({
     });
 
     console.log('mongo connected');
+    initLogger();
 
     afterHook && (await afterHook());
   } catch (error) {
+    global.mongodb.disconnect();
     console.log('error->', 'mongo connect error', error);
     global.mongodb = undefined;
   }

packages/service/common/mongo/sessionRun.ts

@@ -1,39 +1,21 @@
-import mongoose from './index';
+import mongoose, { connectionMongo } from './index';
 
-export class MongoSession {
-  tasks: (() => Promise<any>)[] = [];
-  session: mongoose.mongo.ClientSession | null = null;
-  opts: {
-    session: mongoose.mongo.ClientSession;
-    new: boolean;
-  } | null = null;
+export async function mongoSessionTask(
+  fn: (session: mongoose.mongo.ClientSession) => Promise<any>
+) {
+  const session = await connectionMongo.startSession();
 
-  constructor() {}
-  async init() {
-    this.session = await mongoose.startSession();
-    this.opts = { session: this.session, new: true };
-  }
-  push(
-    tasks: ((opts: {
-      session: mongoose.mongo.ClientSession;
-      new: boolean;
-    }) => () => Promise<any>)[] = []
-  ) {
-    if (!this.opts) return;
-    // this.tasks = this.tasks.concat(tasks.map((item) => item(this.opts)));
-  }
-  async run() {
-    if (!this.session || !this.opts) return;
-    try {
-      this.session.startTransaction();
+  try {
+    session.startTransaction();
 
-      const opts = { session: this.session, new: true };
+    await fn(session);
 
-      await this.session.commitTransaction();
-    } catch (error) {
-      await this.session.abortTransaction();
-      console.error(error);
-    }
-    this.session.endSession();
+    await session.commitTransaction();
+    await session.endSession();
+  } catch (error) {
+    await session.abortTransaction();
+    await session.endSession();
+    console.error(error);
+    return Promise.reject(error);
   }
 }

packages/service/common/pg/index.ts

@@ -0,0 +1,187 @@
+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) {
+    return global.pgClient;
+  }
+
+  global.pgClient = new Pool({
+    connectionString: process.env.PG_URL,
+    max: Number(process.env.DB_MAX_LINK || 5),
+    keepAlive: true,
+    idleTimeoutMillis: 60000,
+    connectionTimeoutMillis: 20000
+  });
+
+  global.pgClient.on('error', (err) => {
+    console.log(err);
+    global.pgClient?.end();
+    global.pgClient = null;
+    connectPg();
+  });
+
+  try {
+    await global.pgClient.connect();
+    console.log('pg connected');
+    return global.pgClient;
+  } catch (error) {
+    global.pgClient = null;
+    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();
+    return pg.query<T>(sql);
+  }
+}
+
+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,
+          q TEXT NOT NULL,
+          a TEXT
+      );
+      CREATE INDEX IF NOT EXISTS vector_index ON ${PgDatasetTableName} USING hnsw (vector vector_ip_ops) WITH (m = 24, 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;

packages/service/common/pg/type.d.ts

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

packages/service/common/response/constant.ts

@@ -0,0 +1,6 @@
+export enum sseResponseEventEnum {
+  error = 'error',
+  answer = 'answer',
+  moduleStatus = 'moduleStatus',
+  appStreamResponse = 'appStreamResponse' // sse response request
+}

packages/service/common/response/index.ts

@@ -1,4 +1,98 @@
 import type { NextApiResponse } from 'next';
+import { sseResponseEventEnum } from './constant';
+import { proxyError, ERROR_RESPONSE, ERROR_ENUM } from '@fastgpt/global/common/error/errorCode';
+import { addLog } from '../mongo/controller';
+import { clearCookie } from '../../support/permission/controller';
+
+export interface ResponseType<T = any> {
+  code: number;
+  message: string;
+  data: T;
+}
+
+export const jsonRes = <T = any>(
+  res: NextApiResponse,
+  props?: {
+    code?: number;
+    message?: string;
+    data?: T;
+    error?: any;
+  }
+) => {
+  const { code = 200, message = '', data = null, error } = props || {};
+
+  const errResponseKey = typeof error === 'string' ? error : error?.message;
+  // Specified error
+  if (ERROR_RESPONSE[errResponseKey]) {
+    // login is expired
+    if (errResponseKey === ERROR_ENUM.unAuthorization) {
+      clearCookie(res);
+    }
+
+    return res.json(ERROR_RESPONSE[errResponseKey]);
+  }
+
+  // another error
+  let msg = '';
+  if ((code < 200 || code >= 400) && !message) {
+    msg = error?.response?.statusText || error?.message || '请求错误';
+    if (typeof error === 'string') {
+      msg = error;
+    } else if (proxyError[error?.code]) {
+      msg = '网络连接异常';
+    } else if (error?.response?.data?.error?.message) {
+      msg = error?.response?.data?.error?.message;
+    } else if (error?.error?.message) {
+      msg = error?.error?.message;
+    }
+
+    addLog.error(`response error: ${msg}`, error);
+  }
+
+  res.status(code).json({
+    code,
+    statusText: '',
+    message: message || msg,
+    data: data !== undefined ? data : null
+  });
+};
+
+export const sseErrRes = (res: NextApiResponse, error: any) => {
+  const errResponseKey = typeof error === 'string' ? error : error?.message;
+
+  // Specified error
+  if (ERROR_RESPONSE[errResponseKey]) {
+    // login is expired
+    if (errResponseKey === ERROR_ENUM.unAuthorization) {
+      clearCookie(res);
+    }
+
+    return responseWrite({
+      res,
+      event: sseResponseEventEnum.error,
+      data: JSON.stringify(ERROR_RESPONSE[errResponseKey])
+    });
+  }
+
+  let msg = error?.response?.statusText || error?.message || '请求错误';
+  if (typeof error === 'string') {
+    msg = error;
+  } else if (proxyError[error?.code]) {
+    msg = '网络连接异常';
+  } else if (error?.response?.data?.error?.message) {
+    msg = error?.response?.data?.error?.message;
+  } else if (error?.error?.message) {
+    msg = error?.error?.message;
+  }
+
+  addLog.error(`sse error: ${msg}`, error);
+
+  responseWrite({
+    res,
+    event: sseResponseEventEnum.error,
+    data: JSON.stringify({ message: msg })
+  });
+};
 
 export function responseWriteController({
   res,