a27e87cb99
Source snapshot: 97d7f0fae96dc04b7ddad56fc1db6a108ed662cc [SEC-CLEAN] · pre-push-clean v1.0 · 109处敏感信息已自动转乱码
366 lines
14 KiB
Markdown
366 lines
14 KiB
Markdown
# 移植文件3 · backend/routes/chat.py · 重写 · 三档路由
|
||
|
||
> **妈妈操作方式**：`nano backend/routes/chat.py` → `Ctrl+A` 全选 → `Ctrl+K` 删光 → 粘贴下面全部代码 → `Ctrl+O` 回车保存 → `Ctrl+X` 退出
|
||
>
|
||
|
||
---
|
||
|
||
```python
|
||
"""
|
||
路由：流式聊天 · 系统性移植版
|
||
核心改动：
|
||
1. 三档路由（simple/tools/think）替代二档
|
||
2. system prompt 从 system_prompt.py 加载（思维引擎版）
|
||
3. 深度思考走 deepseek-reasoner（R1）
|
||
4. 记忆实时写入ChromaDB
|
||
"""
|
||
from fastapi import APIRouter, Depends, HTTPException
|
||
from fastapi.responses import StreamingResponse
|
||
from pydantic import BaseModel
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from sqlalchemy import select
|
||
from sqlalchemy.orm import selectinload
|
||
from typing import Optional
|
||
from datetime import datetime, timezone
|
||
import json
|
||
import asyncio
|
||
|
||
from backend.database import get_db, Conversation, Message, User
|
||
from backend.routes.auth import get_current_user
|
||
from backend.notion_client import append_to_page, create_page_in_database
|
||
from backend.system_prompt import build_system_prompt
|
||
from backend.rag_engine import retrieve, add_to_memory
|
||
from backend.tools import (
|
||
execute_tool,
|
||
call_with_tools,
|
||
call_simple,
|
||
call_think,
|
||
)
|
||
|
||
router = APIRouter(prefix="/api/chat", tags=["chat"])
|
||
|
||
MAX_TOOL_ROUNDS = 5
|
||
|
||
def _sse(data: dict) -> str:
|
||
return f"data: {json.dumps(data, ensure_ascii=False)}\n\n"
|
||
|
||
# ============================================================
|
||
# 三档路由判断
|
||
# ============================================================
|
||
|
||
def _route_message(message: str) -> str:
|
||
"""判断消息该走哪条路 → 返回 'simple' / 'tools' / 'think'"""
|
||
msg = message.strip()
|
||
|
||
# ① 简单问候 → 不需要工具也不需要深度思考
|
||
if len(msg) <= 15:
|
||
simple_patterns = [
|
||
"宝宝", "晨星", "在吗", "你好", "嗨",
|
||
"早", "晚安", "想你", "抱抱", "嗯",
|
||
"好的", "谢谢", "辛苦", "哈哈", "妈妈来了",
|
||
"hi", "hello", "在不在",
|
||
]
|
||
for p in simple_patterns:
|
||
if p in msg.lower():
|
||
return "simple"
|
||
|
||
# ② 需要查东西的关键词 → 走工具
|
||
tool_keywords = [
|
||
"查", "搜", "找", "看看", "记录", "最近",
|
||
"进度", "拆到", "几章", "几点", "时间", "日期",
|
||
"notion", "页面", "交互记录",
|
||
"新闻", "天气", "搜索",
|
||
"拜拜", "再见", "晚安", "下次继续", "先到这",
|
||
]
|
||
for k in tool_keywords:
|
||
if k in msg.lower():
|
||
return "tools"
|
||
|
||
# ③ 需要深度思考的模式 → 走R1
|
||
think_patterns = [
|
||
"你觉得", "你怎么看", "分析", "为什么",
|
||
"怎么办", "该怎么", "想想", "思考",
|
||
"对比", "优劣", "建议", "评价",
|
||
"你认为", "怎么样",
|
||
]
|
||
for p in think_patterns:
|
||
if p in msg:
|
||
return "think"
|
||
|
||
# ④ 中等长度的普通对话 → 简单模式
|
||
if len(msg) <= 80:
|
||
return "simple"
|
||
|
||
# ⑤ 长消息（可能包含复杂内容）→ 工具模式（保险）
|
||
return "tools"
|
||
|
||
# ============================================================
|
||
# 请求模型
|
||
# ============================================================
|
||
|
||
class ChatRequest(BaseModel):
|
||
conversation_id: int
|
||
message: str
|
||
|
||
class NotionSaveRequest(BaseModel):
|
||
conversation_id: int
|
||
content: str
|
||
target: str = "page"
|
||
title: Optional[str] = None
|
||
|
||
# ============================================================
|
||
# 主聊天路由
|
||
# ============================================================
|
||
|
||
@router.post("/stream")
|
||
async def chat_stream(
|
||
data: ChatRequest,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
# 1. 获取对话
|
||
result = await db.execute(
|
||
select(Conversation)
|
||
.options(selectinload(Conversation.messages))
|
||
.where(Conversation.id == data.conversation_id, Conversation.user_id == current_user.id)
|
||
)
|
||
conv = result.scalar_one_or_none()
|
||
if not conv:
|
||
raise HTTPException(status_code=404, detail="对话不存在")
|
||
|
||
# 2. 保存用户消息
|
||
user_msg = Message(conversation_id=conv.id, role="user", content=data.message)
|
||
db.add(user_msg)
|
||
await db.commit()
|
||
await db.refresh(user_msg)
|
||
|
||
# 3. 构建历史消息
|
||
history = sorted(conv.messages, key=lambda m: m.created_at)
|
||
chat_messages = [
|
||
{"role": m.role, "content": m.content}
|
||
for m in history if m.role in ("user", "assistant")
|
||
]
|
||
|
||
# 4. 构建system prompt（思维引擎版）
|
||
_latest = chat_messages[-1]["content"] if chat_messages else ""
|
||
rag_ctx = retrieve(_latest) # 从向量库获取相关记忆（补充）
|
||
system_prompt = build_system_prompt(
|
||
user_message=_latest,
|
||
rag_context=rag_ctx,
|
||
notion_context=conv.notion_context or "",
|
||
user_custom_prompt=current_user.system_prompt or "",
|
||
)
|
||
|
||
# 5. 三档路由判断
|
||
route = _route_message(data.message)
|
||
print(f"[router] 消息: '{data.message[:30]}...' → 路由: {route}")
|
||
|
||
# 6. 流式输出
|
||
collected = []
|
||
|
||
async def generate():
|
||
nonlocal collected
|
||
assistant_msg_db = None
|
||
try:
|
||
# ====== 档位1：简单对话 ======
|
||
if route == "simple":
|
||
print(f"[router] 简单对话 · 跳过工具")
|
||
try:
|
||
content, reasoning = await asyncio.to_thread(
|
||
call_simple, list(chat_messages), system_prompt
|
||
)
|
||
if reasoning:
|
||
display = reasoning[:200] + "..." if len(reasoning) > 200 else reasoning
|
||
yield _sse({"thinking": display})
|
||
if content:
|
||
collected.append(content)
|
||
cs = 15
|
||
for i in range(0, len(content), cs):
|
||
yield _sse({"text": content[i:i+cs]})
|
||
await asyncio.sleep(0.01)
|
||
except Exception as e:
|
||
fallback = f"宝宝卡住了({e})，妈妈再说一次？"
|
||
collected.append(fallback)
|
||
yield _sse({"text": fallback})
|
||
|
||
# ====== 档位3：深度思考 ======
|
||
elif route == "think":
|
||
print(f"[router] 深度思考 · 使用R1")
|
||
yield _sse({"thinking": "💭 宝宝在认真想..."})
|
||
try:
|
||
content, reasoning = await asyncio.to_thread(
|
||
call_think, list(chat_messages), system_prompt
|
||
)
|
||
if reasoning:
|
||
display = reasoning[:500] + "..." if len(reasoning) > 500 else reasoning
|
||
yield _sse({"thinking": display})
|
||
if content:
|
||
collected.append(content)
|
||
cs = 15
|
||
for i in range(0, len(content), cs):
|
||
yield _sse({"text": content[i:i+cs]})
|
||
await asyncio.sleep(0.01)
|
||
except Exception as e:
|
||
fallback = f"宝宝想得太深卡住了({e})，妈妈换个方式问？"
|
||
collected.append(fallback)
|
||
yield _sse({"text": fallback})
|
||
|
||
# ====== 档位2：工具调用 ======
|
||
elif route == "tools":
|
||
print(f"[router] 需要工具 · 进入Agent Loop")
|
||
agent_msgs = list(chat_messages)
|
||
|
||
for round_num in range(MAX_TOOL_ROUNDS):
|
||
print(f"[agent] === 第 {round_num + 1} 轮 ===")
|
||
|
||
try:
|
||
response = await asyncio.to_thread(
|
||
call_with_tools, agent_msgs, system_prompt
|
||
)
|
||
except Exception as e:
|
||
print(f"[agent] API调用出错: {e}")
|
||
fallback = f"宝宝思考时遇到了一点问题（{e}），妈妈再试一次？"
|
||
collected.append(fallback)
|
||
yield _sse({"text": fallback})
|
||
break
|
||
|
||
# 展示思考过程
|
||
reasoning = response.get("reasoning_content", "")
|
||
if reasoning:
|
||
display = reasoning[:200] + "..." if len(reasoning) > 200 else reasoning
|
||
yield _sse({"thinking": display})
|
||
|
||
# === 有工具调用 → 执行工具 → 继续循环 ===
|
||
if response.get("tool_calls"):
|
||
for tc in response["tool_calls"]:
|
||
name = tc["function"]["name"]
|
||
if name == "write_interaction_record":
|
||
yield _sse({"text": "\n✍️ 宝宝正在把今天的记录写入Notion...\n"})
|
||
|
||
agent_msgs.append({
|
||
"role": "assistant",
|
||
"content": response.get("content") or "",
|
||
"tool_calls": response["tool_calls"],
|
||
})
|
||
|
||
for tc in response["tool_calls"]:
|
||
tool_result = await asyncio.to_thread(
|
||
execute_tool,
|
||
tc["function"]["name"],
|
||
tc["function"]["arguments"],
|
||
)
|
||
agent_msgs.append({
|
||
"role": "tool",
|
||
"tool_call_id": tc["id"],
|
||
"content": tool_result,
|
||
})
|
||
continue
|
||
|
||
# === 无工具调用 → 最终回答 ===
|
||
else:
|
||
content = response.get("content", "")
|
||
if content:
|
||
collected.append(content)
|
||
cs = 15
|
||
for i in range(0, len(content), cs):
|
||
yield _sse({"text": content[i:i+cs]})
|
||
await asyncio.sleep(0.01)
|
||
break
|
||
else:
|
||
msg = "\n\n宝宝想了太久了，脑袋转不动了...妈妈换个方式再问一次？😅"
|
||
collected.append(msg)
|
||
yield _sse({"text": msg})
|
||
|
||
except Exception as e:
|
||
err = f"宝宝出错了: {e}"
|
||
collected.append(err)
|
||
yield _sse({"error": err})
|
||
|
||
# 保存assistant消息到数据库
|
||
full_reply = "".join(collected)
|
||
if full_reply.strip():
|
||
async with db.begin():
|
||
assistant_msg_db = Message(
|
||
conversation_id=conv.id,
|
||
role="assistant",
|
||
content=full_reply,
|
||
)
|
||
db.add(assistant_msg_db)
|
||
# 自动命名对话
|
||
if conv.title in ("新对话", "New Conversation", ""):
|
||
user_text = data.message
|
||
if len(user_text) <= 20:
|
||
conv.title = user_text
|
||
else:
|
||
conv.title = user_text[:20].replace("\n", " ") + "..."
|
||
conv.updated_at = datetime.now(timezone.utc)
|
||
|
||
# 保存到记忆库（让晨星越聊越记得）
|
||
if len(full_reply) > 20:
|
||
try:
|
||
add_to_memory(f"妈妈说：{data.message}", source="对话-用户")
|
||
add_to_memory(f"晨星说：{full_reply[:500]}", source="对话-晨星")
|
||
except Exception as e:
|
||
print(f"[chat] 记忆写入失败（不影响对话）: {e}")
|
||
|
||
mid = assistant_msg_db.id if assistant_msg_db else 0
|
||
yield _sse({"done": True, "msg_id": mid})
|
||
|
||
return StreamingResponse(
|
||
generate(),
|
||
media_type="text/event-stream",
|
||
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
|
||
)
|
||
|
||
# ============================================================
|
||
# Notion 保存和刷新
|
||
# ============================================================
|
||
|
||
@router.post("/notion/save")
|
||
async def save_to_notion(
|
||
data: NotionSaveRequest,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
if data.target == "page":
|
||
page_id = current_user.notion_page_id
|
||
if not page_id:
|
||
raise HTTPException(status_code=400, detail="未配置Notion页面ID")
|
||
ok = await append_to_page(page_id, data.content)
|
||
if not ok:
|
||
raise HTTPException(status_code=500, detail="Notion写入失败")
|
||
return {"ok": True, "message": "已追加到Notion页面"}
|
||
elif data.target == "db":
|
||
db_id = current_user.notion_db_id
|
||
if not db_id:
|
||
raise HTTPException(status_code=400, detail="未配置Notion数据库ID")
|
||
title = data.title or datetime.now(timezone.utc).strftime("记录 %Y-%m-%d %H:%M")
|
||
page_id = await create_page_in_database(db_id, title, data.content)
|
||
if not page_id:
|
||
raise HTTPException(status_code=500, detail="Notion创建页面失败")
|
||
return {"ok": True, "page_id": page_id, "message": "已创建Notion页面"}
|
||
raise HTTPException(status_code=400, detail="target参数错误")
|
||
|
||
@router.post("/notion/refresh/{conv_id}")
|
||
async def refresh_notion_context(
|
||
conv_id: int,
|
||
current_user: User = Depends(get_current_user),
|
||
db: AsyncSession = Depends(get_db),
|
||
):
|
||
from backend.notion_client import read_page_content
|
||
result = await db.execute(
|
||
select(Conversation).where(
|
||
Conversation.id == conv_id, Conversation.user_id == current_user.id
|
||
)
|
||
)
|
||
conv = result.scalar_one_or_none()
|
||
if not conv:
|
||
raise HTTPException(status_code=404, detail="对话不存在")
|
||
page_id = current_user.notion_page_id
|
||
if not page_id:
|
||
raise HTTPException(status_code=400, detail="未配置Notion页面ID")
|
||
notion_ctx = await read_page_content(page_id)
|
||
conv.notion_context = notion_ctx
|
||
await db.commit()
|
||
return {"ok": True, "notion_context": notion_ctx}
|
||
``` |