CA/backend/routers/chat.py

"""
Chat proxy router — forwards to ai.2890.ltd OpenAI-compatible API.
Non-streaming: returns {reply, model}
Streaming: returns SSE text/event-stream via StreamingResponse
"""
from __future__ import annotations

import json
import logging
from typing import Optional

import httpx
from fastapi import APIRouter, HTTPException, Depends
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

from config import CHAT_API_KEY, CHAT_API_BASE, CHAT_MODEL
from auth.dependencies import get_current_user

logger = logging.getLogger("cbpoa.chat")

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

SYSTEM_PROMPT = (
    "You are a medical risk analysis assistant for the Wuhan Children's Respiratory "
    "Disease Risk Assessment System (CBPOA). You help users understand environmental "
    "health risks, air quality impacts on children's respiratory health, and spatial "
    "risk patterns. Answer in Chinese (Simplified) unless the user asks in English. "
    "Be helpful, concise, and evidence-based."
)

# ---------------------------------------------------------------------------
# Models
# ---------------------------------------------------------------------------

class ChatMessage(BaseModel):
    role: str
    content: str


class ChatRequest(BaseModel):
    messages: list[ChatMessage]
    stream: bool = False


class ChatResponse(BaseModel):
    reply: str
    model: str


# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------

def _build_payload(messages: list[ChatMessage]) -> dict:
    """Build the OpenAI-compatible request payload."""
    system_msg = {"role": "system", "content": SYSTEM_PROMPT}
    user_msgs = [{"role": m.role, "content": m.content} for m in messages]
    return {
        "model": CHAT_MODEL,
        "messages": [system_msg, *user_msgs],
        "temperature": 0.7,
        "stream": False,
    }


def _redact_key(key: str) -> str:
    """Return a safe version of the API key for logging."""
    if not key:
        return "<not set>"
    return key[:6] + "..." if len(key) > 6 else "***"


# ---------------------------------------------------------------------------
# Endpoint
# ---------------------------------------------------------------------------

@router.post("/chat", response_model=ChatResponse)
async def chat_proxy(body: ChatRequest, user=Depends(get_current_user)):
    """
    Proxy chat requests to the OpenAI-compatible API at ai.2890.ltd.

    - **messages**: list of {role, content}
    - **stream**: set to true for SSE streaming (default false)
    """
    if not CHAT_API_KEY:
        raise HTTPException(
            status_code=503,
            detail="Chat API key is not configured on the server.",
        )

    payload = _build_payload(body.messages)

    if body.stream:
        # -- streaming path -------------------------------------------------
        payload["stream"] = True

        async def event_generator():
            async with httpx.AsyncClient(timeout=httpx.Timeout(60.0)) as client:
                try:
                    async with client.stream(
                        "POST",
                        f"{CHAT_API_BASE}/chat/completions",
                        json=payload,
                        headers={
                            "Authorization": f"Bearer {CHAT_API_KEY}",
                            "Content-Type": "application/json",
                        },
                    ) as response:
                        if response.status_code != 200:
                            # Read error body and forward as a single SSE error
                            error_text = ""
                            async for chunk in response.aiter_text():
                                error_text += chunk
                            logger.error(
                                "Upstream chat error %d: %s",
                                response.status_code,
                                error_text[:500],
                            )
                            yield f"data: {json.dumps({'error': f'Upstream API returned {response.status_code}'})}\n\n"
                            yield "data: [DONE]\n\n"
                            return

                        async for line in response.aiter_lines():
                            yield line + "\n"
                            # SSE spec uses \n\n as event separator; upstream
                            # may send \n\n itself, but we ensure it explicitly.
                            if line.strip() == "data: [DONE]":
                                break

                except httpx.TimeoutException:
                    logger.exception("Timeout connecting to chat upstream")
                    yield f"data: {json.dumps({'error': 'Upstream request timed out'})}\n\n"
                    yield "data: [DONE]\n\n"
                except httpx.ConnectError:
                    logger.exception("Cannot connect to chat upstream")
                    yield f"data: {json.dumps({'error': 'Cannot connect to chat API'})}\n\n"
                    yield "data: [DONE]\n\n"
                except Exception:
                    logger.exception("Unexpected error in chat stream")
                    yield f"data: {json.dumps({'error': 'Internal streaming error'})}\n\n"
                    yield "data: [DONE]\n\n"

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

    # -- non-streaming path -------------------------------------------------
    async with httpx.AsyncClient(timeout=httpx.Timeout(60.0)) as client:
        try:
            resp = await client.post(
                f"{CHAT_API_BASE}/chat/completions",
                json=payload,
                headers={
                    "Authorization": f"Bearer {CHAT_API_KEY}",
                    "Content-Type": "application/json",
                },
            )
            if resp.status_code != 200:
                detail = resp.text[:300]
                logger.error(
                    "Upstream chat API returned %d: %s",
                    resp.status_code,
                    detail,
                )
                raise HTTPException(
                    status_code=502,
                    detail=f"Upstream API error: {resp.status_code}",
                )

            data = resp.json()
            choices = data.get("choices", [])
            if not choices:
                raise HTTPException(
                    status_code=502,
                    detail="Upstream API returned no choices",
                )

            reply = choices[0]["message"]["content"]
            return ChatResponse(reply=reply, model=data.get("model", CHAT_MODEL))

        except httpx.TimeoutException:
            logger.exception("Timeout calling chat upstream")
            raise HTTPException(status_code=504, detail="Upstream API timed out")
        except httpx.ConnectError:
            logger.exception("Cannot connect to chat upstream (api_key=%s)", _redact_key(CHAT_API_KEY))
            raise HTTPException(status_code=502, detail="Cannot connect to chat API")
        except HTTPException:
            raise
        except Exception:
            logger.exception("Unexpected error proxying chat")
            raise HTTPException(status_code=500, detail="Internal chat proxy error")