Fix: update core/agent.py

core/agent.py

@@ -1,6 +1,6 @@
 """
-Agent Core — Uses LLMRouter for ALL LLM calls (no direct API calls)
-Planner + Executor + Self-Heal Loop
+Agent Core — Planner + Executor + Self-Heal Loop
+LLM-powered with OpenAI/Anthropic support, streaming tokens
 """
 
 import asyncio
@@ -9,18 +9,21 @@ import os
 import time
 from typing import Any, Dict, List, Optional
 
+import httpx
 import structlog
 
 from core.models import TaskPlan, TaskStep
 from api.websocket_manager import WebSocketManager
 from memory.db import save_memory, get_history, search_memory
-from ai_router.router import LLMRouter
 
 log = structlog.get_logger()
 
-DEFAULT_MODEL  = os.environ.get("DEFAULT_MODEL", "meta-llama/llama-4-scout-17b-16e-instruct")
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
+ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_API_KEY", "")
+DEFAULT_MODEL = os.environ.get("DEFAULT_MODEL", "gpt-4o")
 OPENAI_BASE_URL = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
 
+
 SYSTEM_PROMPT = """You are an elite autonomous AI software engineer — like Devin or Manus.
 You can plan, code, debug, refactor, test, and deploy software autonomously.
 You think step-by-step, write production-quality code, and self-heal on errors.
@@ -30,18 +33,18 @@ Always respond in structured JSON when asked for plans or structured output.
 PLANNER_PROMPT = """You are a senior software architect. Given a goal, produce a detailed execution plan.
 
 Respond ONLY with valid JSON:
-{{
+{
   "steps": [
-    {{
+    {
       "name": "Step name",
       "description": "What this step does",
       "tool": "code|shell|file|browser|github|memory|search|test|none",
       "estimated_seconds": 10
-    }}
+    }
   ],
   "estimated_duration": 60,
   "tools_needed": ["code", "shell"]
-}}
+}
 
 Goal: {goal}
 Context: {context}
@@ -51,10 +54,9 @@ Context: {context}
 class AgentCore:
     def __init__(self, ws_manager: WebSocketManager):
         self.ws = ws_manager
-        self.router = LLMRouter(ws_manager)
         self.model = DEFAULT_MODEL
 
-    # ─── LLM Call (delegates to router) ───────────────────────────────────────
+    # ─── LLM Call (with streaming) ─────────────────────────────────────────────
 
     async def llm_stream(
         self,
@@ -65,35 +67,180 @@ class AgentCore:
         temperature: float = 0.7,
         max_tokens: int = 4096,
     ) -> str:
-        """All LLM calls go through LLMRouter — no direct API calls."""
-        return await self.router.ask(
-            messages=messages,
-            task_id=task_id,
-            session_id=session_id,
-            temperature=temperature,
-            max_tokens=max_tokens,
-            model=model or self.model,
-            stream=True,
+        """Stream LLM tokens, emitting llm_chunk events via WebSocket."""
+        model = model or self.model
+        full_text = ""
+
+        if OPENAI_API_KEY:
+            full_text = await self._openai_stream(
+                messages, task_id, session_id, model, temperature, max_tokens
+            )
+        elif ANTHROPIC_API_KEY:
+            full_text = await self._anthropic_stream(
+                messages, task_id, session_id, temperature, max_tokens
+            )
+        else:
+            # Demo mode — simulate streaming
+            full_text = await self._demo_stream(messages, task_id, session_id)
+
+        return full_text
+
+    async def _openai_stream(
+        self, messages, task_id, session_id, model, temperature, max_tokens
+    ) -> str:
+        full_text = ""
+        headers = {
+            "Authorization": f"Bearer {OPENAI_API_KEY}",
+            "Content-Type": "application/json",
+        }
+        payload = {
+            "model": model,
+            "messages": messages,
+            "stream": True,
+            "temperature": temperature,
+            "max_tokens": max_tokens,
+        }
+        async with httpx.AsyncClient(timeout=120) as client:
+            async with client.stream(
+                "POST", f"{OPENAI_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
+                            if task_id:
+                                await self.ws.emit(task_id, "llm_chunk", {
+                                    "chunk": delta,
+                                    "accumulated": len(full_text),
+                                }, session_id=session_id)
+                            if session_id and not task_id:
+                                await self.ws.emit_chat(session_id, "llm_chunk", {
+                                    "chunk": delta,
+                                })
+                    except Exception:
+                        pass
+        return full_text
+
+    async def _anthropic_stream(
+        self, messages, task_id, session_id, temperature, max_tokens
+    ) -> str:
+        full_text = ""
+        system = ""
+        filtered = []
+        for m in messages:
+            if m["role"] == "system":
+                system = m["content"]
+            else:
+                filtered.append(m)
+        headers = {
+            "x-api-key": ANTHROPIC_API_KEY,
+            "anthropic-version": "2023-06-01",
+            "Content-Type": "application/json",
+        }
+        payload = {
+            "model": "claude-3-5-sonnet-20241022",
+            "max_tokens": max_tokens,
+            "messages": filtered,
+            "stream": True,
+        }
+        if system:
+            payload["system"] = system
+        async with httpx.AsyncClient(timeout=120) as client:
+            async with client.stream(
+                "POST", "https://api.anthropic.com/v1/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
+                                if task_id:
+                                    await self.ws.emit(task_id, "llm_chunk", {
+                                        "chunk": delta,
+                                    }, session_id=session_id)
+                                if session_id and not task_id:
+                                    await self.ws.emit_chat(session_id, "llm_chunk", {
+                                        "chunk": delta,
+                                    })
+                    except Exception:
+                        pass
+        return full_text
+
+    async def _demo_stream(self, messages, task_id, session_id) -> str:
+        """Demo mode — simulate LLM streaming without API key."""
+        last_user = next(
+            (m["content"] for m in reversed(messages) if m["role"] == "user"), "Hello"
+        )
+        response = (
+            f"🤖 **Devin Agent** (Demo Mode)\n\n"
+            f"I received your request: *{last_user[:100]}*\n\n"
+            f"To enable real AI responses, set `OPENAI_API_KEY` or `ANTHROPIC_API_KEY` in your environment.\n\n"
+            f"**What I can do with a real API key:**\n"
+            f"- 📋 Generate detailed execution plans\n"
+            f"- 💻 Write and execute code autonomously\n"
+            f"- 🔧 Debug and self-heal on errors\n"
+            f"- 🐙 Manage GitHub repos autonomously\n"
+            f"- 🧠 Remember long-running project context\n"
+            f"- 🚀 Deploy applications automatically\n"
         )
+        full_text = ""
+        for word in response.split():
+            chunk = word + " "
+            full_text += chunk
+            await asyncio.sleep(0.03)
+            if task_id:
+                await self.ws.emit(task_id, "llm_chunk", {
+                    "chunk": chunk,
+                    "demo": True,
+                }, session_id=session_id)
+            if session_id and not task_id:
+                await self.ws.emit_chat(session_id, "llm_chunk", {
+                    "chunk": chunk,
+                    "demo": True,
+                })
+        return full_text
 
     # ─── Planning ──────────────────────────────────────────────────────────────
 
     async def plan(self, goal: str, task_id: str, session_id: str = "") -> TaskPlan:
-        """Generate a structured execution plan via LLMRouter."""
+        """Generate a structured execution plan."""
+        # Get context from memory
         memories = await search_memory(goal[:50], session_id=session_id)
         context = "\n".join([m["content"][:200] for m in memories[:3]])
 
         prompt = PLANNER_PROMPT.format(goal=goal, context=context or "No prior context")
+
         messages = [
             {"role": "system", "content": SYSTEM_PROMPT},
-            {"role": "user",   "content": prompt},
+            {"role": "user", "content": prompt},
         ]
 
+        if not OPENAI_API_KEY and not ANTHROPIC_API_KEY:
+            # Demo plan
+            return self._demo_plan(goal)
+
         raw = await self.llm_stream(messages, task_id=task_id, session_id=session_id)
 
+        # Extract JSON from response
         try:
+            # Find JSON block
             start = raw.find("{")
-            end   = raw.rfind("}") + 1
+            end = raw.rfind("}") + 1
             if start >= 0 and end > start:
                 data = json.loads(raw[start:end])
             else:
@@ -118,14 +265,20 @@ class AgentCore:
             return self._demo_plan(goal)
 
     def _demo_plan(self, goal: str) -> TaskPlan:
+        """Fallback plan for demo mode."""
         steps = [
             TaskStep(name="Analyze Requirements", description=f"Analyze: {goal[:60]}", tool="none"),
-            TaskStep(name="Design Solution",      description="Design the solution architecture", tool="none"),
-            TaskStep(name="Implement",            description="Write the implementation code", tool="code"),
-            TaskStep(name="Test",                 description="Test the implementation", tool="test"),
-            TaskStep(name="Document",             description="Write documentation", tool="none"),
+            TaskStep(name="Design Solution", description="Design the solution architecture", tool="none"),
+            TaskStep(name="Implement", description="Write the implementation code", tool="code"),
+            TaskStep(name="Test", description="Test the implementation", tool="test"),
+            TaskStep(name="Document", description="Write documentation", tool="none"),
         ]
-        return TaskPlan(goal=goal, steps=steps, estimated_duration=120, tools_needed=["code", "test"])
+        return TaskPlan(
+            goal=goal,
+            steps=steps,
+            estimated_duration=120,
+            tools_needed=["code", "test"],
+        )
 
     # ─── Step Execution ────────────────────────────────────────────────────────
 
@@ -136,6 +289,7 @@ class AgentCore:
         session_id: str = "",
         context: Dict = {},
     ) -> str:
+        """Execute a single step using the appropriate tool."""
         from tools.executor import ToolExecutor
         executor = ToolExecutor(self.ws)
 
@@ -155,14 +309,18 @@ class AgentCore:
                 session_id=session_id,
             )
             await self.ws.emit(task_id, "tool_result", {
-                "tool": step.tool, "step": step.name,
-                "result": str(result)[:500], "success": True,
+                "tool": step.tool,
+                "step": step.name,
+                "result": str(result)[:500],
+                "success": True,
             }, session_id=session_id)
             return result
         except Exception as e:
             await self.ws.emit(task_id, "tool_result", {
-                "tool": step.tool, "step": step.name,
-                "error": str(e), "success": False,
+                "tool": step.tool,
+                "step": step.name,
+                "error": str(e),
+                "success": False,
             }, session_id=session_id)
             return f"Error in {step.name}: {str(e)}"
 
@@ -176,11 +334,15 @@ class AgentCore:
         task_id: str,
         session_id: str = "",
     ) -> str:
-        steps_summary = "\n".join([f"- {s.name}: {r[:200]}" for s, r in zip(steps, results)])
+        """Compile final result summary."""
+        steps_summary = "\n".join([
+            f"- {s.name}: {r[:200]}" for s, r in zip(steps, results)
+        ])
         messages = [
             {"role": "system", "content": SYSTEM_PROMPT},
             {"role": "user", "content": (
-                f"Summarize the completion of this goal:\nGoal: {goal}\n\n"
+                f"Summarize the completion of this goal:\n"
+                f"Goal: {goal}\n\n"
                 f"Steps completed:\n{steps_summary}\n\n"
                 f"Write a concise success summary with key outcomes."
             )},
@@ -191,21 +353,40 @@ class AgentCore:
     # ─── Chat ──────────────────────────────────────────────────────────────────
 
     async def stream_chat(self, session_id: str, user_message: str):
-        await save_memory(content=user_message, memory_type="conversation",
-                          session_id=session_id, key="user_message")
+        """Stream a conversational chat response."""
+        # Save user message to memory
+        await save_memory(
+            content=user_message,
+            memory_type="conversation",
+            session_id=session_id,
+            key="user_message",
+        )
 
-        history  = await get_history(session_id, limit=10)
+        # Get conversation history
+        history = await get_history(session_id, limit=10)
         messages = [{"role": "system", "content": SYSTEM_PROMPT}]
         for h in reversed(history[-10:]):
             messages.append({"role": "user", "content": h["content"]})
+
         messages.append({"role": "user", "content": user_message})
 
-        await self.ws.emit_chat(session_id, "stream_start", {"status": "generating"})
+        await self.ws.emit_chat(session_id, "stream_start", {
+            "status": "generating",
+        })
+
         response = await self.llm_stream(messages, session_id=session_id)
 
-        await save_memory(content=response, memory_type="conversation",
-                          session_id=session_id, key="assistant_response")
+        # Save assistant response to memory
+        await save_memory(
+            content=response,
+            memory_type="conversation",
+            session_id=session_id,
+            key="assistant_response",
+        )
+
         await self.ws.emit_chat(session_id, "stream_end", {
-            "full_response": response, "status": "complete"
+            "full_response": response,
+            "status": "complete",
         })
+
         return response