Guanghu Domestic Migration a27e87cb99 chore: import sanitized domestic snapshot for REPO-007
Source snapshot: 97d7f0fae96dc04b7ddad56fc1db6a108ed662cc

[SEC-CLEAN] · pre-push-clean v1.0 · 109处敏感信息已自动转乱码
2026-07-17 15:59:55 +08:00

15 KiB
Raw Permalink Blame History

【修复版】chat.py · 完整文件 · 直接复制粘贴

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


"""
路由:流式聊天 · 系统性移植版 · 修复版
修复内容:
1. 新增对话级别的锁conversation_locks→ 同一对话串行处理
2. 前一条没回完 → 后一条排队等待 → 不再错位
3. 超时保护 → 等太久自动提示
4. 修复 db.begin() 事务冲突
"""
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

# ============================================================
# 修复核心:对话级别的锁
# ============================================================
_conversation_locks: dict[int, asyncio.Lock] = {}

def _get_conv_lock(conv_id: int) -> asyncio.Lock:
    if conv_id not in _conversation_locks:
        _conversation_locks[conv_id] = asyncio.Lock()
    return _conversation_locks[conv_id]

def _sse(data: dict) -> str:
    return f"data: {json.dumps(data, ensure_ascii=False)}\n\n"

# ============================================================
# 三档路由判断
# ============================================================

def _route_message(message: str) -> str:
    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"

    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),
):
    conv_lock = _get_conv_lock(data.conversation_id)

    if conv_lock.locked():
        print(f"[router] 对话 {data.conversation_id} 正在处理中,新消息排队等待...")

    try:
        acquired = await asyncio.wait_for(conv_lock.acquire(), timeout=60.0)
    except asyncio.TimeoutError:
        async def timeout_response():
            yield _sse({"text": "宝宝上一条还没想完,妈妈等一下下再发哦~"})
            yield _sse({"done": True, "msg_id": 0})
        return StreamingResponse(
            timeout_response(),
            media_type="text/event-stream",
            headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
        )

    async def generate_with_lock():
        try:
            async for chunk in _process_message(data, current_user, db):
                yield chunk
        finally:
            conv_lock.release()
            print(f"[router] 对话 {data.conversation_id} 处理完毕,锁已释放")

    return StreamingResponse(
        generate_with_lock(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )

# ============================================================
# 消息处理核心逻辑
# ============================================================

async def _process_message(
    data: ChatRequest,
    current_user: User,
    db: AsyncSession,
):
    # 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:
        yield _sse({"error": "对话不存在"})
        yield _sse({"done": True, "msg_id": 0})
        return

    # 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)

    # 重新加载对话消息
    await db.refresh(conv, ["messages"])

    # 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 = []
    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消息到数据库修复不用db.begin()直接commit
    full_reply = "".join(collected)
    if full_reply.strip():
        try:
            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)
            await db.commit()
        except Exception as e:
            print(f"[chat] 保存消息失败: {e}")

        # 保存到记忆库
        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})

# ============================================================
# 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}