199 lines
7.4 KiB
Python
199 lines
7.4 KiB
Python
|
"""
|
||
|
|
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")
|