gangyan/langchain-chat/server/chat/gen_title.py

from fastapi import Body, HTTPException
from configs import LLM_MODELS, TEMPERATURE, MAX_TOKENS
from configs.model_config import DEEPSEEK_MODELS
from server.chat.policy_fun_iast import get_llm_model_response
from typing import List, Dict, Optional, Union
import logging
from pydantic import BaseModel
from server.chat.utils import History

# 初始化日志
logger = logging.getLogger(__name__)


async def gen_title(
            history: List[History] = Body([],
                                        description="历史对话，设为一个整数可以从数据库中读取历史消息",
                                        examples=[[
                                            {"role": "user",
                                            "content": "我们来玩成语接龙，我先来，生龙活虎"},
                                            {"role": "assistant", "content": "虎头虎脑"}]]
                                        ),
            model_name: Optional[str] = Body(LLM_MODELS[0], description="LLM 模型名称。"),
):
    """
    根据一轮对话历史生成简洁标题\n
    :param history:第一轮对话历史\n
    :param model_name:模型名称，非必传\n
    :return: 返回一个字典，包含title字段，值为生成的标题。
    """
    history = [History.to_msg_tuple(h) for h in history]

    if model_name == "R1-70B":
        model_name = DEEPSEEK_MODELS[1]
    elif model_name == "QIANWEN":
        model_name = LLM_MODELS[0]
    else:
        model_name = model_name

    try:
        logger.info("开始生成对话标题...")

        # 调用模型生成标题
        conv_title = get_llm_model_response(
            strategy_name="gen_title",
            llm_model_name=model_name,
            template_prompt_name="gen_title",
            prompt_param_dict={"history": history},
            temperature=TEMPERATURE,
            max_tokens=MAX_TOKENS
        )

        # 清理返回结果
        clean_title = conv_title.strip("“”\"'")
        logger.info(f"成功生成对话标题: {clean_title}")

        return {
            "status_code": 200,
            "data": {"title": clean_title},
        }

    except HTTPException:
        raise  # 主动抛出的HTTP异常直接传递
    except Exception as e:
        logger.exception("标题生成失败")
        return {
            "status_code": 500,
            "data": {"title": ""},
            "error": f"标题生成失败: {str(e)}"
        }