proxy_core.py

"""
代理核心模块 - 提供路由层共用的请求解析、错误格式化和代理流式响应构建逻辑。

消除 routes_codex.py 和 routes_claude.py 中的重复代码。
"""

from __future__ import annotations

import logging
from typing import Any, Optional, Type

from fastapi import Request, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from codebuddy_client import CodeBuddyClient, CodeBuddyClientError

logger = logging.getLogger(__name__)


# ═══════════════════════════════════════════════════════════════
# 统一错误格式
# ═══════════════════════════════════════════════════════════════


def make_error(code: str, message: str) -> dict:
    """构造统一错误响应体

    所有路由和全局异常处理器统一使用此格式：
        {"error": {"code": "...", "message": "..."}}
    """
    return {"error": {"code": code, "message": message}}


def raise_upstream_error(e: CodeBuddyClientError) -> None:
    """将 CodeBuddyClientError 转换为统一的 HTTPException"""
    raise HTTPException(
        status_code=e.status_code or 502,
        detail=make_error("UPSTREAM_ERROR", str(e)),
    )


def raise_internal_error(e: Exception) -> None:
    """将内部异常转换为统一的 HTTPException"""
    raise HTTPException(
        status_code=500,
        detail=make_error("INTERNAL_ERROR", str(e)),
    )


# ═══════════════════════════════════════════════════════════════
# 请求解析
# ═══════════════════════════════════════════════════════════════


async def parse_request_body(request: Request) -> dict:
    """读取并解析请求 JSON 体"""
    try:
        return await request.json()
    except Exception as e:
        logger.error("Failed to parse request JSON: %s", e)
        raise HTTPException(
            status_code=400,
            detail=make_error("BAD_REQUEST", f"Invalid JSON: {e}"),
        )


def validate_request(raw_body: dict, model_class: Type[BaseModel]) -> BaseModel:
    """用 Pydantic 模型校验请求体"""
    try:
        return model_class.model_validate(raw_body)
    except Exception as e:
        logger.error(
            "Request validation failed for %s: %s\nRaw keys: %s",
            model_class.__name__,
            e,
            list(raw_body.keys()) if isinstance(raw_body, dict) else type(raw_body),
        )
        raise HTTPException(
            status_code=422,
            detail=make_error("VALIDATION_ERROR", f"Request validation failed: {e}"),
        )


# ═══════════════════════════════════════════════════════════════
# 流式代理响应构建
# ═══════════════════════════════════════════════════════════════


async def build_streaming_response(
    client: CodeBuddyClient,
    prompt_text: str,
    model: str,
    stream_translator,
    extra_headers: Optional[dict[str, str]] = None,
) -> StreamingResponse:
    """构建流式 SSE 代理响应的通用逻辑。

    1. 创建 CodeBuddy run
    2. 订阅 SSE 流
    3. 翻译为客户端格式
    4. 包装为 StreamingResponse

    Args:
        client: CodeBuddy 客户端
        prompt_text: 翻译后的 prompt 文本
        model: 客户端指定的模型名
        stream_translator: 翻译函数，签名为 (stream: AsyncGenerator, model: str) -> AsyncGenerator
        extra_headers: 额外的响应头

    Returns:
        StreamingResponse 对象
    """
    try:
        run_id = await client.post_run(text=prompt_text)
        codebuddy_stream = client.stream_run(run_id)
        translated_stream = stream_translator(codebuddy_stream, model=model)

        headers = {
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",
        }
        if extra_headers:
            headers.update(extra_headers)

        return StreamingResponse(
            translated_stream,
            media_type="text/event-stream",
            headers=headers,
        )

    except CodeBuddyClientError as e:
        logger.error("CodeBuddy error in proxy: %s", e)
        raise_upstream_error(e)
    except Exception as e:
        logger.exception("Unexpected error in proxy")
        raise_internal_error(e)