# 移植文件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}
```