import os import time from fastapi import FastAPI, Request, HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import Response, JSONResponse, StreamingResponse import httpx from bs4 import BeautifulSoup from typing import List, Dict import asyncio app = FastAPI() app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["GET", "POST"], allow_headers=["*"], ) OLLAMA_LIBRARY_URL = "https://ollama.com/library" RATE_LIMIT = 25 WINDOW_SECONDS = 60 * 60 * 24 ip_store = {} # { ip: { "count": int, "reset": timestamp } } def check_rate_limit(ip: str): now = time.time() if ip not in ip_store: ip_store[ip] = {"count": 0, "reset": now + WINDOW_SECONDS} entry = ip_store[ip] if now > entry["reset"]: entry["count"] = 0 entry["reset"] = now + WINDOW_SECONDS if entry["count"] >= RATE_LIMIT: raise HTTPException( status_code=429, detail="Daily limit reached: 25 images per IP" ) entry["count"] += 1 PKEY = os.getenv("POLLINATIONS_KEY", "") CHAT_RATE_LIMIT = 50 CHAT_WINDOW_SECONDS = 60 * 60 chat_ip_store = {} GROQ_TOOL_MODELS = [ "openai/gpt-oss-120b", "openai/gpt-oss-20b", "meta-llama/llama-4-scout-17b-16e-instruct", "qwen/qwen3-32b", "moonshotai/kimi-k2-instruct", ] GROQ_NORMAL_MODELS = [ "llama-3.1-8b-instant", "llama-3.3-70b-versatile", "meta-llama/llama-4-maverick-17b-128e-instruct", "meta-llama/llama-guard-4-12b", "openai/gpt-oss-safeguard-20b", "qwen/qwen3-32b", ] CEREBRAS_MODELS = [ "gpt-oss-120b", "llama3.1-8b", "qwen-3-235b-a22b-instruct-2507", "zai-glm-4.7", ] def check_chat_rate_limit(ip: str): now = time.time() if ip not in chat_ip_store: chat_ip_store[ip] = { "count": 0, "reset": now + CHAT_WINDOW_SECONDS } entry = chat_ip_store[ip] if now > entry["reset"]: entry["count"] = 0 entry["reset"] = now + CHAT_WINDOW_SECONDS if entry["count"] >= CHAT_RATE_LIMIT: raise HTTPException( status_code=429, detail="Chat rate limit exceeded" ) entry["count"] += 1 return entry["count"] def detect_tool_use(messages: list) -> bool: """ Detect if the request uses tools. We check for: - presence of "tool_calls" - messages containing function_call-like structures """ for m in messages: if "tool_calls" in m: return True if "function_call" in m: return True return False def choose_model(messages: list, msg_count: int): uses_tools = detect_tool_use(messages) if uses_tools: if msg_count > 20: return "openai/gpt-oss-120b", "groq" return "openai/gpt-oss-20b", "groq" if msg_count > 20: return "gpt-oss-120b", "cerebras" return "llama-3.1-8b-instant", "groq" @app.get("/genimg/{prompt}") async def generate_image(prompt: str, request: Request): client_ip = request.client.host check_rate_limit(client_ip) url = f"https://gen.pollinations.ai/image/{prompt}?model=zimage&key={PKEY}" async with httpx.AsyncClient() as client: response = await client.get(url) if response.status_code != 200: raise HTTPException( status_code=500, detail=f"Pollinations error: {response.status_code}" ) return Response( content=response.content, media_type="image/jpeg" ) @app.get("/models") async def get_models() -> List[Dict]: async with httpx.AsyncClient() as client: response = await client.get(OLLAMA_LIBRARY_URL) html = response.text soup = BeautifulSoup(html, "html.parser") items = soup.select("li[x-test-model]") models = [] for item in items: name = item.select_one("[x-test-model-title] span") description = item.select_one("p.max-w-lg") sizes = [el.get_text(strip=True) for el in item.select("[x-test-size]")] pulls = item.select_one("[x-test-pull-count]") tags = [t.get_text(strip=True) for t in item.select('span[class*="text-blue-600"]')] updated = item.select_one("[x-test-updated]") link = item.select_one("a") models.append({ "name": name.get_text(strip=True) if name else "", "description": description.get_text(strip=True) if description else "No description", "sizes": sizes, "pulls": pulls.get_text(strip=True) if pulls else "Unknown", "tags": tags, "updated": updated.get_text(strip=True) if updated else "Unknown", "link": link.get("href") if link else None, }) return models @app.post("/gen/chat/completions") async def generate_text(request: Request): body = await request.json() messages = body.get("messages", []) if not isinstance(messages, list) or len(messages) == 0: raise HTTPException(400, "messages[] is required") ip = request.client.host msg_count = check_chat_rate_limit(ip) uses_tools = ( "tools" in body and isinstance(body["tools"], list) and len(body["tools"]) > 0 ) or ("tool_choice" in body and body["tool_choice"] not in [None, "none"]) requested_model = body.get("model") if uses_tools: if msg_count > 20: chosen_model = "openai/gpt-oss-120b" else: chosen_model = "openai/gpt-oss-20b" provider = "groq" else: if msg_count > 20: chosen_model = "gpt-oss-120b" provider = "cerebras" else: chosen_model = "llama-3.1-8b-instant" provider = "groq" body["model"] = chosen_model stream = body.get("stream", False) if provider == "groq": API_KEY = os.getenv("GROQ_KEY", "") if not API_KEY: raise HTTPException(500, "Missing GROQ_KEY") url = "https://api.groq.com/openai/v1/chat/completions" elif provider == "cerebras": API_KEY = os.getenv("CER_KEY", "") if not API_KEY: raise HTTPException(500, "Missing CER_KEY") url = "https://api.cerebras.ai/v1/chat/completions" else: raise HTTPException(500, "Unknown provider routing error") headers = {"Authorization": f"Bearer {API_KEY}"} if stream: async def event_generator(): async with httpx.AsyncClient(timeout=None) as client: async with client.stream("POST", url, json=body, headers=headers) as r: async for chunk in r.aiter_raw(): yield chunk return StreamingResponse( event_generator(), media_type="text/event-stream", ) else: async with httpx.AsyncClient(timeout=None) as client: r = await client.post(url, json=body, headers=headers) return JSONResponse( status_code=r.status_code, content=r.json() ) raise HTTPException(500, "Unknown provider routing error")