main.py

import getopt
import logging
import os
import socket
import sys
import time
import uuid
from datetime import datetime

from flask import Flask, jsonify
from flask import request
from flask_cors import CORS
from data import conversation_service as conv_service
from data import set_database
from data.conversation_service import save, update_by_id, del_by_id
from data.database import ChatI
from util.Result import err, log, ok, create_request_content, stop_stream, get_chat_api_

app = Flask(__name__)
CORS(app)
# 设置内置日志记录器等级
logger = logging.getLogger('my_logger')

logger.setLevel(logging.INFO)  # 设置日志级别，可以是 DEBUG, INFO, WARNING, ERROR, CRITICAL


@app.route("/get_chat_list", methods=["GET", "POST"])
def get_chat_list():
    message = get_chat_api_()

    return ok(message)


@app.route("/generate/id", methods=['POST'])
def generate_id():
    return ok(str(uuid.uuid4()))


@app.route("/chat/character", methods=['POST'])
def character_info():
    record = conv_service.get_character()
    return ok(record)


@app.route("/chat/repeat/<cid>", methods=['POST', 'GET'])
def chat_repeat(cid):
    logger.info(log(f"id:{cid}"))
    character_id = request.json.get('character')
    model = request.json.get('model')
    prompt = request.json.get('prompt')
    if not prompt:
        return jsonify({"error": "请求参数缺失！"}), 400
    logger.info(log(f"id:{cid}\nprompt:{prompt}\ncharacter:{character_id}"))
    record = conv_service.get_by_id(cid)
    conv_list = []
    human_msg = {
        "speaker": "human",
        "speech": prompt,
        "createTime": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    }
    if record is None:
        record = ChatI(id=cid, conversation={"_id": cid, "title": prompt, "convs": []}, create_time=datetime.now(),
                       update_time=None)
        conv_list.append(human_msg)
        record.conversation['convs'] = conv_list
        save(record)
    else:
        conv_list = record['conversation']['convs']
        conv_list.append(human_msg)
        update_by_id(cid, human_msg)
    try:
        generate = create_request_content(cid, conv_list, character_id, False, model)
        return app.response_class(generate(), mimetype='application/json')
    except Exception as e:
        logger.error(f"Error in chat function: {e}")
        return jsonify({"error": "内部服务器错误"}), 500


@app.route("/stop/chat/<cid>", methods=['POST', 'GET'])
def chat_stop(cid):
    logger.info(log(f"关闭流id:{cid}"))
    stop_stream(cid)
    return str("success")


@app.route("/chat/clear", methods=['POST', 'GET'])
def chat_clear():
    ids = request.json.get('ids')
    if ids is not None:
        try:
            conv_service.del_by_id(ids)
        except Exception as e:
            logger.error(f"Error in chat function: {e}")
    return str("success")


@app.route("/chat/title/<cid>", methods=['POST', 'GET'])
def chat_title(cid):
    logger.info(log(f"id:{cid}"))
    record = conv_service.get_by_id(cid)
    if record is None:
        return err("对话不存在")
    result = record.conversation['title']
    return str(result)


@app.route("/conv/<cid>", methods=['GET'])
def conv(cid):
    logger.info(log(f"id:{cid}"))
    record = conv_service.get_by_id(cid)
    return ok(record)


@app.route("/chat/<cid>", methods=['POST', 'GET'])
def chat(cid):
    # 从 URL 查询参数中获取 prompt
    prompt = request.json.get('prompt')
    character_id = request.json.get('character')
    model = request.json.get('model')
    if not prompt:
        return jsonify({"error": "请求参数缺失！"}), 400
    logger.info(log(f"id:{cid}\nprompt:{prompt}\ncharacter:{character_id}"))
    record = conv_service.get_by_id(cid)
    conv_list = []
    human_msg = {
        "speaker": "human",
        "speech": prompt,
        "createTime": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    }
    if record is None:
        record = ChatI(id=cid, conversation={"_id": cid, "title": prompt, "convs": []}, create_time=datetime.now(),
                       update_time=None)
        conv_list.append(human_msg)
        record.conversation['convs'] = conv_list
        save(record)
    else:
        conv_list = record['conversation']['convs']
        conv_list.append(human_msg)
        update_by_id(cid, human_msg)
    try:
        generate = create_request_content(cid, conv_list, character_id, False, model)
        return app.response_class(generate(), mimetype='application/json')
    except Exception as e:
        logger.error(f"Error in chat function: {e}")
        return jsonify({"error": "内部服务器错误"}), 500


def init_logging():
    hostname = socket.gethostname()
    logfile_path = '/code/%s' % hostname
    os.makedirs(logfile_path, exist_ok=True)
    logfile_name = logfile_path + '/cat.log'
    # 创建控制台处理器并设置日志级别
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.DEBUG)

    # 创建日志格式
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    console_handler.setFormatter(formatter)
    logging.basicConfig(level=logging.DEBUG, filename=logfile_name,
                        format="%(asctime)s - [%(levelname)s] %(filename)s$%(funcName)s:%(lineno)d\t"
                               "%(message)s",
                        datefmt="%F %T")
    # 获取根日志记录器
    logger = logging.getLogger()

    # 将控制台处理器添加到根日志记录器
    logger.addHandler(console_handler)


def init_database():
    logging.info(sys.argv[1:])
    opts, args = getopt.getopt(sys.argv[1:], 'u:p:', ["host=", "port=", "databaseName="])
    set_database(opts)


if __name__ == '__main__':
    init_logging()
    init_database()
    app.run(host="0.0.0.0", port=8383, debug=False)