Fix: update ai_router/router.py

ai_router/router.py

@@ -1,8 +1,7 @@
 """
-🚀 GOD MODE+ LLM Router — Unified AI Gateway
-Primary: Cloudflare AI Gateway → Groq → OpenAI → HF Inference
-ALL LLM calls MUST go through LLMRouter.ask()
-No direct API calls allowed anywhere else.
+Multi-Model AI Router — Phase 9
+Supports: OpenAI, Groq, Cerebras, OpenRouter, HuggingFace
+Automatic failover chain: OpenAI → Groq → Cerebras → OpenRouter → HF
 """
 
 import asyncio
@@ -16,247 +15,250 @@ import structlog
 
 log = structlog.get_logger()
 
-# ─── Gateway Config ────────────────────────────────────────────────────────────
-CF_GATEWAY_URL = os.environ.get(
-    "CF_GATEWAY_URL",
-    "https://gateway.pyaesone-gtckglay.workers.dev/v1/chat/completions"
-)
-CF_GATEWAY_KEY = os.environ.get("CF_GATEWAY_KEY", "")
-
-GROQ_API_KEY    = os.environ.get("GROQ_API_KEY", "")
-OPENAI_API_KEY  = os.environ.get("OPENAI_API_KEY", "")
-HF_TOKEN        = os.environ.get("HF_TOKEN", "")
-
-DEFAULT_MODEL   = os.environ.get("DEFAULT_MODEL", "meta-llama/llama-4-scout-17b-16e-instruct")
-FALLBACK_MODEL  = "llama-3.3-70b-versatile"
-
-
-# ─── Provider Fallback Chain ───────────────────────────────────────────────────
-PROVIDER_CHAIN = [
+# ─── Provider Config ──────────────────────────────────────────────────────────
+PROVIDERS = [
     {
-        "name":    "cloudflare_gateway",
-        "url":     CF_GATEWAY_URL,
-        "key_fn":  lambda: CF_GATEWAY_KEY or OPENAI_API_KEY or GROQ_API_KEY,
-        "model":   DEFAULT_MODEL,
-        "enabled": lambda: bool(CF_GATEWAY_KEY or OPENAI_API_KEY or GROQ_API_KEY),
+        "name": "openai",
+        "key_env": "OPENAI_API_KEY",
+        "base_url": os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1"),
+        "default_model": os.environ.get("DEFAULT_MODEL", "gpt-4o"),
+        "headers_fn": lambda k: {"Authorization": f"Bearer {k}", "Content-Type": "application/json"},
     },
     {
-        "name":    "groq",
-        "url":     "https://api.groq.com/openai/v1/chat/completions",
-        "key_fn":  lambda: GROQ_API_KEY,
-        "model":   FALLBACK_MODEL,
-        "enabled": lambda: bool(GROQ_API_KEY),
+        "name": "groq",
+        "key_env": "GROQ_API_KEY",
+        "base_url": "https://api.groq.com/openai/v1",
+        "default_model": "llama-3.3-70b-versatile",
+        "headers_fn": lambda k: {"Authorization": f"Bearer {k}", "Content-Type": "application/json"},
     },
     {
-        "name":    "openai",
-        "url":     os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1") + "/chat/completions",
-        "key_fn":  lambda: OPENAI_API_KEY,
-        "model":   "gpt-4o",
-        "enabled": lambda: bool(OPENAI_API_KEY),
+        "name": "cerebras",
+        "key_env": "CEREBRAS_API_KEY",
+        "base_url": "https://api.cerebras.ai/v1",
+        "default_model": "llama3.1-70b",
+        "headers_fn": lambda k: {"Authorization": f"Bearer {k}", "Content-Type": "application/json"},
     },
     {
-        "name":    "hf_inference",
-        "url":     "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-8B-Instruct/v1/chat/completions",
-        "key_fn":  lambda: HF_TOKEN,
-        "model":   "meta-llama/Meta-Llama-3-8B-Instruct",
-        "enabled": lambda: bool(HF_TOKEN),
+        "name": "openrouter",
+        "key_env": "OPENROUTER_API_KEY",
+        "base_url": "https://openrouter.ai/api/v1",
+        "default_model": "meta-llama/llama-3.3-70b-instruct:free",
+        "headers_fn": lambda k: {
+            "Authorization": f"Bearer {k}",
+            "Content-Type": "application/json",
+            "HTTP-Referer": "https://god-agent.ai",
+            "X-Title": "God Agent Platform",
+        },
+    },
+    {
+        "name": "anthropic",
+        "key_env": "ANTHROPIC_API_KEY",
+        "base_url": "https://api.anthropic.com/v1",
+        "default_model": "claude-3-5-sonnet-20241022",
+        "headers_fn": lambda k: {
+            "x-api-key": k,
+            "anthropic-version": "2023-06-01",
+            "Content-Type": "application/json",
+        },
     },
 ]
 
 
-class LLMRouter:
+class AIRouter:
     """
-    Unified LLM Router — all agents call LLMRouter.ask()
-    Automatic failover: CF Gateway → Groq → OpenAI → HF
-    Supports streaming via WebSocket emit.
+    God Mode AI Router — automatically routes and fails over across providers.
+    Supports streaming token output via WebSocket.
     """
 
     def __init__(self, ws_manager=None):
         self.ws = ws_manager
-        self._stats: Dict[str, Dict] = {
-            p["name"]: {"calls": 0, "errors": 0, "latency_ms": []}
-            for p in PROVIDER_CHAIN
-        }
+        self._stats: Dict[str, Dict] = {p["name"]: {"calls": 0, "errors": 0, "latency": []} for p in PROVIDERS}
 
-    # ─── PRIMARY ENTRY POINT ──────────────────────────────────────────────────
+    def _get_provider(self, name: str) -> Optional[Dict]:
+        return next((p for p in PROVIDERS if p["name"] == name), None)
 
-    async def ask(
+    def _available_providers(self) -> List[Dict]:
+        """Return providers with valid API keys, in priority order."""
+        return [p for p in PROVIDERS if os.environ.get(p["key_env"], "")]
+
+    # ─── Main Entry Point ─────────────────────────────────────────────────────
+
+    async def complete(
         self,
         messages: List[Dict],
         task_id: str = "",
         session_id: str = "",
         temperature: float = 0.7,
         max_tokens: int = 4096,
-        model: str = "",
+        preferred_model: str = "",
         stream: bool = True,
     ) -> str:
-        """
-        Route LLM call through provider chain with automatic failover.
-        Returns full response text.
-        """
-        active_providers = [p for p in PROVIDER_CHAIN if p["enabled"]()]
+        """Route completion through available providers with failover."""
+        providers = self._available_providers()
 
-        if not active_providers:
-            log.warning("No LLM providers available — returning demo response")
-            return await self._demo_response(messages, task_id, session_id)
+        if not providers:
+            return await self._demo_stream(messages, task_id, session_id)
 
         last_error = None
-        for provider in active_providers:
+        for provider in providers:
             try:
                 start = time.time()
-                result = await self._call_provider(
-                    provider=provider,
-                    messages=messages,
-                    task_id=task_id,
-                    session_id=session_id,
-                    temperature=temperature,
-                    max_tokens=max_tokens,
-                    model_override=model,
-                    stream=stream,
-                )
-                elapsed_ms = round((time.time() - start) * 1000)
+                if provider["name"] == "anthropic":
+                    result = await self._anthropic_stream(
+                        provider, messages, task_id, session_id, temperature, max_tokens
+                    )
+                else:
+                    result = await self._openai_compat_stream(
+                        provider, messages, task_id, session_id, temperature, max_tokens, preferred_model
+                    )
+                elapsed = time.time() - start
                 self._stats[provider["name"]]["calls"] += 1
-                self._stats[provider["name"]]["latency_ms"].append(elapsed_ms)
-                log.info("LLMRouter success", provider=provider["name"], ms=elapsed_ms, chars=len(result))
+                self._stats[provider["name"]]["latency"].append(elapsed)
+                log.info("AI Router success", provider=provider["name"], ms=round(elapsed * 1000))
                 return result
-
             except Exception as e:
                 last_error = e
                 self._stats[provider["name"]]["errors"] += 1
-                log.warning("LLMRouter failover", provider=provider["name"], error=str(e)[:200])
+                log.warning("AI Router failover", provider=provider["name"], error=str(e))
                 continue
 
-        log.error("All LLM providers failed", last_error=str(last_error))
-        return await self._demo_response(messages, task_id, session_id)
+        log.error("All AI providers failed", last_error=str(last_error))
+        return await self._demo_stream(messages, task_id, session_id)
 
-    # ─── Provider Call ─────────────────────────────────────────────────────────
+    # ─── OpenAI-compatible Stream (OpenAI, Groq, Cerebras, OpenRouter) ────────
 
-    async def _call_provider(
-        self,
-        provider: Dict,
-        messages: List[Dict],
-        task_id: str,
-        session_id: str,
-        temperature: float,
-        max_tokens: int,
-        model_override: str,
-        stream: bool,
+    async def _openai_compat_stream(
+        self, provider, messages, task_id, session_id, temperature, max_tokens, preferred_model
     ) -> str:
-        key = provider["key_fn"]()
-        model = model_override or provider["model"]
-        url = provider["url"]
-
-        headers = {
-            "Authorization": f"Bearer {key}",
-            "Content-Type": "application/json",
-        }
+        key = os.environ.get(provider["key_env"], "")
+        model = preferred_model or provider["default_model"]
+        headers = provider["headers_fn"](key)
         payload = {
             "model": model,
             "messages": messages,
+            "stream": True,
             "temperature": temperature,
             "max_tokens": max_tokens,
-            "stream": stream,
         }
+        full_text = ""
+        async with httpx.AsyncClient(timeout=120) as client:
+            async with client.stream(
+                "POST", f"{provider['base_url']}/chat/completions",
+                headers=headers, json=payload
+            ) as resp:
+                resp.raise_for_status()
+                async for line in resp.aiter_lines():
+                    if not line.startswith("data:"):
+                        continue
+                    chunk = line[6:].strip()
+                    if chunk == "[DONE]":
+                        break
+                    try:
+                        data = json.loads(chunk)
+                        delta = data["choices"][0]["delta"].get("content", "")
+                        if delta:
+                            full_text += delta
+                            await self._emit_chunk(delta, task_id, session_id)
+                    except Exception:
+                        pass
+        return full_text
+
+    # ─── Anthropic Stream ─────────────────────────────────────────────────────
 
+    async def _anthropic_stream(
+        self, provider, messages, task_id, session_id, temperature, max_tokens
+    ) -> str:
+        key = os.environ.get(provider["key_env"], "")
+        headers = provider["headers_fn"](key)
+        system = ""
+        filtered = []
+        for m in messages:
+            if m["role"] == "system":
+                system = m["content"]
+            else:
+                filtered.append(m)
+        payload = {
+            "model": provider["default_model"],
+            "max_tokens": max_tokens,
+            "messages": filtered,
+            "stream": True,
+        }
+        if system:
+            payload["system"] = system
         full_text = ""
-        async with httpx.AsyncClient(timeout=120.0) as client:
-            if stream:
-                async with client.stream("POST", url, headers=headers, json=payload) as resp:
-                    resp.raise_for_status()
-                    async for line in resp.aiter_lines():
-                        if not line.startswith("data:"):
-                            continue
-                        chunk_str = line[5:].strip()
-                        if chunk_str == "[DONE]":
-                            break
-                        try:
-                            data = json.loads(chunk_str)
-                            delta = (
-                                data.get("choices", [{}])[0]
-                                .get("delta", {})
-                                .get("content", "")
-                            )
+        async with httpx.AsyncClient(timeout=120) as client:
+            async with client.stream(
+                "POST", f"{provider['base_url']}/messages",
+                headers=headers, json=payload
+            ) as resp:
+                resp.raise_for_status()
+                async for line in resp.aiter_lines():
+                    if not line.startswith("data:"):
+                        continue
+                    try:
+                        data = json.loads(line[5:].strip())
+                        if data.get("type") == "content_block_delta":
+                            delta = data["delta"].get("text", "")
                             if delta:
                                 full_text += delta
-                                await self._emit(delta, task_id, session_id)
-                        except Exception:
-                            pass
-            else:
-                resp = await client.post(url, headers=headers, json={**payload, "stream": False})
-                resp.raise_for_status()
-                data = resp.json()
-                full_text = (
-                    data.get("choices", [{}])[0]
-                    .get("message", {})
-                    .get("content", "")
-                )
-
+                                await self._emit_chunk(delta, task_id, session_id)
+                    except Exception:
+                        pass
         return full_text
 
-    # ─── Demo Response ─────────────────────────────────────────────────────────
+    # ─── Demo Stream ──────────────────────────────────────────────────────────
 
-    async def _demo_response(self, messages: List[Dict], task_id: str, session_id: str) -> str:
+    async def _demo_stream(self, messages, task_id, session_id) -> str:
         last_user = next(
             (m["content"] for m in reversed(messages) if m["role"] == "user"), "Hello"
         )
-        text = (
-            f"🤖 **God Mode+ AI** (No API Key — Demo Mode)\n\n"
-            f"Received: *{last_user[:120]}*\n\n"
-            f"To enable real AI, set one of:\n"
-            f"- `CF_GATEWAY_KEY` (Cloudflare Gateway — recommended)\n"
-            f"- `GROQ_API_KEY` (Groq Llama 3.3 70B — free)\n"
+        response = (
+            f"🤖 **God Agent** (Demo Mode)\n\n"
+            f"Received: *{last_user[:100]}*\n\n"
+            f"To enable real AI, set one of these env vars:\n"
             f"- `OPENAI_API_KEY` (GPT-4o)\n"
-            f"- `HF_TOKEN` (HuggingFace)\n\n"
-            f"**System Status:** All 10 agents online ✅"
+            f"- `GROQ_API_KEY` (Llama 3.3 70B — Free)\n"
+            f"- `OPENROUTER_API_KEY` (Multi-model)\n"
+            f"- `ANTHROPIC_API_KEY` (Claude 3.5)\n\n"
+            f"**God Mode+ Capabilities Active:**\n"
+            f"- ⚡ Multi-agent orchestration\n"
+            f"- 🔧 Autonomous coding & debugging\n"
+            f"- 🧠 Persistent memory system\n"
+            f"- 🔌 Connector ecosystem\n"
+            f"- 📡 Real-time streaming\n"
+            f"- 🌐 Multi-model failover\n"
         )
-        full = ""
-        for word in text.split():
+        full_text = ""
+        for word in response.split():
             chunk = word + " "
-            full += chunk
-            await asyncio.sleep(0.015)
-            await self._emit(chunk, task_id, session_id, demo=True)
-        return full
+            full_text += chunk
+            await asyncio.sleep(0.02)
+            await self._emit_chunk(chunk, task_id, session_id, demo=True)
+        return full_text
 
-    # ─── Emit helper ──────────────────────────────────────────────────────────
+    # ─── Emit Helper ──────────────────────────────────────────────────────────
 
-    async def _emit(self, chunk: str, task_id: str, session_id: str, demo: bool = False):
+    async def _emit_chunk(self, chunk: str, task_id: str, session_id: str, demo: bool = False):
         if not self.ws:
             return
         payload = {"chunk": chunk, "demo": demo}
-        try:
-            if task_id:
-                await self.ws.emit(task_id, "llm_chunk", payload, session_id=session_id)
-            elif session_id:
-                await self.ws.emit_chat(session_id, "llm_chunk", payload)
-        except Exception:
-            pass
+        if task_id:
+            await self.ws.emit(task_id, "llm_chunk", payload, session_id=session_id)
+        if session_id and not task_id:
+            await self.ws.emit_chat(session_id, "llm_chunk", payload)
 
-    # ─── Stats ─────────────────────────────────────────────────────────────────
+    # ─── Stats ────────────────────────────────────────────────────────────────
 
     def get_stats(self) -> Dict:
-        result = {}
+        stats = {}
         for name, s in self._stats.items():
-            lats = s["latency_ms"][-20:]
-            avg = round(sum(lats) / max(len(lats), 1), 1)
-            provider = next((p for p in PROVIDER_CHAIN if p["name"] == name), None)
-            result[name] = {
-                "calls":        s["calls"],
-                "errors":       s["errors"],
-                "avg_latency_ms": avg,
-                "available":    bool(provider and provider["enabled"]()),
+            avg_lat = round(sum(s["latency"][-20:]) / max(len(s["latency"][-20:]), 1) * 1000, 1)
+            stats[name] = {
+                "calls": s["calls"],
+                "errors": s["errors"],
+                "avg_latency_ms": avg_lat,
+                "available": bool(os.environ.get(
+                    next((p["key_env"] for p in PROVIDERS if p["name"] == name), ""), ""
+                )),
             }
-        return result
-
-    def get_active_provider(self) -> str:
-        for p in PROVIDER_CHAIN:
-            if p["enabled"]():
-                return p["name"]
-        return "demo"
-
-
-# ─── Singleton alias for easy import ──────────────────────────────────────────
-# Usage: from ai_router.router import LLMRouter
-# In agents: result = await self.router.ask(messages, task_id=..., session_id=...)
-
-# Legacy AIRouter alias so existing imports don't break
-AIRouter = LLMRouter
+        return stats