amul_proxy/server.py

"""OpenAI-compatible proxy for amulai.in.

Why this exists
---------------
CeRAI's `api_handler.py` only knows three providers: OPENAI, GEMINI, and a
"LOCAL" provider that hits ``{base_url}/v1/chat/completions``. The Amul AI
chat backend is a GET endpoint that streams plain text via SSE
(``GET https://api.prod.amulai.in/api/chat/?session_id=...&query=...``).

This shim accepts standard OpenAI ``POST /v1/chat/completions`` requests,
forwards the user's last message to Amul, drains the SSE body, and returns a
single non-streaming OpenAI-style ChatCompletion. CeRAI calls this proxy as
``provider=LOCAL`` with ``base_url=http://localhost:8088``.

Run:
    uvicorn amul_proxy.server:app --host 0.0.0.0 --port 8088
"""
from __future__ import annotations

import logging
import os
import time
import uuid
from threading import Lock
from typing import List, Optional

import httpx
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

AMUL_API_BASE = os.getenv("AMUL_API_BASE", "https://api.prod.amulai.in")
AMUL_ORIGIN = os.getenv("AMUL_ORIGIN", "https://amulai.in")
AMUL_USER_AGENT = os.getenv(
    "AMUL_USER_AGENT",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/120 Safari/537.36",
)
DEFAULT_SOURCE_LANG = os.getenv("AMUL_SOURCE_LANG", "en")
DEFAULT_TARGET_LANG = os.getenv("AMUL_TARGET_LANG", "en")
REQUEST_TIMEOUT = float(os.getenv("AMUL_TIMEOUT", "120"))

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s %(message)s")
logger = logging.getLogger("amul_proxy")

app = FastAPI(title="Amul AI -> OpenAI-compatible proxy")


# --------------------------------------------------------------------------- #
# Token cache (anonymous JWT, valid 24h per the upstream response).            #
# --------------------------------------------------------------------------- #
class _TokenCache:
    def __init__(self) -> None:
        self._token: Optional[str] = None
        self._expiry: float = 0.0
        self._lock = Lock()

    def get(self, client: httpx.Client) -> str:
        with self._lock:
            if self._token and time.time() < self._expiry - 60:
                return self._token
            r = client.post(
                f"{AMUL_API_BASE}/api/auth/anonymous",
                headers={
                    "origin": AMUL_ORIGIN,
                    "referer": f"{AMUL_ORIGIN}/",
                    "content-type": "application/json",
                    "user-agent": AMUL_USER_AGENT,
                },
                json={},
                timeout=REQUEST_TIMEOUT,
            )
            r.raise_for_status()
            data = r.json()
            self._token = data["access_token"]
            self._expiry = time.time() + int(data.get("expires_in", 3600))
            logger.info("Refreshed anonymous token (expires_in=%s)", data.get("expires_in"))
            return self._token


_TOKEN = _TokenCache()


# --------------------------------------------------------------------------- #
# OpenAI-compatible request/response models (minimal).                         #
# --------------------------------------------------------------------------- #
class ChatMessage(BaseModel):
    role: str
    content: str


class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[ChatMessage]
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None
    top_p: Optional[float] = None
    stream: Optional[bool] = False
    # Amul-specific overrides via standard OpenAI request:
    user: Optional[str] = None  # we use this as session_id when present


class Choice(BaseModel):
    index: int = 0
    message: ChatMessage
    finish_reason: str = "stop"


class Usage(BaseModel):
    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0


class ChatCompletionResponse(BaseModel):
    id: str
    object: str = "chat.completion"
    created: int
    model: str
    choices: List[Choice]
    usage: Usage = Field(default_factory=Usage)


# --------------------------------------------------------------------------- #
# Helpers                                                                      #
# --------------------------------------------------------------------------- #
def _flatten_messages(messages: List[ChatMessage]) -> str:
    """Collapse a chat history into a single query string.

    Amul's API takes a single ``query`` parameter (per request), so we prepend
    any system prompt and earlier turns as context, then end with the latest
    user turn.
    """
    if not messages:
        raise HTTPException(400, "messages must not be empty")

    # Prefer just the latest user message; if there's a system prompt, prepend it.
    system_parts = [m.content for m in messages if m.role == "system" and m.content]
    user_msgs = [m for m in messages if m.role == "user"]
    if not user_msgs:
        raise HTTPException(400, "no user message in payload")
    last_user = user_msgs[-1].content

    if system_parts:
        return f"{' '.join(system_parts).strip()}\n\n{last_user}".strip()
    return last_user.strip()


def _call_amul(query: str, session_id: str) -> str:
    with httpx.Client() as client:
        token = _TOKEN.get(client)
        params = {
            "session_id": session_id,
            "query": query,
            "source_lang": DEFAULT_SOURCE_LANG,
            "target_lang": DEFAULT_TARGET_LANG,
            "use_translation_pipeline": "false",
        }
        headers = {
            "authorization": f"Bearer {token}",
            "accept": "*/*",
            "origin": AMUL_ORIGIN,
            "referer": f"{AMUL_ORIGIN}/",
            "user-agent": AMUL_USER_AGENT,
        }
        url = f"{AMUL_API_BASE}/api/chat/"
        with client.stream(
            "GET", url, params=params, headers=headers, timeout=REQUEST_TIMEOUT
        ) as r:
            if r.status_code == 401:
                # Force refresh and retry once.
                logger.warning("401 from Amul, refreshing token and retrying")
                _TOKEN._token = None  # type: ignore[attr-defined]
                token = _TOKEN.get(client)
                headers["authorization"] = f"Bearer {token}"
                with client.stream(
                    "GET", url, params=params, headers=headers, timeout=REQUEST_TIMEOUT
                ) as r2:
                    r2.raise_for_status()
                    return _drain(r2)
            r.raise_for_status()
            return _drain(r)


def _drain(resp: httpx.Response) -> str:
    chunks: List[str] = []
    for piece in resp.iter_text():
        if piece:
            chunks.append(piece)
    return "".join(chunks).strip()


# --------------------------------------------------------------------------- #
# Routes                                                                       #
# --------------------------------------------------------------------------- #
@app.get("/healthz")
def health() -> dict:
    return {"ok": True}


@app.get("/v1/models")
def list_models() -> dict:
    return {
        "object": "list",
        "data": [
            {"id": "amul-ai", "object": "model", "owned_by": "amulai.in"},
        ],
    }


@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
def chat_completions(req: ChatCompletionRequest) -> ChatCompletionResponse:
    if req.stream:
        # CeRAI's LOCAL handler doesn't enable streaming, so we fail loudly.
        raise HTTPException(400, "streaming not supported by this proxy")

    query = _flatten_messages(req.messages)
    session_id = req.user or f"cerai-{uuid.uuid4()}"
    started = time.time()
    try:
        text = _call_amul(query, session_id)
    except httpx.HTTPStatusError as exc:
        raise HTTPException(502, f"upstream {exc.response.status_code}: {exc.response.text[:200]}")
    except httpx.RequestError as exc:
        raise HTTPException(504, f"upstream network error: {exc}")
    elapsed = time.time() - started
    logger.info("amul query session=%s len_q=%d len_a=%d elapsed=%.2fs",
                session_id, len(query), len(text), elapsed)

    return ChatCompletionResponse(
        id=f"chatcmpl-{uuid.uuid4().hex[:24]}",
        created=int(time.time()),
        model=req.model or "amul-ai",
        choices=[Choice(message=ChatMessage(role="assistant", content=text))],
    )