backend/model_router.py

import os, random, aiohttp, asyncio, json
from typing import List, Dict, AsyncGenerator
from openai import AsyncOpenAI
import httpx

groq_keys = os.environ.get("GROQ_API_KEYS", "").split(",")
openrouter_keys = os.environ.get("OPENROUTER_KEYS", "").split(",")
ollama_base = os.environ.get("OLLAMA_BASE_URL", "http://host.docker.internal:11434")
POLLINATIONS_URL = "https://text.pollinations.ai/openai"

# ---------- client helpers ----------
def _make_client(base_url: str, api_key: str) -> AsyncOpenAI:
    http_client = httpx.AsyncClient(trust_env=False)
    return AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=http_client)

def get_groq_client():
    return _make_client("https://api.groq.com/openai/v1", random.choice(groq_keys))

def get_openrouter_client():
    return _make_client("https://openrouter.ai/api/v1", random.choice(openrouter_keys))

def get_ollama_client():
    return _make_client(f"{ollama_base}/v1", "ollama")

# ---------- dynamic model cache ----------
_cached_models = []

async def _fetch_groq_models() -> list:
    try:
        client = get_groq_client()
        resp = await client.models.list()
        # filter out audio and guard models
        return [m.id for m in resp.data if "whisper" not in m.id and "guard" not in m.id]
    except:
        return []

async def _fetch_openrouter_models() -> list:
    try:
        client = get_openrouter_client()
        resp = await client.models.list()
        # filter out image models
        return [m.id for m in resp.data if "flux" not in m.id][:80]
    except:
        return []

async def refresh_model_cache():
    global _cached_models
    models = []

    # Groq
    for m in await _fetch_groq_models():
        models.append({"id": f"groq-{m}", "name": f"Groq {m}", "free": False})

    # OpenRouter
    for m in await _fetch_openrouter_models():
        models.append({"id": f"openrouter-{m}", "name": f"OR {m}", "free": False})

    # Free Pollinations
    models.append({"id": "free-pollinations", "name": "Free Pollinations (GPT-4o-mini)", "free": True})

    # Ollama
    try:
        import requests
        resp = requests.get(f"{ollama_base}/api/tags", timeout=2)
        if resp.status_code == 200:
            for model in resp.json().get("models", []):
                models.append({"id": f"ollama-{model['name']}", "name": f"Ollama {model['name']}", "free": True})
    except:
        pass

    # Extra models from env (user override / additions)
    extra = json.loads(os.environ.get("EXTRA_MODELS", "[]"))
    models.extend(extra)

    _cached_models = models

def get_available_models_sync():
    return _cached_models

# ---------- chat routing ----------
async def route_chat(model: str, messages: List[Dict]) -> AsyncGenerator[str, None]:
    if model.startswith("groq-"):
        client = get_groq_client()
        model_id = model[5:]
        stream = await client.chat.completions.create(messages=messages, model=model_id, stream=True)
        async for chunk in stream:
            if chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content
    elif model.startswith("openrouter-"):
        client = get_openrouter_client()
        model_id = model[11:]
        stream = await client.chat.completions.create(messages=messages, model=model_id, stream=True)
        async for chunk in stream:
            if chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content
    elif model.startswith("ollama-"):
        client = get_ollama_client()
        model_id = model[7:]
        stream = await client.chat.completions.create(messages=messages, model=model_id, stream=True)
        async for chunk in stream:
            if chunk.choices[0].delta.content:
                yield chunk.choices[0].delta.content
    elif model == "free-pollinations":
        async with aiohttp.ClientSession() as session:
            async with session.post(POLLINATIONS_URL, json={"messages": messages, "model": "openai"}, headers={"Content-Type": "application/json"}) as resp:
                if resp.status == 200:
                    data = await resp.json()
                    yield data.get("choices", [{}])[0].get("message", {}).get("content", "")
                    return
                raise Exception(f"Pollinations error {resp.status}")
    else:
        raise ValueError(f"Unknown model: {model}")

async def route_autocomplete(model: str, prefix: str, suffix: str, max_tokens: int) -> str:
    client = get_groq_client()
    model_id = "llama-3.1-8b-instant"
    prompt = f"Complete the following code:\n```\n{prefix}\n```"
    if suffix:
        prompt += f"\nThe code should continue until before: {suffix}"
    response = await client.chat.completions.create(
        model=model_id,
        messages=[{"role": "user", "content": prompt}],
        max_tokens=max_tokens, temperature=0.1,
    )
    return response.choices[0].message.content