proxy.py

#!/usr/bin/env python3

from typing import Any, Dict

import httpx
import uvicorn
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse

MODEL_NAME = "nomic-ai/nomic-embed-text-v1.5"
UPSTREAM_MODEL = "embed"
PORT = 5055

UPSTREAM_EMBEDDINGS_URL = "http://127.0.0.1:8005/v1/embeddings"
UPSTREAM_MODELS_URL = "http://127.0.0.1:8005/v1/models"

app = FastAPI(title="Embedding Proxy", version="1.0.0")


@app.get("/health")
async def health() -> Dict[str, str]:
    return {"status": "ok"}


@app.get("/health/model")
async def health_model() -> Dict[str, Any]:
    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            response = await client.get(UPSTREAM_MODELS_URL)
    except httpx.RequestError as exc:
        raise HTTPException(
            status_code=502,
            detail={
                "status": "error",
                "message": "Could not reach upstream model server",
                "upstream_url": UPSTREAM_MODELS_URL,
                "error": str(exc),
            },
        )

    if response.status_code != 200:
        raise HTTPException(
            status_code=502,
            detail={
                "status": "error",
                "message": "Upstream model server returned non-200 status",
                "upstream_status_code": response.status_code,
                "upstream_body": response.text,
            },
        )

    try:
        data = response.json()
    except Exception as exc:
        raise HTTPException(
            status_code=502,
            detail={
                "status": "error",
                "message": "Upstream model server returned invalid JSON",
                "error": str(exc),
                "upstream_body": response.text,
            },
        )

    models = data.get("data", [])
    model_ids = {
        model.get("id")
        for model in models
        if isinstance(model, dict)
    }

    upstream_available = UPSTREAM_MODEL in model_ids

    return {
        "status": "ok" if upstream_available else "degraded",
        "public_model": MODEL_NAME,
        "upstream_model": UPSTREAM_MODEL,
        "upstream_url": UPSTREAM_MODELS_URL,
        "upstream_available": upstream_available,
        "available_upstream_models": sorted(model_ids),
    }

@app.get("/v1/models")
async def list_models() -> Dict[str, Any]:
    return {
        "object": "list",
        "data": [
            {
                "id": MODEL_NAME,
                "object": "model",
                "owned_by": "local",
            }
        ],
    }


@app.post("/v1/embeddings")
async def proxy_embeddings(request: Request) -> JSONResponse:
    try:
        payload = await request.json()
    except Exception as exc:
        raise HTTPException(status_code=400, detail=f"Invalid JSON: {exc}")

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

    requested_model = payload.get("model")

    if requested_model not in {MODEL_NAME, UPSTREAM_MODEL, None}:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported model '{requested_model}'. Expected '{MODEL_NAME}'.",
        )

    upstream_payload = dict(payload)
    upstream_payload["model"] = UPSTREAM_MODEL

    headers = {}

    auth = request.headers.get("authorization")
    if auth:
        headers["authorization"] = auth

    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            upstream_response = await client.post(
                UPSTREAM_EMBEDDINGS_URL,
                json=upstream_payload,
                headers=headers,
            )
    except httpx.RequestError as exc:
        raise HTTPException(
            status_code=502,
            detail=f"Could not reach upstream embedding server: {exc}",
        )

    try:
        data = upstream_response.json()
    except Exception:
        return JSONResponse(
            status_code=upstream_response.status_code,
            content={
                "error": {
                    "message": upstream_response.text,
                    "type": "upstream_non_json_response",
                }
            },
        )

    return JSONResponse(
        status_code=upstream_response.status_code,
        content=data,
    )


@app.post("/embeddings")
async def proxy_embeddings_short(request: Request) -> JSONResponse:
    return await proxy_embeddings(request)


if __name__ == "__main__":
    uvicorn.run(
        "proxy:app",
        host="0.0.0.0",
        port=PORT,
        reload=False,
    )