"""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))], )