AlphaPulse/backend/src/alphapulse/api/chat.py

"""Chat：单一主会话 + CONTEXT.md。生成过程以 SSE 推送；浏览器断开即停止生成"""

import asyncio
import json
import logging
from collections.abc import AsyncIterator
from datetime import UTC, datetime
from typing import Annotated, Literal

import anyio
from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.sse import EventSourceResponse, ServerSentEvent
from pydantic import BaseModel, Field
from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

from alphapulse.ai import llm
from alphapulse.ai.prompt import build_system_prompt, current_context, estimate_tokens
from alphapulse.ai.tools import ToolOutcome, run_tool, tool_specs
from alphapulse.core.auth import require_admin
from alphapulse.core.db import SessionLocal, get_session
from alphapulse.models import AiProvider, ChatMessage, ChatSession, ContextVersion

router = APIRouter(prefix="/api/chat", tags=["chat"], dependencies=[Depends(require_admin)])

# 一轮回答中最多调用工具的轮数，超出后不再提供工具，让模型给出当前结论
MAX_TOOL_ROUNDS = 8
# 发给前端展示的工具结果长度
TOOL_PREVIEW_CHARS = 2000
STOPPED_RESULT = json.dumps({"error": "用户已停止生成，未完成的工具结果尚未确认，请核对后再重试"}, ensure_ascii=False)

# api 只有一个进程：同一时间只允许一轮生成，Reset 也要等生成结束
generating = asyncio.Lock()
logger = logging.getLogger("alphapulse.api.chat")


class MessageOut(BaseModel):
    id: int
    role: Literal["user", "assistant", "tool"]
    content: dict
    provider_id: int | None
    usage: dict | None
    created_at: datetime


class SessionOut(BaseModel):
    id: int
    created_at: datetime
    messages: list[MessageOut]
    # 下一轮的 prompt 估算（不含新输入），前端据此提示是否接近上下文上限
    estimated_tokens: int
    generating: bool


class SendIn(BaseModel):
    text: Annotated[str, Field(min_length=1, max_length=20000)]
    provider_id: int | None = None
    # 本轮思考强度；为空时用模型配置里的默认值
    reasoning_effort: llm.Effort | None = None


class ResetIn(BaseModel):
    clear_context: bool = False


class ContextOut(BaseModel):
    id: int | None
    content: str
    source: str | None
    created_at: datetime | None


class ContextIn(BaseModel):
    content: str


def event(name: str, data: dict) -> ServerSentEvent:
    return ServerSentEvent(raw_data=json.dumps(data, ensure_ascii=False, default=str), event=name)


def neutral(m: ChatMessage) -> dict:
    return {"role": m.role, **m.content}


async def active_session(session: AsyncSession) -> ChatSession:
    current = await session.scalar(select(ChatSession).where(ChatSession.ended_at.is_(None)).order_by(ChatSession.id.desc()).limit(1))
    if current is None:
        current = ChatSession()
        session.add(current)
        await session.commit()
    return current


async def session_messages(session: AsyncSession, session_id: int) -> list[ChatMessage]:
    return list((await session.scalars(select(ChatMessage).where(ChatMessage.session_id == session_id).order_by(ChatMessage.id))).all())


async def save_messages(*messages: ChatMessage) -> None:
    async with SessionLocal() as session:
        session.add_all(messages)
        await session.commit()


async def save_shielded(*messages: ChatMessage) -> None:
    """停止生成（请求被取消）时仍要落库，保证历史完整：屏蔽取消再写"""
    with anyio.CancelScope(shield=True):
        await save_messages(*messages)


@router.get("/session")
async def get_chat_session(session: Annotated[AsyncSession, Depends(get_session)]) -> SessionOut:
    # 先取生成状态再读消息：反过来的话，读消息后生成恰好结束，会返回「未在生成」却缺最后一条
    is_generating = generating.locked()
    current = await active_session(session)
    messages = await session_messages(session, current.id)
    system = await build_system_prompt(session)
    history = json.dumps([neutral(m) for m in messages], ensure_ascii=False)
    return SessionOut(
        id=current.id, created_at=current.created_at,
        messages=[MessageOut.model_validate(m, from_attributes=True) for m in messages],
        estimated_tokens=estimate_tokens(system) + estimate_tokens(history),
        generating=is_generating,
    )


@router.post("/messages", response_class=EventSourceResponse)
async def send_message(body: SendIn) -> AsyncIterator[ServerSentEvent]:
    """事件：prompt、thinking、delta、tool_call、tool_result、skill_draft、context_updated、usage、error、done"""
    if generating.locked():
        yield event("error", {"message": "上一轮还在生成，请等待完成或先停止"})
        return
    async with generating:
        async with SessionLocal() as session:
            provider = await session.get(AiProvider, body.provider_id) if body.provider_id else await session.scalar(
                select(AiProvider).where(AiProvider.is_default)
            )
            if provider is None:
                yield event("error", {"message": "没有可用的模型配置，请先到「环境变量」页面添加"})
                return
            session.expunge(provider)
            current = await active_session(session)
            system = await build_system_prompt(session)
            history = [neutral(m) for m in await session_messages(session, current.id)]

        user = {"role": "user", "text": body.text}
        messages = [*history, user]
        estimated = estimate_tokens(system) + estimate_tokens(json.dumps(messages, ensure_ascii=False))
        yield event("prompt", {"estimated_tokens": estimated, "context_window": provider.context_window, "provider": provider.name})
        if estimated + provider.max_tokens > provider.context_window:
            yield event("error", {"message": f"上下文约 {estimated} token，加上最大输出 {provider.max_tokens} 超过上限 {provider.context_window}。请 Reset 或精简 skills / CONTEXT.md"})
            return
        await save_messages(ChatMessage(session_id=current.id, role="user", content={"text": body.text}))

        specs = tool_specs()
        effort = body.reasoning_effort or provider.reasoning_effort
        for round_no in range(MAX_TOOL_ROUNDS + 1):
            text = ""
            completion: llm.Completion | None = None
            try:
                async for item in llm.stream(provider, system, messages, specs if round_no < MAX_TOOL_ROUNDS else [], effort):
                    if isinstance(item, llm.TextDelta):
                        text += item.text
                        yield event("delta", {"text": item.text})
                    elif isinstance(item, llm.ThinkingStarted):
                        yield event("thinking", {})
                    else:
                        completion = item
            except llm.LlmError as e:
                if text:
                    await save_messages(ChatMessage(session_id=current.id, role="assistant", content={"text": text, "tool_calls": [], "error": str(e)}, provider_id=provider.id))
                yield event("error", {"message": str(e)})
                return
            except (asyncio.CancelledError, GeneratorExit):
                # 停止生成：保留已输出的部分，不保留未完成的工具调用。
                # 取消发生在 await 时是 CancelledError，发生在 yield 处（生成器被关闭）时是 GeneratorExit
                await save_shielded(ChatMessage(session_id=current.id, role="assistant", content={"text": text, "tool_calls": [], "stopped": True}, provider_id=provider.id))
                raise

            assistant = {"text": completion.text, "tool_calls": completion.tool_calls}
            if completion.thinking_blocks:
                assistant["thinking_blocks"] = completion.thinking_blocks
            truncated_message = f"模型输出达到最大输出上限（{provider.max_tokens} token）被截断，可在「环境变量」页调大「最大输出」"
            if completion.truncated and not completion.tool_calls:
                assistant["error"] = truncated_message
            pending = list(completion.tool_calls)
            assistant_saved = False
            try:
                # 消息提交和本地进度必须一起完成，避免取消恰好发生在提交之后而留下孤立调用。
                with anyio.CancelScope(shield=True):
                    await save_messages(ChatMessage(session_id=current.id, role="assistant", content=assistant, provider_id=provider.id, usage=completion.usage))
                    assistant_saved = True
                messages.append({"role": "assistant", **assistant})
                yield event("usage", completion.usage)
                if not completion.tool_calls:
                    if completion.truncated:
                        yield event("error", {"message": truncated_message})
                    break
                while pending:
                    call = pending[0]
                    yield event("tool_call", {"id": call["id"], "name": call["name"], "arguments": call["arguments"]})
                    # 截断时不执行，把原因告诉模型，让它精简后重试
                    outcome = ToolOutcome(llm.TRUNCATED_RESULT) if completion.truncated else await run_tool(call["name"], call["arguments"])
                    result = {"tool_call_id": call["id"], "name": call["name"], "result": outcome.result}
                    with anyio.CancelScope(shield=True):
                        await save_messages(ChatMessage(session_id=current.id, role="tool", content=result))
                        pending.pop(0)
                    messages.append({"role": "tool", **result})
                    yield event("tool_result", {"id": call["id"], "name": call["name"], "preview": outcome.result[:TOOL_PREVIEW_CHARS]})
                    if outcome.event:
                        yield event(*outcome.event)
            except (asyncio.CancelledError, GeneratorExit):
                # 每个工具调用都要有结果，否则下一轮发给模型的历史不合法
                if assistant_saved and pending:
                    await save_shielded(*(
                        ChatMessage(session_id=current.id, role="tool", content={"tool_call_id": c["id"], "name": c["name"], "result": STOPPED_RESULT})
                        for c in pending
                    ))
                raise
            except Exception:
                logger.exception("Chat 工具轮次失败")
                if assistant_saved and pending:
                    failed = json.dumps({"error": "工具轮次异常中止，执行结果未确认；后续工具未执行"}, ensure_ascii=False)
                    await save_shielded(*(
                        ChatMessage(session_id=current.id, role="tool", content={"tool_call_id": c["id"], "name": c["name"], "result": failed})
                        for c in pending
                    ))
                yield event("error", {"message": "工具轮次执行失败，请查看服务日志"})
                return
        yield event("done", {})


@router.post("/reset")
async def reset(body: ResetIn, session: Annotated[AsyncSession, Depends(get_session)]) -> None:
    """结束当前会话（消息保留在库里，不再展示），可选同时清空 CONTEXT.md（写入空版本，可找回）"""
    if generating.locked():
        raise HTTPException(status.HTTP_409_CONFLICT, "正在生成，请先停止")
    await session.execute(update(ChatSession).where(ChatSession.ended_at.is_(None)).values(ended_at=datetime.now(UTC)))
    if body.clear_context:
        session.add(ContextVersion(content="", source="reset"))
    await session.commit()


@router.get("/context")
async def get_context(session: Annotated[AsyncSession, Depends(get_session)]) -> ContextOut:
    current = await current_context(session)
    if current is None:
        return ContextOut(id=None, content="", source=None, created_at=None)
    return ContextOut.model_validate(current, from_attributes=True)


@router.put("/context")
async def put_context(body: ContextIn, session: Annotated[AsyncSession, Depends(get_session)]) -> ContextOut:
    version = ContextVersion(content=body.content, source="user")
    session.add(version)
    await session.commit()
    await session.refresh(version)
    return ContextOut.model_validate(version, from_attributes=True)


@router.get("/context/versions")
async def context_versions(session: Annotated[AsyncSession, Depends(get_session)]) -> list[ContextOut]:
    rows = await session.scalars(select(ContextVersion).order_by(ContextVersion.id.desc()).limit(100))
    return [ContextOut.model_validate(v, from_attributes=True) for v in rows]