app/api/pseudo_stream.py

"""伪流式：/v1/xtc/chat/pseudo/start 与 /v1/xtc/chat/pseudo/poll。

替代原 Netlify Blobs 实现，用 SQLite pseudo_sessions/pseudo_events 表。
"""
from __future__ import annotations

import asyncio
from typing import Any, Optional

from fastapi import APIRouter, Depends, Header, Request
from fastapi.responses import JSONResponse

from ..adapters import gemini_api, openai as openai_adapter
from ..auth import require_access_key
from ..config import GEMINI_DEFAULT_MODEL
from ..errors import HttpError
from ..media.imagefix import fix_images_in_messages, infer_fix_mode
from ..services import pseudo_store
from ..services.config_store import load_config
from ..utils.thought import extract_thought_and_answer
from ._common import (
    CORS_HEADERS,
    normalize_chat_body,
    ok_with_cors,
    record_usage,
    select_provider_and_key,
)
from .xtc import _parse_multipart

router = APIRouter(prefix="/v1/xtc/chat/pseudo", tags=["pseudo-stream"])


@router.post("/start")
async def pseudo_start(
    request: Request,
    _key: str = Depends(require_access_key),
    x_provider: Optional[str] = Header(default=None, alias="x-provider"),
):
    """启动伪流式会话：立即返回 session_id，后台异步执行对话。

    支持 JSON 与 multipart/form-data 两种请求体。multipart 用于带文件/图片
    的场景：传统做法是带文件走 /v1/xtc/chat 的长连接（stream=false，后端等
    上游完整返回），但手表平台的 request.upload 默认超时较短（约 30s），
    上游 LLM 生成耗时超过该值会被平台单方面掐断并报 999。让带文件的请求也
    走伪流式（立即返回 session_id，后台处理），可避开长连接超时。
    """
    content_type = (request.headers.get("content-type") or "").lower()
    if "multipart/form-data" in content_type:
        body = await _parse_multipart(request)
    else:
        body = await request.json()
    if not isinstance(body, dict):
        raise HttpError("invalid body", status=400, code="bad_request")

    # 统一归一化：兼容 input+images 简化形态 与 messages OpenAI 形态
    messages, model_from_body, provider_id = normalize_chat_body(body)
    if not messages:
        raise HttpError("input or messages is required", status=400, code="bad_request")

    model = model_from_body or GEMINI_DEFAULT_MODEL
    provider_id = provider_id or x_provider

    # 图片修正
    image_fix_mode = body.get("image_fix")
    camera_facing = body.get("image_camera_facing")
    if not image_fix_mode and camera_facing:
        image_fix_mode = infer_fix_mode(camera_facing)
    if image_fix_mode:
        await fix_images_in_messages(messages, image_fix_mode)

    config = await load_config()
    provider, api_key, clean_model = await select_provider_and_key(
        config=config, provider_id=provider_id, model=model,
        fallback_model=GEMINI_DEFAULT_MODEL,
    )

    # 构造上游 body
    upstream_body = {
        "model": clean_model,
        "messages": messages,
    }
    for k in ("temperature", "top_p", "top_k", "max_tokens", "max_completion_tokens",
              "presence_penalty", "frequency_penalty", "seed", "stop", "n",
              "response_format", "tools", "tool_choice", "reasoning_effort",
              "extra_body", "google"):
        if k in body and body[k] is not None:
            upstream_body[k] = body[k]
    # 让上游在最后一个流式 chunk 里返回 usage，否则 token 计数永远为 0。
    so = upstream_body.get("stream_options")
    if not isinstance(so, dict):
        so = {}
    so = {**so, "include_usage": True}
    upstream_body["stream_options"] = so

    session_id = pseudo_store.create_session(
        payload={"provider": provider.id, "model": clean_model}
    )

    # 后台执行
    asyncio.create_task(
        _run_pseudo_session(
            session_id, provider, api_key, clean_model, upstream_body, messages, _key
        )
    )

    return ok_with_cors(
        {
            "session_id": session_id,
            "provider": provider.id,
            "model": clean_model,
            "poll_after_ms": 600,
            "expires_in_sec": 600,
        },
        extra_headers={"x-xtc-provider": provider.id, "x-xtc-model": clean_model},
    )


@router.api_route("/poll", methods=["GET", "POST"])
async def pseudo_poll(
    request: Request,
    _key: str = Depends(require_access_key),
):
    if request.method == "POST":
        body = await request.json()
    else:
        body = dict(request.query_params)
    session_id = body.get("session_id") or body.get("sessionId")
    cursor = int(body.get("cursor") or 0)
    if not session_id:
        raise HttpError("session_id is required", status=400, code="bad_request")
    result = pseudo_store.poll(session_id, cursor)
    if result is None:
        raise HttpError(
            f"session not found or expired: {session_id}",
            status=404,
            code="not_found",
        )
    return ok_with_cors(result)


async def _run_pseudo_session(
    session_id: str,
    provider,
    api_key: str,
    clean_model: str,
    upstream_body: dict,
    messages: list,
    access_key: str,
) -> None:
    """后台执行对话：流式请求上游，收完整后一次性写入 session。

    伪流式精髓：后端用流式请求大模型（尽快拿到首字节、可中断、不阻塞），
    但对前端是“一次转发”——收完完整回复后把 text/thought 一次性写入 session
    字段，只 emit 一个 done 事件。前端轮询到 done 直接读 body.text，无需
    自己拼接 deltas 增量。

    不再边收边 emit thought/text 增量，原因：
    1. 前端 extractTextFromPseudoBody 只收集 type:"text"/"raw" 的 delta，
       而上游推理模型把内容放在 reasoning_content，会被标成 type:"thought"，
       前端拼不出回复。
    2. 用户需求是后端流式取完整回复后一次返回，不要把流式字符暴露给前端。
    """
    try:
        if provider.type == "gemini":
            chunk_iter = await gemini_api.chat_completions(
                api_key=api_key, model=clean_model, messages=messages,
                body=upstream_body, stream=True,
            )
        else:
            # OpenAI 厂商：读原始 SSE 转 chunk
            chunk_iter = _openai_passthrough_chunks(
                provider.base_url, api_key, upstream_body
            )

        full_text_parts: list[str] = []
        thought_parts: list[str] = []
        last_finish: Optional[str] = None
        captured_usage: Optional[dict] = None
        async for chunk in chunk_iter:
            # usage 通常出现在最后一个 chunk（choices 为空、带 usage 字段），
            # 之前的实现因为 ``if not choices: continue`` 直接跳过了它。
            if isinstance(chunk, dict) and chunk.get("usage"):
                captured_usage = chunk["usage"]
            choices = chunk.get("choices") or []
            if not choices:
                continue
            choice = choices[0]
            delta = choice.get("delta") or {}
            content = delta.get("content")
            # OpenAI 兼容推理模型把思考链放在 reasoning_content / reasoning
            reasoning = delta.get("reasoning_content") or delta.get("reasoning")
            finish_reason = choice.get("finish_reason")
            if reasoning:
                thought_parts.append(reasoning)
            if content:
                full_text_parts.append(content)
            if finish_reason:
                last_finish = finish_reason

        full_text = "".join(full_text_parts)
        # 优先用流式累积的 reasoning_content；为空则从 <thought> 标签抽取（Gemini）
        streamed_thought = "".join(thought_parts)
        tag_thought, text = extract_thought_and_answer(full_text)
        thought = streamed_thought or tag_thought
        # 一次性写入完整回复 + emit done。前端轮询到 done 直接读 body.text。
        pseudo_store.update_session_state(
            session_id,
            done=True,
            thought=thought,
            text=text,
        )
        pseudo_store.append_event(
            session_id, type="done", delta="", finish_reason=last_finish or "stop",
        )
        record_usage(
            access_key=access_key, provider=provider.id, model=clean_model,
            usage=captured_usage, ok=True,
        )
    except HttpError as e:
        pseudo_store.update_session_state(session_id, done=True, error=e.message)
        pseudo_store.append_event(
            session_id, type="error",
            extra={"code": e.code, "message": e.message, "status": e.status},
        )
        record_usage(
            access_key=access_key, provider=provider.id, model=clean_model,
            usage=None, ok=False, error_code=e.code,
        )
    except Exception as e:
        import traceback
        traceback.print_exc()
        pseudo_store.update_session_state(session_id, done=True, error=str(e))
        pseudo_store.append_event(session_id, type="error", extra={"message": str(e)})
        record_usage(
            access_key=access_key, provider=provider.id, model=clean_model,
            usage=None, ok=False, error_code="internal_error",
        )


async def _openai_passthrough_chunks(base_url: str, api_key: str, body: dict):
    """把 OpenAI 厂商的原始 SSE 字节流解析为 OpenAI chunk dict。

    必须用 parse_sse_stream 按 ``\\n\\n`` 帧分隔符跨 chunk 缓冲解析。
    httpx 的 aiter_bytes() 字节块边界与 SSE 帧边界无关，若按每个 chunk
    独立 split('\\n') 解析，跨块切断的 ``data: {...}`` 行会 json.loads
    失败被静默丢弃，导致回复少字、换行被破坏。
    """
    from ..utils.sse import parse_sse_stream

    chunk_iter = await openai_adapter.chat_completions(
        base_url=base_url, api_key=api_key, body=body, stream=True
    )
    async for sse in parse_sse_stream(chunk_iter):
        if sse.get("__done__"):
            return
        if sse.get("__raw__") is not None:
            # 非 JSON 帧，消费方只关心 OpenAI chunk 结构，跳过
            continue
        yield sse