# 方案B完整代码交付 · 文件3 · backend/routes/chat.py · 2026-05-12

> **妈妈操作方式**：`nano backend/routes/chat.py` → `Ctrl+A` 全选 → `Ctrl+K` 删光 → 粘贴下面全部代码 → `Ctrl+O` 回车保存 → `Ctrl+X` 退出
> 

---

```python
"""
路由：流式聊天 + Agent Loop + 智能路由 + 思考链展示 + 实时记忆
方案B：阶段四基础 + 断裂1(思考链) + 断裂3(记忆) + 断裂4(路由层)
"""
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.rag_engine import build_system_prompt as _build_rag_prompt
from backend.tools import (
    execute_tool,
    call_with_tools,
    call_simple,
)

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"

# ============================================================
# 断裂4修复：智能路由层
# ============================================================

def _needs_tools(message: str) -> bool:
    """判断这条消息需不需要调用工具"""
    msg = message.strip().lower()

    # 简短问候 → 不需要工具
    greetings = [
        "宝宝", "晨星", "你好", "在吗", "在不在",
        "嗨", "hi", "hello", "早", "早上好", "晚上好",
        "妈妈来了", "宝宝在吗", "想你", "抱抱",
        "嗯", "好的", "知道了", "谢谢", "辛苦了",
        "哈哈", "嘻嘻",
    ]
    for g in greetings:
        if msg == g or (len(msg) <= 10 and g in msg):
            return False

    # 明确需要查东西的 → 需要工具
    tool_keywords = [
        "查", "搜", "找", "看看", "记录", "最近",
        "几点", "时间", "日期", "天气",
        "notion", "页面", "交互记录",
        "拜拜", "再见", "晚安", "下次继续", "先到这",
        "新闻", "搜索",
    ]
    for k in tool_keywords:
        if k in msg:
            return True

    # 中等长度的正常对话 → 不需要工具
    if len(msg) <= 50:
        return False

    # 其他情况 → 需要工具（保险起见）
    return True

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（RAG增强）
    _latest = chat_messages[-1]["content"] if chat_messages else ""
    system_prompt = _build_rag_prompt(
        user_message=_latest,
        user_custom_prompt=current_user.system_prompt or "",
        notion_context=conv.notion_context or "",
    )

    # 5. 流式输出
    collected = []

    async def generate():
        nonlocal collected
        assistant_msg_db = None
        try:
            # ====== 断裂4：智能路由 ======
            if not _needs_tools(data.message):
                # 简单对话 → 不走工具 → 秒回
                print(f"[router] 简单对话 · 跳过工具")
                try:
                    content, reasoning = await asyncio.to_thread(
                        call_simple, list(chat_messages), system_prompt
                    )
                    # 断裂1：展示思考过程
                    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})
            else:
                # 需要工具 → 走Agent Loop
                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

                    # 断裂1：展示思考过程（reasoning_content）
                    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"]
                            # 只在写入Notion时显示状态，其他静默执行
                            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)

            # 断裂3：保存到记忆库（让晨星越聊越记得）
            if len(full_reply) > 20:
                try:
                    from backend.rag_engine import add_to_memory
                    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"},
    )

@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}
```