proxy/router.py

import json
import logging
from typing import Any

from fastapi import APIRouter, Depends, FastAPI, Request
from fastapi.responses import JSONResponse, Response
from starlette.exceptions import HTTPException as StarletteHTTPException

from config import HOST, PORT, PROVIDER_KEYS
from proxy.auth import get_real_key, resolve_provider
from proxy.client import forward_request
from proxy.providers import PROVIDERS

logger = logging.getLogger(__name__)
api = APIRouter(prefix="/v1")



MODELS = {
    "nvidia": [
        "deepseek-ai/deepseek-v4-pro",
        "deepseek-ai/deepseek-v4-flash",
        "moonshotai/kimi-k2.6",
        "moonshotai/kimi-k2.5",
        "google/gemma-4-31b-it",
    ],
}


@api.get("/models")
async def list_models() -> JSONResponse:
    data = {
        "object": "list",
        "data": [
            {
                "id": model,
                "object": "model",
                "owned_by": provider,
            }
            for provider, models in MODELS.items()
            for model in models
        ],
    }
    return JSONResponse(content=data)


@api.post("/chat/completions")
async def chat_completions(request: Request) -> Response:
    provider = resolve_provider(request)
    real_key = get_real_key(provider)

    try:
        body = await request.json()
    except Exception as exc:
        raise StarletteHTTPException(status_code=400, detail=f"Invalid JSON body: {exc}") from exc

    if not isinstance(body, dict):
        raise StarletteHTTPException(status_code=400, detail="Request body must be a JSON object")

    model = body.get("model")
    allowed = MODELS.get(provider.name, [])
    if model and model not in allowed:
        raise StarletteHTTPException(
            status_code=400,
            detail=f"Model '{model}' is not available for provider '{provider.name}'. Available: {', '.join(allowed)}",
        )

    logger.info("Proxying request to provider=%s model=%s", provider.name, model or "default")

    result = await forward_request(provider, real_key, request, body)

    if isinstance(result, Response):
        return result

    return Response(
        content=result["content"],
        status_code=result["status_code"],
        headers={k: v for k, v in result["headers"].items() if k.lower() not in ("content-encoding", "transfer-encoding", "content-length")},
    )


@api.get("/health")
async def health() -> JSONResponse:
    return JSONResponse({"status": "ok", "providers": list(PROVIDERS.keys())})


@api.get("/keys")
async def proxy_keys() -> JSONResponse:
    masked = {}
    for name, key in PROVIDER_KEYS.items():
        if key:
            visible = key[:8] + "..." if len(key) > 11 else key
            masked[name] = visible
    return JSONResponse(masked)


def create_app() -> FastAPI:
    app = FastAPI(title="OpenAI-Compatible API Proxy", version="1.0.0")
    app.include_router(api)

    @app.get("/")
    async def root():
        return {"message": "OpenAI API Proxy", "docs": "/v1/health"}

    @app.exception_handler(StarletteHTTPException)
    async def http_exception_handler(_request: Request, exc: StarletteHTTPException) -> JSONResponse:
        return JSONResponse(status_code=exc.status_code, content={"error": {"message": exc.detail, "type": "proxy_error"}})

    return app