model

server/asr/asr_server.py

@@ -0,0 +1,58 @@
+from fastapi import FastAPI
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import PlainTextResponse
+from loguru import logger
+from pydantic import BaseModel
+
+from ..web_configs import WEB_CONFIGS
+from .asr_worker import load_asr_model, process_asr
+
+app = FastAPI()
+
+if WEB_CONFIGS.ENABLE_ASR:
+    ASR_HANDLER = load_asr_model()
+else:
+    ASR_HANDLER = None
+
+
+class ASRItem(BaseModel):
+    user_id: int  # User 识别号，用于区分不用的用户调用
+    request_id: str  # 请求 ID，用于生成 TTS & 数字人
+    wav_path: str  # wav 文件路径
+
+
+@app.post("/asr")
+async def get_asr(asr_item: ASRItem):
+    # 语音转文字
+    result = ""
+    status = "success"
+    if ASR_HANDLER is None:
+        result = "ASR not enable in sever"
+        status = "fail"
+        logger.error(f"ASR not enable...")
+    else:
+        result = process_asr(ASR_HANDLER, asr_item.wav_path)
+    logger.info(f"ASR res for id {asr_item.request_id}, res = {result}")
+
+    return {"user_id": asr_item.user_id, "request_id": asr_item.request_id, "status": status, "result": result}
+
+
+@app.get("/asr/check")
+async def check_server():
+    return {"message": "server enabled"}
+
+
+@app.exception_handler(RequestValidationError)
+async def validation_exception_handler(request, exc):
+    """调 API 入参错误的回调接口
+
+    Args:
+        request (_type_): _description_
+        exc (_type_): _description_
+
+    Returns:
+        _type_: _description_
+    """
+    logger.info(request)
+    logger.info(exc)
+    return PlainTextResponse(str(exc), status_code=400)

server/asr/asr_worker.py

@@ -0,0 +1,54 @@
+import datetime
+
+from funasr import AutoModel
+from funasr.download.name_maps_from_hub import name_maps_ms as NAME_MAPS_MS
+from modelscope import snapshot_download
+from modelscope.utils.constant import Invoke, ThirdParty
+
+from ..web_configs import WEB_CONFIGS
+
+
+def load_asr_model():
+
+    # 模型下载
+    model_path_info = dict()
+    for model_name in ["paraformer-zh", "fsmn-vad", "ct-punc"]:
+        print(f"downloading asr model : {NAME_MAPS_MS[model_name]}")
+        mode_dir = snapshot_download(
+            NAME_MAPS_MS[model_name],
+            revision="master",
+            user_agent={Invoke.KEY: Invoke.PIPELINE, ThirdParty.KEY: "funasr"},
+            cache_dir=WEB_CONFIGS.ASR_MODEL_DIR,
+        )
+        model_path_info[model_name] = mode_dir
+        NAME_MAPS_MS[model_name] = mode_dir  # 更新权重路径环境变量
+
+    print(f"ASR model path info = {model_path_info}")
+    # paraformer-zh is a multi-functional asr model
+    # use vad, punc, spk or not as you need
+    model = AutoModel(
+        model="paraformer-zh",  # 语音识别，带时间戳输出，非实时
+        vad_model="fsmn-vad",  # 语音端点检测，实时
+        punc_model="ct-punc",  # 标点恢复
+        # spk_model="cam++" # 说话人确认/分割
+        model_path=model_path_info["paraformer-zh"],
+        vad_kwargs={"model_path": model_path_info["fsmn-vad"]},
+        punc_kwargs={"model_path": model_path_info["ct-punc"]},
+    )
+    return model
+
+
+def process_asr(model: AutoModel, wav_path):
+    # https://github.com/modelscope/FunASR/blob/main/README_zh.md#%E5%AE%9E%E6%97%B6%E8%AF%AD%E9%9F%B3%E8%AF%86%E5%88%AB
+    f_start_time = datetime.datetime.now()
+    res = model.generate(input=wav_path, batch_size_s=50, hotword="魔搭")
+    delta_time = datetime.datetime.now() - f_start_time
+
+    try:
+        print(f"ASR using time {delta_time}s, text: ", res[0]["text"])
+        res_str = res[0]["text"]
+    except Exception as e:
+        print("ASR 解析失败，无法获取到文字")
+        return ""
+
+    return res_str

server/base/base_server.py

@@ -0,0 +1,206 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   base_server.py
+@Time    :   2024/09/02
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   中台服务入口文件
+"""
+
+import time
+import uuid
+from contextlib import asynccontextmanager
+from pathlib import Path
+
+from fastapi import Depends, FastAPI, File, HTTPException, Response, UploadFile
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import PlainTextResponse
+from fastapi.staticfiles import StaticFiles
+from loguru import logger
+
+from ..web_configs import API_CONFIG, WEB_CONFIGS
+from .database.init_db import create_db_and_tables
+from .routers import digital_human, llm, products, streamer_info, streaming_room, users
+from .server_info import SERVER_PLUGINS_INFO
+from .utils import ChatItem, ResultCode, gen_default_data, make_return_data, streamer_sales_process
+
+swagger_description = """
+
+## 项目地址
+
+[销冠 —— 卖货主播大模型 && 后台管理系统](https://github.com/PeterH0323/Streamer-Sales)
+
+## 功能点
+
+1. 📜 **主播文案一键生成**
+2. 🚀 KV cache + Turbomind **推理加速**
+3. 📚 RAG **检索增强生成**
+4. 🔊 TTS **文字转语音**
+5. 🦸 **数字人生成**
+6. 🌐 **Agent 网络查询**
+7. 🎙️ **ASR 语音转文字**
+8. 🍍 **Vue + pinia + element-plus **搭建的前端，可自由扩展快速开发
+9. 🗝️ 后端采用 FastAPI + Uvicorn + PostgreSQL，**高性能，高效编码，生产可用，同时具有 JWT 身份验证**
+10. 🐋 采用 Docker-compose 部署，**一键实现分布式部署**
+
+"""
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """服务生命周期函数"""
+    # 启动
+    create_db_and_tables()  # 创建数据库和数据表
+
+    # 新服务，生成默认数据，可以自行注释 or 修改
+    gen_default_data()
+
+    if WEB_CONFIGS.ENABLE_RAG:
+        from .modules.rag.rag_worker import load_rag_model
+
+        # 生成 rag 数据库
+        await load_rag_model(user_id=1)
+
+    yield
+
+    # 结束
+    logger.info("Base server stopped.")
+
+
+app = FastAPI(
+    title="销冠 —— 卖货主播大模型 && 后台管理系统",
+    description=swagger_description,
+    summary="一个能够根据给定的商品特点从激发用户购买意愿角度出发进行商品解说的卖货主播大模型。",
+    version="1.0.0",
+    license_info={
+        "name": "AGPL-3.0 license",
+        "url": "https://github.com/PeterH0323/Streamer-Sales/blob/main/LICENSE",
+    },
+    root_path=API_CONFIG.API_V1_STR,
+    lifespan=lifespan,
+)
+
+# 注册路由
+app.include_router(users.router)
+app.include_router(products.router)
+app.include_router(llm.router)
+app.include_router(streamer_info.router)
+app.include_router(streaming_room.router)
+app.include_router(digital_human.router)
+
+
+# 挂载静态文件目录，以便访问上传的文件
+WEB_CONFIGS.SERVER_FILE_ROOT = str(Path(WEB_CONFIGS.SERVER_FILE_ROOT).absolute())
+Path(WEB_CONFIGS.SERVER_FILE_ROOT).mkdir(parents=True, exist_ok=True)
+logger.info(f"上传文件挂载路径: {WEB_CONFIGS.SERVER_FILE_ROOT}")
+logger.info(f"上传文件访问 URL: {API_CONFIG.REQUEST_FILES_URL}")
+app.mount(
+    f"/{API_CONFIG.REQUEST_FILES_URL.split('/')[-1]}",
+    StaticFiles(directory=WEB_CONFIGS.SERVER_FILE_ROOT),
+    name=API_CONFIG.REQUEST_FILES_URL.split("/")[-1],
+)
+
+
+@app.get("/")
+async def hello():
+    return {"message": "Hello Streamer-Sales"}
+
+
+@app.exception_handler(RequestValidationError)
+async def validation_exception_handler(request, exc):
+    """调 API 入参错误的回调接口
+
+    Args:
+        request (_type_): _description_
+        exc (_type_): _description_
+
+    Returns:
+        _type_: _description_
+    """
+    logger.info(request.headers)
+    logger.info(exc)
+    return PlainTextResponse(str(exc), status_code=400)
+
+
+@app.get("/dashboard", tags=["base"], summary="获取主页信息接口")
+async def get_dashboard_info():
+    """首页展示数据"""
+    fake_dashboard_data = {
+        "registeredBrandNum": 98431,  # 入驻品牌方
+        "productNum": 49132,  # 商品数
+        "dailyActivity": 68431,  # 日活
+        "todayOrder": 8461321,  # 订单量
+        "totalSales": 245578131857,  # 销售额
+        "conversionRate": 90.0,  # 转化率
+        # 折线图
+        "orderNumList": [46813, 68461, 99561, 138131, 233812, 84613, 846122],  # 订单量
+        "totalSalesList": [46813, 68461, 99561, 138131, 23383, 84613, 841213],  # 销售额
+        "newUserList": [3215, 65131, 6513, 6815, 2338, 84614, 84213],  # 新增用户
+        "activityUserList": [132, 684, 59431, 4618, 31354, 68431, 88431],  # 活跃用户
+        # 柱状图
+        "knowledgeBasesNum": 12,  # 知识库数量
+        "digitalHumanNum": 3,  # 数字人数量
+        "LiveRoomNum": 5,  # 直播间数量
+    }
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", fake_dashboard_data)
+
+
+@app.get("/plugins_info", tags=["base"], summary="获取组件信息接口")
+async def get_plugins_info():
+
+    plugins_info = SERVER_PLUGINS_INFO.get_status()
+    return make_return_data(True, ResultCode.SUCCESS, "成功", plugins_info)
+
+
+@app.post("/upload/file", tags=["base"], summary="上传文件接口")
+async def upload_product_api(file: UploadFile = File(...), user_id: int = Depends(users.get_current_user_info)):
+
+    file_type = file.filename.split(".")[-1]  # eg. png
+    logger.info(f"upload file type = {file_type}")
+
+    sub_dir_name_map = {
+        "md": WEB_CONFIGS.INSTRUCTIONS_DIR,
+        "png": WEB_CONFIGS.IMAGES_DIR,
+        "jpg": WEB_CONFIGS.IMAGES_DIR,
+        "mp4": WEB_CONFIGS.STREAMER_INFO_FILES_DIR,
+        "wav": WEB_CONFIGS.STREAMER_INFO_FILES_DIR,
+        "webm": WEB_CONFIGS.ASR_FILE_DIR,
+    }
+    if file_type in ["wav", "mp4"]:
+        save_root = WEB_CONFIGS.STREAMER_FILE_DIR
+    elif file_type in ["webm"]:
+        save_root = ""
+    else:
+        save_root = WEB_CONFIGS.PRODUCT_FILE_DIR
+
+    upload_time = str(int(time.time())) + "__" + str(uuid.uuid4().hex)
+
+    sub_dir_name = sub_dir_name_map[file_type]
+    save_path = Path(WEB_CONFIGS.SERVER_FILE_ROOT).joinpath(save_root, sub_dir_name, upload_time + "." + file_type)
+    save_path.parent.mkdir(exist_ok=True, parents=True)
+    logger.info(f"save path = {save_path}")
+
+    # 使用流式处理接收文件
+    with open(save_path, "wb") as buffer:
+        while chunk := await file.read(1024 * 1024 * 5):  # 每次读取 5MB 的数据块
+            buffer.write(chunk)
+
+    split_dir_name = Path(WEB_CONFIGS.SERVER_FILE_ROOT).name  # 保存文件夹根目录名字
+    file_url = f"{API_CONFIG.REQUEST_FILES_URL}{str(save_path).split(split_dir_name)[-1]}"
+
+    # TODO 文件归属记录表
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", file_url)
+
+
+@app.post("/streamer-sales/chat", tags=["base"], summary="对话接口", deprecated=True)
+async def streamer_sales_chat(chat_item: ChatItem, response: Response):
+    from sse_starlette import EventSourceResponse
+
+    # 对话总接口
+    response.headers["Content-Type"] = "text/event-stream"
+    response.headers["Cache-Control"] = "no-cache"
+    return EventSourceResponse(streamer_sales_process(chat_item))

server/base/database/init_db.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   init_db.py
+@Time    :   2024/09/06
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   数据库初始化
+"""
+
+from loguru import logger
+from pydantic import PostgresDsn
+from pydantic_core import MultiHostUrl
+from sqlmodel import SQLModel, create_engine
+
+from ...web_configs import WEB_CONFIGS
+
+ECHO_DB_MESG = True  # 数据库执行中是否回显，for debug
+
+
+def sqlalchemy_db_url() -> PostgresDsn:
+    """生成数据库 URL
+
+    Returns:
+        PostgresDsn: 数据库地址
+    """
+    return MultiHostUrl.build(
+        scheme="postgresql+psycopg",
+        username=WEB_CONFIGS.POSTGRES_USER,
+        password=WEB_CONFIGS.POSTGRES_PASSWORD,
+        host=WEB_CONFIGS.POSTGRES_SERVER,
+        port=WEB_CONFIGS.POSTGRES_PORT,
+        path=WEB_CONFIGS.POSTGRES_DB,
+    )
+
+
+logger.info(f"connecting to db: {str(sqlalchemy_db_url())}")
+DB_ENGINE = create_engine(str(sqlalchemy_db_url()), echo=ECHO_DB_MESG)
+
+
+def create_db_and_tables():
+    """创建所有数据库和对应的表，有则跳过"""
+    SQLModel.metadata.create_all(DB_ENGINE)

server/base/database/llm_db.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   llm_db.py
+@Time    :   2024/09/01
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   大模型对话数据库交互
+"""
+
+import yaml
+
+from ...web_configs import WEB_CONFIGS
+
+
+async def get_llm_product_prompt_base_info():
+    # 加载对话配置文件
+    with open(WEB_CONFIGS.CONVERSATION_CFG_YAML_PATH, "r", encoding="utf-8") as f:
+        dataset_yaml = yaml.safe_load(f)
+
+    return dataset_yaml

server/base/database/product_db.py

@@ -0,0 +1,186 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   product_db.py
+@Time    :   2024/08/30
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   商品数据表文件读写
+"""
+
+from typing import List, Tuple
+
+from loguru import logger
+from sqlalchemy import func
+from sqlmodel import Session, and_, not_, select
+
+from ...web_configs import API_CONFIG
+from ..models.product_model import ProductInfo
+from .init_db import DB_ENGINE
+
+
+async def get_db_product_info(
+    user_id: int,
+    current_page: int = -1,
+    page_size: int = 10,
+    product_name: str | None = None,
+    product_id: int | None = None,
+    exclude_list: List[int] | None = None,
+) -> Tuple[List[ProductInfo], int]:
+    """查询数据库中的商品信息
+
+    Args:
+        user_id (int): 用户 ID
+        current_page (int, optional): 页数. Defaults to -1.
+        page_size (int, optional): 每页的大小. Defaults to 10.
+        product_name (str | None, optional): 商品名称，模糊搜索. Defaults to None.
+        product_id (int | None, optional): 商品 ID，用户获取特定商品信息. Defaults to None.
+
+    Returns:
+         List[ProductInfo]: 商品信息
+         int : 该用户持有的总商品数，已剔除删除的
+    """
+
+    assert current_page != 0
+    assert page_size != 0
+
+    # 查询条件
+    query_condiction = and_(ProductInfo.user_id == user_id, ProductInfo.delete == False)
+
+    # 获取总数
+    with Session(DB_ENGINE) as session:
+        # 获得该用户所有商品的总数
+        total_product_num = session.scalar(select(func.count(ProductInfo.product_id)).where(query_condiction))
+
+        if product_name is not None:
+            # 查询条件更改为商品名称模糊搜索
+            query_condiction = and_(
+                ProductInfo.user_id == user_id, ProductInfo.delete == False, ProductInfo.product_name.ilike(f"%{product_name}%")
+            )
+
+        elif product_id is not None:
+            # 查询条件更改为查找特定 ID
+            query_condiction = and_(
+                ProductInfo.user_id == user_id, ProductInfo.delete == False, ProductInfo.product_id == product_id
+            )
+
+        elif exclude_list is not None:
+            # 排除查询
+            query_condiction = and_(
+                ProductInfo.user_id == user_id, ProductInfo.delete == False, not_(ProductInfo.product_id.in_(exclude_list))
+            )
+
+        # 查询获取商品
+        if current_page < 0:
+            # 全部查询
+            product_list = session.exec(select(ProductInfo).where(query_condiction).order_by(ProductInfo.product_id)).all()
+        else:
+            # 分页查询
+            offset_idx = (current_page - 1) * page_size
+            product_list = session.exec(
+                select(ProductInfo).where(query_condiction).offset(offset_idx).limit(page_size).order_by(ProductInfo.product_id)
+            ).all()
+
+    if product_list is None:
+        logger.warning("nothing to find in db...")
+        product_list = []
+
+    # 将路径换成服务器路径
+    for product in product_list:
+        product.image_path = API_CONFIG.REQUEST_FILES_URL + product.image_path
+        product.instruction = API_CONFIG.REQUEST_FILES_URL + product.instruction
+
+    logger.info(product_list)
+    logger.info(f"len {len(product_list)}")
+
+    return product_list, total_product_num
+
+
+async def delete_product_id(product_id: int, user_id: int) -> bool:
+    """删除特定的商品 ID
+
+    Args:
+        product_id (int): 商品 ID
+        user_id (int): 用户 ID，用于防止其他用户恶意删除
+
+    Returns:
+        bool: 是否删除成功
+    """
+
+    delete_success = True
+
+    try:
+        # 获取总数
+        with Session(DB_ENGINE) as session:
+            # 查找特定 ID
+            product_info = session.exec(
+                select(ProductInfo).where(and_(ProductInfo.product_id == product_id, ProductInfo.user_id == user_id))
+            ).one()
+
+            if product_info is None:
+                logger.error("Delete by other ID !!!")
+                return False
+
+            product_info.delete = True  # 设置为删除
+            session.add(product_info)
+            session.commit()  # 提交
+    except Exception:
+        delete_success = False
+
+    return delete_success
+
+
+def create_or_update_db_product_by_id(product_id: int, new_info: ProductInfo, user_id: int) -> bool:
+    """新增 or 编辑商品信息
+
+    Args:
+        product_id (int): 商品 ID
+        new_info (ProductInfo): 新的信息
+        user_id (int): 用户 ID，用于防止其他用户恶意修改
+
+    Returns:
+        bool: 说明书是否变化
+    """
+
+    instruction_updated = False
+
+    # 去掉服务器地址
+    new_info.image_path = new_info.image_path.replace(API_CONFIG.REQUEST_FILES_URL, "")
+    new_info.instruction = new_info.instruction.replace(API_CONFIG.REQUEST_FILES_URL, "")
+
+    with Session(DB_ENGINE) as session:
+
+        if product_id > 0:
+            # 更新特定 ID
+            product_info = session.exec(
+                select(ProductInfo).where(and_(ProductInfo.product_id == product_id, ProductInfo.user_id == user_id))
+            ).one()
+
+            if product_info is None:
+                logger.error("Edit by other ID !!!")
+                return False
+
+            if product_info.instruction != new_info.instruction:
+                # 判断说明书是否变化了
+                instruction_updated = True
+
+            # 更新对应的值
+            product_info.product_name = new_info.product_name
+            product_info.product_class = new_info.product_class
+            product_info.heighlights = new_info.heighlights
+            product_info.image_path = new_info.image_path
+            product_info.instruction = new_info.instruction
+            product_info.departure_place = new_info.departure_place
+            product_info.delivery_company = new_info.delivery_company
+            product_info.selling_price = new_info.selling_price
+            product_info.amount = new_info.amount
+
+            session.add(product_info)
+        else:
+            # 新增，直接添加即可
+            session.add(new_info)
+            instruction_updated = True
+
+        session.commit()  # 提交
+    return instruction_updated

server/base/database/streamer_info_db.py

@@ -0,0 +1,152 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   streamer_info_db.py
+@Time    :   2024/08/30
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   主播信息数据库操作
+"""
+
+
+from typing import List
+
+from loguru import logger
+from sqlmodel import Session, and_, select
+
+from ...web_configs import API_CONFIG
+from ..models.streamer_info_model import StreamerInfo
+from .init_db import DB_ENGINE
+
+
+async def get_db_streamer_info(user_id: int, streamer_id: int | None = None) -> List[StreamerInfo] | None:
+    """查询数据库中的主播信息
+
+    Args:
+        user_id (int): 用户 ID
+        streamer_id (int | None, optional): 主播 ID，用户获取特定主播信息. Defaults to None.
+
+    Returns:
+        List[StreamerInfo] | StreamerInfo | None: 主播信息，如果获取全部则返回 list，如果获取单个则返回单个，如果查不到返回 None
+    """
+
+    # 查询条件
+    query_condiction = and_(StreamerInfo.user_id == user_id, StreamerInfo.delete == False)
+
+    # 获取总数
+    with Session(DB_ENGINE) as session:
+        # 获得该用户所有主播的总数
+        # total_product_num = session.scalar(select(func.count(StreamerInfo.product_id)).where(query_condiction))
+
+        if streamer_id is not None:
+            # 查询条件更改为查找特定 ID
+            query_condiction = and_(
+                StreamerInfo.user_id == user_id, StreamerInfo.delete == False, StreamerInfo.streamer_id == streamer_id
+            )
+
+        # 查询主播商品，并根据 ID 进行排序
+        try:
+            streamer_list = session.exec(select(StreamerInfo).where(query_condiction).order_by(StreamerInfo.streamer_id)).all()
+        except Exception as e:
+            streamer_list = None
+
+    if streamer_list is None:
+        logger.warning("nothing to find in db...")
+        streamer_list = []
+
+    # 将路径换成服务器路径
+    for streamer in streamer_list:
+        streamer.avatar = API_CONFIG.REQUEST_FILES_URL + streamer.avatar
+        streamer.tts_reference_audio = API_CONFIG.REQUEST_FILES_URL + streamer.tts_reference_audio
+        streamer.poster_image = API_CONFIG.REQUEST_FILES_URL + streamer.poster_image
+        streamer.base_mp4_path = API_CONFIG.REQUEST_FILES_URL + streamer.base_mp4_path
+
+    logger.info(streamer_list)
+    logger.info(f"len {len(streamer_list)}")
+
+    return streamer_list
+
+
+async def delete_streamer_id(streamer_id: int, user_id: int) -> bool:
+    """删除特定的主播 ID
+
+    Args:
+        streamer_id (int): 主播 ID
+        user_id (int): 用户 ID，用于防止其他用户恶意删除
+
+    Returns:
+        bool: 是否删除成功
+    """
+
+    delete_success = True
+
+    try:
+        # 获取总数
+        with Session(DB_ENGINE) as session:
+            # 查找特定 ID
+            streamer_info = session.exec(
+                select(StreamerInfo).where(and_(StreamerInfo.streamer_id == streamer_id, StreamerInfo.user_id == user_id))
+            ).one()
+
+            if streamer_info is None:
+                logger.error("Delete by other ID !!!")
+                return False
+
+            streamer_info.delete = True  # 设置为删除
+            session.add(streamer_info)
+            session.commit()  # 提交
+    except Exception:
+        delete_success = False
+
+    return delete_success
+
+
+def create_or_update_db_streamer_by_id(streamer_id: int, new_info: StreamerInfo, user_id: int) -> int:
+    """新增 or 编辑主播信息
+
+    Args:
+        product_id (int): 商品 ID
+        new_info (ProductInfo): 新的信息
+        user_id (int): 用户 ID，用于防止其他用户恶意修改
+
+    Returns:
+        int: 主播 ID
+    """
+
+    # 去掉服务器地址
+    new_info.avatar = new_info.avatar.replace(API_CONFIG.REQUEST_FILES_URL, "")
+    new_info.tts_reference_audio = new_info.tts_reference_audio.replace(API_CONFIG.REQUEST_FILES_URL, "")
+    new_info.poster_image = new_info.poster_image.replace(API_CONFIG.REQUEST_FILES_URL, "")
+    new_info.base_mp4_path = new_info.base_mp4_path.replace(API_CONFIG.REQUEST_FILES_URL, "")
+
+    with Session(DB_ENGINE) as session:
+
+        if streamer_id > 0:
+            # 更新特定 ID
+            streamer_info = session.exec(
+                select(StreamerInfo).where(and_(StreamerInfo.streamer_id == streamer_id, StreamerInfo.user_id == user_id))
+            ).one()
+
+            if streamer_info is None:
+                logger.error("Edit by other ID !!!")
+                return -1
+        else:
+            # 新增，直接添加即可
+            streamer_info = StreamerInfo(user_id=user_id)
+
+        # 更新对应的值
+        streamer_info.name = new_info.name
+        streamer_info.character = new_info.character
+        streamer_info.avatar = new_info.avatar
+        streamer_info.tts_weight_tag = new_info.tts_weight_tag
+        streamer_info.tts_reference_sentence = new_info.tts_reference_sentence
+        streamer_info.tts_reference_audio = new_info.tts_reference_audio
+        streamer_info.poster_image = new_info.poster_image
+        streamer_info.base_mp4_path = new_info.base_mp4_path
+
+        session.add(streamer_info)
+        session.commit()  # 提交
+        session.refresh(streamer_info)
+
+        return int(streamer_info.streamer_id)

server/base/database/streamer_room_db.py

@@ -0,0 +1,415 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   streamer_room_db.py
+@Time    :   2024/08/31
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   直播间信息数据库操作
+"""
+
+
+from datetime import datetime
+from typing import List
+
+from loguru import logger
+from sqlmodel import Session, and_, not_, select
+
+from ...web_configs import API_CONFIG
+from ..models.streamer_room_model import ChatMessageInfo, OnAirRoomStatusItem, SalesDocAndVideoInfo, StreamRoomInfo
+from .init_db import DB_ENGINE
+
+
+async def get_db_streaming_room_info(user_id: int, room_id: int | None = None) -> List[StreamRoomInfo] | None:
+    """查询数据库中的商品信息
+
+    Args:
+        user_id (int): 用户 ID
+        streamer_id (int | None, optional): 主播 ID，用户获取特定主播信息. Defaults to None.
+
+    Returns:
+        List[StreamRoomInfo] | None: 直播间信息
+    """
+
+    # 查询条件
+    query_condiction = and_(StreamRoomInfo.user_id == user_id, StreamRoomInfo.delete == False)
+
+    # 获取总数
+    with Session(DB_ENGINE) as session:
+        if room_id is not None:
+            # 查询条件更改为查找特定 ID
+            query_condiction = and_(
+                StreamRoomInfo.user_id == user_id, StreamRoomInfo.delete == False, StreamRoomInfo.room_id == room_id
+            )
+
+        # 查询获取直播间信息
+        stream_room_list = session.exec(select(StreamRoomInfo).where(query_condiction).order_by(StreamRoomInfo.room_id)).all()
+
+    if stream_room_list is None:
+        logger.warning("nothing to find in db...")
+        stream_room_list = []
+
+    # 将路径换成服务器路径
+    for stream_room in stream_room_list:
+        # 主播信息
+        stream_room.streamer_info.avatar = API_CONFIG.REQUEST_FILES_URL + stream_room.streamer_info.avatar
+        stream_room.streamer_info.tts_reference_audio = (
+            API_CONFIG.REQUEST_FILES_URL + stream_room.streamer_info.tts_reference_audio
+        )
+        stream_room.streamer_info.poster_image = API_CONFIG.REQUEST_FILES_URL + stream_room.streamer_info.poster_image
+        stream_room.streamer_info.base_mp4_path = API_CONFIG.REQUEST_FILES_URL + stream_room.streamer_info.base_mp4_path
+
+        # 商品信息
+        for idx, product in enumerate(stream_room.product_list):
+            stream_room.product_list[idx].product_info.image_path = API_CONFIG.REQUEST_FILES_URL + product.product_info.image_path
+            stream_room.product_list[idx].product_info.instruction = (
+                API_CONFIG.REQUEST_FILES_URL + product.product_info.instruction
+            )
+
+    logger.info(stream_room_list)
+    logger.info(f"len {len(stream_room_list)}")
+
+    return stream_room_list
+
+
+async def delete_room_id(room_id: int, user_id: int) -> bool:
+    """删除特定的主播间 ID
+
+    Args:
+        room_id (int): 直播间 ID
+        user_id (int): 用户 ID，用于防止其他用户恶意删除
+
+    Returns:
+        bool: 是否删除成功
+    """
+
+    delete_success = True
+
+    try:
+        # 获取总数
+        with Session(DB_ENGINE) as session:
+            # 查找特定 ID
+            room_info = session.exec(
+                select(StreamRoomInfo).where(and_(StreamRoomInfo.room_id == room_id, StreamRoomInfo.user_id == user_id))
+            ).one()
+
+            if room_info is None:
+                logger.error("Delete by other ID !!!")
+                return False
+
+            room_info.delete = True  # 设置为删除
+            session.add(room_info)
+            session.commit()  # 提交
+    except Exception:
+        delete_success = False
+
+    return delete_success
+
+
+def create_or_update_db_room_by_id(room_id: int, new_info: StreamRoomInfo, user_id: int):
+    """新增 or 编辑直播间信息
+
+    Args:
+        room_id (int): 直播间 ID
+        new_info (StreamRoomInfo): 新的信息
+        user_id (int): 用户 ID，用于防止其他用户恶意修改
+    """
+
+    with Session(DB_ENGINE) as session:
+
+        # 更新 status 内容
+        if new_info.status_id is not None:
+            status_info = session.exec(
+                select(OnAirRoomStatusItem).where(OnAirRoomStatusItem.status_id == new_info.status_id)
+            ).one()
+        else:
+            status_info = OnAirRoomStatusItem()
+
+        status_info.streaming_video_path = new_info.status.streaming_video_path.replace(API_CONFIG.REQUEST_FILES_URL, "")
+        status_info.live_status = new_info.status.live_status
+        session.add(status_info)
+        session.commit()
+        session.refresh(status_info)
+
+        if room_id > 0:
+
+            # 更新主播间其他信息
+            room_info = session.exec(
+                select(StreamRoomInfo).where(and_(StreamRoomInfo.room_id == room_id, StreamRoomInfo.user_id == user_id))
+            ).one()
+
+            if room_info is None:
+                logger.error("Edit by other ID !!!")
+                return
+
+        else:
+            room_info = StreamRoomInfo(status_id=status_info.status_id, user_id=user_id)
+
+        # 更新直播间基础信息
+        room_info.name = new_info.name
+        room_info.prohibited_words_id = new_info.prohibited_words_id
+        room_info.room_poster = new_info.room_poster.replace(API_CONFIG.REQUEST_FILES_URL, "")
+        room_info.background_image = new_info.background_image.replace(API_CONFIG.REQUEST_FILES_URL, "")
+        room_info.streamer_id = new_info.streamer_id
+
+        session.add(room_info)
+        session.commit()  # 提交
+        session.refresh(room_info)
+
+        # 更新商品信息
+        if len(new_info.product_list) > 0:
+            selected_id_list = [product.product_id for product in new_info.product_list]
+            for product in new_info.product_list:
+                if product.sales_info_id is not None:
+                    # 更新
+                    sales_info = session.exec(
+                        select(SalesDocAndVideoInfo).where(
+                            and_(
+                                SalesDocAndVideoInfo.room_id == room_info.room_id,
+                                SalesDocAndVideoInfo.product_id == product.product_id,
+                                SalesDocAndVideoInfo.sales_info_id == product.sales_info_id,
+                            )
+                        )
+                    ).one()
+                else:
+                    # 新建
+                    sales_info = SalesDocAndVideoInfo()
+
+                sales_info.product_id = product.product_id
+                sales_info.sales_doc = product.sales_doc
+                sales_info.start_time = product.start_time
+                sales_info.start_video = product.start_video.replace(API_CONFIG.REQUEST_FILES_URL, "")
+                sales_info.selected = True
+                sales_info.room_id = room_info.room_id
+                session.add(sales_info)
+                session.commit()
+
+            # 删除没选上的
+            if len(selected_id_list) > 0:
+                cancel_select_sales_info = session.exec(
+                    select(SalesDocAndVideoInfo).where(
+                        and_(
+                            SalesDocAndVideoInfo.room_id == room_info.room_id,
+                            not_(SalesDocAndVideoInfo.product_id.in_(selected_id_list)),
+                        )
+                    )
+                ).all()
+
+                if cancel_select_sales_info is not None:
+                    for cancel_select in cancel_select_sales_info:
+                        session.delete(cancel_select)
+                        session.commit()
+
+        return room_info.room_id
+
+
+def init_conversation(db_session, sales_info_id: int, streamer_id: int, sales_doc: str):
+    """新建直播间对话，一般触发于点击 开始直播 or 下一个商品
+
+    Args:
+        db_session (it): 数据库句柄
+        sales_info_id (int): 销售 ID
+        streamer_id (int): 主播 ID
+        sales_doc (str): 主播文案
+    """
+    message_info = ChatMessageInfo(
+        sales_info_id=sales_info_id, streamer_id=streamer_id, role="streamer", message=sales_doc, send_time=datetime.now()
+    )
+    db_session.add(message_info)
+
+
+def update_message_info(sales_info_id: int, role_id: int, role: str, message: str):
+    """新增对话记录
+
+    Args:
+        sales_info_id (int): 销售 ID
+        role_id (int): 角色 ID
+        role (str): 角色类型："streamer", "user"
+        message (str): 插入的消息
+    """
+
+    assert role in ["streamer", "user"]
+
+    with Session(DB_ENGINE) as session:
+
+        role_key = "streamer_id" if role == "streamer" else "user_id"
+        role_id_info = {role_key: role_id}
+
+        message_info = ChatMessageInfo(
+            **role_id_info, sales_info_id=sales_info_id, role=role, message=message, send_time=datetime.now()
+        )
+        session.add(message_info)
+        session.commit()
+
+
+def update_db_room_status(room_id: int, user_id: int, process_type: str):
+    """编辑直播间状态信息
+
+    Args:
+        room_id (int): 直播间 ID
+        new_status_info (OnAirRoomStatusItem): 新的信息
+        user_id (int): 用户 ID，用于防止其他用户恶意修改
+    """
+
+    with Session(DB_ENGINE) as session:
+
+        # 更新主播间其他信息
+        room_info = session.exec(
+            select(StreamRoomInfo).where(and_(StreamRoomInfo.room_id == room_id, StreamRoomInfo.user_id == user_id))
+        ).one()
+
+        if room_info is None:
+            logger.error("Edit by other ID !!!")
+            return
+
+        # 更新 status 内容
+        if room_info.status_id is not None:
+            status_info = session.exec(
+                select(OnAirRoomStatusItem).where(OnAirRoomStatusItem.status_id == room_info.status_id)
+            ).one()
+
+        if status_info is None:
+            logger.error("status_info is None !!!")
+            return
+
+        if process_type in ["online", "next-product"]:
+
+            if process_type == "online":
+                status_info.live_status = 1
+                status_info.start_time = datetime.now()
+                status_info.end_time = None
+                status_info.current_product_index = 0
+
+            elif process_type == "next-product":
+                status_info.current_product_index += 1
+
+            current_idx = status_info.current_product_index
+
+            status_info.streaming_video_path = room_info.product_list[current_idx].start_video
+            status_info.sales_info_id = room_info.product_list[current_idx].sales_info_id
+
+            sales_info = session.exec(
+                select(SalesDocAndVideoInfo).where(
+                    SalesDocAndVideoInfo.sales_info_id == room_info.product_list[current_idx].sales_info_id
+                )
+            ).one()
+
+            sales_info.start_time = datetime.now()
+            session.add(sales_info)
+
+            # 新建对话
+            init_conversation(
+                session, status_info.sales_info_id, room_info.streamer_id, room_info.product_list[current_idx].sales_doc
+            )
+
+        elif process_type == "offline":
+            status_info.streaming_video_path = ""
+            status_info.live_status = 2
+            status_info.end_time = datetime.now()
+
+        else:
+            raise NotImplemented("process type error !!")
+
+        session.add(status_info)
+        session.commit()
+
+
+def get_message_list(sales_info_id: int) -> List[ChatMessageInfo]:
+    """根据销售 ID 获取全部对话
+
+    Args:
+        sales_info_id (int): 销售 ID
+
+    Returns:
+        List[ChatMessageInfo]: 对话列表
+    """
+    with Session(DB_ENGINE) as session:
+
+        message_info = session.exec(
+            select(ChatMessageInfo)
+            .where(and_(ChatMessageInfo.sales_info_id == sales_info_id))
+            .order_by(ChatMessageInfo.message_id)
+        ).all()
+
+        if message_info is None:
+            return []
+
+    formate_message_list = []
+    for message_ in message_info:
+        chat_item = {
+            "role": message_.role,
+            "avatar": message_.user_info.avatar if message_.role == "user" else message_.streamer_info.avatar,
+            "userName": message_.user_info.username if message_.role == "user" else message_.streamer_info.name,
+            "message": message_.message,
+            "datetime": message_.send_time,
+        }
+
+        chat_item["avatar"] = API_CONFIG.REQUEST_FILES_URL + chat_item["avatar"]
+        formate_message_list.append(chat_item)
+
+    return formate_message_list
+
+
+def update_room_video_path(status_id: int, news_video_server_path: str):
+    """数据库更新 status 主播视频
+
+    Args:
+        status_id (int): 主播间 status ID
+        news_video_server_path (str): 需要更新的主播视频 服务器地址
+
+    """
+    with Session(DB_ENGINE) as session:
+        # 更新 status 内容
+        status_info = session.exec(select(OnAirRoomStatusItem).where(OnAirRoomStatusItem.status_id == status_id)).one()
+
+        status_info.streaming_video_path = news_video_server_path.replace(API_CONFIG.REQUEST_FILES_URL, "")
+        session.add(status_info)
+        session.commit()
+
+
+async def get_live_room_info(user_id: int, room_id: int):
+    """获取直播间的开播实时信息
+
+    Args:
+        user_id (int): 用户 ID
+        room_id (int): 直播间 ID
+
+    Returns:
+        dict: 直播间实时信息
+    """
+
+    # 根据直播间 ID 获取信息
+    streaming_room_info = await get_db_streaming_room_info(user_id, room_id)
+    streaming_room_info = streaming_room_info[0]
+
+    # 主播信息
+    streamer_info = streaming_room_info.streamer_info
+
+    # 商品索引
+    prodcut_index = streaming_room_info.status.current_product_index
+
+    # 是否为最后的商品
+    final_procut = True if len(streaming_room_info.product_list) - 1 == prodcut_index else False
+
+    # 对话信息
+    conversation_list = get_message_list(streaming_room_info.status.sales_info_id)
+
+    # 视频转换为服务器地址
+    video_path = API_CONFIG.REQUEST_FILES_URL + streaming_room_info.status.streaming_video_path
+
+    # 返回报文
+    res_data = {
+        "streamerInfo": streamer_info,
+        "conversation": conversation_list,
+        "currentProductInfo": streaming_room_info.product_list[prodcut_index].product_info,
+        "currentStreamerVideo": video_path,
+        "currentProductIndex": streaming_room_info.status.current_product_index,
+        "startTime": streaming_room_info.status.start_time,
+        "currentPoductStartTime": streaming_room_info.product_list[prodcut_index].start_time,
+        "finalProduct": final_procut,
+    }
+
+    logger.info(res_data)
+
+    return res_data

server/base/database/user_db.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   user_db.py
+@Time    :   2024/08/31
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   用户信息数据库操作
+"""
+
+from sqlmodel import Session, select
+
+from ...web_configs import API_CONFIG
+from ..models.user_model import UserBaseInfo, UserInfo
+from .init_db import DB_ENGINE
+
+
+def get_db_user_info(id: int = -1, username: str = "", all_info: bool = False) -> UserBaseInfo | UserInfo | None:
+    """查询数据库获取用户信息
+
+    Args:
+        id (int): 用户 ID
+        username (str): 用户名
+        all_info (bool): 是否返回含有密码串的敏感信息
+
+    Returns:
+        UserInfo | None: 用户信息，没有查到返回 None
+    """
+
+    if username == "":
+        # 使用 ID 的方式进行查询
+        query = select(UserInfo).where(UserInfo.user_id == id)
+    else:
+        query = select(UserInfo).where(UserInfo.username == username)
+
+    # 查询数据库
+    with Session(DB_ENGINE) as session:
+        results = session.exec(query).first()
+
+    # 返回服务器地址
+    results.avatar = API_CONFIG.REQUEST_FILES_URL + results.avatar
+
+    if results is not None and all_info is False:
+        # 返回不含用户敏感信息的基本信息
+        results = UserBaseInfo(**results.model_dump())
+
+    return results

server/base/models/llm_model.py

@@ -0,0 +1,17 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   llm_model.py
+@Time    :   2024/09/01
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   大模型对话数据结构
+"""
+
+from pydantic import BaseModel
+
+
+class GenProductItem(BaseModel):
+    gen_type: str
+    instruction: str

server/base/models/product_model.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   product_model.py
+@Time    :   2024/08/30
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   商品数据类型定义
+"""
+
+from datetime import datetime
+from typing import List
+from pydantic import BaseModel
+from sqlmodel import Field, Relationship, SQLModel
+
+
+# =======================================================
+#                      数据库模型
+# =======================================================
+
+
+class ProductInfo(SQLModel, table=True):
+    """商品信息"""
+
+    __tablename__ = "product_info"
+
+    product_id: int | None = Field(default=None, primary_key=True, unique=True)
+    product_name: str = Field(index=True, unique=True)
+    product_class: str
+    heighlights: str
+    image_path: str
+    instruction: str
+    departure_place: str
+    delivery_company: str
+    selling_price: float
+    amount: int
+    upload_date: datetime = datetime.now()
+    delete: bool = False
+
+    user_id: int | None = Field(default=None, foreign_key="user_info.user_id")
+
+    sales_info: list["SalesDocAndVideoInfo"] = Relationship(back_populates="product_info")
+
+
+# =======================================================
+#                      基本模型
+# =======================================================
+
+
+class ProductPageItem(BaseModel):
+    product_list: List[ProductInfo] = []
+    currentPage: int = 0  # 当前页数
+    pageSize: int = 0  # 页面的组件数量
+    totalSize: int = 0  # 总大小
+
+
+class ProductQueryItem(BaseModel):
+    instructionPath: str = ""  # 商品说明书路径，用于获取说明书内容

server/base/models/streamer_info_model.py

@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   streamer_info_model.py
+@Time    :   2024/08/30
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   主播信息数据结构
+"""
+
+from typing import Optional
+from sqlmodel import Field, Relationship, SQLModel
+
+
+# =======================================================
+#                      数据库模型
+# =======================================================
+class StreamerInfo(SQLModel, table=True):
+    __tablename__ = "streamer_info"
+
+    streamer_id: int | None = Field(default=None, primary_key=True, unique=True)
+    name: str = Field(index=True, unique=True)
+    character: str = ""
+    avatar: str = ""  # 头像
+
+    tts_weight_tag: str = ""  # 艾丝妲
+    tts_reference_sentence: str = ""
+    tts_reference_audio: str = ""
+
+    poster_image: str = ""
+    base_mp4_path: str = ""
+
+    delete: bool = False
+
+    user_id: int | None = Field(default=None, foreign_key="user_info.user_id")
+
+    room_info: Optional["StreamRoomInfo"] | None = Relationship(
+        back_populates="streamer_info", sa_relationship_kwargs={"lazy": "selectin"}
+    )

server/base/models/streamer_room_model.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   streamer_room_model.py
+@Time    :   2024/08/31
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   直播间信息数据结构定义
+"""
+
+from datetime import datetime
+from typing import Optional
+
+from pydantic import BaseModel
+from sqlmodel import Field, Relationship, SQLModel
+
+from ..models.user_model import UserInfo
+from ..models.product_model import ProductInfo
+from ..models.streamer_info_model import StreamerInfo
+
+
+class RoomChatItem(BaseModel):
+    roomId: int
+    message: str = ""
+    asrFileUrl: str = ""
+
+
+# =======================================================
+#                      直播间数据库模型
+# =======================================================
+
+
+class SalesDocAndVideoInfo(SQLModel, table=True):
+    """直播间 文案 和 数字人介绍视频数据结构"""
+
+    __tablename__ = "sales_doc_and_video_info"
+
+    sales_info_id: int | None = Field(default=None, primary_key=True, unique=True)
+
+    sales_doc: str = ""  # 讲解文案
+    start_video: str = ""  # 开播时候第一个讲解视频
+    start_time: datetime | None = None  # 当前商品开始时间
+    selected: bool = True
+
+    product_id: int | None = Field(default=None, foreign_key="product_info.product_id")
+    product_info: ProductInfo | None = Relationship(back_populates="sales_info", sa_relationship_kwargs={"lazy": "selectin"})
+
+    room_id: int | None = Field(default=None, foreign_key="stream_room_info.room_id")
+    stream_room: Optional["StreamRoomInfo"] | None = Relationship(back_populates="product_list")
+
+
+class OnAirRoomStatusItem(SQLModel, table=True):
+    """直播间状态信息"""
+
+    __tablename__ = "on_air_room_status_item"
+
+    status_id: int | None = Field(default=None, primary_key=True, unique=True)  # 直播间 ID
+
+    sales_info_id: int | None = Field(default=None, foreign_key="sales_doc_and_video_info.sales_info_id")
+
+    current_product_index: int = 0  # 目前讲解的商品列表索引
+    streaming_video_path: str = ""  # 目前介绍使用的视频
+
+    live_status: int = 0  # 直播间状态 0 未开播，1 正在直播，2 下播了
+    start_time: datetime | None = None  # 直播开始时间
+    end_time: datetime | None = None  # 直播下播时间
+
+    room_info: Optional["StreamRoomInfo"] | None = Relationship(
+        back_populates="status", sa_relationship_kwargs={"lazy": "selectin"}
+    )
+
+    """直播间信息，数据库保存时的数据结构"""
+
+
+class StreamRoomInfo(SQLModel, table=True):
+
+    __tablename__ = "stream_room_info"
+
+    room_id: int | None = Field(default=None, primary_key=True, unique=True)  # 直播间 ID
+
+    name: str = ""  # 直播间名字
+
+    product_list: list[SalesDocAndVideoInfo] = Relationship(
+        back_populates="stream_room",
+        sa_relationship_kwargs={"lazy": "selectin", "order_by": "asc(SalesDocAndVideoInfo.product_id)"},
+    )  # 商品列表，查找的时候加上 order_by 自动排序，desc -> 降序; asc -> 升序
+
+    prohibited_words_id: int = 0  # 违禁词表 ID
+    room_poster: str = ""  # 海报图
+    background_image: str = ""  # 主播背景图
+
+    delete: bool = False  # 是否删除
+
+    status_id: int | None = Field(default=None, foreign_key="on_air_room_status_item.status_id")
+    status: OnAirRoomStatusItem | None = Relationship(back_populates="room_info", sa_relationship_kwargs={"lazy": "selectin"})
+
+    streamer_id: int | None = Field(default=None, foreign_key="streamer_info.streamer_id")  # 主播 ID
+    streamer_info: StreamerInfo | None = Relationship(back_populates="room_info", sa_relationship_kwargs={"lazy": "selectin"})
+
+    user_id: int | None = Field(default=None, foreign_key="user_info.user_id")
+
+
+# =======================================================
+#                    直播对话数据库模型
+# =======================================================
+
+
+class ChatMessageInfo(SQLModel, table=True):
+    """直播页面对话数据结构"""
+
+    __tablename__ = "chat_message_info"
+
+    message_id: int | None = Field(default=None, primary_key=True, unique=True)  # 消息 ID
+
+    sales_info_id: int | None = Field(default=None, foreign_key="sales_doc_and_video_info.sales_info_id")
+    sales_info: SalesDocAndVideoInfo | None = Relationship(sa_relationship_kwargs={"lazy": "selectin"})
+
+    user_id: int | None = Field(default=None, foreign_key="user_info.user_id")
+    user_info: UserInfo | None = Relationship(sa_relationship_kwargs={"lazy": "selectin"})
+
+    streamer_id: int | None = Field(default=None, foreign_key="streamer_info.streamer_id")
+    streamer_info: StreamerInfo | None = Relationship(sa_relationship_kwargs={"lazy": "selectin"})
+
+    role: str
+    message: str
+    send_time: datetime | None = None

server/base/models/user_model.py

@@ -0,0 +1,43 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   user_model.py
+@Time    :   2024/08/31
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   用户信息数据结构
+"""
+
+from datetime import datetime
+from ipaddress import IPv4Address
+from pydantic import BaseModel
+from sqlmodel import Field, SQLModel
+
+
+# =======================================================
+#                      基本模型
+# =======================================================
+class TokenItem(BaseModel):
+    access_token: str
+    token_type: str
+
+
+class UserBaseInfo(BaseModel):
+    user_id: int | None = Field(default=None, primary_key=True, unique=True)
+    username: str = Field(index=True, unique=True)
+    email: str | None = None
+    avatar: str | None = None
+    create_time: datetime = datetime.now()
+
+
+# =======================================================
+#                      数据库模型
+# =======================================================
+class UserInfo(UserBaseInfo, SQLModel, table=True):
+
+    __tablename__ = "user_info"
+
+    hashed_password: str
+    ip_address: IPv4Address | None = None
+    delete: bool = False

server/base/modules/agent/agent_worker.py

@@ -0,0 +1,200 @@
+import json
+import logging
+
+from lagent.actions import ActionExecutor
+from lagent.agents.internlm2_agent import Internlm2Protocol
+from lagent.schema import ActionReturn, AgentReturn
+from loguru import logger
+
+from .delivery_time_query import DeliveryTimeQueryAction
+
+
+def init_handlers(departure_place, delivery_company_name):
+    META_CN = "当开启工具以及代码时，根据需求选择合适的工具进行调用"
+
+    INTERPRETER_CN = (
+        "你现在已经能够在一个有状态的 Jupyter 笔记本环境中运行 Python 代码。"
+        "当你向 python 发送含有 Python 代码的消息时，它将在该环境中执行。"
+        "这个工具适用于多种场景，如数据分析或处理（包括数据操作、统计分析、图表绘制），"
+        "复杂的计算问题（解决数学和物理难题），编程示例（理解编程概念或特性），"
+        "文本处理和分析（比如文本解析和自然语言处理），"
+        "机器学习和数据科学（用于展示模型训练和数据可视化），"
+        "以及文件操作和数据导入（处理CSV、JSON等格式的文件）。"
+    )
+
+    PLUGIN_CN = (
+        "你可以使用如下工具："
+        "\n{prompt}\n"
+        "如果你已经获得足够信息，请直接给出答案. 避免不必要的工具调用! "
+        "同时注意你可以使用的工具，不要随意捏造！"
+    )
+
+    protocol_handler = Internlm2Protocol(
+        meta_prompt=META_CN,
+        interpreter_prompt=INTERPRETER_CN,
+        plugin_prompt=PLUGIN_CN,
+        tool=dict(
+            begin="{start_token}{name}\n",
+            start_token="<|action_start|>",
+            name_map=dict(plugin="<|plugin|>", interpreter="<|interpreter|>"),
+            belong="assistant",
+            end="<|action_end|>\n",
+        ),
+    )
+    action_list = [
+        DeliveryTimeQueryAction(
+            departure_place=departure_place,
+            delivery_company_name=delivery_company_name,
+        ),
+    ]
+    plugin_map = {action.name: action for action in action_list}
+    plugin_name = [action.name for action in action_list]
+    plugin_action = [plugin_map[name] for name in plugin_name]
+    action_executor = ActionExecutor(actions=plugin_action)
+
+    return action_executor, protocol_handler
+
+
+def get_agent_result(llm_model_handler, prompt_input, departure_place, delivery_company_name):
+
+    action_executor, protocol_handler = init_handlers(departure_place, delivery_company_name)
+
+    # 第一次将 prompt 生成 agent 形式的 prompt
+    # [{'role': 'system', 'content': '当开启工具以及代码时，根据需求选择合适的工具进行调用'},
+    # {'role': 'system', 'content': '你可以使用如下工具：\n[\n    {\n        "name": "ArxivSearch.get_arxiv_article_information",\n
+    #                                                                       "description": "This is the subfunction for tool \'ArxivSearch\', you can use this tool. The description of this function is: \\nRun Arxiv search and get the article meta information.",\n
+    #                                                                       "parameters": [\n            {\n                "name": "query",\n                "type": "STRING",\n                "description": "the content of search query"\n            }\n        ],\n        "required": [\n            "query"\n        ],\n        "return_data": [\n            {\n                "name": "content",\n                "description": "a list of 3 arxiv search papers",\n                "type": "STRING"\n            }\n        ],\n        "parameter_description": "If you call this tool, you must pass arguments in the JSON format {key: value}, where the key is the parameter name."\n    }\n]\n
+    #                                                       如果你已经获得足够信息，请直接给出答案. 避免不必要的工具调用! 同时注意你可以使用的工具，不要随意捏造！',
+    #                                       'name': 'plugin'},
+    # {'role': 'user', 'content': '帮我搜索 InternLM2 Technical Report'}]
+
+    # 推理得出：'<|action_start|><|plugin|>\n{"name": "ArxivSearch.get_arxiv_article_information", "parameters": {"query": "InternLM2 Technical Report"}}<|action_end|>\n'
+    # 放入 assient 中
+
+    # 使用 ArxivSearch.get_arxiv_article_information 方法得出结果，放到 envrinment 里面，结果是：
+    # [{'role': 'system', 'content': '当开启工具以及代码时，根据需求选择合适的工具进行调用'},
+    # {'role': 'system', 'content': '你可以使用如下工具：\n[\n    {\n        "name": "ArxivSearch.get_arxiv_article_information",\n        "description": "This is the subfunction for tool \'ArxivSearch\', you can use this tool. The description of this function is: \\nRun Arxiv search and get the article meta information.",\n        "parameters": [\n            {\n                "name": "query",\n                "type": "STRING",\n                "description": "the content of search query"\n            }\n        ],\n        "required": [\n            "query"\n        ],\n        "return_data": [\n            {\n                "name": "content",\n                "description": "a list of 3 arxiv search papers",\n                "type": "STRING"\n            }\n        ],\n        "parameter_description": "If you call this tool, you must pass arguments in the JSON format {key: value}, where the key is the parameter name."\n    }\n]\n如果你已经获得足够信息，请直接给出答案. 避免不必要的工具调用! 同时注意你可以使用的工具，不要随意捏造！', 'name': 'plugin'},
+    # {'role': 'user', 'content': '帮我搜索 InternLM2 Technical Report'},
+    # {'role': 'assistant', 'content': '<|action_start|><|plugin|>\n{"name": "ArxivSearch.get_arxiv_article_information", "parameters": {"query": "InternLM2 Technical Report"}}<|action_end|>\n'},
+    # {'role': 'environment', 'content': '{"content": "Published: 2024-03-26\\nTitle: InternLM2 Technical Report\\nAuthors: Zheng Cai, Maosong Cao, Haojiong Chen, Kai Chen, Keyu Chen, Xin Chen, Xun Chen, Zehui Chen, Zhi Chen, Pei Chu, Xiaoyi Dong, Haodong Duan, Qi Fan, Zhaoye Fei, Yang Gao, Jiaye Ge, Chenya Gu, Yuzhe Gu, Tao Gui, Aijia Guo, Qipeng Guo, Conghui He, Yingfan Hu, Ting Huang, Tao Jiang, Penglong Jiao, Zhenjiang Jin, Zhikai Lei, Jiaxing Li, Jingwen Li, Linyang Li, Shuaibin Li, Wei Li, Yining Li, Hongwei Liu, Jiangning Liu, Jiawei Hong, Kaiwen Liu, Kuikun Liu, Xiaoran Liu, Chengqi Lv, Haijun Lv, Kai Lv, Li Ma, Runyuan Ma, Zerun Ma, Wenchang Ning, Linke Ouyang, Jiantao Qiu, Yuan Qu, Fukai Shang, Yunfan Shao, Demin Song, Zifan Song, Zhihao Sui, Peng Sun, Yu Sun, Huanze Tang, Bin Wang, Guoteng Wang, Jiaqi Wang, Jiayu Wang, Rui Wang, Yudong Wang, Ziyi Wang, Xingjian Wei, Qizhen Weng, Fan Wu, Yingtong Xiong, Chao Xu, Ruiliang Xu, Hang Yan, Yirong Yan, Xiaogui Yang, Haochen Ye, Huaiyuan Ying, Jia Yu, Jing Yu, Yuhang Zang, Chuyu Zhang, Li Zhang, Pan Zhang, Peng Zhang, Ruijie Zhang, Shuo Zhang, Songyang Zhang, Wenjian Zhang, Wenwei Zhang, Xingcheng Zhang, Xinyue Zhang, Hui Zhao, Qian Zhao, Xiaomeng Zhao, Fengzhe Zhou, Zaida Zhou, Jingming Zhuo, Yicheng Zou, Xipeng Qiu, Yu Qiao, Dahua Lin\\nSummary: The evolution of Large Language Models (LLMs) like ChatGPT and GPT-4 has\\nsparked discussions on the advent of Artificial General Intelligence (AGI).\\nHowever, replicating such advancements in open-source models has been\\nchallenging. This paper introduces InternLM2, an open-source LLM that\\noutperforms its predecessors in comprehensive evaluations across 6 dimensions\\nand 30 benchmarks, long-context modeling, and open-ended subjective evaluations\\nthrough innovative pre-training and optimization techniques. The pre-training\\nprocess of InternLM2 is meticulously detailed, highlighting the preparation of\\ndiverse data types including text, code, and long-context data. InternLM2\\nefficiently captures long-term dependencies, initially trained on 4k tokens\\nbefore advancing to 32k tokens in pre-training and fine-tuning stages,\\nexhibiting remarkable performance on the 200k ``Needle-in-a-Haystack\\" test.\\nInternLM2 is further aligned using Supervised Fine-Tuning (SFT) and a novel\\nConditional Online Reinforcement Learning from Human Feedback (COOL RLHF)\\nstrategy that addresses conflicting human preferences and reward hacking. By\\nreleasing InternLM2 models in different training stages and model sizes, we\\nprovide the community with insights into the model\'s evolution.\\n\\nPublished: 2017-07-27\\nTitle: Cumulative Reports of the SoNDe Project July 2017\\nAuthors: Sebastian Jaksch, Ralf Engels, Günter Kemmerling, Codin Gheorghe, Philip Pahlsson, Sylvain Désert, Frederic Ott\\nSummary: This are the cumulated reports of the SoNDe detector Project as of July 2017.\\nThe contained reports are: - Report on the 1x1 module technical demonstrator -\\nReport on used materials - Report on radiation hardness of components - Report\\non potential additional applications - Report on the 2x2 module technical\\ndemonstrator - Report on test results of the 2x2 technical demonstrator\\n\\nPublished: 2023-03-12\\nTitle: Banach Couples. I. Elementary Theory\\nAuthors: Jaak Peetre, Per Nilsson\\nSummary: This note is an (exact) copy of the report of Jaak Peetre, \\"Banach Couples.\\nI. Elementary Theory\\". Published as Technical Report, Lund (1971). Some more\\nrecent general references have been added and some references updated though"}', 'name': 'plugin'}]
+
+    # 然后调用大模型推理总结，stream 输出
+
+    # 判断 name is None ，跳出循环
+    inner_history = [{"role": "user", "content": prompt_input}]
+    interpreter_executor = None
+    max_turn = 7
+    for _ in range(max_turn):
+
+        prompt = protocol_handler.format(  # 生成 agent prompt
+            inner_step=inner_history,
+            plugin_executor=action_executor,
+            interpreter_executor=interpreter_executor,
+        )
+        cur_response = ""
+
+        agent_return = AgentReturn()
+
+        # 根据 tokenizer_config.json 中查找到特殊的 token ：
+        # token_map = {
+        #     92538: "<|plugin|>",
+        #     92539: "<|interpreter|>",
+        #     92540: "<|action_end|>",
+        #     92541: "<|action_start|>",
+        # }
+
+        # 将 prompt 给模型
+        # [{'role': 'system', 'content': '当开启工具以及代码时，根据需求选择合适的工具进行调用'},
+        # {'role': 'system', 'content': '你可以使用如下工具：\n[\n    {\n        "name": "ArxivSearch.get_arxiv_article_information",\n
+        #                                                                       "description": "This is the subfunction for tool \'ArxivSearch\', you can use this tool. The description of this function is: \\nRun Arxiv search and get the article meta information.",\n
+        #                                                                       "parameters": [\n            {\n                "name": "query",\n                "type": "STRING",\n                "description": "the content of search query"\n            }\n        ],\n        "required": [\n            "query"\n        ],\n        "return_data": [\n            {\n                "name": "content",\n                "description": "a list of 3 arxiv search papers",\n                "type": "STRING"\n            }\n        ],\n        "parameter_description": "If you call this tool, you must pass arguments in the JSON format {key: value}, where the key is the parameter name."\n    }\n]\n
+        #                                                       如果你已经获得足够信息，请直接给出答案. 避免不必要的工具调用! 同时注意你可以使用的工具，不要随意捏造！',
+        #                                       'name': 'plugin'},
+        # {'role': 'user', 'content': '帮我搜索 InternLM2 Technical Report'}]
+
+        # skip_special_tokens = False 输出 <|action_start|> <|plugin|> 等特殊字符
+        # for item in model_pipe.stream_infer(prompt, gen_config=prepare_generation_config(skip_special_tokens=False)):
+
+        logger.info(f"agent input for llm: {prompt}")
+
+        model_name = llm_model_handler.available_models[0]
+        for item in llm_model_handler.chat_completions_v1(
+            model=model_name, messages=prompt, stream=True, skip_special_tokens=False
+        ):
+            # 从 prompt 推理结果例子：
+            # '<|action_start|><|plugin|>\n{"name": "ArxivSearch.get_arxiv_article_information", "parameters": {"query": "InternLM2 Technical Report"}}<|action_end|>\n'
+
+            logger.info(f"agent return = {item}")
+            if "content" not in item["choices"][0]["delta"]:
+                continue
+            current_res = item["choices"][0]["delta"]["content"]
+
+            if "~" in current_res:
+                current_res = item.text.replace("~", "。").replace("。。", "。")
+
+            cur_response += current_res
+
+            logger.info(f"agent return = {item}")
+
+            name, language, action = protocol_handler.parse(
+                message=cur_response,
+                plugin_executor=action_executor,
+                interpreter_executor=interpreter_executor,
+            )
+            if name:  # "plugin"
+                if name == "plugin":
+                    if action_executor:
+                        executor = action_executor
+                    else:
+                        logging.info(msg="No plugin is instantiated!")
+                        continue
+                    try:
+                        action = json.loads(action)
+                    except Exception as e:
+                        logging.info(msg=f"Invaild action {e}")
+                        continue
+                elif name == "interpreter":
+                    if interpreter_executor:
+                        executor = interpreter_executor
+                    else:
+                        logging.info(msg="No interpreter is instantiated!")
+                        continue
+                # agent_return.state = agent_state
+                agent_return.response = action
+
+        print(f"Agent response: {cur_response}")
+
+        if name:
+            print(f"Agent action: {action}")
+            action_return: ActionReturn = executor(action["name"], action["parameters"])
+            # action_return.thought = language
+            # agent_return.actions.append(action_return)
+            try:
+                return_str = action_return.result[0]["content"]
+                return return_str
+            except Exception as e:
+                return ""
+
+            # agent_return_list.append(dict(role='assistant', name=name, content=action))
+            # agent_return_list.append(protocol_handler.format_response(action_return, name=name))
+
+        # inner_history.append(dict(role="language", content=language))
+
+        if not name:
+            agent_return.response = language
+            break
+        # elif action_return.type == executor.finish_action.name:
+        #     try:
+        #         response = action_return.args["text"]["response"]
+        #     except Exception:
+        #         logging.info(msg="Unable to parse FinishAction.")
+        #         response = ""
+        #     agent_return.response = response
+        #     break
+        # else:
+        #     inner_history.append(dict(role="tool", content=action, name=name))
+        #     inner_history.append(protocol_handler.format_response(action_return, name=name))
+        #     # agent_state += 1
+        #     # agent_return.state = agent_state
+        #     # yield agent_return
+    return ""

server/base/modules/agent/delivery_time_query.py

@@ -0,0 +1,300 @@
+import base64
+from datetime import datetime
+import hashlib
+import json
+from typing import Optional, Type
+
+import jionlp as jio
+import requests
+from lagent.actions.base_action import BaseAction, tool_api
+from lagent.actions.parser import BaseParser, JsonParser
+from lagent.schema import ActionReturn, ActionStatusCode
+
+from ....web_configs import WEB_CONFIGS
+
+
+class DeliveryTimeQueryAction(BaseAction):
+    """快递时效查询插件，用于根据用户提出的收货地址查询到达期限"""
+
+    def __init__(
+        self,
+        departure_place: str,
+        delivery_company_name: str,
+        description: Optional[dict] = None,
+        parser: Type[BaseParser] = JsonParser,
+        enable: bool = True,
+    ) -> None:
+        super().__init__(description, parser, enable)
+        self.departure_place = departure_place  # 发货地
+
+        # 天气查询
+        self.weather_query_handler = WeatherQuery(departure_place, WEB_CONFIGS.AGENT_WEATHER_API_KEY)
+        self.delivery_time_handler = DeliveryTimeQuery(delivery_company_name, WEB_CONFIGS.AGENT_DELIVERY_TIME_API_KEY)
+
+    @tool_api
+    def run(self, query: str) -> ActionReturn:
+        """一个到货时间查询API。可以根据城市名查询到货时间信息。
+
+        Args:
+            query (:class:`str`): 需要查询的城市名。
+        """
+
+        # 获取文本中收货地，发货地后台设置
+        # 防止 LLM 将城市识别错误，进行兜底
+        city_info = jio.parse_location(query, town_village=True)
+        city_name = city_info["city"]
+
+        # 获取收货地代号 -> 天气
+        destination_weather = self.weather_query_handler(city_name)
+
+        # 获取发货地代号 -> 天气
+        departure_weather = self.weather_query_handler(self.departure_place)
+
+        # 获取到达时间
+        delivery_time = self.delivery_time_handler(self.departure_place, city_name)
+
+        final_str = (
+            f"今天日期：{datetime.now().strftime('%m月%d日')}\n"
+            f"收货地天气：{destination_weather.result[0]['content']}\n"
+            f"发货地天气：{departure_weather.result[0]['content']}\n"
+            f"物流信息：{delivery_time.result[0]['content']}\n"
+            "回答突出“预计送达时间”和“收货地天气”，如果收货地或者发货地遇到暴雨暴雪等极端天气，须告知用户快递到达时间会有所增加。"
+        )
+
+        tool_return = ActionReturn(type=self.name)
+        tool_return.result = [dict(type="text", content=final_str)]
+        return tool_return
+
+
+class WeatherQuery:
+    """快递时效查询插件，用于根据用户提出的收货地址查询到达期限"""
+
+    def __init__(
+        self,
+        departure_place: str,
+        api_key: Optional[str] = None,
+    ) -> None:
+        self.departure_place = departure_place  # 发货地
+
+        # 天气查询
+        # api_key = os.environ.get("WEATHER_API_KEY", key)
+        if api_key is None:
+            raise ValueError("Please set Weather API key either in the environment as WEATHER_API_KEY")
+        self.api_key = api_key
+        self.location_query_url = "https://geoapi.qweather.com/v2/city/lookup"
+        self.weather_query_url = "https://devapi.qweather.com/v7/weather/now"
+
+    def parse_results(self, city_name: str, results: dict) -> str:
+        """解析 API 返回的信息
+
+        Args:
+            results (dict): JSON 格式的 API 报文。
+
+        Returns:
+            str: 解析后的结果。
+        """
+        now = results["now"]
+        data = (
+            # f'数据观测时间: {now["obsTime"]}；'
+            f"城市名: {city_name}；"
+            f'温度: {now["temp"]}°C；'
+            f'体感温度: {now["feelsLike"]}°C；'
+            f'天气: {now["text"]}；'
+            # f'风向: {now["windDir"]}，角度为 {now["wind360"]}°；'
+            f'风力等级: {now["windScale"]}，风速为 {now["windSpeed"]} km/h；'
+            f'相对湿度: {now["humidity"]}；'
+            f'当前小时累计降水量: {now["precip"]} mm；'
+            # f'大气压强: {now["pressure"]} 百帕；'
+            f'能见度: {now["vis"]} km。'
+        )
+        return data
+
+    def __call__(self, query):
+        tool_return = ActionReturn()
+        status_code, response = self.search_weather_with_city(query)
+        if status_code == -1:
+            tool_return.errmsg = response
+            tool_return.state = ActionStatusCode.HTTP_ERROR
+        elif status_code == 200:
+            parsed_res = self.parse_results(query, response)
+            tool_return.result = [dict(type="text", content=str(parsed_res))]
+            tool_return.state = ActionStatusCode.SUCCESS
+        else:
+            tool_return.errmsg = str(status_code)
+            tool_return.state = ActionStatusCode.API_ERROR
+        return tool_return
+
+    def search_weather_with_city(self, query: str):
+        """根据城市名获取城市代号，然后进行天气���询
+
+        Args:
+            query (str): 城市名
+
+        Returns:
+            int: 天气接口调用状态码
+            dict: 天气接口返回信息
+        """
+
+        # 获取城市代号
+        try:
+            city_code_response = requests.get(self.location_query_url, params={"key": self.api_key, "location": query})
+        except Exception as e:
+            return -1, str(e)
+
+        if city_code_response.status_code != 200:
+            return city_code_response.status_code, city_code_response.json()
+        city_code_response = city_code_response.json()
+        if len(city_code_response["location"]) == 0:
+            return -1, "未查询到城市"
+        city_code = city_code_response["location"][0]["id"]
+
+        # 获取天气
+        try:
+            weather_response = requests.get(self.weather_query_url, params={"key": self.api_key, "location": city_code})
+        except Exception as e:
+            return -1, str(e)
+        return weather_response.status_code, weather_response.json()
+
+
+class DeliveryTimeQuery:
+    def __init__(
+        self,
+        delivery_company_name: Optional[str] = "中通",
+        api_key: Optional[str] = None,
+    ) -> None:
+
+        # 快递时效查询
+        # api_key = os.environ.get("DELIVERY_TIME_API_KEY", key)
+        if api_key is None or "," not in api_key:
+            raise ValueError(
+                'Please set Delivery time API key either in the environment as DELIVERY_TIME_API_KEY="${e_business_id},${api_key}"'
+            )
+        self.e_business_id = api_key.split(",")[0]
+        self.api_key = api_key.split(",")[1]
+        self.api_url = "http://api.kdniao.com/api/dist"  # 快递鸟
+        self.china_location = jio.china_location_loader()
+        # 快递鸟对应的
+        DELIVERY_COMPANY_MAP = {
+            "德邦": "DBL",
+            "邮政": "EMS",
+            "京东": "JD",
+            "极兔速递": "JTSD",
+            "顺丰": "SF",
+            "申通": "STO",
+            "韵达": "YD",
+            "圆通": "YTO",
+            "中通": "ZTO",
+        }
+        self.delivery_company_name = delivery_company_name
+        self.delivery_company_id = DELIVERY_COMPANY_MAP[delivery_company_name]
+
+    @staticmethod
+    def data_md5(n):
+        # md5加密
+        md5 = hashlib.md5()
+        md5.update(str(n).encode("utf-8"))
+        return md5.hexdigest()
+
+    def get_data_sign(self, n):
+        # 签名
+        md5Data = self.data_md5(json.dumps(n) + self.api_key)
+        res = str(base64.b64encode(md5Data.encode("utf-8")), "utf-8")
+        return res
+
+    def get_city_detail(self, name):
+        # 如果是城市名，使用第一个区名
+        city_info = jio.parse_location(name, town_village=True)
+        # china_location = jio.china_location_loader()
+
+        county_name = ""
+        for i in self.china_location[city_info["province"]][city_info["city"]].keys():
+            if "区" == i[-1]:
+                county_name = i
+                break
+
+        return {
+            "province": city_info["province"],
+            "city": city_info["city"],
+            "county": county_name,
+        }
+
+    def get_params(self, send_city, receive_city):
+
+        # 根据市查出省份和区名称
+        send_city_info = self.get_city_detail(send_city)
+        receive_city_info = self.get_city_detail(receive_city)
+
+        # 预计送达时间接口文档；https://www.yuque.com/kdnjishuzhichi/dfcrg1/ynkmts0e5owsnpvu
+        # 请求接口指令
+        RequestType = "6004"
+        # 组装应用级参数
+        RequestData = {
+            "ShipperCode": self.delivery_company_id,
+            "ReceiveArea": receive_city_info["county"],
+            "ReceiveCity": receive_city_info["city"],
+            "ReceiveProvince": receive_city_info["province"],
+            "SendArea": send_city_info["county"],
+            "SendCity": send_city_info["city"],
+            "SendProvince": send_city_info["province"],
+        }
+        # 组装系统级参数
+        data = {
+            "RequestData": json.dumps(RequestData),
+            "RequestType": RequestType,
+            "EBusinessID": self.e_business_id,
+            "DataSign": self.get_data_sign(RequestData),
+            "DataType": 2,
+        }
+        return data
+
+    def parse_results(self, response):
+
+        # 返回例子：
+        # {
+        # "EBusinessID" : "1000000",
+        # "Data" : {
+        #     "DeliveryTime" : "06月15日下午可达",
+        #     "SendAddress" : null,
+        #     "ReceiveArea" : "芙蓉区",
+        #     "SendProvince" : "广东省",
+        #     "ReceiveProvince" : "湖南省",
+        #     "ShipperCode" : "DBL",
+        #     "Hour" : "52h",
+        #     "SendArea" : "白云区",
+        #     "ReceiveAddress" : null,
+        #     "SendCity" : "广州市",
+        #     "ReceiveCity" : "长沙市"
+        # },
+        # "ResultCode" : "100",
+        # "Success" : true
+        # }
+
+        response = response["Data"]
+        data = (
+            f'发货地点: {response["SendProvince"]} {response["SendCity"]}；'
+            f'收货地点: {response["ReceiveProvince"]} {response["ReceiveCity"]}；'
+            f'预计送达时间: {response["DeliveryTime"]}；'
+            f"快递公司: {self.delivery_company_name}；"
+            f'预计时效: {response["Hour"]}。'
+        )
+        return data
+
+    def __call__(self, send_city, receive_city):
+        tool_return = ActionReturn()
+        try:
+            res = requests.post(self.api_url, self.get_params(send_city, receive_city))
+            status_code = res.status_code
+            response = res.json()
+        except Exception as e:
+            tool_return.errmsg = str(e)
+            tool_return.state = ActionStatusCode.API_ERROR
+            return tool_return
+
+        if status_code == 200:
+            parsed_res = self.parse_results(response)
+            tool_return.result = [dict(type="text", content=str(parsed_res))]
+            tool_return.state = ActionStatusCode.SUCCESS
+        else:
+            tool_return.errmsg = str(status_code)
+            tool_return.state = ActionStatusCode.API_ERROR
+        return tool_return

server/base/modules/rag/feature_store.py

@@ -0,0 +1,545 @@
+"""extract feature and search with user query."""
+
+import argparse
+import json
+import os
+import re
+import shutil
+from multiprocessing import Pool
+from pathlib import Path
+from typing import Any, List, Optional
+
+import yaml
+
+# 解决 Warning：huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks…
+# os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+from BCEmbedding.tools.langchain import BCERerank
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.text_splitter import (MarkdownHeaderTextSplitter,
+                                     MarkdownTextSplitter,
+                                     RecursiveCharacterTextSplitter)
+from langchain.vectorstores.faiss import FAISS as Vectorstore
+from langchain_core.documents import Document
+from loguru import logger
+from torch.cuda import empty_cache
+
+from .file_operation import FileName, FileOperation
+from .retriever import CacheRetriever, Retriever
+
+
+def read_and_save(file: FileName):
+    if os.path.exists(file.copypath):
+        # already exists, return
+        logger.info("already exist, skip load")
+        return
+    file_opr = FileOperation()
+    logger.info("reading {}, would save to {}".format(file.origin, file.copypath))
+    content, error = file_opr.read(file.origin)
+    if error is not None:
+        logger.error("{} load error: {}".format(file.origin, str(error)))
+        return
+
+    if content is None or len(content) < 1:
+        logger.warning("{} empty, skip save".format(file.origin))
+        return
+
+    with open(file.copypath, "w") as f:
+        f.write(content)
+
+
+def _split_text_with_regex_from_end(text: str, separator: str, keep_separator: bool) -> List[str]:
+    # Now that we have the separator, split the text
+    if separator:
+        if keep_separator:
+            # The parentheses in the pattern keep the delimiters in the result.
+            _splits = re.split(f"({separator})", text)
+            splits = ["".join(i) for i in zip(_splits[0::2], _splits[1::2])]
+            if len(_splits) % 2 == 1:
+                splits += _splits[-1:]
+            # splits = [_splits[0]] + splits
+        else:
+            splits = re.split(separator, text)
+    else:
+        splits = list(text)
+    return [s for s in splits if s != ""]
+
+
+# copy from https://github.com/chatchat-space/Langchain-Chatchat/blob/master/text_splitter/chinese_recursive_text_splitter.py
+class ChineseRecursiveTextSplitter(RecursiveCharacterTextSplitter):
+
+    def __init__(
+        self,
+        separators: Optional[List[str]] = None,
+        keep_separator: bool = True,
+        is_separator_regex: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        """Create a new TextSplitter."""
+        super().__init__(keep_separator=keep_separator, **kwargs)
+        self._separators = separators or ["\n\n", "\n", "。|！|？", "\.\s|\!\s|\?\s", "；|;\s", "，|,\s"]
+        self._is_separator_regex = is_separator_regex
+
+    def _split_text(self, text: str, separators: List[str]) -> List[str]:
+        """Split incoming text and return chunks."""
+        final_chunks = []
+        # Get appropriate separator to use
+        separator = separators[-1]
+        new_separators = []
+        for i, _s in enumerate(separators):
+            _separator = _s if self._is_separator_regex else re.escape(_s)
+            if _s == "":
+                separator = _s
+                break
+            if re.search(_separator, text):
+                separator = _s
+                new_separators = separators[i + 1 :]
+                break
+
+        _separator = separator if self._is_separator_regex else re.escape(separator)
+        splits = _split_text_with_regex_from_end(text, _separator, self._keep_separator)
+
+        # Now go merging things, recursively splitting longer texts.
+        _good_splits = []
+        _separator = "" if self._keep_separator else separator
+        for s in splits:
+            if self._length_function(s) < self._chunk_size:
+                _good_splits.append(s)
+            else:
+                if _good_splits:
+                    merged_text = self._merge_splits(_good_splits, _separator)
+                    final_chunks.extend(merged_text)
+                    _good_splits = []
+                if not new_separators:
+                    final_chunks.append(s)
+                else:
+                    other_info = self._split_text(s, new_separators)
+                    final_chunks.extend(other_info)
+        if _good_splits:
+            merged_text = self._merge_splits(_good_splits, _separator)
+            final_chunks.extend(merged_text)
+        return [re.sub(r"\n{2,}", "\n", chunk.strip()) for chunk in final_chunks if chunk.strip() != ""]
+
+
+class FeatureStore:
+    """Tokenize and extract features from the project's documents, for use in
+    the reject pipeline and response pipeline."""
+
+    def __init__(
+        self, embeddings: HuggingFaceEmbeddings, reranker: BCERerank, config_path: str = "config.ini", language: str = "zh"
+    ) -> None:
+        """Init with model device type and config."""
+        self.config_path = config_path
+        self.reject_throttle = -1
+        self.language = language
+        with open(config_path, "r", encoding="utf-8") as f:
+            config = yaml.safe_load(f)["feature_store"]
+            self.reject_throttle = config["reject_throttle"]
+
+        logger.warning(
+            "!!! If your feature generated by `text2vec-large-chinese` before 20240208, please rerun `python3 -m huixiangdou.service.feature_store`"  # noqa E501
+        )
+
+        logger.debug("loading text2vec model..")
+        self.embeddings = embeddings
+        self.reranker = reranker
+        self.compression_retriever = None
+        self.rejecter = None
+        self.retriever = None
+        self.md_splitter = MarkdownTextSplitter(chunk_size=768, chunk_overlap=32)
+
+        if language == "zh":
+            self.text_splitter = ChineseRecursiveTextSplitter(
+                keep_separator=True, is_separator_regex=True, chunk_size=768, chunk_overlap=32
+            )
+        else:
+            self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=768, chunk_overlap=32)
+
+        self.head_splitter = MarkdownHeaderTextSplitter(
+            headers_to_split_on=[
+                ("#", "Header 1"),
+                ("##", "Header 2"),
+                ("###", "Header 3"),
+            ]
+        )
+
+    def split_md(self, text: str, source: None):
+        """Split the markdown document in a nested way, first extracting the
+        header.
+
+        If the extraction result exceeds 1024, split it again according to
+        length.
+        """
+        docs = self.head_splitter.split_text(text)
+
+        final = []
+        for doc in docs:
+            header = ""
+            if len(doc.metadata) > 0:
+                if "Header 1" in doc.metadata:
+                    header += doc.metadata["Header 1"]
+                if "Header 2" in doc.metadata:
+                    header += " "
+                    header += doc.metadata["Header 2"]
+                if "Header 3" in doc.metadata:
+                    header += " "
+                    header += doc.metadata["Header 3"]
+
+            if len(doc.page_content) >= 1024:
+                subdocs = self.md_splitter.create_documents([doc.page_content])
+                for subdoc in subdocs:
+                    if len(subdoc.page_content) >= 10:
+                        final.append("{} {}".format(header, subdoc.page_content.lower()))
+            elif len(doc.page_content) >= 10:
+                final.append("{} {}".format(header, doc.page_content.lower()))  # noqa E501
+
+        for item in final:
+            if len(item) >= 1024:
+                logger.debug("source {} split length {}".format(source, len(item)))
+        return final
+
+    def clean_md(self, text: str):
+        """Remove parts of the markdown document that do not contain the key
+        question words, such as code blocks, URL links, etc."""
+        # remove ref
+        pattern_ref = r"\[(.*?)\]\(.*?\)"
+        new_text = re.sub(pattern_ref, r"\1", text)
+
+        # remove code block
+        pattern_code = r"```.*?```"
+        new_text = re.sub(pattern_code, "", new_text, flags=re.DOTALL)
+
+        # remove underline
+        new_text = re.sub("_{5,}", "", new_text)
+
+        # remove table
+        # new_text = re.sub('\|.*?\|\n\| *\:.*\: *\|.*\n(\|.*\|.*\n)*', '', new_text, flags=re.DOTALL)   # noqa E501
+
+        # use lower
+        new_text = new_text.lower()
+        return new_text
+
+    def get_md_documents(self, file: FileName):
+        documents = []
+        length = 0
+        text = ""
+        with open(file.copypath, encoding="utf8") as f:
+            text = f.read()
+        text = file.prefix + "\n" + self.clean_md(text)
+        if len(text) <= 1:
+            return [], length
+
+        chunks = self.split_md(text=text, source=os.path.abspath(file.copypath))
+        for chunk in chunks:
+            new_doc = Document(page_content=chunk, metadata={"source": file.basename, "read": file.copypath})
+            length += len(chunk)
+            documents.append(new_doc)
+        return documents, length
+
+    def get_text_documents(self, text: str, file: FileName):
+        if len(text) <= 1:
+            return []
+        chunks = self.text_splitter.create_documents([text])
+        documents = []
+        for chunk in chunks:
+            # `source` is for return references
+            # `read` is for LLM response
+            chunk.metadata = {"source": file.basename, "read": file.copypath}
+            documents.append(chunk)
+        return documents
+
+    def ingress_response(self, files: list, work_dir: str):
+        """Extract the features required for the response pipeline based on the
+        document."""
+        feature_dir = os.path.join(work_dir, "db_response")
+        if not os.path.exists(feature_dir):
+            os.makedirs(feature_dir)
+
+        # logger.info('glob {} in dir {}'.format(files, file_dir))
+        file_opr = FileOperation()
+        documents = []
+
+        for i, file in enumerate(files):
+            logger.debug("{}/{}.. {}".format(i + 1, len(files), file.basename))
+            if not file.state:
+                continue
+
+            if file._type == "md":
+                md_documents, md_length = self.get_md_documents(file)
+                documents += md_documents
+                logger.info("{} content length {}".format(file._type, md_length))
+                file.reason = str(md_length)
+
+            else:
+                # now read pdf/word/excel/ppt text
+                text, error = file_opr.read(file.copypath)
+                if error is not None:
+                    file.state = False
+                    file.reason = str(error)
+                    continue
+                file.reason = str(len(text))
+                logger.info("{} content length {}".format(file._type, len(text)))
+                text = file.prefix + text
+                documents += self.get_text_documents(text, file)
+
+        if len(documents) < 1:
+            return
+        vs = Vectorstore.from_documents(documents, self.embeddings)
+        vs.save_local(feature_dir)
+
+    def ingress_reject(self, files: list, work_dir: str):
+        """Extract the features required for the reject pipeline based on
+        documents."""
+        feature_dir = os.path.join(work_dir, "db_reject")
+        if not os.path.exists(feature_dir):
+            os.makedirs(feature_dir)
+
+        documents = []
+        file_opr = FileOperation()
+
+        logger.debug("ingress reject..")
+        for i, file in enumerate(files):
+            if not file.state:
+                continue
+
+            if file._type == "md":
+                # reject base not clean md
+                text = file.basename + "\n"
+                with open(file.copypath, encoding="utf8") as f:
+                    text += f.read()
+                if len(text) <= 1:
+                    continue
+
+                chunks = self.split_md(text=text, source=os.path.abspath(file.copypath))
+                for chunk in chunks:
+                    new_doc = Document(page_content=chunk, metadata={"source": file.basename, "read": file.copypath})
+                    documents.append(new_doc)
+
+            else:
+                text, error = file_opr.read(file.copypath)
+                if error is not None:
+                    continue
+                text = file.basename + text
+                documents += self.get_text_documents(text, file)
+
+        if len(documents) < 1:
+            return
+        vs = Vectorstore.from_documents(documents, self.embeddings)
+        vs.save_local(feature_dir)
+
+    def preprocess(self, files: list, work_dir: str):
+        """Preprocesses files in a given directory. Copies each file to
+        'preprocess' with new name formed by joining all subdirectories with
+        '_'.
+
+        Args:
+            files (list): original file list.
+            work_dir (str): Working directory where preprocessed files will be stored.  # noqa E501
+
+        Returns:
+            str: Path to the directory where preprocessed markdown files are saved.
+
+        Raises:
+            Exception: Raise an exception if no markdown files are found in the provided repository directory.  # noqa E501
+        """
+        preproc_dir = os.path.join(work_dir, "preprocess")
+        if not os.path.exists(preproc_dir):
+            os.makedirs(preproc_dir)
+
+        pool = Pool(processes=16)
+        file_opr = FileOperation()
+        for idx, file in enumerate(files):
+            if not os.path.exists(file.origin):
+                file.state = False
+                file.reason = "skip not exist"
+                continue
+
+            if file._type == "image":
+                file.state = False
+                file.reason = "skip image"
+
+            elif file._type in ["pdf", "word", "excel", "ppt", "html"]:
+                # read pdf/word/excel file and save to text format
+                md5 = file_opr.md5(file.origin)
+                file.copypath = os.path.join(preproc_dir, "{}.text".format(md5))
+                pool.apply_async(read_and_save, (file,))
+
+            elif file._type in ["md", "text"]:
+                # rename text files to new dir
+                md5 = file_opr.md5(file.origin)
+                file.copypath = os.path.join(preproc_dir, file.origin.replace("/", "_")[-84:])
+                try:
+                    shutil.copy(file.origin, file.copypath)
+                    file.state = True
+                    file.reason = "preprocessed"
+                except Exception as e:
+                    file.state = False
+                    file.reason = str(e)
+
+            else:
+                file.state = False
+                file.reason = "skip unknown format"
+        pool.close()
+        logger.debug("waiting for preprocess read finish..")
+        pool.join()
+
+        # check process result
+        for file in files:
+            if file._type in ["pdf", "word", "excel"]:
+                if os.path.exists(file.copypath):
+                    file.state = True
+                    file.reason = "preprocessed"
+                else:
+                    file.state = False
+                    file.reason = "read error"
+
+    def initialize(self, files: list, work_dir: str):
+        """Initializes response and reject feature store.
+
+        Only needs to be called once. Also calculates the optimal threshold
+        based on provided good and bad question examples, and saves it in the
+        configuration file.
+        """
+        logger.info("initialize response and reject feature store, you only need call this once.")  # noqa E501
+        self.preprocess(files=files, work_dir=work_dir)
+        self.ingress_response(files=files, work_dir=work_dir)
+        self.ingress_reject(files=files, work_dir=work_dir)
+
+
+def parse_args():
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(description="Feature store for processing directories.")
+    parser.add_argument("--work_dir", type=str, default="work_dir", help="Working directory.")
+    parser.add_argument("--repo_dir", type=str, default="repodir", help="Root directory where the repositories are located.")
+    parser.add_argument(
+        "--config_path", default="config.ini", help="Feature store configuration path. Default value is config.ini"
+    )
+    parser.add_argument(
+        "--good_questions",
+        default="resource/good_questions.json",
+        help="Positive examples in the dataset. Default value is resource/good_questions.json",  # noqa E251  # noqa E501
+    )
+    parser.add_argument(
+        "--bad_questions",
+        default="resource/bad_questions.json",
+        help="Negative examples json path. Default value is resource/bad_questions.json",  # noqa E251  # noqa E501
+    )
+    parser.add_argument("--sample", help="Input an json file, save reject and search output.")
+    args = parser.parse_args()
+    return args
+
+
+def test_reject(retriever: Retriever, sample: str = None):
+    """Simple test reject pipeline."""
+    if sample is None:
+        real_questions = [
+            "SAM 10个T 的训练集，怎么比比较公平呢~？速度上还有缺陷吧？",
+            "想问下，如果只是推理的话，amp的fp16是不会省显存么，我看parameter仍然是float32，开和不开推理的显存占用都是一样的。能不能直接用把数据和model都 .half() 代替呢，相比之下amp好在哪里",  # noqa E501
+            "mmdeploy支持ncnn vulkan部署么，我只找到了ncnn cpu 版本",
+            "大佬们，如果我想在高空检测安全帽，我应该用 mmdetection 还是 mmrotate",
+            "请问 ncnn 全称是什么",
+            "有啥中文的 text to speech 模型吗?",
+            "今天中午吃什么？",
+            "huixiangdou 是什么？",
+            "mmpose 如何安装？",
+            "使用科研仪器需要注意什么？",
+        ]
+    else:
+        with open(sample) as f:
+            real_questions = json.load(f)
+
+    for example in real_questions:
+        reject, _ = retriever.is_reject(example)
+
+        if reject:
+            logger.error(f"reject query: {example}")
+        else:
+            logger.warning(f"process query: {example}")
+
+        if sample is not None:
+            if reject:
+                with open("workdir/negative.txt", "a+") as f:
+                    f.write(example)
+                    f.write("\n")
+            else:
+                with open("workdir/positive.txt", "a+") as f:
+                    f.write(example)
+                    f.write("\n")
+
+    empty_cache()
+
+
+def test_query(retriever: Retriever, sample: str = None):
+    """Simple test response pipeline."""
+    if sample is not None:
+        with open(sample) as f:
+            real_questions = json.load(f)
+        logger.add("logs/feature_store_query.log", rotation="4MB")
+    else:
+        real_questions = ["mmpose installation", "how to use std::vector ?"]
+
+    for example in real_questions:
+        example = example[0:400]
+        print(retriever.query(example))
+        empty_cache()
+
+    empty_cache()
+
+
+def fix_system_error():
+    """
+    Fix `No module named 'faiss.swigfaiss_avx2`
+    """
+    import os
+    from pathlib import Path
+
+    import faiss
+
+    if Path(faiss.__file__).parent.joinpath("swigfaiss_avx2.py").exists():
+        return
+
+    print("Fixing faiss error...")
+    os.system(f"cd {Path(faiss.__file__).parent} && ln -s swigfaiss.py swigfaiss_avx2.py")
+
+
+def gen_vector_db(config_path, source_dir, work_dir, test_mode=False, update_reject=False):
+
+    # 解决 faiss 导入问题
+    fix_system_error()
+
+    # 必须是绝对路径，否则加载会有问题
+    work_dir = str(Path(work_dir).absolute())
+
+    cache = CacheRetriever(config_path=config_path)
+
+    # 生成向量数据库
+    fs_init = FeatureStore(embeddings=cache.embeddings, reranker=cache.reranker, config_path=config_path)
+
+    # walk all files in repo dir
+    file_opr = FileOperation()
+    files = file_opr.scan_dir(repo_dir=source_dir)
+    fs_init.initialize(files=files, work_dir=work_dir)
+    file_opr.summarize(files)
+    del fs_init
+
+    # update reject throttle
+    if update_reject:
+        retriever = cache.get(config_path=config_path, work_dir=work_dir)
+        with open(os.path.join("resource", "good_questions.json")) as f:
+            good_questions = json.load(f)
+        with open(os.path.join("resource", "bad_questions.json")) as f:
+            bad_questions = json.load(f)
+        retriever.update_throttle(config_path=config_path, good_questions=good_questions, bad_questions=bad_questions)
+
+        cache.pop("default")
+
+    if test_mode:
+        # test
+        retriever = cache.get(config_path=config_path, work_dir=work_dir)
+        # test_reject(retriever, args.sample)
+        test_query(retriever, args.sample)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    gen_vector_db(args.config_path, args.repo_dir, args.work_dir, test_mode=True)

server/base/modules/rag/file_operation.py

@@ -0,0 +1,228 @@
+import hashlib
+import os
+
+import pandas as pd
+from bs4 import BeautifulSoup
+from loguru import logger
+
+
+class FileName:
+    """Record file original name, state and copied filepath with text
+    format."""
+
+    def __init__(self, root: str, filename: str, _type: str):
+        self.root = root
+        self.prefix = filename.replace("/", "_")
+        self.basename = os.path.basename(filename)
+        self.origin = os.path.join(root, filename)
+        self.copypath = ""
+        self._type = _type
+        self.state = True
+        self.reason = ""
+
+    def __str__(self):
+        return "{},{},{},{}\n".format(self.basename, self.copypath, self.state, self.reason)
+
+
+class FileOperation:
+    """Encapsulate all file reading operations."""
+
+    def __init__(self):
+        self.image_suffix = [".jpg", ".jpeg", ".png", ".bmp"]
+        self.md_suffix = ".md"
+        self.text_suffix = [".txt", ".text"]
+        self.excel_suffix = [".xlsx", ".xls", ".csv"]
+        self.pdf_suffix = ".pdf"
+        self.ppt_suffix = ".pptx"
+        self.html_suffix = [".html", ".htm", ".shtml", ".xhtml"]
+        self.word_suffix = [".docx", ".doc"]
+        self.normal_suffix = (
+            [self.md_suffix]
+            + self.text_suffix
+            + self.excel_suffix
+            + [self.pdf_suffix]
+            + self.word_suffix
+            + [self.ppt_suffix]
+            + self.html_suffix
+        )
+
+    def get_type(self, filepath: str):
+        filepath = filepath.lower()
+        if filepath.endswith(self.pdf_suffix):
+            return "pdf"
+
+        if filepath.endswith(self.md_suffix):
+            return "md"
+
+        if filepath.endswith(self.ppt_suffix):
+            return "ppt"
+
+        for suffix in self.image_suffix:
+            if filepath.endswith(suffix):
+                return "image"
+
+        for suffix in self.text_suffix:
+            if filepath.endswith(suffix):
+                return "text"
+
+        for suffix in self.word_suffix:
+            if filepath.endswith(suffix):
+                return "word"
+
+        for suffix in self.excel_suffix:
+            if filepath.endswith(suffix):
+                return "excel"
+
+        for suffix in self.html_suffix:
+            if filepath.endswith(suffix):
+                return "html"
+        return None
+
+    def md5(self, filepath: str):
+        hash_object = hashlib.sha256()
+        with open(filepath, "rb") as file:
+            chunk_size = 8192
+            while chunk := file.read(chunk_size):
+                hash_object.update(chunk)
+
+        return hash_object.hexdigest()[0:8]
+
+    def summarize(self, files: list):
+        success = 0
+        skip = 0
+        failed = 0
+
+        for file in files:
+            if file.state:
+                success += 1
+            elif file.reason == "skip":
+                skip += 1
+            else:
+                logger.info("{} {}".format(file.origin, file.reason))
+                failed += 1
+
+            logger.info("{} {}".format(file.reason, file.copypath))
+        logger.info("累计{}文件，成功{}个，跳过{}个，异常{}个".format(len(files), success, skip, failed))
+
+    def scan_dir(self, repo_dir: str):
+        files = []
+        for root, _, filenames in os.walk(repo_dir):
+            for filename in filenames:
+                _type = self.get_type(filename)
+                if _type is not None:
+                    files.append(FileName(root=root, filename=filename, _type=_type))
+        return files
+
+    def read_pdf(self, filepath: str):
+        # load pdf and serialize table
+
+        # TODO fitz 安装有些不兼容，后续按需完善
+        import fitz
+
+        text = ""
+        with fitz.open(filepath) as pages:
+            for page in pages:
+                text += page.get_text()
+                tables = page.find_tables()
+                for table in tables:
+                    tablename = "_".join(filter(lambda x: x is not None and "Col" not in x, table.header.names))
+                    pan = table.to_pandas()
+                    json_text = pan.dropna(axis=1).to_json(force_ascii=False)
+                    text += tablename
+                    text += "\n"
+                    text += json_text
+                    text += "\n"
+        return text
+
+    def read_excel(self, filepath: str):
+        table = None
+        if filepath.endswith(".csv"):
+            table = pd.read_csv(filepath)
+        else:
+            table = pd.read_excel(filepath)
+        if table is None:
+            return ""
+        json_text = table.dropna(axis=1).to_json(force_ascii=False)
+        return json_text
+
+    def read(self, filepath: str):
+        file_type = self.get_type(filepath)
+
+        text = ""
+
+        if not os.path.exists(filepath):
+            return text, None
+
+        try:
+
+            if file_type == "md" or file_type == "text":
+                with open(filepath) as f:
+                    text = f.read()
+
+            elif file_type == "pdf":
+                text += self.read_pdf(filepath)
+
+            elif file_type == "excel":
+                text += self.read_excel(filepath)
+
+            elif file_type == "word" or file_type == "ppt":
+                # https://stackoverflow.com/questions/36001482/read-doc-file-with-python
+                # https://textract.readthedocs.io/en/latest/installation.html
+
+                # TODO textract 在 pip 高于 24.1 后安装不了，因为其库自身原因，后续按需进行完善
+                # 可自行安装 pip install textract==1.6.5
+                import textract  # for word and ppt
+
+                text = textract.process(filepath).decode("utf8")
+                if file_type == "ppt":
+                    text = text.replace("\n", " ")
+
+            elif file_type == "html":
+                with open(filepath) as f:
+                    soup = BeautifulSoup(f.read(), "html.parser")
+                    text += soup.text
+
+        except Exception as e:
+            logger.error((filepath, str(e)))
+            return "", e
+        text = text.replace("\n\n", "\n")
+        text = text.replace("\n\n", "\n")
+        text = text.replace("\n\n", "\n")
+        text = text.replace("  ", " ")
+        text = text.replace("  ", " ")
+        text = text.replace("  ", " ")
+        return text, None
+
+
+if __name__ == "__main__":
+
+    def get_pdf_files(directory):
+        pdf_files = []
+        # 遍历目录
+        for root, dirs, files in os.walk(directory):
+            for file in files:
+                # 检查文件扩展名是否为.pdf
+                if file.lower().endswith(".pdf"):
+                    # 将完整路径添加到列表中
+                    pdf_files.append(os.path.abspath(os.path.join(root, file)))
+        return pdf_files
+
+    # 将你想要搜索的目录替换为下面的路径
+    pdf_list = get_pdf_files("/home/khj/huixiangdou-web-online-data/hxd-bad-file")
+
+    # 打印所有找到的PDF文件的绝对路径
+
+    opr = FileOperation()
+    for pdf_path in pdf_list:
+        text, error = opr.read(pdf_path)
+        print("processing {}".format(pdf_path))
+        if error is not None:
+            # pdb.set_trace()
+            print("")
+
+        else:
+            if text is not None:
+                print(len(text))
+            else:
+                # pdb.set_trace()
+                print("")

server/base/modules/rag/rag_worker.py

@@ -0,0 +1,122 @@
+import shutil
+from pathlib import Path
+
+import torch
+from loguru import logger
+
+from ....web_configs import WEB_CONFIGS
+from ...database.product_db import get_db_product_info
+from .feature_store import gen_vector_db
+from .retriever import CacheRetriever
+
+# 基础配置
+CONTEXT_MAX_LENGTH = 3000  # 上下文最大长度
+GENERATE_TEMPLATE = "这是说明书：“{}”\n 客户的问题：“{}” \n 请阅读说明并运用你的性格进行解答。"  # RAG prompt 模板
+
+# RAG 实例句柄
+RAG_RETRIEVER = None
+
+
+def build_rag_prompt(rag_retriever: CacheRetriever, product_name, prompt):
+
+    real_retriever = rag_retriever.get(fs_id="default")
+
+    if isinstance(real_retriever, tuple):
+        logger.info(f" @@@ GOT real_retriever == tuple : {real_retriever}")
+        return ""
+
+    chunk, db_context, references = real_retriever.query(
+        f"商品名：{product_name}。{prompt}", context_max_length=CONTEXT_MAX_LENGTH - 2 * len(GENERATE_TEMPLATE)
+    )
+    logger.info(f"db_context = {db_context}")
+
+    if db_context is not None and len(db_context) > 1:
+        prompt_rag = GENERATE_TEMPLATE.format(db_context, prompt)
+    else:
+        logger.info("db_context get error")
+        prompt_rag = prompt
+
+    logger.info(f"RAG reference = {references}")
+    logger.info("=" * 20)
+
+    return prompt_rag
+
+
+def init_rag_retriever(rag_config: str, db_path: str):
+    torch.cuda.empty_cache()
+
+    retriever = CacheRetriever(config_path=rag_config)
+
+    # 初始化
+    retriever.get(fs_id="default", config_path=rag_config, work_dir=db_path)
+
+    return retriever
+
+
+async def gen_rag_db(user_id, force_gen=False):
+    """
+    生成向量数据库。
+
+    参数:
+    force_gen - 布尔值，当设置为 True 时，即使数据库已存在也会重新生成数据库。
+    """
+
+    # 检查数据库目录是否存在，如果存在且force_gen为False，则不执行生成操作
+    if Path(WEB_CONFIGS.RAG_VECTOR_DB_DIR).exists() and not force_gen:
+        return
+
+    if force_gen and Path(WEB_CONFIGS.RAG_VECTOR_DB_DIR).exists():
+        shutil.rmtree(WEB_CONFIGS.RAG_VECTOR_DB_DIR)
+
+    # 仅仅遍历 instructions 字段里面的文件
+    if Path(WEB_CONFIGS.PRODUCT_INSTRUCTION_DIR_GEN_DB_TMP).exists():
+        shutil.rmtree(WEB_CONFIGS.PRODUCT_INSTRUCTION_DIR_GEN_DB_TMP)
+    Path(WEB_CONFIGS.PRODUCT_INSTRUCTION_DIR_GEN_DB_TMP).mkdir(exist_ok=True, parents=True)
+
+    # 读取 yaml 文件，获取所有说明书路径，并移动到 tmp 目录
+    product_list, _ = await get_db_product_info(user_id)
+
+    for info in product_list:
+
+        shutil.copyfile(
+            Path(
+                WEB_CONFIGS.SERVER_FILE_ROOT,
+                WEB_CONFIGS.PRODUCT_FILE_DIR,
+                WEB_CONFIGS.INSTRUCTIONS_DIR,
+                Path(info.instruction).name,
+            ),
+            Path(WEB_CONFIGS.PRODUCT_INSTRUCTION_DIR_GEN_DB_TMP).joinpath(Path(info.instruction).name),
+        )
+
+    logger.info("Generating rag database, pls wait ...")
+    # 调用函数生成向量数据库
+    gen_vector_db(
+        WEB_CONFIGS.RAG_CONFIG_PATH,
+        str(Path(WEB_CONFIGS.PRODUCT_INSTRUCTION_DIR_GEN_DB_TMP).absolute()),
+        WEB_CONFIGS.RAG_VECTOR_DB_DIR,
+    )
+
+    # 删除过程文件
+    shutil.rmtree(WEB_CONFIGS.PRODUCT_INSTRUCTION_DIR_GEN_DB_TMP)
+
+
+async def load_rag_model(user_id):
+
+    global RAG_RETRIEVER
+
+    # 重新生成 RAG 向量数据库
+    await gen_rag_db(user_id)
+
+    # 加载 rag 模型
+    RAG_RETRIEVER = init_rag_retriever(rag_config=WEB_CONFIGS.RAG_CONFIG_PATH, db_path=WEB_CONFIGS.RAG_VECTOR_DB_DIR)
+    logger.info("load rag model done !...")
+
+
+async def rebuild_rag_db(user_id, db_name="default"):
+
+    # 重新生成 RAG 向量数据库
+    await gen_rag_db(user_id, force_gen=True)
+
+    # 重新加载 retriever
+    RAG_RETRIEVER.pop(db_name)
+    RAG_RETRIEVER.get(fs_id=db_name, config_path=WEB_CONFIGS.RAG_CONFIG_PATH, work_dir=WEB_CONFIGS.RAG_VECTOR_DB_DIR)

server/base/modules/rag/retriever.py

@@ -0,0 +1,244 @@
+"""extract feature and search with user query."""
+
+import os
+import time
+
+import numpy as np
+import yaml
+from BCEmbedding.tools.langchain import BCERerank
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.vectorstores.faiss import FAISS as Vectorstore
+from langchain_community.vectorstores.utils import DistanceStrategy
+from loguru import logger
+from modelscope import snapshot_download
+from sklearn.metrics import precision_recall_curve
+
+from ....web_configs import WEB_CONFIGS
+from .file_operation import FileOperation
+
+
+class Retriever:
+    """Tokenize and extract features from the project's documents, for use in
+    the reject pipeline and response pipeline."""
+
+    def __init__(self, embeddings, reranker, work_dir: str, reject_throttle: float) -> None:
+        """Init with model device type and config."""
+        self.reject_throttle = reject_throttle
+        self.rejecter = Vectorstore.load_local(
+            os.path.join(work_dir, "db_reject"), embeddings=embeddings, allow_dangerous_deserialization=True
+        )
+        self.retriever = Vectorstore.load_local(
+            os.path.join(work_dir, "db_response"),
+            embeddings=embeddings,
+            allow_dangerous_deserialization=True,
+            distance_strategy=DistanceStrategy.MAX_INNER_PRODUCT,
+        ).as_retriever(search_type="similarity", search_kwargs={"score_threshold": 0.15, "k": 30})
+        self.compression_retriever = ContextualCompressionRetriever(base_compressor=reranker, base_retriever=self.retriever)
+
+    def is_reject(self, question, k=30, disable_throttle=False):
+        """If no search results below the threshold can be found from the
+        database, reject this query."""
+        if disable_throttle:
+            # for searching throttle during update sample
+            docs_with_score = self.rejecter.similarity_search_with_relevance_scores(question, k=1)
+            if len(docs_with_score) < 1:
+                return True, docs_with_score
+            return False, docs_with_score
+        else:
+            # for retrieve result
+            # if no chunk passed the throttle, give the max
+            docs_with_score = self.rejecter.similarity_search_with_relevance_scores(question, k=k)
+            ret = []
+            max_score = -1
+            top1 = None
+            for doc, score in docs_with_score:
+                if score >= self.reject_throttle:
+                    ret.append(doc)
+                if score > max_score:
+                    max_score = score
+                    top1 = (doc, score)
+            reject = False if len(ret) > 0 else True
+            return reject, [top1]
+
+    def update_throttle(self, config_path: str = "config.yaml", good_questions=[], bad_questions=[]):
+        """Update reject throttle based on positive and negative examples."""
+
+        if len(good_questions) == 0 or len(bad_questions) == 0:
+            raise Exception("good and bad question examples cat not be empty.")
+        questions = good_questions + bad_questions
+        predictions = []
+        for question in questions:
+            self.reject_throttle = -1
+            _, docs = self.is_reject(question=question, disable_throttle=True)
+            score = docs[0][1]
+            predictions.append(max(0, score))
+
+        labels = [1 for _ in range(len(good_questions))] + [0 for _ in range(len(bad_questions))]
+        precision, recall, thresholds = precision_recall_curve(labels, predictions)
+
+        # get the best index for sum(precision, recall)
+        sum_precision_recall = precision[:-1] + recall[:-1]
+        index_max = np.argmax(sum_precision_recall)
+        optimal_threshold = max(thresholds[index_max], 0.0)
+
+        with open(config_path, "r", encoding="utf-8") as f:
+            config = yaml.safe_load(f)
+        config["feature_store"]["reject_throttle"] = float(optimal_threshold)
+        with open(config_path, "w", encoding="utf8") as f:
+            yaml.dump(config, f)
+        logger.info(f"The optimal threshold is: {optimal_threshold}, saved it to {config_path}")  # noqa E501
+
+    def query(self, question: str, context_max_length: int = 16000):  # , tracker: QueryTracker = None):
+        """Processes a query and returns the best match from the vector store
+        database. If the question is rejected, returns None.
+
+        Args:
+            question (str): The question asked by the user.
+
+        Returns:
+            str: The best matching chunk, or None.
+            str: The best matching text, or None
+        """
+        print(f"DEBUG -1: enter query")
+
+        if question is None or len(question) < 1:
+            print(f"DEBUG 0: len error")
+
+            return None, None, []
+
+        if len(question) > 512:
+            logger.warning("input too long, truncate to 512")
+            question = question[0:512]
+
+        # reject, docs = self.is_reject(question=question)
+        # assert (len(docs) > 0)
+        # if reject:
+        # return None, None, [docs[0][0].metadata['source']]
+
+        docs = self.compression_retriever.get_relevant_documents(question)
+
+        print(f"DEBUG 1: {docs}")
+
+        # if tracker is not None:
+        #     tracker.log('retrieve', [doc.metadata['source'] for doc in docs])
+        chunks = []
+        context = ""
+        references = []
+
+        # add file text to context, until exceed `context_max_length`
+
+        file_opr = FileOperation()
+        for idx, doc in enumerate(docs):
+            chunk = doc.page_content
+            chunks.append(chunk)
+
+            if "read" not in doc.metadata:
+                logger.error(
+                    "If you are using the version before 20240319, please rerun `python3 -m huixiangdou.service.feature_store`"
+                )
+                raise Exception("huixiangdou version mismatch")
+            file_text, error = file_opr.read(doc.metadata["read"])
+            if error is not None:
+                # read file failed, skip
+                print(f"DEBUG 2: error")
+
+                continue
+
+            source = doc.metadata["source"]
+            logger.info("target {} file length {}".format(source, len(file_text)))
+
+            print(f"DEBUG 3: target {source}, file length {len(file_text)}")
+
+            if len(file_text) + len(context) > context_max_length:
+                if source in references:
+                    continue
+                references.append(source)
+                # add and break
+                add_len = context_max_length - len(context)
+                if add_len <= 0:
+                    break
+                chunk_index = file_text.find(chunk)
+                if chunk_index == -1:
+                    # chunk not in file_text
+                    context += chunk
+                    context += "\n"
+                    context += file_text[0 : add_len - len(chunk) - 1]
+                else:
+                    start_index = max(0, chunk_index - (add_len - len(chunk)))
+                    context += file_text[start_index : start_index + add_len]
+                break
+
+            if source not in references:
+                context += file_text
+                context += "\n"
+                references.append(source)
+
+        context = context[0:context_max_length]
+        logger.debug("query:{} top1 file:{}".format(question, references[0]))
+        return "\n".join(chunks), context, [os.path.basename(r) for r in references]
+
+
+class CacheRetriever:
+
+    def __init__(self, config_path: str, max_len: int = 4):
+        self.cache = dict()
+        self.max_len = max_len
+        with open(config_path, "r", encoding="utf-8") as f:
+            config = yaml.safe_load(f)["feature_store"]
+            embedding_model_path = config["embedding_model_path"]
+            reranker_model_path = config["reranker_model_path"]
+
+        embedding_model_path = snapshot_download(embedding_model_path, cache_dir=WEB_CONFIGS.RAG_MODEL_DIR)
+        reranker_model_path = snapshot_download(reranker_model_path, cache_dir=WEB_CONFIGS.RAG_MODEL_DIR)
+
+        # load text2vec and rerank model
+        logger.info("loading test2vec and rerank models")
+        self.embeddings = HuggingFaceEmbeddings(
+            model_name=embedding_model_path,
+            model_kwargs={"device": "cuda"},
+            encode_kwargs={"batch_size": 1, "normalize_embeddings": True},
+        )
+        self.embeddings.client = self.embeddings.client.half()
+        reranker_args = {"model": reranker_model_path, "top_n": 7, "device": "cuda", "use_fp16": True}
+        self.reranker = BCERerank(**reranker_args)
+
+    def get(self, fs_id: str = "default", config_path="config.yaml", work_dir="workdir"):
+        if fs_id in self.cache:
+            self.cache[fs_id]["time"] = time.time()
+            return self.cache[fs_id]["retriever"]
+
+        if not os.path.exists(work_dir) or not os.path.exists(config_path):
+            return None, "workdir or config.yaml not exist"
+
+        with open(config_path, "r", encoding="utf-8") as f:
+            reject_throttle = yaml.safe_load(f)["feature_store"]["reject_throttle"]
+
+        if len(self.cache) >= self.max_len:
+            # drop the oldest one
+            del_key = None
+            min_time = time.time()
+            for key, value in self.cache.items():
+                cur_time = value["time"]
+                if cur_time < min_time:
+                    min_time = cur_time
+                    del_key = key
+
+            if del_key is not None:
+                del_value = self.cache[del_key]
+                self.cache.pop(del_key)
+                del del_value["retriever"]
+
+        retriever = Retriever(
+            embeddings=self.embeddings, reranker=self.reranker, work_dir=work_dir, reject_throttle=reject_throttle
+        )
+        self.cache[fs_id] = {"retriever": retriever, "time": time.time()}
+        return retriever
+
+    def pop(self, fs_id: str):
+        if fs_id not in self.cache:
+            return
+        del_value = self.cache[fs_id]
+        self.cache.pop(fs_id)
+        # manually free memory
+        del del_value

server/base/modules/rag/test_queries.json

@@ -0,0 +1,4 @@
+[
+  "我的商品是牛肉。饲养天数",
+  "我的商品是唇膏。净含量是多少"
+]

server/base/queue_thread.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   queue_thread.py
+@Time    :   2024/09/02
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   队列调取相关逻辑(半废弃状态)
+"""
+
+
+from loguru import logger
+import requests
+import multiprocessing
+
+from ..web_configs import API_CONFIG
+from .server_info import SERVER_PLUGINS_INFO
+
+
+def process_tts(tts_text_queue):
+
+    while True:
+        try:
+            text_chunk = tts_text_queue.get(block=True, timeout=1)
+        except Exception as e:
+            # logger.info(f"### {e}")
+            continue
+        logger.info(f"Get tts quene: {type(text_chunk)} , {text_chunk}")
+        res = requests.post(API_CONFIG.TTS_URL, json=text_chunk)
+
+        # # tts 推理成功，放入数字人队列进行推理
+        # res_json = res.json()
+        # tts_request_dict = {
+        #     "user_id": "123",
+        #     "request_id": text_chunk["request_id"],
+        #     "chunk_id": text_chunk["chunk_id"],
+        #     "tts_path": res_json["wav_path"],
+        # }
+
+        # DIGITAL_HUMAN_QUENE.put(tts_request_dict)
+
+        logger.info(f"tts res = {res}")
+
+
+def process_digital_human(digital_human_queue):
+
+    while True:
+        try:
+            text_chunk = digital_human_queue.get(block=True, timeout=1)
+        except Exception as e:
+            # logger.info(f"### {e}")
+            continue
+        logger.info(f"Get digital human quene: {type(text_chunk)} , {text_chunk}")
+        res = requests.post(API_CONFIG.DIGITAL_HUMAN_URL, json=text_chunk)
+        logger.info(f"digital human res = {res}")
+
+
+if SERVER_PLUGINS_INFO.tts_server_enabled:
+    TTS_TEXT_QUENE = multiprocessing.Queue(maxsize=100)
+    tts_thread = multiprocessing.Process(target=process_tts, args=(TTS_TEXT_QUENE,), name="tts_processer")
+    tts_thread.start()
+else:
+    TTS_TEXT_QUENE = None
+
+if SERVER_PLUGINS_INFO.digital_human_server_enabled:
+    DIGITAL_HUMAN_QUENE = multiprocessing.Queue(maxsize=100)
+    digital_human_thread = multiprocessing.Process(
+        target=process_digital_human, args=(DIGITAL_HUMAN_QUENE,), name="digital_human_processer"
+    )
+    digital_human_thread.start()
+else:
+    DIGITAL_HUMAN_QUENE = None

server/base/routers/digital_human.py

@@ -0,0 +1,85 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   digital_human.py
+@Time    :   2024/09/02
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   数字人接口
+"""
+
+
+from pathlib import Path
+import uuid
+import requests
+from fastapi import APIRouter
+from loguru import logger
+from pydantic import BaseModel
+
+from ...web_configs import API_CONFIG, WEB_CONFIGS
+from ..utils import ResultCode, make_return_data
+
+router = APIRouter(
+    prefix="/digital-human",
+    tags=["digital-human"],
+    responses={404: {"description": "Not found"}},
+)
+
+
+class GenDigitalHumanVideoItem(BaseModel):
+    streamerId: int
+    salesDoc: str
+
+
+async def gen_tts_and_digital_human_video_app(streamer_id: int, sales_doc: str):
+    logger.info(sales_doc)
+
+    request_id = str(uuid.uuid1())
+    sentence_id = 1  # 直接推理，所以设置成 1
+    user_id = "123"
+
+    # 生成 TTS wav
+    tts_json = {
+        "user_id": user_id,
+        "request_id": request_id,
+        "sentence": sales_doc,
+        "chunk_id": sentence_id,
+        # "wav_save_name": chat_item.request_id + f"{str(sentence_id).zfill(8)}.wav",
+    }
+    tts_save_path = Path(WEB_CONFIGS.TTS_WAV_GEN_PATH, request_id + f"-{str(1).zfill(8)}.wav")
+    logger.info(f"waiting for wav generating done: {tts_save_path}")
+    _ = requests.post(API_CONFIG.TTS_URL, json=tts_json)
+
+    # 生成数字人视频
+    digital_human_gen_info = {
+        "user_id": user_id,
+        "request_id": request_id,
+        "chunk_id": 0,
+        "tts_path": str(tts_save_path),
+        "streamer_id": str(streamer_id),
+    }
+    video_path = Path(WEB_CONFIGS.DIGITAL_HUMAN_VIDEO_OUTPUT_PATH).joinpath(request_id + ".mp4")
+    logger.info(f"Generating digital human: {video_path}")
+    _ = requests.post(API_CONFIG.DIGITAL_HUMAN_URL, json=digital_human_gen_info)
+
+    # 删除过程文件
+    tts_save_path.unlink()
+
+    server_video_path = f"{API_CONFIG.REQUEST_FILES_URL}/{WEB_CONFIGS.STREAMER_FILE_DIR}/vid_output/{request_id}.mp4"
+    logger.info(server_video_path)
+
+    return server_video_path
+
+
+@router.post("/gen")
+async def get_digital_human_according_doc_api(gen_item: GenDigitalHumanVideoItem):
+    """根据口播文案生成数字人介绍视频
+
+    Args:
+        gen_item (GenDigitalHumanVideoItem): _description_
+
+    """
+    server_video_path = await gen_tts_and_digital_human_video_app(gen_item.streamerId, gen_item.salesDoc)
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", server_video_path)

server/base/routers/llm.py

@@ -0,0 +1,187 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   llm.py
+@Time    :   2024/09/02
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   大模型接口
+"""
+
+
+from typing import Dict, List
+
+from fastapi import APIRouter, Depends
+from loguru import logger
+
+from ..database.llm_db import get_llm_product_prompt_base_info
+from ..database.product_db import get_db_product_info
+from ..database.streamer_info_db import get_db_streamer_info
+from ..models.product_model import ProductInfo
+from ..models.streamer_info_model import StreamerInfo
+from ..modules.agent.agent_worker import get_agent_result
+from ..server_info import SERVER_PLUGINS_INFO
+from ..utils import LLM_MODEL_HANDLER, ResultCode, make_return_data
+from .users import get_current_user_info
+
+router = APIRouter(
+    prefix="/llm",
+    tags=["llm"],
+    responses={404: {"description": "Not found"}},
+)
+
+
+def combine_history(prompt: list, history_msg: list):
+    """生成对话历史 prompt
+
+    Args:
+        prompt (_type_): _description_
+        history_msg (_type_): _description_. Defaults to None.
+
+    Returns:
+        _type_: _description_
+    """
+    # 角色映射表
+    role_map = {"streamer": "assistant", "user": "user"}
+
+    # 生成历史对话信息
+    for message in history_msg:
+        prompt.append({"role": role_map[message["role"]], "content": message["message"]})
+
+    return prompt
+
+
+async def gen_poduct_base_prompt(
+    user_id: int,
+    streamer_id: int = -1,
+    product_id: int = -1,
+    streamer_info: StreamerInfo | None = None,
+    product_info: ProductInfo | None = None,
+) -> List[Dict[str, str]]:
+    """生成商品介绍的 prompt
+
+    Args:
+        user_id (int): 用户 ID
+        streamer_id (int): 主播 ID
+        product_id (int): 商品 ID
+        streamer_info (StreamerInfo, optional): 主播信息，如果为空则根据 streamer_id 查表
+        product_info (ProductInfo, optional): 商品信息，如果为空则根据 product_id 查表
+
+    Returns:
+        List[Dict[str,str]]: 生成的 promot
+    """
+
+    assert (streamer_id == -1 and streamer_info is not None) or (streamer_id != -1 and streamer_info is None)
+    assert (product_id == -1 and product_info is not None) or (product_id != -1 and product_info is None)
+
+    # 加载对话配置文件
+    dataset_yaml = await get_llm_product_prompt_base_info()
+
+    # 从配置中提取对话设置相关的信息
+    # system_str: 系统词，针对销售角色定制
+    # first_input_template: 对话开始时的第一个输入模板
+    # product_info_struct_template: 产品信息结构模板
+    system = dataset_yaml["conversation_setting"]["system"]
+    first_input_template = dataset_yaml["conversation_setting"]["first_input"]
+    product_info_struct_template = dataset_yaml["product_info_struct"]
+
+    # 根据 ID 获取主播信息
+    if streamer_info is None:
+        streamer_info = await get_db_streamer_info(user_id, streamer_id)
+        streamer_info = streamer_info[0]
+
+    # 将销售角色名和角色信息插入到 system prompt
+    character_str = streamer_info.character.replace(";", "、")
+    system_str = system.replace("{role_type}", streamer_info.name).replace("{character}", character_str)
+
+    # 根据 ID 获取商品信息
+    if product_info is None:
+        product_list, _ = await get_db_product_info(user_id, product_id=product_id)
+        product_info = product_list[0]
+
+    heighlights_str = product_info.heighlights.replace(";", "、")
+    product_info_str = product_info_struct_template[0].replace("{name}", product_info.product_name)
+    product_info_str += product_info_struct_template[1].replace("{highlights}", heighlights_str)
+
+    # 生成商品文案 prompt
+    sales_doc_prompt = first_input_template.replace("{product_info}", product_info_str)
+
+    prompt = [{"role": "system", "content": system_str}, {"role": "user", "content": sales_doc_prompt}]
+    logger.info(prompt)
+
+    return prompt
+
+
+async def get_agent_res(prompt, departure_place, delivery_company):
+    """调用 Agent 能力"""
+    agent_response = ""
+
+    if not SERVER_PLUGINS_INFO.agent_enabled:
+        # 如果不开启则直接返回空
+        return ""
+
+    GENERATE_AGENT_TEMPLATE = (
+        "这是网上获取到的信息：“{}”\n 客户的问题：“{}” \n 请认真阅读信息并运用你的性格进行解答。"  # Agent prompt 模板
+    )
+    input_prompt = prompt[-1]["content"]
+    agent_response = get_agent_result(LLM_MODEL_HANDLER, input_prompt, departure_place, delivery_company)
+    if agent_response != "":
+        agent_response = GENERATE_AGENT_TEMPLATE.format(agent_response, input_prompt)
+        logger.info(f"Agent response: {agent_response}")
+
+    return agent_response
+
+
+async def get_llm_res(prompt):
+    """获取 LLM 推理返回
+
+    Args:
+        prompt (str): _description_
+
+    Returns:
+        _type_: _description_
+    """
+
+    logger.info(prompt)
+    model_name = LLM_MODEL_HANDLER.available_models[0]
+
+    res_data = ""
+    for item in LLM_MODEL_HANDLER.chat_completions_v1(model=model_name, messages=prompt):
+        res_data = item["choices"][0]["message"]["content"]
+
+    return res_data
+
+
+@router.get("/gen_sales_doc", summary="生成主播文案接口")
+async def get_product_info_api(streamer_id: int, product_id: int, user_id: int = Depends(get_current_user_info)):
+    """生成口播文案
+
+    Args:
+        streamer_id (int): 主播 ID，用于获取性格等信息
+        product_id (int): 商品 ID
+    """
+
+    prompt = await gen_poduct_base_prompt(user_id, streamer_id, product_id)
+
+    res_data = await get_llm_res(prompt)
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", res_data)
+
+
+@router.get("/gen_product_info")
+async def get_product_info_api(product_id: int, user_id: int = Depends(get_current_user_info)):
+    """TODO 根据说明书内容生成商品信息
+
+    Args:
+        gen_product_item (GenProductItem): _description_
+    """
+
+    raise NotImplemented()
+    instruction_str = ""
+    prompt = [{"system": "现在你是一个文档小助手，你可以从文档里面总结出我需要的信息", "input": ""}]
+
+    res_data = ""
+    model_name = LLM_MODEL_HANDLER.available_models[0]
+    for item in LLM_MODEL_HANDLER.chat_completions_v1(model=model_name, messages=prompt):
+        res_data += item

server/base/routers/products.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   products.py
+@Time    :   2024/08/30
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   商品信息接口
+"""
+
+from pathlib import Path
+
+from fastapi import APIRouter, Depends
+
+from ...web_configs import WEB_CONFIGS
+from ..database.product_db import (
+    create_or_update_db_product_by_id,
+    delete_product_id,
+    get_db_product_info,
+)
+from ..models.product_model import ProductInfo, ProductPageItem, ProductQueryItem
+from ..modules.rag.rag_worker import rebuild_rag_db
+from ..utils import ResultCode, make_return_data
+from .users import get_current_user_info
+
+router = APIRouter(
+    prefix="/products",
+    tags=["products"],
+    responses={404: {"description": "Not found"}},
+)
+
+
+@router.get("/list", summary="获取分页商品信息接口")
+async def get_product_info_api(
+    currentPage: int = 1, pageSize: int = 5, productName: str | None = None, user_id: int = Depends(get_current_user_info)
+):
+    product_list, db_product_size = await get_db_product_info(
+        user_id=user_id,
+        current_page=currentPage,
+        page_size=pageSize,
+        product_name=productName,
+    )
+
+    res_data = ProductPageItem(product_list=product_list, currentPage=currentPage, pageSize=pageSize, totalSize=db_product_size)
+    return make_return_data(True, ResultCode.SUCCESS, "成功", res_data)
+
+
+@router.get("/info/{productId}", summary="获取特定商品 ID 的详细信息接口")
+async def get_product_id_info_api(productId: int, user_id: int = Depends(get_current_user_info)):
+    product_list, _ = await get_db_product_info(user_id=user_id, product_id=productId)
+
+    if len(product_list) == 1:
+        product_list = product_list[0]
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", product_list)
+
+
+@router.post("/create", summary="新增商品接口")
+async def upload_product_api(upload_product_item: ProductInfo, user_id: int = Depends(get_current_user_info)):
+
+    upload_product_item.user_id = user_id
+    upload_product_item.product_id = None
+
+    rebuild_rag_db_flag = create_or_update_db_product_by_id(0, upload_product_item)
+
+    if WEB_CONFIGS.ENABLE_RAG and rebuild_rag_db_flag:
+        # 重新生成 RAG 向量数据库
+        await rebuild_rag_db(user_id)
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", "")
+
+
+@router.put("/edit/{product_id}", summary="编辑商品接口")
+async def upload_product_api(product_id: int, upload_product_item: ProductInfo, user_id: int = Depends(get_current_user_info)):
+
+    rebuild_rag_db_flag = create_or_update_db_product_by_id(product_id, upload_product_item, user_id)
+
+    if WEB_CONFIGS.ENABLE_RAG and rebuild_rag_db_flag:
+        # 重新生成 RAG 向量数据库
+        await rebuild_rag_db(user_id)
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", "")
+
+
+@router.delete("/delete/{productId}", summary="删除特定商品 ID 接口")
+async def upload_product_api(productId: int, user_id: int = Depends(get_current_user_info)):
+
+    process_success_flag = await delete_product_id(productId, user_id)
+
+    if not process_success_flag:
+        return make_return_data(False, ResultCode.FAIL, "失败", "")
+
+    if WEB_CONFIGS.ENABLE_RAG:
+        # 重新生成 RAG 向量数据库
+        await rebuild_rag_db(user_id)
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", "")
+
+
+@router.post("/instruction", summary="获取对应商品的说明书内容接口", dependencies=[Depends(get_current_user_info)])
+async def get_product_instruction_info_api(instruction_path: ProductQueryItem):
+    """获取对应商品的说明书
+
+    Args:
+        instruction_path (ProductInstructionItem): 说明书路径
+
+    """
+    # TODO 后续改为前端 axios 直接获取
+    loacl_path = Path(WEB_CONFIGS.SERVER_FILE_ROOT).joinpath(
+        WEB_CONFIGS.PRODUCT_FILE_DIR, WEB_CONFIGS.INSTRUCTIONS_DIR, Path(instruction_path.instructionPath).name
+    )
+    if not loacl_path.exists():
+        return make_return_data(False, ResultCode.FAIL, "文件不存在", "")
+
+    with open(loacl_path, "r") as f:
+        instruction_content = f.read()
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", instruction_content)

server/base/routers/streamer_info.py

@@ -0,0 +1,156 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   streamer_info.py
+@Time    :   2024/08/10
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   主播管理信息页面接口
+"""
+
+from typing import Tuple
+import uuid
+from pathlib import Path
+
+import requests
+from fastapi import APIRouter, Depends
+from loguru import logger
+
+from ...web_configs import API_CONFIG, WEB_CONFIGS
+from ..database.streamer_info_db import create_or_update_db_streamer_by_id, delete_streamer_id, get_db_streamer_info
+from ..models.streamer_info_model import StreamerInfo
+from ..utils import ResultCode, make_poster_by_video_first_frame, make_return_data
+from .users import get_current_user_info
+
+router = APIRouter(
+    prefix="/streamer",
+    tags=["streamer"],
+    responses={404: {"description": "Not found"}},
+)
+
+
+async def gen_digital_human(user_id, streamer_id: int, new_streamer_info: StreamerInfo) -> Tuple[str, str]:
+    """生成数字人视频
+
+    Args:
+        user_id (int): 用户 ID
+        streamer_id (int): 主播 ID
+        new_streamer_info (StreamerInfo): 新的主播信息
+
+    Returns:
+        str: 数字人视频地址
+        str: 数字人头像/海报地址
+    """
+
+    streamer_info_db = await get_db_streamer_info(user_id, streamer_id)
+    streamer_info_db = streamer_info_db[0]
+
+    new_base_mp4_path = new_streamer_info.base_mp4_path.replace(API_CONFIG.REQUEST_FILES_URL, "")
+    if streamer_info_db.base_mp4_path.replace(API_CONFIG.REQUEST_FILES_URL, "") == new_base_mp4_path:
+        # 数字人视频没更新，跳过
+        return streamer_info_db.base_mp4_path, streamer_info_db.poster_image
+
+    # 调取接口生成进行数字人预处理
+
+    # new_streamer_info.base_mp4_path 是 服务器地址，需要进行转换
+    video_local_dir = Path(WEB_CONFIGS.SERVER_FILE_ROOT).joinpath(
+        WEB_CONFIGS.STREAMER_FILE_DIR, WEB_CONFIGS.STREAMER_INFO_FILES_DIR
+    )
+
+    digital_human_gen_info = {
+        "user_id": str(user_id),
+        "request_id": str(uuid.uuid1()),
+        "streamer_id": str(new_streamer_info.streamer_id),
+        "video_path": str(video_local_dir.joinpath(Path(new_streamer_info.base_mp4_path).name)),
+    }
+    logger.info(f"Getting digital human preprocessing: {new_streamer_info.streamer_id}")
+    _ = requests.post(API_CONFIG.DIGITAL_HUMAN_PREPROCESS_URL, json=digital_human_gen_info)
+
+    # 根据视频第一帧生成头图
+    poster_save_name = Path(new_streamer_info.base_mp4_path).stem + ".png"
+    make_poster_by_video_first_frame(str(video_local_dir.joinpath(Path(new_streamer_info.base_mp4_path).name)), poster_save_name)
+
+    # 生成头图服务器地址
+    poster_server_url = str(Path(new_streamer_info.base_mp4_path).parent.joinpath(poster_save_name))
+    if "http://" not in poster_server_url and "http:/" in poster_server_url:
+        poster_server_url = poster_server_url.replace("http:/", "http://")
+
+    return new_streamer_info.base_mp4_path, poster_server_url
+
+
+@router.get("/list", summary="获取所有主播信息接口，用于用户进行主播的选择")
+async def get_streamer_info_api(user_id: int = Depends(get_current_user_info)):
+    """获取所有主播信息，用于用户进行主播的选择"""
+    streamer_list = await get_db_streamer_info(user_id)
+    return make_return_data(True, ResultCode.SUCCESS, "成功", streamer_list)
+
+
+@router.get("/info/{streamerId}", summary="用于获取特定主播的信息接口")
+async def get_streamer_info_api(streamerId: int, user_id: int = Depends(get_current_user_info)):
+    """用于获取特定主播的信息"""
+
+    streamer_list = await get_db_streamer_info(user_id, streamerId)
+    if len(streamer_list) == 1:
+        streamer_list = streamer_list[0]
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", streamer_list)
+
+
+@router.post("/create", summary="新增主播信息接口")
+async def create_streamer_info_api(streamerItem: StreamerInfo, user_id: int = Depends(get_current_user_info)):
+    """新增主播信息"""
+    streamer_info = streamerItem
+    streamer_info.user_id = user_id
+    streamer_info.streamer_id = None
+
+    poster_image = streamer_info.poster_image
+    base_mp4_path = streamer_info.base_mp4_path
+
+    streamer_info.poster_image = ""
+    streamer_info.base_mp4_path = ""
+
+    # 更新数据库，才能拿到 stream_id
+    streamer_id = create_or_update_db_streamer_by_id(0, streamer_info, user_id)
+
+    streamer_info.poster_image = poster_image
+    streamer_info.base_mp4_path = base_mp4_path
+    streamer_info.streamer_id = streamer_id
+
+    # 数字人视频对其进行初始化，同时生成头图
+    video_info = await gen_digital_human(user_id, streamer_id, streamer_info)
+
+    streamer_info.base_mp4_path = video_info[0]
+    streamer_info.poster_image = video_info[1]
+    streamer_info.avatar = video_info[1]
+
+    create_or_update_db_streamer_by_id(streamer_id, streamer_info, user_id)
+    return make_return_data(True, ResultCode.SUCCESS, "成功", streamer_id)
+
+
+@router.put("/edit/{streamer_id}", summary="修改主播信息接口")
+async def edit_streamer_info_api(streamer_id: int, streamer_info: StreamerInfo, user_id: int = Depends(get_current_user_info)):
+    """修改主播信息"""
+
+    # 如果更新了数字人视频对其进行初始化，同时生成头图
+    video_info = await gen_digital_human(user_id, streamer_id, streamer_info)
+
+    streamer_info.base_mp4_path = video_info[0]
+    streamer_info.poster_image = video_info[1]
+    streamer_info.avatar = video_info[1]
+
+    # 更新数据库
+    create_or_update_db_streamer_by_id(streamer_id, streamer_info, user_id)
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", streamer_id)
+
+
+@router.delete("/delete/{streamerId}", summary="删除主播接口")
+async def upload_product_api(streamerId: int, user_id: int = Depends(get_current_user_info)):
+
+    process_success_flag = await delete_streamer_id(streamerId, user_id)
+
+    if not process_success_flag:
+        return make_return_data(False, ResultCode.FAIL, "失败", "")
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", "")

server/base/routers/streaming_room.py

@@ -0,0 +1,335 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   streaming_room.py
+@Time    :   2024/08/31
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   主播间信息交互接口
+"""
+
+import uuid
+from copy import deepcopy
+from pathlib import Path
+
+import requests
+from fastapi import APIRouter, Depends
+from loguru import logger
+
+from ...web_configs import API_CONFIG, WEB_CONFIGS
+from ..database.product_db import get_db_product_info
+from ..database.streamer_room_db import (
+    create_or_update_db_room_by_id,
+    get_live_room_info,
+    get_message_list,
+    update_db_room_status,
+    delete_room_id,
+    get_db_streaming_room_info,
+    update_message_info,
+    update_room_video_path,
+)
+from ..models.product_model import ProductInfo
+from ..models.streamer_room_model import OnAirRoomStatusItem, RoomChatItem, SalesDocAndVideoInfo, StreamRoomInfo
+from ..modules.rag.rag_worker import RAG_RETRIEVER, build_rag_prompt
+from ..routers.users import get_current_user_info
+from ..server_info import SERVER_PLUGINS_INFO
+from ..utils import ResultCode, make_return_data
+from .digital_human import gen_tts_and_digital_human_video_app
+from .llm import combine_history, gen_poduct_base_prompt, get_agent_res, get_llm_res
+
+router = APIRouter(
+    prefix="/streaming-room",
+    tags=["streaming-room"],
+    responses={404: {"description": "Not found"}},
+)
+
+
+@router.get("/list", summary="获取所有直播间信息接口")
+async def get_streaming_room_api(user_id: int = Depends(get_current_user_info)):
+    """获取所有直播间信息"""
+    # 加载直播间数据
+    streaming_room_list = await get_db_streaming_room_info(user_id)
+
+    for i in range(len(streaming_room_list)):
+        # 直接返回会导致字段丢失，需要转 dict 确保返回值里面有该字段
+        streaming_room_list[i] = dict(streaming_room_list[i])
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", streaming_room_list)
+
+
+@router.get("/info/{roomId}", summary="获取特定直播间信息接口")
+async def get_streaming_room_id_api(
+    roomId: int, currentPage: int = 1, pageSize: int = 5, user_id: int = Depends(get_current_user_info)
+):
+    """获取特定直播间信息"""
+    # 加载直播间配置文件
+    assert roomId != 0
+
+    # TODO 加入分页
+
+    # 加载直播间数据
+    streaming_room_list = await get_db_streaming_room_info(user_id, room_id=roomId)
+
+    if len(streaming_room_list) == 1:
+        # 直接返回会导致字段丢失，需要转 dict 确保返回值里面有该字段
+        format_product_list = []
+        for db_product in streaming_room_list[0].product_list:
+
+            product_dict = dict(db_product)
+            # 将 start_video 改为服务器地址
+            if product_dict["start_video"] != "":
+                product_dict["start_video"] = API_CONFIG.REQUEST_FILES_URL + product_dict["start_video"]
+
+            format_product_list.append(product_dict)
+        streaming_room_list = dict(streaming_room_list[0])
+        streaming_room_list["product_list"] = format_product_list
+    else:
+        streaming_room_list = []
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", streaming_room_list)
+
+
+@router.get("/product-edit-list/{roomId}", summary="获取直播间商品编辑列表，含有已选中的标识")
+async def get_streaming_room_product_list_api(
+    roomId: int, currentPage: int = 1, pageSize: int = 0, user_id: int = Depends(get_current_user_info)
+):
+    """获取直播间商品编辑列表，含有已选中的标识"""
+
+    # 获取目前直播间商品列表
+    if roomId == 0:
+        # 新的直播间
+        merge_list = []
+        exclude_list = []
+    else:
+        streaming_room_info = await get_db_streaming_room_info(user_id, roomId)
+
+        if len(streaming_room_info) == 0:
+            raise "401"
+
+        streaming_room_info = streaming_room_info[0]
+        # 获取未被选中的商品
+        exclude_list = [product.product_id for product in streaming_room_info.product_list]
+        merge_list = deepcopy(streaming_room_info.product_list)
+
+    # 获取未选中的商品信息
+    not_select_product_list, db_product_size = await get_db_product_info(user_id=user_id, exclude_list=exclude_list)
+
+    # 合并商品信息
+    for not_select_product in not_select_product_list:
+        merge_list.append(
+            SalesDocAndVideoInfo(
+                product_id=not_select_product.product_id,
+                product_info=ProductInfo(**dict(not_select_product)),
+                selected=False,
+            )
+        )
+
+    # TODO 懒加载分页
+
+    # 格式化
+    format_merge_list = []
+    for product in merge_list:
+        # 直接返回会导致字段丢失，需要转 dict 确保返回值里面有该字段
+        dict_info = dict(product)
+        if "stream_room" in dict_info:
+            dict_info.pop("stream_room")
+        format_merge_list.append(dict_info)
+
+    page_info = dict(
+        product_list=format_merge_list,
+        current=currentPage,
+        pageSize=db_product_size,
+        totalSize=db_product_size,
+    )
+    logger.info(page_info)
+    return make_return_data(True, ResultCode.SUCCESS, "成功", page_info)
+
+
+@router.post("/create", summary="新增直播间接口")
+async def streaming_room_edit_api(edit_item: dict, user_id: int = Depends(get_current_user_info)):
+    product_list = edit_item.pop("product_list")
+    status = edit_item.pop("status")
+    edit_item.pop("streamer_info")
+    edit_item.pop("room_id")
+
+    if "status_id" in edit_item:
+        edit_item.pop("status_id")
+
+    formate_product_list = []
+    for product in product_list:
+        if not product["selected"]:
+            continue
+        product.pop("product_info")
+        product_item = SalesDocAndVideoInfo(**product)
+        formate_product_list.append(product_item)
+
+    edit_item["user_id"] = user_id
+    formate_info = StreamRoomInfo(**edit_item, product_list=formate_product_list, status=OnAirRoomStatusItem(**status))
+    room_id = create_or_update_db_room_by_id(0, formate_info, user_id)
+    return make_return_data(True, ResultCode.SUCCESS, "成功", room_id)
+
+
+@router.put("/edit/{room_id}", summary="编辑直播间接口")
+async def streaming_room_edit_api(room_id: int, edit_item: dict, user_id: int = Depends(get_current_user_info)):
+    """编辑直播间接口
+
+    Args:
+        edit_item (StreamRoomInfo): _description_
+    """
+
+    product_list = edit_item.pop("product_list")
+    status = edit_item.pop("status")
+    edit_item.pop("streamer_info")
+
+    formate_product_list = []
+    for product in product_list:
+        if not product["selected"]:
+            continue
+        product.pop("product_info")
+        product_item = SalesDocAndVideoInfo(**product)
+        formate_product_list.append(product_item)
+
+    formate_info = StreamRoomInfo(**edit_item, product_list=formate_product_list, status=OnAirRoomStatusItem(**status))
+    create_or_update_db_room_by_id(room_id, formate_info, user_id)
+    return make_return_data(True, ResultCode.SUCCESS, "成功", room_id)
+
+
+@router.delete("/delete/{roomId}", summary="删除直播间接口")
+async def delete_room_api(roomId: int, user_id: int = Depends(get_current_user_info)):
+
+    process_success_flag = await delete_room_id(roomId, user_id)
+
+    if not process_success_flag:
+        return make_return_data(False, ResultCode.FAIL, "失败", "")
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", "")
+
+
+# ============================================================
+#                          开播接口
+# ============================================================
+
+
+@router.post("/online/{roomId}", summary="直播间开播接口")
+async def offline_api(roomId: int, user_id: int = Depends(get_current_user_info)):
+
+    update_db_room_status(roomId, user_id, "online")
+    return make_return_data(True, ResultCode.SUCCESS, "成功", "")
+
+
+@router.put("/offline/{roomId}", summary="直播间下播接口")
+async def offline_api(roomId: int, user_id: int = Depends(get_current_user_info)):
+
+    update_db_room_status(roomId, user_id, "offline")
+    return make_return_data(True, ResultCode.SUCCESS, "成功", "")
+
+
+@router.post("/next-product/{roomId}", summary="直播间进行下一个商品讲解接口")
+async def on_air_live_room_next_product_api(roomId: int, user_id: int = Depends(get_current_user_info)):
+    """直播间进行下一个商品讲解
+
+    Args:
+        roomId (int): 直播间 ID
+    """
+
+    update_db_room_status(roomId, user_id, "next-product")
+    return make_return_data(True, ResultCode.SUCCESS, "成功", "")
+
+
+@router.get("/live-info/{roomId}", summary="获取正在直播的直播间信息接口")
+async def get_on_air_live_room_api(roomId: int, user_id: int = Depends(get_current_user_info)):
+    """获取正在直播的直播间信息
+
+    1. 主播视频地址
+    2. 商品信息，显示在右下角的商品缩略图
+    3. 对话记录 conversation_list
+
+    Args:
+        roomId (int): 直播间 ID
+    """
+
+    res_data = await get_live_room_info(user_id, roomId)
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", res_data)
+
+
+@router.put("/chat", summary="直播间对话接口")
+async def get_on_air_live_room_api(room_chat: RoomChatItem, user_id: int = Depends(get_current_user_info)):
+    # 根据直播间 ID 获取信息
+    streaming_room_info = await get_db_streaming_room_info(user_id, room_chat.roomId)
+    streaming_room_info = streaming_room_info[0]
+
+    # 商品索引
+    product_detail = streaming_room_info.product_list[streaming_room_info.status.current_product_index].product_info
+
+    # 销售 ID
+    sales_info_id = streaming_room_info.product_list[streaming_room_info.status.current_product_index].sales_info_id
+
+    # 更新对话记录
+    update_message_info(sales_info_id, user_id, role="user", message=room_chat.message)
+
+    # 获取最新的对话记录
+    conversation_list = get_message_list(sales_info_id)
+
+    # 根据对话记录生成 prompt
+    prompt = await gen_poduct_base_prompt(
+        user_id,
+        streamer_info=streaming_room_info.streamer_info,
+        product_info=product_detail,
+    )  # system + 获取商品文案prompt
+
+    prompt = combine_history(prompt, conversation_list)
+
+    # ====================== Agent ======================
+    # 调取 Agent
+    agent_response = await get_agent_res(prompt, product_detail.departure_place, product_detail.delivery_company)
+    if agent_response != "":
+        logger.info("Agent 执行成功，不执行 RAG")
+        prompt[-1]["content"] = agent_response
+
+    # ====================== RAG ======================
+    # 调取 rag
+    elif SERVER_PLUGINS_INFO.rag_enabled:
+        logger.info("Agent 未执行 or 未开启，调用 RAG")
+        # agent 失败，调取 rag, chat_item.plugins.rag 为 True，则使用 RAG 查询数据库
+        rag_res = build_rag_prompt(RAG_RETRIEVER, product_detail.product_name, prompt[-1]["content"])
+        if rag_res != "":
+            prompt[-1]["content"] = rag_res
+
+    # 调取 LLM
+    streamer_res = await get_llm_res(prompt)
+
+    # 生成数字人视频
+    server_video_path = await gen_tts_and_digital_human_video_app(streaming_room_info.streamer_info.streamer_id, streamer_res)
+
+    # 更新直播间数字人视频信息
+    update_room_video_path(streaming_room_info.status_id, server_video_path)
+
+    # 更新对话记录
+    update_message_info(sales_info_id, streaming_room_info.streamer_info.streamer_id, role="streamer", message=streamer_res)
+
+    return make_return_data(True, ResultCode.SUCCESS, "成功", "")
+
+
+@router.post("/asr", summary="直播间调取 ASR 语音转文字 接口")
+async def get_on_air_live_room_api(room_chat: RoomChatItem, user_id: int = Depends(get_current_user_info)):
+
+    # room_chat.asr_file 是 服务器地址，需要进行转换
+    asr_local_path = Path(WEB_CONFIGS.SERVER_FILE_ROOT).joinpath(WEB_CONFIGS.ASR_FILE_DIR, Path(room_chat.asrFileUrl).name)
+
+    # 获取 ASR 结果
+    req_data = {
+        "user_id": user_id,
+        "request_id": str(uuid.uuid1()),
+        "wav_path": str(asr_local_path),
+    }
+    logger.info(req_data)
+
+    res = requests.post(API_CONFIG.ASR_URL, json=req_data).json()
+    asr_str = res["result"]
+    logger.info(f"ASR res = {asr_str}")
+
+    # 删除过程文件
+    asr_local_path.unlink()
+    return make_return_data(True, ResultCode.SUCCESS, "成功", asr_str)

server/base/routers/users.py

@@ -0,0 +1,157 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   users.py
+@Time    :   2024/08/30
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   用户登录和 Token 认证接口
+"""
+
+from datetime import datetime, timedelta, timezone
+
+import jwt
+from fastapi import APIRouter, Depends, HTTPException
+from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
+from loguru import logger
+from passlib.context import CryptContext
+
+from ...web_configs import WEB_CONFIGS
+from ..database.user_db import get_db_user_info
+from ..models.user_model import TokenItem
+from ..utils import ResultCode, make_return_data
+
+router = APIRouter(
+    prefix="/user",
+    tags=["user"],
+    responses={404: {"description": "Not found"}},
+)
+
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="/user/login")
+
+# 密码加解密
+PWD_CONTEXT = CryptContext(schemes=["bcrypt"], deprecated="auto")
+
+
+def verify_password(plain_password: str, hashed_password: str) -> bool:
+    """密码校验
+
+    Args:
+        plain_password (str): 明文密码
+        hashed_password (str): 加密后的密码，用于对比
+
+    Returns:
+        bool: 校验时候通过
+    """
+    logger.info(f"expect password = {PWD_CONTEXT.hash('123456')}")
+    return PWD_CONTEXT.verify(plain_password, hashed_password)
+
+
+def get_password_hash(password: str) -> str:
+    """生成哈希密码
+
+    Args:
+        password (str): 明文密码
+
+    Returns:
+        str: 加密后的哈希密码
+    """
+    return PWD_CONTEXT.hash(password)
+
+
+def authenticate_user(username: str, password: str) -> bool:
+    """对用户名和密码进行校验
+
+    Args:
+        username (str): 用户名
+        password (str): 密码
+
+    Returns:
+        bool: 是否检验通过
+    """
+
+    # 获取用户信息
+    user_info = get_db_user_info(username=username, all_info=True)
+    if not user_info:
+        # 没有找到用户名
+        logger.info(f"Cannot find username = {username}")
+        return False
+
+    # 校验密码
+    if not verify_password(password, user_info.hashed_password):
+        logger.info(f"verify_password fail")
+        # 密码校验失败
+        return False
+
+    return user_info
+
+
+def get_current_user_info(token: str = Depends(oauth2_scheme)):
+    """在 token 中提取 user id
+
+    Args:
+        token (str, optional): token. Defaults to Depends(oauth2_scheme).
+
+    Raises:
+        HTTPException: 401 获取失败
+
+    Returns:
+        int: 用户 ID
+    """
+    logger.info(token)
+    try:
+        token_data = jwt.decode(token, WEB_CONFIGS.TOKEN_JWT_SECURITY_KEY, algorithms=WEB_CONFIGS.TOKEN_JWT_ALGORITHM)
+        logger.info(token_data)
+        user_id = token_data.get("user_id", None)
+    except Exception as e:
+        logger.error(e)
+        raise HTTPException(status_code=401, detail="Could not validate credentials")
+
+    if not user_id:
+        logger.error(f"can not get user_id: {user_id}")
+        raise HTTPException(status_code=401, detail="Could not validate credentials")
+
+    # TODO 超时强制重新登录
+
+    logger.info(f"Got user_id: {user_id}")
+    return user_id
+
+
+@router.post("/login", summary="登录接口")
+async def login(form_data: OAuth2PasswordRequestForm = Depends()):
+
+    # 校验用户名和密码
+    user_info = authenticate_user(form_data.username, form_data.password)
+
+    if not user_info:
+        raise HTTPException(status_code=401, detail="Incorrect username or password", headers={"WWW-Authenticate": "Bearer"})
+
+    # 过期时间
+    token_expires = datetime.now(timezone.utc) + timedelta(days=7)
+
+    # token 生成包含内容，记录 IP 的原因是防止被其他人拿到用户的 token 进行假冒访问
+    token_data = {
+        "user_id": user_info.user_id,
+        "username": user_info.username,
+        "exp": int(token_expires.timestamp()),
+        "ip": user_info.ip_address,
+        "login_time": int(datetime.now(timezone.utc).timestamp()),
+    }
+    logger.info(f"token_data = {token_data}")
+
+    # 生成 token
+    token = jwt.encode(token_data, WEB_CONFIGS.TOKEN_JWT_SECURITY_KEY, algorithm=WEB_CONFIGS.TOKEN_JWT_ALGORITHM)
+
+    # 返回
+    res_json = TokenItem(access_token=token, token_type="bearer")
+    logger.info(f"Got token info = {res_json}")
+    # return make_return_data(True, ResultCode.SUCCESS, "成功", content)
+    return res_json
+
+
+@router.get("/me", summary="获取用户信息")
+async def get_streaming_room_api(user_id: int = Depends(get_current_user_info)):
+    """获取用户信息"""
+    user_info = get_db_user_info(id=user_id)
+    return make_return_data(True, ResultCode.SUCCESS, "成功", user_info)

server/base/server_info.py

@@ -0,0 +1,134 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   server_info.py
+@Time    :   2024/09/02
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   组件信息获取逻辑
+"""
+
+
+import random
+import requests
+from loguru import logger
+
+from ..web_configs import API_CONFIG, WEB_CONFIGS
+
+
+class ServerPluginsInfo:
+
+    def __init__(self) -> None:
+        self.update_info()
+
+    def update_info(self):
+
+        self.tts_server_enabled = self._check_server(API_CONFIG.TTS_URL + "/check")
+        self.digital_human_server_enabled = self._check_server(API_CONFIG.DIGITAL_HUMAN_CHECK_URL)
+        self.asr_server_enabled = self._check_server(API_CONFIG.ASR_URL + "/check")
+        self.llm_enabled = self._check_server(API_CONFIG.LLM_URL)
+
+        if WEB_CONFIGS.AGENT_DELIVERY_TIME_API_KEY is None or WEB_CONFIGS.AGENT_WEATHER_API_KEY is None:
+            self.agent_enabled = False
+        else:
+            self.agent_enabled = True
+
+        self.rag_enabled = WEB_CONFIGS.ENABLE_RAG
+
+        logger.info(
+            "\nself check plugins info : \n"
+            f"| llm            | {self.llm_enabled} |\n"
+            f"| rag            | {self.rag_enabled} |\n"
+            f"| tts            | {self.tts_server_enabled} |\n"
+            f"| digital hunam  | {self.digital_human_server_enabled} |\n"
+            f"| asr            | {self.asr_server_enabled} |\n"
+            f"| agent          | {self.agent_enabled} |\n"
+        )
+
+    @staticmethod
+    def _check_server(url):
+
+        try:
+            res = requests.get(url)
+        except requests.exceptions.ConnectionError:
+            return False
+
+        if res.status_code == 200:
+            return True
+        else:
+            return False
+
+    @staticmethod
+    def _make_color_list(color_num):
+
+        color_list = [
+            "#FF3838",
+            "#FF9D97",
+            "#FF701F",
+            "#FFB21D",
+            "#CFD231",
+            "#48F90A",
+            "#92CC17",
+            "#3DDB86",
+            "#1A9334",
+            "#00D4BB",
+            "#2C99A8",
+            "#00C2FF",
+            "#344593",
+            "#6473FF",
+            "#0018EC",
+            "#8438FF",
+            "#520085",
+            "#CB38FF",
+            "#FF95C8",
+            "#FF37C7",
+        ]
+
+        return random.sample(color_list, color_num)
+
+    def get_status(self):
+        self.update_info()
+
+        info_list = [
+            {
+                "plugin_name": "LLM",
+                "describe": "大语言模型，用于根据客户历史对话，生成对话信息",
+                "enabled": self.llm_enabled,
+            },
+            {
+                "plugin_name": "RAG",
+                "describe": "用于调用知识库实时更新信息",
+                "enabled": self.rag_enabled,
+            },
+            {
+                "plugin_name": "TTS",
+                "describe": "文字转语音，让主播的文字也能听到",
+                "enabled": self.tts_server_enabled,
+            },
+            {
+                "plugin_name": "数字人",
+                "describe": "数字人服务，用于生成数字人，需要和 TTS 一起开启才有效果",
+                "enabled": self.digital_human_server_enabled,
+            },
+            {
+                "plugin_name": "Agent",
+                "describe": "用于根据用户对话，获取网络的实时信息",
+                "enabled": self.agent_enabled,
+            },
+            {
+                "plugin_name": "ASR",
+                "describe": "语音转文字，让用户无需打字就可以和主播进行对话",
+                "enabled": self.asr_server_enabled,
+            },
+        ]
+
+        # 生成图标背景色
+        color_list = self._make_color_list(len(info_list))
+        for idx, color in enumerate(color_list):
+            info_list[idx].update({"avatar_color": color})
+
+        return info_list
+
+
+SERVER_PLUGINS_INFO = ServerPluginsInfo()

server/base/utils.py

@@ -0,0 +1,485 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@File    :   utils.py
+@Time    :   2024/09/02
+@Project :   https://github.com/PeterH0323/Streamer-Sales
+@Author  :   HinGwenWong
+@Version :   1.0
+@Desc    :   工具集合文件
+"""
+
+
+import asyncio
+from ipaddress import IPv4Address
+import json
+import random
+import wave
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List
+
+import cv2
+from lmdeploy.serve.openai.api_client import APIClient
+from loguru import logger
+from pydantic import BaseModel
+from sqlmodel import Session, select
+from tqdm import tqdm
+
+from server.base.models.user_model import UserInfo
+
+from ..tts.tools import SYMBOL_SPLITS, make_text_chunk
+from ..web_configs import API_CONFIG, WEB_CONFIGS
+from .database.init_db import DB_ENGINE
+from .models.product_model import ProductInfo
+from .models.streamer_info_model import StreamerInfo
+from .models.streamer_room_model import OnAirRoomStatusItem, SalesDocAndVideoInfo, StreamRoomInfo
+
+from .modules.agent.agent_worker import get_agent_result
+from .modules.rag.rag_worker import RAG_RETRIEVER, build_rag_prompt
+from .queue_thread import DIGITAL_HUMAN_QUENE, TTS_TEXT_QUENE
+from .server_info import SERVER_PLUGINS_INFO
+
+
+class ChatGenConfig(BaseModel):
+    # LLM 推理配置
+    top_p: float = 0.8
+    temperature: float = 0.7
+    repetition_penalty: float = 1.005
+
+
+class ProductInfoItem(BaseModel):
+    name: str
+    heighlights: str
+    introduce: str  # 生成商品文案 prompt
+
+    image_path: str
+    departure_place: str
+    delivery_company_name: str
+
+
+class PluginsInfo(BaseModel):
+    rag: bool = True
+    agent: bool = True
+    tts: bool = True
+    digital_human: bool = True
+
+
+class ChatItem(BaseModel):
+    user_id: str  # User 识别号，用于区分不用的用户调用
+    request_id: str  # 请求 ID，用于生成 TTS & 数字人
+    prompt: List[Dict[str, str]]  # 本次的 prompt
+    product_info: ProductInfoItem  # 商品信息
+    plugins: PluginsInfo = PluginsInfo()  # 插件信息
+    chat_config: ChatGenConfig = ChatGenConfig()
+
+
+# 加载 LLM 模型
+LLM_MODEL_HANDLER = APIClient(API_CONFIG.LLM_URL)
+
+
+async def streamer_sales_process(chat_item: ChatItem):
+
+    # ====================== Agent ======================
+    # 调取 Agent
+    agent_response = ""
+    if chat_item.plugins.agent and SERVER_PLUGINS_INFO.agent_enabled:
+        GENERATE_AGENT_TEMPLATE = (
+            "这是网上获取到的信息：“{}”\n 客户的问题：“{}” \n 请认真阅读信息并运用你的性格进行解答。"  # Agent prompt 模板
+        )
+        input_prompt = chat_item.prompt[-1]["content"]
+        agent_response = get_agent_result(
+            LLM_MODEL_HANDLER, input_prompt, chat_item.product_info.departure_place, chat_item.product_info.delivery_company_name
+        )
+        if agent_response != "":
+            agent_response = GENERATE_AGENT_TEMPLATE.format(agent_response, input_prompt)
+            print(f"Agent response: {agent_response}")
+            chat_item.prompt[-1]["content"] = agent_response
+
+    # ====================== RAG ======================
+    # 调取 rag
+    if chat_item.plugins.rag and agent_response == "":
+        # 如果 Agent 没有执行，则使用 RAG 查询数据库
+        rag_prompt = chat_item.prompt[-1]["content"]
+        prompt_pro = build_rag_prompt(RAG_RETRIEVER, chat_item.product_info.name, rag_prompt)
+
+        if prompt_pro != "":
+            chat_item.prompt[-1]["content"] = prompt_pro
+
+    # llm 推理流返回
+    logger.info(chat_item.prompt)
+
+    current_predict = ""
+    idx = 0
+    last_text_index = 0
+    sentence_id = 0
+    model_name = LLM_MODEL_HANDLER.available_models[0]
+    for item in LLM_MODEL_HANDLER.chat_completions_v1(model=model_name, messages=chat_item.prompt, stream=True):
+        logger.debug(f"LLM predict: {item}")
+        if "content" not in item["choices"][0]["delta"]:
+            continue
+        current_res = item["choices"][0]["delta"]["content"]
+
+        if "~" in current_res:
+            current_res = current_res.replace("~", "。").replace("。。", "。")
+
+        current_predict += current_res
+        idx += 1
+
+        if chat_item.plugins.tts and SERVER_PLUGINS_INFO.tts_server_enabled:
+            # 切句子
+            sentence = ""
+            for symbol in SYMBOL_SPLITS:
+                if symbol in current_res:
+                    last_text_index, sentence = make_text_chunk(current_predict, last_text_index)
+                    if len(sentence) <= 3:
+                        # 文字太短的情况，不做生成
+                        sentence = ""
+                    break
+
+            if sentence != "":
+                sentence_id += 1
+                logger.info(f"get sentence: {sentence}")
+                tts_request_dict = {
+                    "user_id": chat_item.user_id,
+                    "request_id": chat_item.request_id,
+                    "sentence": sentence,
+                    "chunk_id": sentence_id,
+                    # "wav_save_name": chat_item.request_id + f"{str(sentence_id).zfill(8)}.wav",
+                }
+
+                TTS_TEXT_QUENE.put(tts_request_dict)
+                await asyncio.sleep(0.01)
+
+        yield json.dumps(
+            {
+                "event": "message",
+                "retry": 100,
+                "id": idx,
+                "data": current_predict,
+                "step": "llm",
+                "end_flag": False,
+            },
+            ensure_ascii=False,
+        )
+        await asyncio.sleep(0.01)  # 加个延时避免无法发出 event stream
+
+    if chat_item.plugins.digital_human and SERVER_PLUGINS_INFO.digital_human_server_enabled:
+
+        wav_list = [
+            Path(WEB_CONFIGS.TTS_WAV_GEN_PATH, chat_item.request_id + f"-{str(i).zfill(8)}.wav")
+            for i in range(1, sentence_id + 1)
+        ]
+        while True:
+            # 等待 TTS 生成完成
+            not_exist_count = 0
+            for tts_wav in wav_list:
+                if not tts_wav.exists():
+                    not_exist_count += 1
+
+            logger.info(f"still need to wait for {not_exist_count}/{sentence_id} wav generating...")
+            if not_exist_count == 0:
+                break
+
+            yield json.dumps(
+                {
+                    "event": "message",
+                    "retry": 100,
+                    "id": idx,
+                    "data": current_predict,
+                    "step": "tts",
+                    "end_flag": False,
+                },
+                ensure_ascii=False,
+            )
+            await asyncio.sleep(1)  # 加个延时避免无法发出 event stream
+
+        # 合并 tts
+        tts_save_path = Path(WEB_CONFIGS.TTS_WAV_GEN_PATH, chat_item.request_id + ".wav")
+        all_tts_data = []
+
+        for wav_file in tqdm(wav_list):
+            logger.info(f"Reading wav file {wav_file}...")
+            with wave.open(str(wav_file), "rb") as wf:
+                all_tts_data.append([wf.getparams(), wf.readframes(wf.getnframes())])
+
+        logger.info(f"Merging wav file to {tts_save_path}...")
+        tts_params = max([tts_data[0] for tts_data in all_tts_data])
+        with wave.open(str(tts_save_path), "wb") as wf:
+            wf.setparams(tts_params)  # 使用第一个音频参数
+
+            for wf_data in all_tts_data:
+                wf.writeframes(wf_data[1])
+        logger.info(f"Merged wav file to {tts_save_path} !")
+
+        # 生成数字人视频
+        tts_request_dict = {
+            "user_id": chat_item.user_id,
+            "request_id": chat_item.request_id,
+            "chunk_id": 0,
+            "tts_path": str(tts_save_path),
+        }
+
+        logger.info(f"Generating digital human...")
+        DIGITAL_HUMAN_QUENE.put(tts_request_dict)
+        while True:
+            if (
+                Path(WEB_CONFIGS.DIGITAL_HUMAN_VIDEO_OUTPUT_PATH)
+                .joinpath(Path(tts_save_path).stem + ".mp4")
+                .with_suffix(".txt")
+                .exists()
+            ):
+                break
+            yield json.dumps(
+                {
+                    "event": "message",
+                    "retry": 100,
+                    "id": idx,
+                    "data": current_predict,
+                    "step": "dg",
+                    "end_flag": False,
+                },
+                ensure_ascii=False,
+            )
+            await asyncio.sleep(1)  # 加个延时避免无法发出 event stream
+
+        # 删除过程文件
+        for wav_file in wav_list:
+            wav_file.unlink()
+
+    yield json.dumps(
+        {
+            "event": "message",
+            "retry": 100,
+            "id": idx,
+            "data": current_predict,
+            "step": "all",
+            "end_flag": True,
+        },
+        ensure_ascii=False,
+    )
+
+
+def make_poster_by_video_first_frame(video_path: str, image_output_name: str):
+    """根据视频第一帧生成缩略图
+
+    Args:
+        video_path (str): 视频文件路径
+
+    Returns:
+        str: 第一帧保存的图片路径
+    """
+
+    # 打开视频文件
+    cap = cv2.VideoCapture(video_path)
+
+    # 读取第一帧
+    ret, frame = cap.read()
+
+    # 检查是否成功读取
+    poster_save_path = str(Path(video_path).parent.joinpath(image_output_name))
+    if ret:
+        # 保存图像到文件
+        cv2.imwrite(poster_save_path, frame)
+        logger.info(f"第一帧已保存为 {poster_save_path}")
+    else:
+        logger.error("无法读取视频帧")
+
+    # 释放视频捕获对象
+    cap.release()
+
+    return poster_save_path
+
+
+@dataclass
+class ResultCode:
+    SUCCESS: int = 0000  # 成功
+    FAIL: int = 1000  # 失败
+
+
+def make_return_data(success_flag: bool, code: ResultCode, message: str, data: dict):
+    return {
+        "success": success_flag,
+        "code": code,
+        "message": message,
+        "data": data,
+        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+    }
+
+
+def gen_default_data():
+    """生成默认数据，包括：
+    - 商品数据
+    - 主播数据
+    - 直播间信息以及关联表
+    """
+
+    def create_default_user():
+        """创建默认用户"""
+        admin_user = UserInfo(
+            username="hingwen.wong",
+            ip_address=IPv4Address("127.0.0.1"),
+            email="[email protected]",
+            hashed_password="$2b$12$zXXveodjipHZMoSxJz5ODul7Z9YeRJd0GeSBjpwHdqEtBbAFvEdre",  # 123456 -> 用 get_password_hash 加密后的字符串
+            avatar="/user/user-avatar.png",
+        )
+
+        with Session(DB_ENGINE) as session:
+            session.add(admin_user)
+            session.commit()
+
+    def init_user() -> bool:
+        """判断是否需要创建默认用户
+
+        Returns:
+            bool: 是否执行创建默认用户
+        """
+        with Session(DB_ENGINE) as session:
+            results = session.exec(select(UserInfo).where(UserInfo.user_id == 1)).first()
+
+        if results is None:
+            # 如果数据库为空，创建初始用户
+            create_default_user()
+            logger.info("created default user info")
+            return True
+
+        return False
+
+    def create_default_product_item():
+        """生成商品默认数据库"""
+        delivery_company_list = ["京东", "顺丰", "韵达", "圆通", "中通"]
+        departure_place_list = ["广州", "北京", "武汉", "杭州", "上海", "深圳", "成都"]
+        default_product_list = {
+            "beef": {
+                "product_name": "进口和牛羽下肉",
+                "heighlights": "富含铁质;营养价值高;肌肉纤维好;红白相间纹理;适合烧烤炖煮;草食动物来源",
+                "product_class": "食品",
+            },
+            "elec_toothblush": {
+                "product_name": "声波电动牙刷",
+                "heighlights": "高效清洁;减少手动压力;定时提醒;智能模式调节;无线充电;噪音低",
+                "product_class": "电子",
+            },
+            "lip_stick": {
+                "product_name": "唇膏",
+                "heighlights": "丰富色号;滋润保湿;显色度高;持久不脱色;易于涂抹;便携包装",
+                "product_class": "美妆",
+            },
+            "mask": {
+                "product_name": "光感润颜面膜",
+                "heighlights": "密集滋养;深层补水;急救修复;快速见效;定期护理;多种类型选择",
+                "product_class": "美妆",
+            },
+            "oled_tv": {
+                "product_name": "65英寸OLED电视",
+                "heighlights": "色彩鲜艳;对比度极高;响应速度快;无背光眩光;厚度较薄;自发光无需额外照明",
+                "product_class": "家电",
+            },
+            "pad": {
+                "product_name": "14英寸平板电脑",
+                "heighlights": "轻薄;触控操作;电池续航好;移动办公便利;娱乐性强;适合儿童学习",
+                "product_class": "电子",
+            },
+            "pants": {
+                "product_name": "速干运动裤",
+                "heighlights": "快干;伸缩自如;吸湿排汗;防风保暖;高腰设计;多口袋实用",
+                "product_class": "衣服",
+            },
+            "pen": {
+                "product_name": "墨水钢笔",
+                "heighlights": "耐用性;可书写性;不同颜色和类型;轻便设计;环保材料;易于携带",
+                "product_class": "文具",
+            },
+            "perfume": {
+                "product_name": "薰衣草淡香氛",
+                "heighlights": "浪漫优雅;花香调为主;情感表达;适合各种年龄;瓶身设计精致;提升女性魅力",
+                "product_class": "家居用品",
+            },
+            "shampoo": {
+                "product_name": "本草精华洗发露",
+                "heighlights": "温和配方;深层清洁;滋养头皮;丰富泡沫;易冲洗;适合各种发质",
+                "product_class": "日用品",
+            },
+            "wok": {
+                "product_name": "不粘煎炒锅",
+                "heighlights": "不粘涂层;耐磨耐用;导热快;易清洗;多种烹饪方式;设计人性化",
+                "product_class": "厨具",
+            },
+            "yoga_mat": {
+                "product_name": "瑜伽垫",
+                "heighlights": "防滑材质;吸湿排汗;厚度适中;耐用易清洁;各种瑜伽动作适用;轻巧便携",
+                "product_class": "运动",
+            },
+        }
+
+        with Session(DB_ENGINE) as session:
+            for product_key, product_info in default_product_list.items():
+                add_item = ProductInfo(
+                    **product_info,
+                    image_path=f"/{WEB_CONFIGS.PRODUCT_FILE_DIR}/{WEB_CONFIGS.IMAGES_DIR}/{product_key}.png",
+                    instruction=f"/{WEB_CONFIGS.PRODUCT_FILE_DIR}/{WEB_CONFIGS.INSTRUCTIONS_DIR}/{product_key}.md",
+                    departure_place=random.choice(departure_place_list),
+                    delivery_company=random.choice(delivery_company_list),
+                    selling_price=round(random.uniform(66.6, 1999.9), 2),
+                    amount=random.randint(999, 9999),
+                    user_id=1,
+                )
+                session.add(add_item)
+            session.commit()
+
+        logger.info("created default product info done!")
+
+    def create_default_streamer():
+
+        with Session(DB_ENGINE) as session:
+            streamer_item = StreamerInfo(
+                name="乐乐喵",
+                character="甜美;可爱;熟练使用各种网络热门梗造句;称呼客户为[家人们]",
+                avatar=f"/{WEB_CONFIGS.STREAMER_FILE_DIR}/{WEB_CONFIGS.STREAMER_INFO_FILES_DIR}/lelemiao.png",
+                base_mp4_path=f"/{WEB_CONFIGS.STREAMER_FILE_DIR}/{WEB_CONFIGS.STREAMER_INFO_FILES_DIR}/lelemiao.mp4",
+                poster_image=f"/{WEB_CONFIGS.STREAMER_FILE_DIR}/{WEB_CONFIGS.STREAMER_INFO_FILES_DIR}/lelemiao.png",
+                tts_reference_audio=f"/{WEB_CONFIGS.STREAMER_FILE_DIR}/{WEB_CONFIGS.STREAMER_INFO_FILES_DIR}/lelemiao.wav",
+                tts_reference_sentence="列车巡游银河，我不一定都能帮上忙，但只要是花钱能解决的事，尽管和我说吧。",
+                tts_weight_tag="艾丝妲",
+                user_id=1,
+            )
+            session.add(streamer_item)
+            session.commit()
+
+    def create_default_room():
+
+        with Session(DB_ENGINE) as session:
+
+            product_list = session.exec(
+                select(ProductInfo).where(ProductInfo.user_id == 1).order_by(ProductInfo.product_id)
+            ).all()
+
+            on_air_status = OnAirRoomStatusItem(user_id=1)
+            session.add(on_air_status)
+            session.commit()
+            session.refresh(on_air_status)
+
+            stream_item = StreamRoomInfo(
+                name="001",
+                user_id=1,
+                status_id=on_air_status.status_id,
+                streamer_id=1,
+            )
+            session.add(stream_item)
+            session.commit()
+            session.refresh(stream_item)
+
+            random_list = random.choices(product_list, k=3)
+            for product_random in random_list:
+                add_sales_info = SalesDocAndVideoInfo(product_id=product_random.product_id, room_id=stream_item.room_id)
+                session.add(add_sales_info)
+                session.commit()
+                session.refresh(add_sales_info)
+
+    # 主要逻辑
+    created = init_user()
+    if created:
+        create_default_product_item()  # 商品信息
+        create_default_streamer()  # 主播信息
+        create_default_room()  # 直播间信息

server/digital_human/digital_human_server.py

@@ -0,0 +1,68 @@
+from fastapi import FastAPI
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import PlainTextResponse
+from loguru import logger
+from pydantic import BaseModel
+
+
+from .modules.digital_human_worker import gen_digital_human_video_app, preprocess_digital_human_app
+
+
+app = FastAPI()
+
+
+class DigitalHumanItem(BaseModel):
+    user_id: str  # User 识别号，用于区分不用的用户调用
+    request_id: str  # 请求 ID，用于生成 TTS & 数字人
+    streamer_id: str  # 数字人 ID
+    tts_path: str = ""  # 文本
+    chunk_id: int = 0  # 句子 ID
+
+
+class DigitalHumanPreprocessItem(BaseModel):
+    user_id: str  # User 识别号，用于区分不用的用户调用
+    request_id: str  # 请求 ID，用于生成 TTS & 数字人
+    streamer_id: str  # 数字人 ID
+    video_path: str  # 数字人视频
+
+
+@app.post("/digital_human/gen")
+async def get_digital_human(dg_item: DigitalHumanItem):
+    """生成数字人视频"""
+    save_tag = (
+        dg_item.request_id + ".mp4" if dg_item.chunk_id == 0 else dg_item.request_id + f"-{str(dg_item.chunk_id).zfill(8)}.mp4"
+    )
+    mp4_path = await gen_digital_human_video_app(dg_item.streamer_id, dg_item.tts_path, save_tag)
+    logger.info(f"digital human mp4 path = {mp4_path}")
+    return {"user_id": dg_item.user_id, "request_id": dg_item.request_id, "digital_human_mp4_path": mp4_path}
+
+
+@app.post("/digital_human/preprocess")
+async def preprocess_digital_human(preprocess_item: DigitalHumanPreprocessItem):
+    """数字人视频预处理，用于新增数字人"""
+
+    _ = await preprocess_digital_human_app(str(preprocess_item.streamer_id), preprocess_item.video_path)
+
+    logger.info(f"digital human process for {preprocess_item.streamer_id} done")
+    return {"user_id": preprocess_item.user_id, "request_id": preprocess_item.request_id}
+
+
+@app.exception_handler(RequestValidationError)
+async def validation_exception_handler(request, exc):
+    """调 API 入参错误的回调接口
+
+    Args:
+        request (_type_): _description_
+        exc (_type_): _description_
+
+    Returns:
+        _type_: _description_
+    """
+    logger.info(request)
+    logger.info(exc)
+    return PlainTextResponse(str(exc), status_code=400)
+
+
+@app.get("/digital_human/check")
+async def check_server():
+    return {"message": "server enabled"}

server/digital_human/modules/__init__.py

@@ -0,0 +1,6 @@
+from torch import hub
+from ...web_configs import WEB_CONFIGS
+from pathlib import Path
+
+# 部分模型会使用 torch download 下载，需要设置路径
+hub.set_dir(str(Path(WEB_CONFIGS.DIGITAL_HUMAN_MODEL_DIR).joinpath("face-alignment")))

server/digital_human/modules/digital_human_worker.py

@@ -0,0 +1,33 @@
+from pathlib import Path
+from .realtime_inference import DIGITAL_HUMAN_HANDLER, gen_digital_human_preprocess, gen_digital_human_video
+from ...web_configs import WEB_CONFIGS
+
+
+async def gen_digital_human_video_app(stream_id, audio_path, save_tag):
+    if DIGITAL_HUMAN_HANDLER is None:
+        return None
+
+    save_path = gen_digital_human_video(
+        DIGITAL_HUMAN_HANDLER,
+        stream_id,
+        audio_path,
+        work_dir=str(Path(WEB_CONFIGS.DIGITAL_HUMAN_VIDEO_OUTPUT_PATH).absolute()),
+        video_path=save_tag,
+        fps=DIGITAL_HUMAN_HANDLER.fps,
+    )
+
+    return save_path
+
+
+async def preprocess_digital_human_app(stream_id, video_path):
+    if DIGITAL_HUMAN_HANDLER is None:
+        return None
+
+    res = gen_digital_human_preprocess(
+        DIGITAL_HUMAN_HANDLER,
+        stream_id,
+        work_dir=str(Path(WEB_CONFIGS.DIGITAL_HUMAN_VIDEO_OUTPUT_PATH).absolute()),
+        video_path=video_path,
+    )
+
+    return res

server/digital_human/modules/musetalk/models/unet.py

@@ -0,0 +1,43 @@
+import torch
+import torch.nn as nn
+import math
+import json
+
+from diffusers import UNet2DConditionModel
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model=384, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        b, seq_len, d_model = x.size()
+        pe = self.pe[:, :seq_len, :]
+        x = x + pe.to(x.device)
+        return x
+    
+class UNet():
+    def __init__(self, 
+                 unet_config,
+                 model_path,
+                 use_float16=False,
+        ):
+        with open(unet_config, 'r') as f:
+            unet_config = json.load(f)
+        self.model = UNet2DConditionModel(**unet_config)
+        self.pe = PositionalEncoding(d_model=384)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        weights = torch.load(model_path) if torch.cuda.is_available() else torch.load(model_path, map_location=self.device)
+        self.model.load_state_dict(weights)
+        if use_float16:
+            self.model = self.model.half()
+        self.model.to(self.device)
+    
+if __name__ == "__main__":
+    unet = UNet()

server/digital_human/modules/musetalk/models/vae.py

@@ -0,0 +1,149 @@
+import os
+
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as transforms
+from diffusers import AutoencoderKL
+
+
+class VAE():
+    """
+    VAE (Variational Autoencoder) class for image processing.
+    """
+
+    def __init__(self, model_path="./models/sd-vae-ft-mse/", resized_img=256, use_float16=False):
+        """
+        Initialize the VAE instance.
+
+        :param model_path: Path to the trained model.
+        :param resized_img: The size to which images are resized.
+        :param use_float16: Whether to use float16 precision.
+        """
+        self.model_path = model_path
+        self.vae = AutoencoderKL.from_pretrained(self.model_path)
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.vae.to(self.device)
+
+        if use_float16:
+            self.vae = self.vae.half()
+            self._use_float16 = True
+        else:
+            self._use_float16 = False
+
+        self.scaling_factor = self.vae.config.scaling_factor
+        self.transform = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+        self._resized_img = resized_img
+        self._mask_tensor = self.get_mask_tensor()
+        
+    def get_mask_tensor(self):
+        """
+        Creates a mask tensor for image processing.
+        :return: A mask tensor.
+        """
+        mask_tensor = torch.zeros((self._resized_img,self._resized_img))
+        mask_tensor[:self._resized_img//2,:] = 1
+        mask_tensor[mask_tensor< 0.5] = 0
+        mask_tensor[mask_tensor>= 0.5] = 1
+        return mask_tensor
+            
+    def preprocess_img(self,img_name,half_mask=False):
+        """
+        Preprocess an image for the VAE.
+
+        :param img_name: The image file path or a list of image file paths.
+        :param half_mask: Whether to apply a half mask to the image.
+        :return: A preprocessed image tensor.
+        """
+        window = []
+        if isinstance(img_name, str):
+            window_fnames = [img_name]
+            for fname in window_fnames:
+                img = cv2.imread(fname)
+                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+                img = cv2.resize(img, (self._resized_img, self._resized_img),
+                                     interpolation=cv2.INTER_LANCZOS4)
+                window.append(img)
+        else:
+            img = cv2.cvtColor(img_name, cv2.COLOR_BGR2RGB)
+            window.append(img)
+            
+        x = np.asarray(window) / 255.
+        x = np.transpose(x, (3, 0, 1, 2))
+        x = torch.squeeze(torch.FloatTensor(x))
+        if half_mask:
+            x = x * (self._mask_tensor>0.5)
+        x = self.transform(x)
+        
+        x = x.unsqueeze(0) # [1, 3, 256, 256] torch tensor
+        x = x.to(self.vae.device)
+
+        return x
+
+    def encode_latents(self,image):
+        """
+        Encode an image into latent variables.
+
+        :param image: The image tensor to encode.
+        :return: The encoded latent variables.
+        """
+        with torch.no_grad():
+            init_latent_dist = self.vae.encode(image.to(self.vae.dtype)).latent_dist
+        init_latents = self.scaling_factor * init_latent_dist.sample()
+        return init_latents
+    
+    def decode_latents(self, latents):
+        """
+        Decode latent variables back into an image.
+        :param latents: The latent variables to decode.
+        :return: A NumPy array representing the decoded image.
+        """
+        latents = (1/  self.scaling_factor) * latents
+        image = self.vae.decode(latents.to(self.vae.dtype)).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        image = image.detach().cpu().permute(0, 2, 3, 1).float().numpy()
+        image = (image * 255).round().astype("uint8")
+        image = image[...,::-1] # RGB to BGR
+        return image
+    
+    def get_latents_for_unet(self,img):
+        """
+        Prepare latent variables for a U-Net model.
+        :param img: The image to process.
+        :return: A concatenated tensor of latents for U-Net input.
+        """
+        
+        ref_image = self.preprocess_img(img,half_mask=True) # [1, 3, 256, 256] RGB, torch tensor
+        masked_latents = self.encode_latents(ref_image) # [1, 4, 32, 32], torch tensor
+        ref_image = self.preprocess_img(img,half_mask=False) # [1, 3, 256, 256] RGB, torch tensor
+        ref_latents = self.encode_latents(ref_image) # [1, 4, 32, 32], torch tensor
+        latent_model_input = torch.cat([masked_latents, ref_latents], dim=1)
+        return latent_model_input
+
+if __name__ == "__main__":
+    vae_mode_path = "./models/sd-vae-ft-mse/"
+    vae = VAE(model_path = vae_mode_path,use_float16=False)
+    img_path = "./results/sun001_crop/00000.png"
+    
+    crop_imgs_path = "./results/sun001_crop/"
+    latents_out_path = "./results/latents/"
+    if not os.path.exists(latents_out_path):
+        os.mkdir(latents_out_path)
+
+    files = os.listdir(crop_imgs_path)
+    files.sort()
+    files = [file for file in files if file.split(".")[-1] == "png"]
+
+    for file in files:
+        index = file.split(".")[0]
+        img_path = crop_imgs_path + file
+        latents = vae.get_latents_for_unet(img_path)
+        print(img_path,"latents",latents.size())
+        #torch.save(latents,os.path.join(latents_out_path,index+".pt"))
+        #reload_tensor = torch.load('tensor.pt')
+        #print(reload_tensor.size())
+        
+
+    

server/digital_human/modules/musetalk/utils/__init__.py

@@ -0,0 +1,5 @@
+import sys
+from os.path import abspath, dirname
+current_dir = dirname(abspath(__file__))
+parent_dir = dirname(current_dir)
+sys.path.append(parent_dir+'/utils')

server/digital_human/modules/musetalk/utils/blending.py

@@ -0,0 +1,110 @@
+from PIL import Image
+import numpy as np
+import cv2
+from face_parsing import FaceParsing
+
+
+def init_face_parsing_model(
+    resnet_path="./models/face-parse-bisent/resnet18-5c106cde.pth", face_model_pth="./models/face-parse-bisent/79999_iter.pth"
+):
+    fp_model = FaceParsing(resnet_path, face_model_pth)
+    return fp_model
+
+
+def get_crop_box(box, expand):
+    x, y, x1, y1 = box
+    x_c, y_c = (x + x1) // 2, (y + y1) // 2
+    w, h = x1 - x, y1 - y
+    s = int(max(w, h) // 2 * expand)
+    crop_box = [x_c - s, y_c - s, x_c + s, y_c + s]
+    return crop_box, s
+
+
+def face_seg(image, fp_model):
+    seg_image = fp_model(image)
+    if seg_image is None:
+        print("error, no person_segment")
+        return None
+
+    seg_image = seg_image.resize(image.size)
+    return seg_image
+
+
+def get_image(image, face, face_box, fp_model, upper_boundary_ratio=0.5, expand=1.2):
+    # print(image.shape)
+    # print(face.shape)
+
+    body = Image.fromarray(image[:, :, ::-1])
+    face = Image.fromarray(face[:, :, ::-1])
+
+    x, y, x1, y1 = face_box
+    # print(x1-x,y1-y)
+    crop_box, s = get_crop_box(face_box, expand)
+    x_s, y_s, x_e, y_e = crop_box
+    face_position = (x, y)
+
+    face_large = body.crop(crop_box)
+    ori_shape = face_large.size
+
+    mask_image = face_seg(face_large, fp_model)
+    mask_small = mask_image.crop((x - x_s, y - y_s, x1 - x_s, y1 - y_s))
+    mask_image = Image.new("L", ori_shape, 0)
+    mask_image.paste(mask_small, (x - x_s, y - y_s, x1 - x_s, y1 - y_s))
+
+    # keep upper_boundary_ratio of talking area
+    width, height = mask_image.size
+    top_boundary = int(height * upper_boundary_ratio)
+    modified_mask_image = Image.new("L", ori_shape, 0)
+    modified_mask_image.paste(mask_image.crop((0, top_boundary, width, height)), (0, top_boundary))
+
+    blur_kernel_size = int(0.1 * ori_shape[0] // 2 * 2) + 1
+    mask_array = cv2.GaussianBlur(np.array(modified_mask_image), (blur_kernel_size, blur_kernel_size), 0)
+    mask_image = Image.fromarray(mask_array)
+
+    face_large.paste(face, (x - x_s, y - y_s, x1 - x_s, y1 - y_s))
+    body.paste(face_large, crop_box[:2], mask_image)
+    body = np.array(body)
+    return body[:, :, ::-1]
+
+
+def get_image_prepare_material(image, face_box, fp_model, upper_boundary_ratio=0.5, expand=1.2):
+    body = Image.fromarray(image[:, :, ::-1])
+
+    x, y, x1, y1 = face_box
+    # print(x1-x,y1-y)
+    crop_box, s = get_crop_box(face_box, expand)
+    x_s, y_s, x_e, y_e = crop_box
+
+    face_large = body.crop(crop_box)
+    ori_shape = face_large.size
+
+    mask_image = face_seg(face_large, fp_model)
+    mask_small = mask_image.crop((x - x_s, y - y_s, x1 - x_s, y1 - y_s))
+    mask_image = Image.new("L", ori_shape, 0)
+    mask_image.paste(mask_small, (x - x_s, y - y_s, x1 - x_s, y1 - y_s))
+
+    # keep upper_boundary_ratio of talking area
+    width, height = mask_image.size
+    top_boundary = int(height * upper_boundary_ratio)
+    modified_mask_image = Image.new("L", ori_shape, 0)
+    modified_mask_image.paste(mask_image.crop((0, top_boundary, width, height)), (0, top_boundary))
+
+    blur_kernel_size = int(0.1 * ori_shape[0] // 2 * 2) + 1
+    mask_array = cv2.GaussianBlur(np.array(modified_mask_image), (blur_kernel_size, blur_kernel_size), 0)
+    return mask_array, crop_box
+
+
+def get_image_blending(image, face, face_box, mask_array, crop_box):
+    body = Image.fromarray(image[:, :, ::-1])
+    face = Image.fromarray(face[:, :, ::-1])
+
+    x, y, x1, y1 = face_box
+    x_s, y_s, x_e, y_e = crop_box
+    face_large = body.crop(crop_box)
+
+    mask_image = Image.fromarray(mask_array)
+    mask_image = mask_image.convert("L")
+    face_large.paste(face, (x - x_s, y - y_s, x1 - x_s, y1 - y_s))
+    body.paste(face_large, crop_box[:2], mask_image)
+    body = np.array(body)
+    return body[:, :, ::-1]

server/digital_human/modules/musetalk/utils/dwpose/default_runtime.py

@@ -0,0 +1,54 @@
+default_scope = 'mmpose'
+
+# hooks
+default_hooks = dict(
+    timer=dict(type='IterTimerHook'),
+    logger=dict(type='LoggerHook', interval=50),
+    param_scheduler=dict(type='ParamSchedulerHook'),
+    checkpoint=dict(type='CheckpointHook', interval=10),
+    sampler_seed=dict(type='DistSamplerSeedHook'),
+    visualization=dict(type='PoseVisualizationHook', enable=False),
+    badcase=dict(
+        type='BadCaseAnalysisHook',
+        enable=False,
+        out_dir='badcase',
+        metric_type='loss',
+        badcase_thr=5))
+
+# custom hooks
+custom_hooks = [
+    # Synchronize model buffers such as running_mean and running_var in BN
+    # at the end of each epoch
+    dict(type='SyncBuffersHook')
+]
+
+# multi-processing backend
+env_cfg = dict(
+    cudnn_benchmark=False,
+    mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0),
+    dist_cfg=dict(backend='nccl'),
+)
+
+# visualizer
+vis_backends = [
+    dict(type='LocalVisBackend'),
+    # dict(type='TensorboardVisBackend'),
+    # dict(type='WandbVisBackend'),
+]
+visualizer = dict(
+    type='PoseLocalVisualizer', vis_backends=vis_backends, name='visualizer')
+
+# logger
+log_processor = dict(
+    type='LogProcessor', window_size=50, by_epoch=True, num_digits=6)
+log_level = 'INFO'
+load_from = None
+resume = False
+
+# file I/O backend
+backend_args = dict(backend='local')
+
+# training/validation/testing progress
+train_cfg = dict(by_epoch=True)
+val_cfg = dict()
+test_cfg = dict()

server/digital_human/modules/musetalk/utils/dwpose/rtmpose-l_8xb32-270e_coco-ubody-wholebody-384x288.py

@@ -0,0 +1,257 @@
+#_base_ = ['../../../_base_/default_runtime.py']
+_base_ = ['default_runtime.py']
+
+# runtime
+max_epochs = 270
+stage2_num_epochs = 30
+base_lr = 4e-3
+train_batch_size = 8
+val_batch_size = 8
+
+train_cfg = dict(max_epochs=max_epochs, val_interval=10)
+randomness = dict(seed=21)
+
+# optimizer
+optim_wrapper = dict(
+    type='OptimWrapper',
+    optimizer=dict(type='AdamW', lr=base_lr, weight_decay=0.05),
+    paramwise_cfg=dict(
+        norm_decay_mult=0, bias_decay_mult=0, bypass_duplicate=True))
+
+# learning rate
+param_scheduler = [
+    dict(
+        type='LinearLR',
+        start_factor=1.0e-5,
+        by_epoch=False,
+        begin=0,
+        end=1000),
+    dict(
+        # use cosine lr from 150 to 300 epoch
+        type='CosineAnnealingLR',
+        eta_min=base_lr * 0.05,
+        begin=max_epochs // 2,
+        end=max_epochs,
+        T_max=max_epochs // 2,
+        by_epoch=True,
+        convert_to_iter_based=True),
+]
+
+# automatically scaling LR based on the actual training batch size
+auto_scale_lr = dict(base_batch_size=512)
+
+# codec settings
+codec = dict(
+    type='SimCCLabel',
+    input_size=(288, 384),
+    sigma=(6., 6.93),
+    simcc_split_ratio=2.0,
+    normalize=False,
+    use_dark=False)
+
+# model settings
+model = dict(
+    type='TopdownPoseEstimator',
+    data_preprocessor=dict(
+        type='PoseDataPreprocessor',
+        mean=[123.675, 116.28, 103.53],
+        std=[58.395, 57.12, 57.375],
+        bgr_to_rgb=True),
+    backbone=dict(
+        _scope_='mmdet',
+        type='CSPNeXt',
+        arch='P5',
+        expand_ratio=0.5,
+        deepen_factor=1.,
+        widen_factor=1.,
+        out_indices=(4, ),
+        channel_attention=True,
+        norm_cfg=dict(type='SyncBN'),
+        act_cfg=dict(type='SiLU'),
+        init_cfg=dict(
+            type='Pretrained',
+            prefix='backbone.',
+            checkpoint='https://download.openmmlab.com/mmpose/v1/projects/'
+            'rtmpose/cspnext-l_udp-aic-coco_210e-256x192-273b7631_20230130.pth'  # noqa: E501
+        )),
+    head=dict(
+        type='RTMCCHead',
+        in_channels=1024,
+        out_channels=133,
+        input_size=codec['input_size'],
+        in_featuremap_size=(9, 12),
+        simcc_split_ratio=codec['simcc_split_ratio'],
+        final_layer_kernel_size=7,
+        gau_cfg=dict(
+            hidden_dims=256,
+            s=128,
+            expansion_factor=2,
+            dropout_rate=0.,
+            drop_path=0.,
+            act_fn='SiLU',
+            use_rel_bias=False,
+            pos_enc=False),
+        loss=dict(
+            type='KLDiscretLoss',
+            use_target_weight=True,
+            beta=10.,
+            label_softmax=True),
+        decoder=codec),
+    test_cfg=dict(flip_test=True, ))
+
+# base dataset settings
+dataset_type = 'UBody2dDataset'
+data_mode = 'topdown'
+data_root = 'data/UBody/'
+
+backend_args = dict(backend='local')
+
+scenes = [
+    'Magic_show', 'Entertainment', 'ConductMusic', 'Online_class', 'TalkShow',
+    'Speech', 'Fitness', 'Interview', 'Olympic', 'TVShow', 'Singing',
+    'SignLanguage', 'Movie', 'LiveVlog', 'VideoConference'
+]
+
+train_datasets = [
+    dict(
+        type='CocoWholeBodyDataset',
+        data_root='data/coco/',
+        data_mode=data_mode,
+        ann_file='annotations/coco_wholebody_train_v1.0.json',
+        data_prefix=dict(img='train2017/'),
+        pipeline=[])
+]
+
+for scene in scenes:
+    train_dataset = dict(
+        type=dataset_type,
+        data_root=data_root,
+        data_mode=data_mode,
+        ann_file=f'annotations/{scene}/train_annotations.json',
+        data_prefix=dict(img='images/'),
+        pipeline=[],
+        sample_interval=10)
+    train_datasets.append(train_dataset)
+
+# pipelines
+train_pipeline = [
+    dict(type='LoadImage', backend_args=backend_args),
+    dict(type='GetBBoxCenterScale'),
+    dict(type='RandomFlip', direction='horizontal'),
+    dict(type='RandomHalfBody'),
+    dict(
+        type='RandomBBoxTransform', scale_factor=[0.5, 1.5], rotate_factor=90),
+    dict(type='TopdownAffine', input_size=codec['input_size']),
+    dict(type='mmdet.YOLOXHSVRandomAug'),
+    dict(
+        type='Albumentation',
+        transforms=[
+            dict(type='Blur', p=0.1),
+            dict(type='MedianBlur', p=0.1),
+            dict(
+                type='CoarseDropout',
+                max_holes=1,
+                max_height=0.4,
+                max_width=0.4,
+                min_holes=1,
+                min_height=0.2,
+                min_width=0.2,
+                p=1.0),
+        ]),
+    dict(type='GenerateTarget', encoder=codec),
+    dict(type='PackPoseInputs')
+]
+val_pipeline = [
+    dict(type='LoadImage', backend_args=backend_args),
+    dict(type='GetBBoxCenterScale'),
+    dict(type='TopdownAffine', input_size=codec['input_size']),
+    dict(type='PackPoseInputs')
+]
+
+train_pipeline_stage2 = [
+    dict(type='LoadImage', backend_args=backend_args),
+    dict(type='GetBBoxCenterScale'),
+    dict(type='RandomFlip', direction='horizontal'),
+    dict(type='RandomHalfBody'),
+    dict(
+        type='RandomBBoxTransform',
+        shift_factor=0.,
+        scale_factor=[0.5, 1.5],
+        rotate_factor=90),
+    dict(type='TopdownAffine', input_size=codec['input_size']),
+    dict(type='mmdet.YOLOXHSVRandomAug'),
+    dict(
+        type='Albumentation',
+        transforms=[
+            dict(type='Blur', p=0.1),
+            dict(type='MedianBlur', p=0.1),
+            dict(
+                type='CoarseDropout',
+                max_holes=1,
+                max_height=0.4,
+                max_width=0.4,
+                min_holes=1,
+                min_height=0.2,
+                min_width=0.2,
+                p=0.5),
+        ]),
+    dict(type='GenerateTarget', encoder=codec),
+    dict(type='PackPoseInputs')
+]
+
+# data loaders
+train_dataloader = dict(
+    batch_size=train_batch_size,
+    num_workers=10,
+    persistent_workers=True,
+    sampler=dict(type='DefaultSampler', shuffle=True),
+    dataset=dict(
+        type='CombinedDataset',
+        metainfo=dict(from_file='configs/_base_/datasets/coco_wholebody.py'),
+        datasets=train_datasets,
+        pipeline=train_pipeline,
+        test_mode=False,
+    ))
+
+val_dataloader = dict(
+    batch_size=val_batch_size,
+    num_workers=10,
+    persistent_workers=True,
+    drop_last=False,
+    sampler=dict(type='DefaultSampler', shuffle=False, round_up=False),
+    dataset=dict(
+        type='CocoWholeBodyDataset',
+        data_root=data_root,
+        data_mode=data_mode,
+        ann_file='data/coco/annotations/coco_wholebody_val_v1.0.json',
+        bbox_file='data/coco/person_detection_results/'
+        'COCO_val2017_detections_AP_H_56_person.json',
+        data_prefix=dict(img='coco/val2017/'),
+        test_mode=True,
+        pipeline=val_pipeline,
+    ))
+test_dataloader = val_dataloader
+
+# hooks
+default_hooks = dict(
+    checkpoint=dict(
+        save_best='coco-wholebody/AP', rule='greater', max_keep_ckpts=1))
+
+custom_hooks = [
+    dict(
+        type='EMAHook',
+        ema_type='ExpMomentumEMA',
+        momentum=0.0002,
+        update_buffers=True,
+        priority=49),
+    dict(
+        type='mmdet.PipelineSwitchHook',
+        switch_epoch=max_epochs - stage2_num_epochs,
+        switch_pipeline=train_pipeline_stage2)
+]
+
+# evaluators
+val_evaluator = dict(
+    type='CocoWholeBodyMetric',
+    ann_file='data/coco/annotations/coco_wholebody_val_v1.0.json')
+test_evaluator = val_evaluator

server/digital_human/modules/musetalk/utils/face_detection/README.md

@@ -0,0 +1 @@
+The code for Face Detection in this folder has been taken from the wonderful [face_alignment](https://github.com/1adrianb/face-alignment) repository. This has been modified to take batches of faces at a time.