70 lines
2.4 KiB
Python
70 lines
2.4 KiB
Python
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[1], 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[1]
|
||
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)}"
|
||
}
|
||
|