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, 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] = { "list_outlines": "查询大纲", "create_outline": "创建大纲", "update_outline": "更新大纲", "delete_outline": "删除大纲", "get_world_setting": "查询世界观", "save_world_setting": "保存世界观", } def _build_system_prompt( novel_id: int | None, session: Session, page_context: str | None = None, tools_enabled: bool = False, ) -> str: """根据小说上下文和用户所在页面构建 system prompt""" base = "你是一个专业的小说创作助手。" if not novel_id: return base novel = session.get(Novel, novel_id) if not novel: return base parts = [ f"你是一个专业的小说创作助手，当前正在协助创作《{novel.title}》。", f"小说简介：{novel.description or '暂无'}", ] if page_context: parts.append(f"\n用户当前正在查看：{page_context}。") if tools_enabled: parts.append( "\n你可以使用工具来查看和管理大纲、世界观设定。" "\n当用户要求查看、添加、修改或删除大纲或世界观时，请使用相应的工具。" "\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" @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.novel_id and 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 = {} 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}) 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 ] 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: pass # 正常完成：保存 AI 回复 if full_response: with Session(engine) as s: ai_msg = ChatMessage( novel_id=data.novel_id, role="assistant", content=full_response, ) s.add(ai_msg) s.commit() 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()