companion-assistant/backend/app/api/chat.py

"""对话路由：会话管理 + 消息编辑/反馈/重新生成 + REST 非流式对话。"""
import json

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session

from ..core.deps import get_current_user
from ..core.response import ok
from ..database import get_db
from ..models import Agent, ChatMessage, ChatSession, User
from ..schemas.chat import ChatMessageUpdate, ChatRequest, ChatSessionCreate, FeedbackRequest
from ..services import chat_service

router = APIRouter(prefix="/chat", tags=["对话"])


def _session_out(s: ChatSession) -> dict:
    return {
        "id": s.id, "title": s.title, "agent_id": s.agent_id, "model": s.model,
        "created_at": s.created_at.isoformat(), "updated_at": s.updated_at.isoformat(),
    }


def _message_out(m: ChatMessage) -> dict:
    try:
        suggestions = json.loads(m.suggestions) if m.suggestions else []
    except json.JSONDecodeError:
        suggestions = []
    try:
        file_ids = json.loads(m.file_ids) if m.file_ids else []
    except json.JSONDecodeError:
        file_ids = []
    return {
        "id": m.id, "session_id": m.session_id, "role": m.role, "content": m.content,
        "model": m.model, "file_ids": file_ids, "feedback": m.feedback or "",
        "suggestions": suggestions, "edited": bool(m.edited), "regenerated": m.regenerated or 0,
        "reasoning": m.reasoning_content or "",
        "created_at": m.created_at.isoformat(), "updated_at": m.updated_at.isoformat(),
    }


def _get_owned_session(db: Session, user: User, session_id: int) -> ChatSession:
    s = db.get(ChatSession, session_id)
    if not s or s.user_id != user.id:
        raise HTTPException(status_code=404, detail="会话不存在")
    return s


@router.get("/sessions")
def list_sessions(user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    sessions = (
        db.query(ChatSession).filter(ChatSession.user_id == user.id)
        .order_by(ChatSession.updated_at.desc()).all()
    )
    return ok([_session_out(s) for s in sessions])


@router.post("/sessions")
def create_session(body: ChatSessionCreate, user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    if body.agent_id:
        agent = db.get(Agent, body.agent_id)
        if not agent or (agent.user_id and agent.user_id != user.id and not agent.is_public):
            raise HTTPException(status_code=404, detail="智能体不存在")
    s = ChatSession(user_id=user.id, title=body.title, agent_id=body.agent_id, model=body.model)
    db.add(s)
    db.commit()
    db.refresh(s)
    return ok(_session_out(s))


@router.get("/sessions/{session_id}/messages")
def list_messages(session_id: int, user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    s = _get_owned_session(db, user, session_id)
    messages = db.query(ChatMessage).filter(ChatMessage.session_id == s.id).order_by(ChatMessage.id).all()
    return ok([_message_out(m) for m in messages])


@router.post("/sessions/{session_id}/messages")
async def send_message(session_id: int, body: ChatRequest,
                       user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    s = _get_owned_session(db, user, session_id)
    if not body.content and not body.file_ids:
        raise HTTPException(status_code=400, detail="消息内容不能为空")
    try:
        reply = await chat_service.chat_once(db, user, s, body.content, body.agent_id, body.model, body.file_ids)
    except LookupError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"模型调用失败：{e}")
    return ok(_message_out(reply))


@router.delete("/sessions/{session_id}")
def delete_session(session_id: int, user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    s = _get_owned_session(db, user, session_id)
    db.delete(s)
    db.commit()
    return ok(message="已删除")


# ---------- 消息操作 ----------

def _get_owned_message(db: Session, user: User, message_id: int) -> ChatMessage:
    m = db.get(ChatMessage, message_id)
    if not m:
        raise HTTPException(status_code=404, detail="消息不存在")
    s = db.get(ChatSession, m.session_id)
    if not s or s.user_id != user.id:
        raise HTTPException(status_code=404, detail="消息不存在")
    return m


@router.put("/messages/{message_id}")
def edit_message(message_id: int, body: ChatMessageUpdate,
                 user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    """编辑用户消息（随后前端会触发重新生成）。"""
    m = _get_owned_message(db, user, message_id)
    if m.role != "user":
        raise HTTPException(status_code=400, detail="只能编辑用户消息")
    m.content = body.content
    m.edited = True
    m.updated_at = __import__("datetime").datetime.utcnow()
    db.commit()
    db.refresh(m)
    return ok(_message_out(m))


@router.post("/messages/{message_id}/feedback")
def set_feedback(message_id: int, body: FeedbackRequest,
                 user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    m = _get_owned_message(db, user, message_id)
    if m.role != "assistant":
        raise HTTPException(status_code=400, detail="只能对 AI 回答反馈")
    m.feedback = body.feedback
    db.commit()
    db.refresh(m)
    return ok(_message_out(m))


@router.post("/sessions/{session_id}/regenerate")
async def regenerate(session_id: int,
                     user: User = Depends(get_current_user), db: Session = Depends(get_db)):
    """REST 非流式重新生成最后一条 AI 回答（备用，前端主要走 WS 流式）。"""
    s = _get_owned_session(db, user, session_id)
    try:
        async for _k, _d, message_id in _regenerate_rest(chat_service.regenerate_stream(db, s)):
            pass  # 流式丢弃，最终内容已入库
    except LookupError as e:
        raise HTTPException(status_code=404, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"模型调用失败：{e}")
    m = db.get(ChatMessage, message_id)
    return ok(_message_out(m))


async def _regenerate_rest(agen):
    """适配 regenerate_stream 的 (kind, delta), id 产出结构。"""
    async for (kind, delta), message_id in agen:
        yield kind, delta, message_id