app.py

from fastapi import FastAPI, Request, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Annotated, Any, Optional
import asyncio
from redis.asyncio import Redis
import httpx
import os

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows your Netlify URL to connect
    allow_credentials=False,
    allow_methods=["*"],  # Resolves the 405 OPTIONS error
    allow_headers=["*"],
)


@app.post("/")
def health_check():
    return "Healthy"

password = os.environ.get("PASSWORD")
r = Redis(
    host='redis-15562.c1.us-west-2-2.ec2.cloud.redislabs.com',
    port=15562,
    decode_responses=True,
    username="default",
    password=password,
)



Model_links = {
    "llama3.2": "https://galactromedaNPMAIECOSYSTEM-model23.hf.space/llama3.2", #https://sonuramashishnpm-npmai.hf.space/llama
    "qwen2.5-coder:7b":"https://galactromedaNPMAIECOSYSTEM-model24.hf.space/qwen2.5-coder:7b", #https://sonuramashishnpm-npmai.hf.space/qwen
    "vicuna:7b":"https://galactromedaNPMAIECOSYSTEM-model25.hf.space/vicuna:7b", #https://sonuramashish22028704-vicuna7b.hf.space/vicuna
    "phi3:medium":"https://sonuramashish22028704-phi3medium.hf.space/phi3medium",
    "llama3.2:3b":"https://npmaiecosystem-model5.hf.space/llama3.2:3b",
    "llama3.1:8b":"https://npmaiecosystem-model5.hf.space/llama3.1:8b",
    "falcon:7b-instruct":"https://sonuramashish22028704-falcon7binstruct.hf.space/falcon",
    "meditron":"https://KarachiUniNPMAIECOSYSTEM-model14.hf.space/meditron",
    "granite4.1:3b":"https://karachiuninpmaiecosystem-model14.hf.space/granite4.1:3b",
    "wizard-math":"https://karachiuninpmaiecosystem-model15.hf.space/wizard-math",
    "llava":"https://karachiuninpmaiecosystem-model15.hf.space/llava",
    "olmo2":"https://karachiuninpmaiecosystem-model19.hf.space/olmo2",
    "nemotron-mini":"https://karachiuninpmaiecosystem-model19.hf.space/nemotron-mini",
    "internlm2:7b":"https://sonuramashish22028704-internlm27b.hf.space/internlm",
    "codellama:7b-instruct":"https://sonuramashish22028704-falcon7binstruct.hf.space/codellama",
    "mistral:7b":"https://sonuramashish22028704-mistral7b.hf.space/mistral",
    "qwen3.5:9b":"https://sonuramashish22028704-vicuna7b.hf.space/qwen359gb",
    "gemma2:9b":"https://sonuramashish22028704-internlm27b.hf.space/gemma29b",
    "qwen2.5:7b":"https://npmaiecosystem-model6.hf.space/qwen2.5:7b",
    "llama3.2:1b":"https://npmaiecosystem-model6.hf.space/llama3.2:1b",
    "qwen2-math:7b":"https://npmaiecosystem-model7.hf.space/qwen2-math:7b",
    "qwen2.5vl":"https://npmaiecosystem-model7.hf.space/qwen2.5vl",
    "phi3:3.8b":"https://sonuramashishnpmai-model8.hf.space/phi3:3.8b",
    "llava-phi3":"https://sonuramashishnpmai-model8.hf.space/llava-phi3",
    "gemma2:2b":"https://sonuramashishnpmai-model9.hf.space/gemma2:2b",
    "openhermes":"https://sonuramashishnpmai-model9.hf.space/openhermes",
    "deepseek-coder:6.7b":"https://sonuramashishnpmai-model10.hf.space/deepseek-coder:6.7b",
    "yi":"https://sonuramashishnpmai-model10.hf.space/yi",
    "granite3.3:2b":"https://sonuramashishnpmai-model12.hf.space/granite3.3:2b",
    "smollm:1.7b":"https://sonuramashishnpmai-model12.hf.space/smollm:1.7b",
    "stablelm2":"https://sonuramashishnpmai-model13.hf.space/stablelm2",
    "cniongolo_biomistral":"https://sonuramashishnpm-model13.hf.space/cniongolo_biomistral",
    "moondream":"https://karachiuninpmaiecosystem-model16.hf.space/moondream",
    "openchat":"https://karachiuninpmaiecosystem-model16.hf.space/openchat",
    "deepseek-r1:7b":"https://karachiuninpmaiecosystem-model17.hf.space/deepseek-r1:7b",
    "openbmb_minicpm-v2.6":"https://karachiuninpmaiecosystem-model17.hf.space/openbmb_minicpm-v2.6",
    "llama3:8b":"https://karachiuninpmaiecosystem-model18.hf.space/llama3:8b",
    "gemma3:4b":"https://karachiuninpmaiecosystem-model18.hf.space/gemma3:4b",
    "stable-code":"https://galactromedanpmaiecosystem-model20.hf.space/stable-code",
    "codellama:7b":"https://galactromedanpmaiecosystem-model20.hf.space/codellama:7b",
    "aya:8b":"https://galactromedanpmaiecosystem-model21.hf.space/aya:8b",
    "phi4:14b":"https://galactromedanpmaiecosystem-model20.hf.space/phi4:14b",
    "gemma4:e2b":"https://galactromedanpmaiecosystem-model22.hf.space/gemma4:e2b",
    "qwen3.5:2b":"https://galactromedanpmaiecosystem-model22.hf.space/qwen3.5:2b",
    "llama3.2_fall":"https://sonuramashishnpm-npm-journalist.hf.space/llm_fall_llama",
    "qwen2.5-coder:7b_fall":"https://sonuramashish22028704-mistral7b.hf.space/llm_fall_qwen2",
    "vicuna:7b_fall":"https://sonuramashishnpm-model4.hf.space/llm_fall_vicuna",
    "gemma3:12b_fall":"https://npmaiecosystem-gemma312b_fall.hf.space/llm_fall_gemma312b",
    "internlm2:7b_fall":"https://sonuramashishnpm-model2.hf.space/llm_fall_interlm",
    "falcon:7b-instruct_fall":"https://sonuramashishnpm-model1.hf.space/llm_fall_falcon",
    "codellama:7b-instruct_fall":"https://sonuramashishnpm-model3.hf.space/llm_fall_codellama",
    "mistral:7b_fall":"https://sonuramashishnpm-model2.hf.space/fall_llm_mistral",
    "phi3:medium_fall":"https://sonuramashishnpm-model1.hf.space/llama_fall_phi",
    "qwen3.5:9b_fall":"https://sonuramashishnpm-model4.hf.space/llm_fall_qwen359gb",
    "gemma2:9b_fall":"https://sonuramashishnpm-model3.hf.space/llm_fall_gemma29b",
    "gemma3:12b":"https://npmaiecosystem-gemma312b.hf.space/gemma312b",
}

# Updated Lua Script
LUA_CHECK_AND_INC = """
local key = KEYS[1]
local status = tonumber(redis.call('HGET', key, 'status') or '0')
if status < 1 then
    redis.call('HSET', key, 'status', status + 1)
    return status
end
return -1
"""

LUA_REMOVAL_STATUS = """
local key = KEYS[1]
local status = tonumber(redis.call('HGET', key, 'status') or '0')
if status > 0 then
    redis.call('HSET', key, 'status', status -1)
    return status -1
end
return 0
"""

async def check_cond(model_link: str, fall_model: Optional[list] = None):
    status = await r.eval(LUA_CHECK_AND_INC, 1, model_link)
    if status != -1:
        return {"link": model_link, "statusno": status}

    if fall_model:
        for model in fall_model:
            status = await r.eval(LUA_CHECK_AND_INC, 1, model)
            if status != -1:
                return {"link": model, "statusno": status}
                
    
    for model in Model_links.values():
        status = await r.eval(LUA_CHECK_AND_INC, 1, model)
        if status != -1:
            return {"link": model, "statusno": status}

    return None


class Input(BaseModel):
    model: str
    temperature: float = 0.5
    prompt: str
    change: bool = True
    Models: Optional[list] = None

@app.post("/load_balancer")
async def llm_router(inputs: Input):
    if not inputs.model or not inputs.prompt:
        raise HTTPException(status_code=400, detail="Model name and Prompt are required.")

    if inputs.model not in Model_links:
        raise HTTPException(status_code=444, detail="Model not found.")

    model_link = Model_links[inputs.model]
    fall_links = []

    fall_models = inputs.Models
    if fall_models:
        for m in fall_models:
            model_name = f"{m}_fall"
            if model_name in Model_links.keys():
                link = Model_links[model_name]
                fall_links.append(link)
            else:
                raise HTTPException(status_code=402, detail="Fallback models are not found in Models Dictionary")

    model_cond = await check_cond(model_link=model_link, fall_model=fall_links)
    
    if model_cond and model_cond.get("link") and model_cond.get("statusno") is not None:
        return await router(
            model_url=model_cond["link"],
            prompt=inputs.prompt,
            temp=inputs.temperature
        )
    else:
        raise HTTPException(status_code=503, detail="All model endpoints and fallbacks are busy.")

async def router(model_url, prompt, temp):
    error_log = ""
    process= ""
    payload = {"prompt": prompt, "temperature": temp}
    timeout = httpx.Timeout(connect=100.0, read=560.0, write=300.0, pool=320.0)

    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            response = await client.post(model_url, json=payload)
            response.raise_for_status()
            f_response = response.json()["response"]
            if f_response is not None and str(f_response).strip() != "":
                process += f_response

            else:
                raise ValueError("Empty string or None returned in response from LLM")
    except Exception as e:
        error_log += f"LLM backend error: {str(e)}"
        print(f"Detailed ERROR:-{e}")

    finally:
        await r.eval(LUA_REMOVAL_STATUS, 1, model_url)

    if error_log:
        raise HTTPException(status_code=502, detail=error_log)
        
    else:
        return {"response": process}