412 lines
16 KiB
Python
412 lines
16 KiB
Python
import json
|
||
from datetime import datetime, timezone
|
||
|
||
from fastapi import APIRouter, Depends, Query
|
||
from fastapi.responses import StreamingResponse
|
||
from sqlmodel import Session, select, func
|
||
|
||
from ..database import get_session, engine
|
||
from ..models import AIConfig, ChatMessage, Novel
|
||
from ..schemas import ChatRequest, ChatMessageRead, ChatMessageList
|
||
from ..services.ai_provider import (
|
||
stream_chat, simple_completion,
|
||
build_assistant_message, build_tool_results_messages, ToolCall,
|
||
)
|
||
from ..services.tools import get_openai_tools, get_anthropic_tools, execute_tool
|
||
|
||
router = APIRouter(prefix="/api/chat", tags=["chat"])
|
||
|
||
# 最大工具调用轮次
|
||
MAX_TOOL_ROUNDS = 8
|
||
|
||
# 工具名称到中文描述
|
||
TOOL_LABELS: dict[str, str] = {
|
||
"get_novel_info": "查看小说信息",
|
||
"update_novel_info": "更新小说信息",
|
||
"list_outlines": "查询大纲",
|
||
"create_outline": "创建大纲",
|
||
"update_outline": "更新大纲",
|
||
"delete_outline": "删除大纲",
|
||
"reorder_outlines": "调整大纲排序",
|
||
"get_world_setting": "查询世界观",
|
||
"save_world_setting": "保存世界观",
|
||
"list_characters": "查询角色",
|
||
"create_character": "创建角色",
|
||
"update_character": "更新角色",
|
||
"delete_character": "删除角色",
|
||
"list_chapters": "查询章节",
|
||
"create_chapter": "创建章节",
|
||
"update_chapter": "更新章节",
|
||
"delete_chapter": "删除章节",
|
||
}
|
||
|
||
|
||
def _build_system_prompt(
|
||
novel_id: int | None,
|
||
session: Session,
|
||
page_context: str | None = None,
|
||
tools_enabled: bool = False,
|
||
) -> str:
|
||
"""根据小说上下文和用户所在页面构建 system prompt"""
|
||
parts = [
|
||
"你是 Writing Red Dot 写作助手，专注于小说创作领域。",
|
||
"你擅长：构思情节、塑造角色、打磨文笔、设计世界观、规划大纲结构。",
|
||
]
|
||
|
||
if novel_id:
|
||
novel = session.get(Novel, novel_id)
|
||
if novel:
|
||
parts.append(f"\n当前小说：《{novel.title}》（ID: {novel.id}）")
|
||
parts.append(f"小说简介：{novel.description or '暂无'}")
|
||
|
||
if page_context:
|
||
parts.append(f"\n用户当前正在查看：{page_context}。请结合当前页面场景提供针对性建议。")
|
||
|
||
if tools_enabled:
|
||
parts.append(
|
||
"\n## 你的工具能力\n"
|
||
"你是一个具备工具调用能力的写作助手。你拥有以下专用工具，没有其他任何工具：\n"
|
||
"\n### 小说信息\n"
|
||
"1. get_novel_info — 查看当前小说的标题、简介等基础信息\n"
|
||
"2. update_novel_info — 修改小说标题或简介\n"
|
||
"\n### 大纲管理\n"
|
||
"3. list_outlines — 列出所有大纲节点\n"
|
||
"4. create_outline — 创建新的大纲节点（需要提供摘要，可选详情和父节点ID）\n"
|
||
"5. update_outline — 更新指定大纲节点的摘要或详情\n"
|
||
"6. delete_outline — 删除指定大纲节点\n"
|
||
"\n### 世界观\n"
|
||
"7. get_world_setting — 获取世界观设定内容\n"
|
||
"8. save_world_setting — 保存或更新世界观设定\n"
|
||
"\n### 角色管理\n"
|
||
"9. list_characters — 列出所有角色\n"
|
||
"10. create_character — 创建新角色（需要提供名称，可选描述）\n"
|
||
"11. update_character — 更新指定角色的名称或描述\n"
|
||
"12. delete_character — 删除指定角色\n"
|
||
"\n### 章节管理\n"
|
||
"13. list_chapters — 列出所有章节\n"
|
||
"14. create_chapter — 创建新章节（需要提供标题，可选正文）\n"
|
||
"15. update_chapter — 更新指定章节的标题或正文\n"
|
||
"16. delete_chapter — 删除指定章节\n"
|
||
"\n当用户询问你有哪些工具或能力时，请据实回答上述工具。"
|
||
"\n当用户要求查看或修改小说信息、大纲、世界观、角色或章节时，请主动使用相应工具完成操作，并简要说明结果。"
|
||
)
|
||
if not novel_id:
|
||
parts.append("\n注意：工具操作需要关联到具体小说。如果用户需要使用工具，请提示他们先进入某本小说。")
|
||
|
||
return "\n".join(parts)
|
||
|
||
|
||
def _get_tools(provider: str) -> list[dict]:
|
||
if provider == "anthropic":
|
||
return get_anthropic_tools()
|
||
return get_openai_tools()
|
||
|
||
|
||
def _sse(data: dict) -> str:
|
||
return f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
|
||
|
||
|
||
def _save_assistant_message(novel_id: int | None, full_response: str, tool_call_logs: list[str]):
|
||
"""保存 AI 回复到 DB（含工具调用摘要）。在 finally 中调用，确保断连也能保存。"""
|
||
save_content = full_response
|
||
if tool_call_logs:
|
||
tool_summary = "\n".join(tool_call_logs)
|
||
save_content = f"[以下操作已执行]\n{tool_summary}\n\n{full_response}"
|
||
if save_content.strip():
|
||
with Session(engine) as s:
|
||
ai_msg = ChatMessage(
|
||
novel_id=novel_id,
|
||
role="assistant",
|
||
content=save_content,
|
||
)
|
||
s.add(ai_msg)
|
||
s.commit()
|
||
|
||
|
||
@router.post("/stream")
|
||
async def chat_stream(
|
||
data: ChatRequest,
|
||
session: Session = Depends(get_session),
|
||
):
|
||
"""SSE 流式对话端点，支持 tool calling + 用户审批"""
|
||
config = session.exec(select(AIConfig)).first()
|
||
if not config or not config.api_key:
|
||
return StreamingResponse(
|
||
_error_stream("请先配置AI服务"),
|
||
media_type="text/event-stream",
|
||
)
|
||
|
||
# 保存用户消息（仅首次请求，续传审批时不保存）
|
||
if not data.pending_tool_calls:
|
||
user_msg = data.messages[-1] if data.messages else None
|
||
if user_msg and user_msg.role == "user":
|
||
db_msg = ChatMessage(
|
||
novel_id=data.novel_id,
|
||
role="user",
|
||
content=user_msg.content,
|
||
)
|
||
session.add(db_msg)
|
||
session.commit()
|
||
|
||
# 构建消息列表
|
||
messages = [{"role": m.role, "content": m.content} for m in data.messages]
|
||
|
||
# 构建上下文
|
||
use_tools = bool(data.tools_enabled)
|
||
system_prompt = _build_system_prompt(
|
||
data.novel_id, session, data.page_context, use_tools
|
||
)
|
||
tools = _get_tools(config.provider) if use_tools else None
|
||
|
||
async def generate():
|
||
nonlocal messages
|
||
|
||
full_response = data.assistant_text or ""
|
||
usage_info: dict = {}
|
||
# 记录本次对话中的工具调用摘要（用于保存到 DB）
|
||
tool_call_logs: list[str] = []
|
||
|
||
try:
|
||
# ── 阶段一：如果有待审批的工具调用，先执行它们 ──
|
||
if data.pending_tool_calls:
|
||
tool_calls = [
|
||
ToolCall(id=tc.id, name=tc.name, arguments=tc.arguments)
|
||
for tc in data.pending_tool_calls
|
||
]
|
||
|
||
# 构建 assistant 消息（含工具调用）
|
||
assistant_msg = build_assistant_message(
|
||
config.provider, data.assistant_text or "", tool_calls
|
||
)
|
||
messages.append(assistant_msg)
|
||
|
||
# 执行每个工具，发送结果
|
||
tool_results = []
|
||
for tc in tool_calls:
|
||
label = TOOL_LABELS.get(tc.name, tc.name)
|
||
result = execute_tool(tc.name, tc.arguments, data.novel_id)
|
||
tool_results.append({"id": tc.id, "result": result})
|
||
tool_call_logs.append(f"[工具调用: {label}] 参数: {json.dumps(tc.arguments, ensure_ascii=False)} → 结果: {result}")
|
||
yield _sse({"tool_result": {
|
||
"name": tc.name,
|
||
"label": label,
|
||
"result": result,
|
||
}})
|
||
|
||
# 追加工具结果消息
|
||
result_msgs = build_tool_results_messages(config.provider, tool_results)
|
||
messages.extend(result_msgs)
|
||
|
||
# ── 阶段二：LLM 循环 ──
|
||
for _round in range(MAX_TOOL_ROUNDS):
|
||
round_text = ""
|
||
round_tool_calls = []
|
||
|
||
async for event in stream_chat(config, messages, system_prompt, tools):
|
||
if event.text:
|
||
round_text += event.text
|
||
yield _sse({"content": event.text})
|
||
if event.usage:
|
||
for k, v in event.usage.items():
|
||
if v:
|
||
usage_info[k] = usage_info.get(k, 0) + v
|
||
if event.tool_calls:
|
||
round_tool_calls = event.tool_calls
|
||
|
||
# 没有工具调用 → 完成
|
||
if not round_tool_calls:
|
||
full_response += round_text
|
||
break
|
||
|
||
# 有工具调用
|
||
full_response += round_text
|
||
calls_data = [
|
||
{
|
||
"id": tc.id,
|
||
"name": tc.name,
|
||
"label": TOOL_LABELS.get(tc.name, tc.name),
|
||
"arguments": tc.arguments,
|
||
}
|
||
for tc in round_tool_calls
|
||
]
|
||
|
||
if data.auto_approve_tools:
|
||
# 自动执行模式：直接执行工具，不暂停
|
||
yield _sse({"tool_calls_auto": calls_data})
|
||
|
||
assistant_msg = build_assistant_message(
|
||
config.provider, full_response, round_tool_calls
|
||
)
|
||
messages.append(assistant_msg)
|
||
|
||
tool_results = []
|
||
for tc in round_tool_calls:
|
||
label = TOOL_LABELS.get(tc.name, tc.name)
|
||
result = execute_tool(tc.name, tc.arguments, data.novel_id)
|
||
tool_results.append({"id": tc.id, "result": result})
|
||
tool_call_logs.append(f"[工具调用: {label}] 参数: {json.dumps(tc.arguments, ensure_ascii=False)} → 结果: {result}")
|
||
yield _sse({"tool_result": {
|
||
"name": tc.name,
|
||
"label": label,
|
||
"result": result,
|
||
}})
|
||
|
||
result_msgs = build_tool_results_messages(config.provider, tool_results)
|
||
messages.extend(result_msgs)
|
||
# 继续循环，让 LLM 基于工具结果生成回复
|
||
continue
|
||
else:
|
||
# 审批模式：暂停，发给前端审批
|
||
yield _sse({
|
||
"tool_calls_pending": calls_data,
|
||
"assistant_text": full_response,
|
||
})
|
||
yield _sse({"done": True, "pending": True, "usage": usage_info or None})
|
||
return # 暂停，不保存到 DB
|
||
|
||
except Exception as e:
|
||
yield _sse({"error": str(e)})
|
||
finally:
|
||
# 无论连接是否断开，都保存已有的回复和工具调用到 DB
|
||
# 这确保刷新页面后 loadHistory 能恢复正确状态
|
||
_save_assistant_message(data.novel_id, full_response, tool_call_logs)
|
||
|
||
yield _sse({"done": True, "usage": usage_info or None})
|
||
|
||
return StreamingResponse(
|
||
generate(),
|
||
media_type="text/event-stream",
|
||
headers={
|
||
"Cache-Control": "no-cache",
|
||
"X-Accel-Buffering": "no",
|
||
},
|
||
)
|
||
|
||
|
||
async def _error_stream(message: str):
|
||
yield _sse({"error": message})
|
||
yield _sse({"done": True})
|
||
|
||
|
||
@router.get("/messages", response_model=ChatMessageList)
|
||
def list_messages(
|
||
novel_id: int | None = Query(default=None),
|
||
limit: int = Query(default=50, ge=1, le=200),
|
||
session: Session = Depends(get_session),
|
||
):
|
||
query = select(ChatMessage)
|
||
count_query = select(func.count(ChatMessage.id))
|
||
if novel_id is not None:
|
||
query = query.where(ChatMessage.novel_id == novel_id)
|
||
count_query = count_query.where(ChatMessage.novel_id == novel_id)
|
||
else:
|
||
query = query.where(ChatMessage.novel_id.is_(None)) # type: ignore
|
||
count_query = count_query.where(ChatMessage.novel_id.is_(None)) # type: ignore
|
||
|
||
total = session.exec(count_query).one()
|
||
items = session.exec(
|
||
query.order_by(ChatMessage.created_at.desc()).limit(limit)
|
||
).all()
|
||
items = list(reversed(items))
|
||
return ChatMessageList(
|
||
items=[ChatMessageRead.model_validate(m, from_attributes=True) for m in items],
|
||
total=total,
|
||
)
|
||
|
||
|
||
@router.delete("/messages", status_code=204)
|
||
def clear_messages(
|
||
novel_id: int | None = Query(default=None),
|
||
session: Session = Depends(get_session),
|
||
):
|
||
query = select(ChatMessage)
|
||
if novel_id is not None:
|
||
query = query.where(ChatMessage.novel_id == novel_id)
|
||
else:
|
||
query = query.where(ChatMessage.novel_id.is_(None)) # type: ignore
|
||
messages = session.exec(query).all()
|
||
for msg in messages:
|
||
session.delete(msg)
|
||
session.commit()
|
||
|
||
|
||
# 保留最近对话轮次数（每轮 = 1 user + 1 assistant）
|
||
KEEP_RECENT_ROUNDS = 2
|
||
|
||
|
||
@router.post("/compress")
|
||
async def compress_context(
|
||
novel_id: int | None = Query(default=None),
|
||
session: Session = Depends(get_session),
|
||
):
|
||
"""压缩对话上下文：用 LLM 总结旧消息，仅保留最近几轮对话。"""
|
||
config = session.exec(select(AIConfig)).first()
|
||
if not config or not config.api_key:
|
||
return {"error": "请先配置AI服务"}
|
||
|
||
# 查询所有消息（按时间正序）
|
||
query = select(ChatMessage)
|
||
if novel_id is not None:
|
||
query = query.where(ChatMessage.novel_id == novel_id)
|
||
else:
|
||
query = query.where(ChatMessage.novel_id.is_(None)) # type: ignore
|
||
all_msgs = list(session.exec(query.order_by(ChatMessage.created_at)).all())
|
||
|
||
if len(all_msgs) <= KEEP_RECENT_ROUNDS * 2:
|
||
return {"compressed": False, "reason": "消息太少，无需压缩"}
|
||
|
||
# 分割：旧消息（需要压缩）+ 最近消息（保留原文）
|
||
keep_count = KEEP_RECENT_ROUNDS * 2
|
||
old_msgs = all_msgs[:-keep_count]
|
||
recent_msgs = all_msgs[-keep_count:]
|
||
|
||
# 检查第一条旧消息是否已是摘要（避免重复压缩无效内容）
|
||
# 构建需要总结的对话文本
|
||
conversation_text = ""
|
||
for msg in old_msgs:
|
||
role_label = "用户" if msg.role == "user" else "AI助手"
|
||
content = msg.content
|
||
# 如果是已压缩的摘要，标注出来
|
||
if content.startswith("[对话摘要]"):
|
||
conversation_text += f"[之前的摘要]{content[5:]}\n\n"
|
||
else:
|
||
conversation_text += f"{role_label}: {content}\n\n"
|
||
|
||
# 调用 LLM 生成摘要
|
||
summary_prompt = (
|
||
"你是一个对话摘要助手。请将以下对话历史压缩为简洁的摘要，保留关键信息：\n"
|
||
"1. 用户提出的核心需求和决策\n"
|
||
"2. AI 执行的重要操作及结果（如创建/修改了哪些大纲、角色、章节等）\n"
|
||
"3. 达成的共识和待办事项\n\n"
|
||
"只输出摘要内容，不要加前缀或解释。用简洁的条目列表形式。"
|
||
)
|
||
summary_messages = [
|
||
{"role": "user", "content": f"请压缩以下对话历史：\n\n{conversation_text}"}
|
||
]
|
||
|
||
try:
|
||
summary = await simple_completion(config, summary_messages, summary_prompt)
|
||
except Exception as e:
|
||
return {"error": f"摘要生成失败: {str(e)}"}
|
||
|
||
# 删除旧消息
|
||
for msg in old_msgs:
|
||
session.delete(msg)
|
||
|
||
# 插入摘要消息（时间设为最早保留消息之前）
|
||
summary_msg = ChatMessage(
|
||
novel_id=novel_id,
|
||
role="assistant",
|
||
content=f"[对话摘要]\n{summary}",
|
||
created_at=recent_msgs[0].created_at,
|
||
)
|
||
session.add(summary_msg)
|
||
session.commit()
|
||
|
||
return {
|
||
"compressed": True,
|
||
"removed_count": len(old_msgs),
|
||
"kept_count": len(recent_msgs),
|
||
"summary_length": len(summary),
|
||
}
|