Update backend/model_router.py

backend/model_router.py

@@ -1,14 +1,19 @@
-import os, random, aiohttp
+
+# Fixed model_router.py
+model_router_py = '''import os
+import random
+import aiohttp
 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(",")
+groq_keys = [k.strip() for k in os.environ.get("GROQ_API_KEYS", "").split(",") if k.strip()]
+openrouter_keys = [k.strip() for k in os.environ.get("OPENROUTER_KEYS", "").split(",") if k.strip()]
 ollama_base = os.environ.get("OLLAMA_BASE_URL", "http://host.docker.internal:11434")
 
 POLLINATIONS_URL = "https://text.pollinations.ai/openai"
 
+
 def _make_client(base_url: str, api_key: str) -> AsyncOpenAI:
     http_client = httpx.AsyncClient(trust_env=False)
     return AsyncOpenAI(
@@ -17,17 +22,25 @@ def _make_client(base_url: str, api_key: str) -> AsyncOpenAI:
         http_client=http_client
     )
 
+
 def get_groq_client():
+    if not groq_keys:
+        raise ValueError("No Groq API keys configured. Set GROQ_API_KEYS env var.")
     key = random.choice(groq_keys)
     return _make_client("https://api.groq.com/openai/v1", key)
 
+
 def get_openrouter_client():
+    if not openrouter_keys:
+        raise ValueError("No OpenRouter API keys configured. Set OPENROUTER_KEYS env var.")
     key = random.choice(openrouter_keys)
     return _make_client("https://openrouter.ai/api/v1", key)
 
+
 def get_ollama_client():
     return _make_client(f"{ollama_base}/v1", "ollama")
 
+
 async def route_chat(model: str, messages: List[Dict]) -> AsyncGenerator[str, None]:
     if model.startswith("groq-"):
         client = get_groq_client()
@@ -40,6 +53,7 @@ async def route_chat(model: str, messages: List[Dict]) -> AsyncGenerator[str, No
         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:]
@@ -51,6 +65,7 @@ async def route_chat(model: str, messages: List[Dict]) -> AsyncGenerator[str, No
         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:]
@@ -62,6 +77,7 @@ async def route_chat(model: str, messages: List[Dict]) -> AsyncGenerator[str, No
         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(
@@ -74,20 +90,33 @@ async def route_chat(model: str, messages: List[Dict]) -> AsyncGenerator[str, No
                     raise Exception(f"Pollinations error {resp.status}: {text}")
                 data = await resp.json()
                 content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
-                yield content
+                # Yield word by word to simulate streaming
+                words = content.split(" ")
+                for word in words:
+                    yield word + " "
     else:
         raise ValueError(f"Unknown model: {model}")
 
+
 async def route_autocomplete(model: str, prefix: str, suffix: str, max_tokens: int) -> str:
     if model.startswith("ollama-"):
         client = get_ollama_client()
         model_id = model[7:]
+    elif model.startswith("groq-"):
+        client = get_groq_client()
+        model_id = model[5:]
+    elif model.startswith("openrouter-"):
+        client = get_openrouter_client()
+        model_id = model[11:]
     else:
+        # Default to Groq for autocomplete
         client = get_groq_client()
         model_id = "llama-3.1-8b-instant"
-    prompt = f"Complete the following code:\n```\n{prefix}\n```"
+
+    prompt = f"Complete the following code:\\n```\\n{prefix}\\n```"
     if suffix:
-        prompt += f"\nThe code should continue until before: {suffix}"
+        prompt += f"\\nThe code should continue until before: {suffix}"
+
     response = await client.chat.completions.create(
         model=model_id,
         messages=[{"role": "user", "content": prompt}],
@@ -96,7 +125,8 @@ async def route_autocomplete(model: str, prefix: str, suffix: str, max_tokens: i
     )
     return response.choices[0].message.content
 
-def get_available_models():
+
+async def get_available_models():
     models = [
         {"id": "groq-llama-3.3-70b-versatile", "name": "Groq Llama 3.3 70B", "free": False},
         {"id": "groq-llama-3.1-8b-instant", "name": "Groq Llama 3.1 8B (Instant)", "free": False},
@@ -107,12 +137,25 @@ def get_available_models():
         {"id": "openrouter-google/gemma-2-9b-it", "name": "Gemma 2 9B (OpenRouter)", "free": False},
         {"id": "free-pollinations", "name": "Free Pollinations (GPT-4o-mini)", "free": True},
     ]
+
+    # Try to fetch Ollama models asynchronously
     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:
+        async with aiohttp.ClientSession() as session:
+            async with session.get(f"{ollama_base}/api/tags", timeout=aiohttp.ClientTimeout(total=2)) as resp:
+                if resp.status == 200:
+                    data = await resp.json()
+                    for model in data.get("models", []):
+                        models.append({
+                            "id": f"ollama-{model['name']}",
+                            "name": f"Ollama {model['name']}",
+                            "free": True
+                        })
+    except Exception:
         pass
-    return models
+
+    return models
+'''
+
+with open('/mnt/agents/output/model_router.py', 'w') as f:
+    f.write(model_router_py)
+print("model_router.py written successfully")