app.py

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")