gateway/server.py

# -*- coding: utf-8 -*-
"""
Darwin-60B-DUO Gateway — FastAPI OpenAI-compatible orchestrator.

Exposes a single OpenAI-compatible endpoint ("darwin-60b-duo") that
internally routes to two backends:
  - Darwin-28B-REASON (English reasoning specialist, HF GPQA Diamond #3)
  - AWAXIS-Think-31B  (Korean specialist, K-AI Leaderboard #1)

Hybrid-A strategy (config.json):
  - 70% Route (single backend)
  - 20% Split / Refine (sequential two-model collaboration)
  - 10% Ensemble V_1 (cross-verification tournament for MCQ / short answers)

Run:
    pip install -r requirements.txt
    python server.py --port 8000 \\
        --darwin-url http://127.0.0.1:8021/v1 \\
        --awaxis-url http://127.0.0.1:8022/v1

License: Gemma (combined-license inheritance — see README).
"""
import argparse
import asyncio
import json
import time
import uuid
from typing import Any, Dict, List, Optional

import httpx
from fastapi import FastAPI, HTTPException
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, Field

from router import select_strategy, RouteDecision
from refine import sequential_refine
from ensemble import ensemble_v1


# ---------------------------------------------------------------------------
# Pydantic models — OpenAI Chat Completions API subset
# ---------------------------------------------------------------------------
class ChatMessage(BaseModel):
    role: str
    content: str


class ChatCompletionRequest(BaseModel):
    model: str = "darwin-60b-duo"
    messages: List[ChatMessage]
    temperature: float = 0.7
    top_p: float = 0.95
    max_tokens: int = 4096
    n: int = 1
    stream: bool = False
    # Optional: force a specific strategy ("route_darwin", "route_awaxis",
    # "split_refine", "ensemble_v1", "auto"). Default "auto" = Hybrid-A router.
    duo_strategy: Optional[str] = "auto"


# ---------------------------------------------------------------------------
# Backend HTTP client
# ---------------------------------------------------------------------------
class Backend:
    def __init__(self, name: str, base_url: str, served_name: str):
        self.name = name
        self.base_url = base_url.rstrip("/")
        self.served_name = served_name
        self.client = httpx.AsyncClient(timeout=httpx.Timeout(900.0))

    async def chat(
        self,
        messages: List[Dict[str, str]],
        temperature: float = 0.7,
        max_tokens: int = 4096,
        n: int = 1,
        top_p: float = 0.95,
    ) -> List[str]:
        payload = {
            "model": self.served_name,
            "messages": messages,
            "temperature": temperature,
            "top_p": top_p,
            "max_tokens": max_tokens,
            "n": n,
        }
        r = await self.client.post(
            f"{self.base_url}/chat/completions", json=payload
        )
        r.raise_for_status()
        data = r.json()
        return [c["message"]["content"] for c in data["choices"]]

    async def health(self) -> bool:
        try:
            r = await self.client.get(f"{self.base_url}/models", timeout=5)
            return r.status_code == 200
        except Exception:
            return False


# ---------------------------------------------------------------------------
# FastAPI app
# ---------------------------------------------------------------------------
app = FastAPI(
    title="Darwin-60B-DUO Gateway",
    version="1.0.0",
    description=(
        "Single OpenAI-compatible endpoint for the Darwin-60B-DUO "
        "(Darwin-28B-REASON + AWAXIS-Think-31B). Hybrid-A routing."
    ),
)

# Initialized via CLI args at startup
DARWIN: Optional[Backend] = None
AWAXIS: Optional[Backend] = None


@app.get("/v1/models")
async def list_models():
    """Expose only the aggregate model to external callers."""
    return {
        "object": "list",
        "data": [
            {
                "id": "darwin-60b-duo",
                "object": "model",
                "owned_by": "FINAL-Bench",
                "created": int(time.time()),
            }
        ],
    }


@app.get("/health")
async def health():
    d_ok = await DARWIN.health() if DARWIN else False
    a_ok = await AWAXIS.health() if AWAXIS else False
    status = "ok" if (d_ok and a_ok) else "degraded"
    return {
        "status": status,
        "backends": {
            "darwin-28r": d_ok,
            "awaxis-31b": a_ok,
        },
        "gateway_version": "1.0.0",
    }


def _build_response(content: str, route_meta: Dict[str, Any]) -> Dict[str, Any]:
    """Build an OpenAI-compatible Chat Completion response with route metadata."""
    return {
        "id": f"chatcmpl-{uuid.uuid4().hex[:24]}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": "darwin-60b-duo",
        "choices": [
            {
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": content,
                },
                "finish_reason": "stop",
            }
        ],
        "usage": {
            "prompt_tokens": -1,  # Aggregate gateway does not track tokens
            "completion_tokens": -1,
            "total_tokens": -1,
        },
        # Non-standard metadata for transparency / debugging
        "_duo_route": route_meta,
    }


@app.post("/v1/chat/completions")
async def chat_completions(req: ChatCompletionRequest):
    if not req.messages:
        raise HTTPException(400, "messages must not be empty")

    user_text = req.messages[-1].content
    messages_dict = [m.dict() for m in req.messages]

    # ----- Strategy selection -----
    if req.duo_strategy and req.duo_strategy != "auto":
        decision = RouteDecision(strategy=req.duo_strategy, reason="user_forced")
    else:
        decision = select_strategy(user_text)

    t0 = time.time()

    # ----- Execute -----
    try:
        if decision.strategy == "route_darwin":
            outputs = await DARWIN.chat(
                messages_dict,
                temperature=req.temperature,
                max_tokens=req.max_tokens,
                top_p=req.top_p,
            )
            content = outputs[0]

        elif decision.strategy == "route_awaxis":
            outputs = await AWAXIS.chat(
                messages_dict,
                temperature=req.temperature,
                max_tokens=req.max_tokens,
                top_p=req.top_p,
            )
            content = outputs[0]

        elif decision.strategy == "split_refine":
            # Darwin reasons in English → AWAXIS polishes in Korean
            content = await sequential_refine(
                DARWIN, AWAXIS, messages_dict,
                temperature=req.temperature, max_tokens=req.max_tokens
            )

        elif decision.strategy == "split_refine_reverse":
            # AWAXIS retrieves Korean context → Darwin polishes in English
            content = await sequential_refine(
                AWAXIS, DARWIN, messages_dict,
                temperature=req.temperature, max_tokens=req.max_tokens
            )

        elif decision.strategy == "ensemble_v1":
            # MCQ / short answer: MAJ@N per model + cross-verify if mismatched
            content = await ensemble_v1(
                DARWIN, AWAXIS, messages_dict,
                temperature=req.temperature, max_tokens=req.max_tokens,
                n_rsa=8,
            )

        else:
            # Fallback: AWAXIS (default for ambiguous / mixed)
            outputs = await AWAXIS.chat(
                messages_dict, temperature=req.temperature,
                max_tokens=req.max_tokens, top_p=req.top_p,
            )
            content = outputs[0]
            decision.strategy = "fallback_awaxis"

    except httpx.HTTPError as e:
        raise HTTPException(503, f"backend error: {type(e).__name__}: {e}")

    elapsed = time.time() - t0
    route_meta = {
        "strategy": decision.strategy,
        "reason": decision.reason,
        "elapsed_s": round(elapsed, 2),
        "language_ratio": decision.korean_ratio,
    }

    return JSONResponse(_build_response(content, route_meta))


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def main():
    p = argparse.ArgumentParser()
    p.add_argument("--host", default="0.0.0.0")
    p.add_argument("--port", type=int, default=8000)
    p.add_argument(
        "--darwin-url", default="http://127.0.0.1:8021/v1",
        help="Darwin-28B-REASON vLLM endpoint",
    )
    p.add_argument(
        "--awaxis-url", default="http://127.0.0.1:8022/v1",
        help="AWAXIS-Think-31B vLLM endpoint",
    )
    p.add_argument("--darwin-served-name", default="darwin-28r")
    p.add_argument("--awaxis-served-name", default="awaxis-31b")
    args = p.parse_args()

    global DARWIN, AWAXIS
    DARWIN = Backend("darwin-28r", args.darwin_url, args.darwin_served_name)
    AWAXIS = Backend("awaxis-31b", args.awaxis_url, args.awaxis_served_name)

    import uvicorn
    uvicorn.run(app, host=args.host, port=args.port, log_level="info")


if __name__ == "__main__":
    main()