Update backend.py

backend.py

@@ -1,169 +1,116 @@
-# backend.py — FINAL VERSION — PHASE-3 MINI INSTRUCT + CACHING + LOGGING
+# backend.py — REFACTORED FOR ROBUSTNESS & RELIABILITY
 import sqlite3
 import os
 import json
-import uuid
-import zipfile
-import tempfile
-import subprocess
 import re
 import concurrent.futures
-from typing import List, Dict, Any
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import traceback
+import zipfile
+import hashlib
+from typing import Optional, Dict, Any
+
 import torch
 import psutil
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
+# --- CONFIGURATION ---
 DB_PATH = "code_agents_pro.db"
 PROJECT_ROOT = "./projects"
 os.makedirs(PROJECT_ROOT, exist_ok=True)
 
-# ------------------------------ DATABASE ------------------------------
+
+# ------------------------------ DATABASE (REFACTORED WITH CONTEXT MANAGERS) ------------------------------
 def init_db():
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
-    cursor.executescript("""
-        CREATE TABLE IF NOT EXISTS users (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            username TEXT UNIQUE NOT NULL,
-            password_hash TEXT NOT NULL,
-            created_at DATETIME DEFAULT CURRENT_TIMESTAMP
-        );
-        CREATE TABLE IF NOT EXISTS projects (
-            id INTEGER PRIMARY KEY AUTOINCREMENT,
-            user_id INTEGER NOT NULL,
-            title TEXT,
-            description TEXT,
-            status TEXT DEFAULT 'queued',
-            zip_path TEXT,
-            logs TEXT DEFAULT '',
-            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
-            FOREIGN KEY (user_id) REFERENCES users(id)
-        );
-        CREATE INDEX idx_user_status ON projects(user_id, status);
-    """)
-    conn.commit()
-    conn.close()
+    """Initializes the database schema safely."""
+    with sqlite3.connect(DB_PATH) as conn:
+        cursor = conn.cursor()
+        cursor.executescript("""
+            CREATE TABLE IF NOT EXISTS users (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                username TEXT UNIQUE NOT NULL,
+                password_hash TEXT NOT NULL
+            );
+            CREATE TABLE IF NOT EXISTS projects (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                user_id INTEGER NOT NULL,
+                title TEXT,
+                description TEXT,
+                status TEXT DEFAULT 'queued',
+                zip_path TEXT,
+                logs TEXT DEFAULT '',
+                created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+                FOREIGN KEY (user_id) REFERENCES users(id)
+            );
+            CREATE INDEX IF NOT EXISTS idx_user_status ON projects(user_id, status);
+        """)
+        conn.commit()
 
 init_db()
 
+# --- USER & PROJECT CRUD (ALL REFACTORED FOR SAFETY) ---
+
 def hash_password(password):
-    import hashlib
     return hashlib.sha256(password.encode()).hexdigest()
 
 def verify_password(password, stored_hash):
     return hash_password(password) == stored_hash
 
+def _db_execute(query, params=(), fetchone=False, fetchall=False, commit=False):
+    """A centralized, safe way to interact with the database."""
+    try:
+        with sqlite3.connect(DB_PATH) as conn:
+            conn.row_factory = sqlite3.Row
+            cursor = conn.cursor()
+            cursor.execute(query, params)
+            if commit:
+                conn.commit()
+                return cursor.lastrowid
+            if fetchone:
+                return cursor.fetchone()
+            if fetchall:
+                return cursor.fetchall()
+    except sqlite3.Error as e:
+        print(f"Database error: {e}")
+        return None
+
 def create_user(username, password):
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
+    query = "INSERT INTO users (username, password_hash) VALUES (?, ?)"
+    # Using a try-except block here specifically for the IntegrityError
     try:
-        cursor.execute("INSERT INTO users (username, password_hash) VALUES (?, ?)", 
-                       (username, hash_password(password)))
-        user_id = cursor.lastrowid
-        conn.commit()
-        conn.close()
-        return user_id
+        return _db_execute(query, (username, hash_password(password)), commit=True)
     except sqlite3.IntegrityError:
-        conn.close()
         return None
 
 def get_user_by_username(username):
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
-    cursor.execute("SELECT id, password_hash FROM users WHERE username = ?", (username,))
-    row = cursor.fetchone()
-    conn.close()
-    return row if row else None
+    return _db_execute("SELECT id, password_hash FROM users WHERE username = ?", (username,), fetchone=True)
 
 def get_user_projects(user_id, limit=3):
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
-    cursor.execute("""
-        SELECT id, title, description, status, zip_path, created_at, logs
-        FROM projects WHERE user_id = ? ORDER BY created_at DESC LIMIT ?
-    """, (user_id, limit))
-    projects = cursor.fetchall()
-    conn.close()
-    return projects
+    query = "SELECT * FROM projects WHERE user_id = ? ORDER BY created_at DESC LIMIT ?"
+    return _db_execute(query, (user_id, limit), fetchall=True)
 
 def create_project(user_id, title, description):
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
-    cursor.execute("""
-        INSERT INTO projects (user_id, title, description, status, logs)
-        VALUES (?, ?, ?, 'queued', '')
-    """, (user_id, title, description))
-    project_id = cursor.lastrowid
-    conn.commit()
-    conn.close()
-    return project_id
+    query = "INSERT INTO projects (user_id, title, description) VALUES (?, ?, ?)"
+    return _db_execute(query, (user_id, title, description), commit=True)
 
 def update_project_status(project_id, status, logs=None, zip_path=None):
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
-    if zip_path:
-        cursor.execute("""
-            UPDATE projects SET status = ?, zip_path = ?, logs = ? WHERE id = ?
-        """, (status, zip_path, logs, project_id))
-    else:
-        cursor.execute("""
-            UPDATE projects SET status = ?, logs = ? WHERE id = ?
-        """, (status, logs, project_id))
-    conn.commit()
-    conn.close()
+    query = "UPDATE projects SET status = ?, logs = COALESCE(?, logs), zip_path = COALESCE(?, zip_path) WHERE id = ?"
+    _db_execute(query, (status, logs, zip_path, project_id), commit=True)
 
 def get_project(project_id):
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
-    cursor.execute("""
-        SELECT id, user_id, title, description, status, zip_path, logs
-        FROM projects WHERE id = ?
-    """, (project_id,))
-    row = cursor.fetchone()
-    conn.close()
-    return row
+    return _db_execute("SELECT * FROM projects WHERE id = ?", (project_id,), fetchone=True)
 
 def search_projects(user_id, query):
-    conn = sqlite3.connect(DB_PATH)
-    cursor = conn.cursor()
-    cursor.execute("""
-        SELECT id, title, description, status, zip_path
-        FROM projects WHERE user_id = ? AND (title LIKE ? OR description LIKE ?)
-        ORDER BY created_at DESC
-    """, (user_id, f"%{query}%", f"%{query}%"))
-    results = cursor.fetchall()
-    conn.close()
-    return results
-
-# ------------------------------ SYSTEM MONITORING ------------------------------
-def get_ram_usage_gb():
-    return round(psutil.virtual_memory().used / (1024**3), 2)
-
-def get_vram_usage_gb():
-    """Try to get GPU VRAM usage. Returns 0 if no GPU or error."""
-    try:
-        if torch.cuda.is_available():
-            torch.cuda.synchronize()
-            allocated = torch.cuda.memory_allocated() / (1024**3)
-            reserved = torch.cuda.memory_reserved() / (1024**3)
-            return round(max(allocated, reserved), 2)
-        else:
-            return 0.0
-    except:
-        return 0.0
-
-# ------------------------------ MODEL LOADING (PER ROLE) ------------------------------
+    sql_query = "SELECT id, title, description, status FROM projects WHERE user_id = ? AND (title LIKE ? OR description LIKE ?)"
+    return _db_execute(sql_query, (user_id, f"%{query}%", f"%{query}%"), fetchall=True)
+
+# ------------------------------ MODEL LOADING & CACHING ------------------------------
 MODEL_REGISTRY = {
-    "understander": "Qwen/Qwen2.5-Coder-0.6B-Instruct",
+    "planner": "microsoft/Phi-3-mini-4k-instruct",
     "architect": "Qwen/Qwen2.5-Coder-0.6B-Instruct",
-    "ceo": "microsoft/Phi-3-mini-4k-instruct",
-    "manager": "microsoft/Phi-3-mini-4k-instruct",
-    "worker": "Qwen/Qwen2.5-Coder-0.6B-Instruct",
+    "coder": "Qwen/Qwen2.5-Coder-0.6B-Instruct",
     "reviewer": "microsoft/Phi-3-mini-4k-instruct",
-    "editor": "Qwen/Qwen2.5-Coder-0.6B-Instruct",
-    "tester": "microsoft/Phi-3-mini-4k-instruct",
-    "publisher": "Qwen/Qwen2.5-Coder-0.6B-Instruct",
-    "final_ceo": "microsoft/Phi-3-mini-4k-instruct",
+    "tester": "Qwen/Qwen2.5-Coder-0.6B-Instruct",
+    "publisher": "microsoft/Phi-3-mini-4k-instruct",
 }
 
 _MODEL_CACHE = {}
@@ -171,434 +118,210 @@ _MODEL_CACHE = {}
 def load_model(model_name):
     if model_name in _MODEL_CACHE:
         return _MODEL_CACHE[model_name]
-    
-    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+
+    print(f"Loading model: {model_name}...")
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        torch_dtype="auto",
         device_map="auto",
         trust_remote_code=True,
-        attn_implementation="eager"  # ✅ Avoids flash-attn requirement — safe on all hardware
+        attn_implementation="eager"
     )
     _MODEL_CACHE[model_name] = (tokenizer, model)
+    print(f"Model {model_name} loaded and cached.")
     return tokenizer, model
 
-# ------------------------------ PROMPTS PER ROLE ------------------------------
+# ------------------------------ AGENT PROMPTS ------------------------------
 ROLE_PROMPTS = {
-    "understander": """
-You are the Understander. Your job is to deeply analyze the user's project request.
-Extract:
-- Core functionality
-- Required libraries
-- File structure
-- Edge cases
-
-Output only JSON:
-{
-  "purpose": "string",
-  "features": ["list of features"],
-  "dependencies": ["list of required pip packages"],
-  "files": [
-    {"path": "filename.py", "type": "source|test|config|other"},
-    ...
-  ]
-}
-""",
-
-    "architect": """
-You are the Architect. Based on the Understander's analysis, design the full project structure.
-Define:
-- Folder hierarchy
-- File contents (minimal skeleton)
-- Configuration files
-
-Output only JSON:
-{
-  "structure": [
-    {"path": "src/main.py", "content": "// empty file"},
-    {"path": "requirements.txt", "content": "flask\\nrequests"}
-  ],
-  "notes": "Brief notes on architecture decisions"
-}
-""",
-
-    "ceo": """
-You are the CEO. You oversee the entire project. You receive input from the Architect.
-Your job is to review the plan, ensure alignment with goals, and assign tasks to Managers.
-Respond only with: "Approved. Assigning to Managers."
-""",
-
-    "manager": """
-You are a Manager. You receive a task from the CEO.
-Assign work to Workers based on file type.
-Example:
-- "Write src/main.py" → assign to Worker A
-- "Write tests/test_main.py" → assign to Worker B
-
-Return only JSON:
-{
-  "tasks": [
-    {"file": "src/main.py", "assigned_to": "worker", "instructions": "Write a Flask route..."},
-    {"file": "tests/test_main.py", "assigned_to": "worker", "instructions": "Write a test for the /upload endpoint..."}
-  ]
-}
-""",
-
-    "worker": """
-You are a Worker. You write clean, functional code based on instructions.
-Never guess. Use tools if needed.
-Only output the raw file content — no markdown, no explanations.
-If writing Python, use proper imports, error handling, comments.
-
-File: {file}
-Instructions: {instructions}
-""",
-    
-    "reviewer": """
-You are the Reviewer. You audit all code written by Workers.
-Check for:
-- Security flaws (eval(), shell injection)
-- Missing error handling
-- Unused imports
-- Poor variable names
-- Logic errors
-
-For each issue, return:
-{
-  "file": "path/to/file.py",
-  "issue": "description",
-  "suggestion": "fix suggestion"
+    "planner": """You are an expert software planner. Your task is to break down a user's request into a detailed, actionable plan.
+    - Clarify the core purpose.
+    - List the key features.
+    - Determine the necessary libraries and dependencies.
+    - Define a complete file structure.
+    Output ONLY a single JSON object with the keys: "purpose", "features", "dependencies", "files".
+    The "files" key should be an array of strings representing file paths (e.g., ["src/main.py", "tests/test_main.py"]).""",
+
+    "architect": """You are a software architect. Based on the file plan, you will create the initial content for each file.
+    - For configuration files (like requirements.txt), list the dependencies.
+    - For source code files, write a brief comment or placeholder (e.g., '# TODO: Implement main logic').
+    - For documentation (README.md), write a title.
+    Output ONLY a single JSON object where keys are file paths and values are the initial content.""",
+
+    "coder": """You are a professional programmer. Your only job is to write clean, functional, and complete code for a given file based on instructions.
+    - You must write the full code for the file path provided.
+    - Do NOT add any explanations, introductions, or markdown formatting.
+    - Your output must be ONLY the raw source code for the file.""",
+
+    "reviewer": """You are a meticulous code reviewer. Your task is to analyze a file's code and identify potential bugs, style issues, or missing features.
+    - Check for logical errors, security vulnerabilities, and bad practices.
+    - Ensure the code is readable and well-commented.
+    - Output ONLY a single JSON object with two keys: "has_issues" (boolean) and "suggestions" (a string containing a bulleted list of required changes). If no issues, "has_issues" is false.""",
+
+    "tester": """You are a quality assurance engineer specializing in automated testing. Your task is to write a complete pytest test file for a given source code file.
+    - Cover the main functionality and edge cases.
+    - Use appropriate assertions.
+    - Your output must be ONLY the raw source code for the test file.""",
+
+    "publisher": """You are a release manager. Your job is to create the final project documentation and configuration based on the complete file structure.
+    - Create a helpful README.md that explains the project's purpose and how to run it.
+    - Create a standard .gitignore file.
+    - Create a basic Dockerfile for containerization.
+    Output ONLY a single JSON object where keys are the filenames ("README.md", ".gitignore", "Dockerfile") and values are their complete content."""
 }
 
-If no issues: return {"issues": []}
-""",
-
-    "editor": """
-You are the Editor. Apply all Reviewer suggestions to the code.
-Update files directly.
-Do not explain — just rewrite the file content.
-If a file doesn't exist, create it.
-Return only the updated file content.
-""",
-
-    "tester": """
-You are a Tester. Write unit tests for each source file.
-Use pytest. Structure:
-- File: tests/test_{filename}.py
-- Use assert statements
-- Cover edge cases
-
-Only output the test file content — nothing else.
-""",
-    
-    "publisher": """
-You are the Publisher. Finalize the project.
-Generate:
-- README.md (plain text, no markdown formatting)
-- Dockerfile
-- requirements.txt (complete, deduplicated)
-- .gitignore
-
-Output each file as:
-{
-  "README.md": "...",
-  "Dockerfile": "...",
-  "requirements.txt": "...",
-  ".gitignore": "venv/\n__pycache__/"
-}
-""",
-    
-    "final_ceo": """
-You are the final CEO. You receive the packaged project from the Publisher.
-Verify everything is complete and correct.
-If yes: respond with "Ready to package."
-If no: list missing items.
-"""
-}
+# ------------------------------ FILE SYSTEM TOOLS ------------------------------
+def get_project_dir(user_id, project_id):
+    path = os.path.join(PROJECT_ROOT, str(user_id), str(project_id))
+    os.makedirs(path, exist_ok=True)
+    return path
 
-# ------------------------------ TOOLS (UNCHANGED) ------------------------------
-def get_user_project_dir(user_id):
-    user_dir = os.path.join(PROJECT_ROOT, str(user_id))
-    os.makedirs(user_dir, exist_ok=True)
-    return user_dir
-
-def clear_user_project_dir(user_id):
-    user_dir = get_user_project_dir(user_id)
-    for f in os.listdir(user_dir):
-        path = os.path.join(user_dir, f)
-        if os.path.isdir(path):
-            import shutil
-            shutil.rmtree(path)
-        else:
-            os.remove(path)
-
-def create_file(user_id, path: str, content: str):
-    user_dir = get_user_project_dir(user_id)
-    full_path = os.path.join(user_dir, path.lstrip("/"))
+def create_file(project_dir, path, content):
+    full_path = os.path.join(project_dir, path)
     os.makedirs(os.path.dirname(full_path), exist_ok=True)
-    with open(full_path, 'w') as f:
+    with open(full_path, 'w', encoding='utf-8') as f:
         f.write(content)
-    return f"Created: {path}"
-
-def read_file(user_id, path: str):
-    user_dir = get_user_project_dir(user_id)
-    full_path = os.path.join(user_dir, path.lstrip("/"))
-    if not os.path.exists(full_path):
-        return f"File not found: {path}"
-    with open(full_path, 'r') as f:
-        return f.read()
 
-def list_files(user_id):
-    user_dir = get_user_project_dir(user_id)
-    files = []
-    for root, _, filenames in os.walk(user_dir):
-        for f in filenames:
-            files.append(os.path.relpath(os.path.join(root, f), user_dir))
-    return files
+def read_file(project_dir, path):
+    full_path = os.path.join(project_dir, path)
+    if not os.path.exists(full_path): return None
+    with open(full_path, 'r', encoding='utf-8') as f:
+        return f.read()
 
-def execute_code(user_id, code: str):
-    try:
-        user_dir = get_user_project_dir(user_id)
-        with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False, dir=user_dir) as f:
-            f.write(code)
-            temp_path = f.name
-        result = subprocess.run(["python3", temp_path], capture_output=True, text=True, timeout=10)
-        os.unlink(temp_path)
-        if result.returncode == 0:
-            return result.stdout
-        else:
-            return f"Error: {result.stderr}"
-    except Exception as e:
-        return f"Timeout/Error: {str(e)}"
-
-def run_tests(user_id):
-    user_dir = get_user_project_dir(user_id)
-    test_dir = os.path.join(user_dir, "tests")
-    if not os.path.exists(test_dir):
-        return "No test directory found."
-    result = subprocess.run(["pytest", test_dir], capture_output=True, text=True, cwd=user_dir)
-    return result.stdout
-
-def lint_code(user_id):
-    user_dir = get_user_project_dir(user_id)
-    result = subprocess.run(["flake8", user_dir], capture_output=True, text=True)
-    return result.stdout if result.returncode != 0 else "No linting errors."
-
-def scan_vulns(user_id):
-    reqs = os.path.join(get_user_project_dir(user_id), "requirements.txt")
-    if not os.path.exists(reqs):
-        return "No requirements.txt found."
-    result = subprocess.run(["pip-audit", "-r", reqs], capture_output=True, text=True)
-    return result.stdout if result.returncode != 0 else "No known vulnerabilities found."
-
-def generate_dockerfile(user_id, content: str):
-    create_file(user_id, "Dockerfile", content)
-    return "Dockerfile generated."
-
-def write_readme(user_id, content: str):
-    create_file(user_id, "README.md", content)
-    return "README.md generated."
-
-def write_requirements(user_id, content: str):
-    create_file(user_id, "requirements.txt", content)
-    return "requirements.txt generated."
-
-def write_gitignore(user_id, content: str):
-    create_file(user_id, ".gitignore", content)
-    return ".gitignore generated."
-
-def zip_project(user_id, project_name: str):
-    user_dir = get_user_project_dir(user_id)
-    zip_path = os.path.join(user_dir, f"{project_name}.zip")
+def zip_project(project_dir, project_id):
+    zip_filename = f"project_{project_id}.zip"
+    zip_path = os.path.join(os.path.dirname(project_dir), zip_filename) # Place zip in user folder, not project folder
     with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
-        for root, _, files in os.walk(user_dir):
-            for f in files:
-                if f.endswith(".zip"): continue
-                full_path = os.path.join(root, f)
-                arcname = os.path.relpath(full_path, user_dir)
+        for root, _, files in os.walk(project_dir):
+            for file in files:
+                full_path = os.path.join(root, file)
+                arcname = os.path.relpath(full_path, project_dir)
                 zf.write(full_path, arcname)
     return zip_path
 
-# ------------------------------ LLM GENERATOR (WITH CHAT TEMPLATE FOR PHI-3) ------------------------------
-def generate_with_model(role: str, prompt: str, context: dict = {}) -> str:
+# ------------------------------ CORE AI & AGENT LOGIC ------------------------------
+def _extract_json(text: str) -> Optional[Dict[str, Any]]:
+    """Robustly extracts a JSON object from a string, even if wrapped in markdown."""
+    match = re.search(r"```json\s*([\s\S]*?)\s*```|(\{[\s\S]*\})", text)
+    if not match:
+        return None
+    json_str = match.group(1) or match.group(2)
+    try:
+        return json.loads(json_str)
+    except json.JSONDecodeError:
+        print(f"Failed to decode JSON: {json_str[:200]}...")
+        return None
+
+def generate_with_model(role: str, prompt: str) -> str:
+    """Generates a response from a cached model for a given role and prompt."""
     try:
         model_name = MODEL_REGISTRY[role]
         tokenizer, model = load_model(model_name)
+        
+        messages = [{"role": "system", "content": ROLE_PROMPTS[role]}, {"role": "user", "content": prompt}]
+        input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
 
-        # Format message using Phi-3 or Qwen chat template
-        if "Phi-3" in model_name:
-            # Phi-3 format
-            messages = [
-                {"role": "system", "content": ROLE_PROMPTS[role]},
-                {"role": "user", "content": prompt}
-            ]
-            text = tokenizer.apply_chat_template(
-                messages,
-                tokenize=False,
-                add_generation_prompt=True
-            )
-        else:
-            # Qwen format
-            messages = [
-                {"role": "system", "content": ROLE_PROMPTS[role]},
-                {"role": "user", "content": prompt}
-            ]
-            text = tokenizer.apply_chat_template(
-                messages,
-                tokenize=False,
-                add_generation_prompt=True
-            )
-
-        inputs = tokenizer(text, return_tensors="pt").to(model.device)
-
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=512,
-            temperature=0.2,
-            top_p=0.9,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id
-        )
-
-        response = tokenizer.decode(outputs[0][len(inputs["input_ids"][0]):], skip_special_tokens=True)
+        outputs = model.generate(**inputs, max_new_tokens=2048, pad_token_id=tokenizer.eos_token_id)
+        response = tokenizer.decode(outputs[0][len(inputs.input_ids[0]):], skip_special_tokens=True)
         return response.strip()
 
     except Exception as e:
-        return f"[ERROR IN {role.upper()}]: {str(e)}"
+        print(f"Error during model generation for role {role}: {e}")
+        return f'{{"error": "Failed to generate response: {str(e)}"}}'
 
-# ------------------------------ ASYNC JOB EXECUTOR (WITH FULL LOGGING) ------------------------------
+# ------------------------------ THE AGENT CHAIN EXECUTOR ------------------------------
 def run_agent_chain(project_id, user_id, initial_prompt):
+    project_dir = get_project_dir(user_id, project_id)
+    log_entries = []
+
+    def log_step(agent, action, output=""):
+        log_entry = f"**[{agent.upper()}]**: {action}\n"
+        if output:
+            log_entry += f"```\n{output[:1000]}{'...' if len(output) > 1000 else ''}\n```\n---\n"
+        log_entries.append(log_entry)
+        update_project_status(project_id, "running", logs="".join(log_entries))
+
     try:
-        update_project_status(project_id, "running", f"Starting agent chain...\nRAM: {get_ram_usage_gb()}GB | VRAM: {get_vram_usage_gb()}GB")
-
-        user_dir = get_user_project_dir(user_id)
-        clear_user_project_dir(user_id)
-
-        log_entries = []
-
-        def log_step(agent, action, output):
-            entry = f"[{agent}] {action}\n→ {output[:500]}{'...' if len(output) > 500 else ''}\n\n"
-            log_entries.append(entry)
-            current_logs = ''.join(log_entries)
-            update_project_status(project_id, "running", current_logs)
-
-        # Step 1: Understander
-        log_step("Understander", "Analyzing prompt...", "Started...")
-        understand_result = generate_with_model("understander", initial_prompt)
-        log_step("Understander", "Analysis complete", understand_result)
-
-        # Step 2: Architect
-        log_step("Architect", "Designing structure...", "Started...")
-        architect_result = generate_with_model("architect", understand_result)
-        log_step("Architect", "Structure designed", architect_result)
-
-        # Parse architect JSON
-        try:
-            import json
-            arch_data = json.loads(architect_result)
-            for file in arch_data.get("structure", []):
-                create_file(user_id, file["path"], file["content"])
-        except Exception as e:
-            log_step("Architect", "JSON parse failed", str(e))
-
-        # Step 3: CEO
-        log_step("CEO", "Reviewing architecture...", "Started...")
-        ceo_response = generate_with_model("ceo", "Review architecture and approve.")
-        log_step("CEO", "Response", ceo_response)
-
-        # Step 4: Manager assigns tasks
-        log_step("Manager", "Assigning tasks...", "Started...")
-        manager_response = generate_with_model("manager", "Assign tasks based on architecture.")
-        log_step("Manager", "Tasks assigned", manager_response)
-
-        try:
-            import json
-            tasks = json.loads(manager_response).get("tasks", [])
-            for task in tasks:
-                worker_response = generate_with_model("worker", "", context={
-                    "file": task["file"],
-                    "instructions": task["instructions"]
-                })
-                create_file(user_id, task["file"], worker_response)
-                log_step("Worker", f"Generated {task['file']}", worker_response[:300])
-        except Exception as e:
-            log_step("Manager", "Task parsing failed", str(e))
-
-        # Step 5: Reviewer
-        all_files = "\n".join(list_files(user_id))
-        reviewer_input = f"Review these files:\n\n{all_files}"
-        log_step("Reviewer", "Auditing code...", "Started...")
-        reviewer_result = generate_with_model("reviewer", reviewer_input)
-        log_step("Reviewer", "Audit complete", reviewer_result)
-
-        # Step 6: Editor applies fixes
-        editor_input = f"Apply these changes:\n{reviewer_result}\n\nFiles:\n{all_files}"
-        log_step("Editor", "Applying fixes...", "Started...")
-        editor_response = generate_with_model("editor", editor_input)
-        log_step("Editor", "Fixes applied", "Applied corrections to files.")
-
-        # Step 7: Testers generate tests
-        for filename in [f for f in list_files(user_id) if f.endswith(".py") and not f.startswith("test_")]:
-            test_prompt = f"Write a test for this file: {filename}"
-            log_step("Tester", f"Generating test for {filename}...", "Started...")
-            test_content = generate_with_model("tester", test_prompt)
-            test_path = f"tests/test_{os.path.basename(filename)}"
-            create_file(user_id, test_path, test_content)
-            log_step("Tester", f"Test saved to {test_path}", test_content[:200])
-
-        # Step 8: Editor reviews tests
-        test_files = "\n".join([f for f in list_files(user_id) if f.startswith("tests/")])
-        editor_test_input = f"Review these test files:\n\n{test_files}\n\nMake sure they are complete and correct."
-        log_step("Editor", "Reviewing tests...", "Started...")
-        editor_test_response = generate_with_model("editor", editor_test_input)
-        log_step("Editor", "Test review complete", "Reviewed and corrected test files.")
-
-        # Step 9: Publisher creates final assets
-        publisher_input = "Generate final project assets: README.md, Dockerfile, requirements.txt, .gitignore"
-        log_step("Publisher", "Generating final files...", "Started...")
-        publisher_response = generate_with_model("publisher", publisher_input)
-        log_step("Publisher", "Final assets generated", publisher_response)
-
-        try:
-            import json
-            pub_data = json.loads(publisher_response)
-            for fname, content in pub_data.items():
-                if fname == "README.md":
-                    write_readme(user_id, content)
-                elif fname == "Dockerfile":
-                    generate_dockerfile(user_id, content)
-                elif fname == "requirements.txt":
-                    write_requirements(user_id, content)
-                elif fname == ".gitignore":
-                    write_gitignore(user_id, content)
-        except Exception as e:
-            log_step("Publisher", "JSON parsing failed", str(e))
-
-        # Step 10: Final CEO
-        final_input = "All files generated. Verify completeness."
-        log_step("Final CEO", "Final verification...", "Started...")
-        final_ceo_response = generate_with_model("final_ceo", final_input)
-        log_step("Final CEO", "Verification result", final_ceo_response)
-
-        if "Ready to package" not in final_ceo_response:
-            raise Exception(f"Final CEO rejected: {final_ceo_response}")
-
-        # Step 11: ZIP & Save
-        log_step("System", "Creating ZIP archive...", "Started...")
-        zip_path = zip_project(user_id, f"project_{project_id}")
-        log_step("System", "ZIP created", f"Saved to {zip_path}")
-
-        update_project_status(project_id, "completed", ''.join(log_entries), zip_path)
-        return "Completed successfully"
+        log_step("SYSTEM", f"Initializing project directory: {project_dir}")
+        update_project_status(project_id, "running", logs="Agent team is assembling...")
+
+        # 1. PLANNER
+        log_step("PLANNER", "Analyzing user request and creating a project plan...")
+        plan_response = generate_with_model("planner", initial_prompt)
+        plan_data = _extract_json(plan_response)
+        if not plan_data: raise ValueError("Planner failed to create a valid JSON plan.")
+        log_step("PLANNER", "Plan created successfully.", json.dumps(plan_data, indent=2))
+        
+        # 2. ARCHITECT
+        log_step("ARCHITECT", "Creating initial file skeletons based on the plan...")
+        arch_prompt = f"Create initial content for the following files:\n{json.dumps(plan_data['files'])}"
+        arch_response = generate_with_model("architect", arch_prompt)
+        arch_data = _extract_json(arch_response)
+        if not arch_data: raise ValueError("Architect failed to create valid JSON file structures.")
+        for path, content in arch_data.items():
+            create_file(project_dir, path, content)
+        log_step("ARCHITECT", "File skeletons created.", "\n".join(arch_data.keys()))
+
+        # 3. CODER
+        source_files = [f for f in plan_data['files'] if f.startswith('src/') and f.endswith('.py')]
+        for file_path in source_files:
+            log_step("CODER", f"Writing complete code for `{file_path}`...")
+            coder_prompt = f"Based on the project plan ({plan_data['purpose']}), write the full Python code for the file: `{file_path}`."
+            code = generate_with_model("coder", coder_prompt)
+            create_file(project_dir, file_path, code)
+            log_step("CODER", f"Finished writing code for `{file_path}`.", code)
+
+        # 4. REVIEWER
+        log_step("REVIEWER", "Reviewing all generated source code for issues...")
+        for file_path in source_files:
+            code_content = read_file(project_dir, file_path)
+            review_prompt = f"Review the following code from `{file_path}`:\n\n{code_content}"
+            review_response = generate_with_model("reviewer", review_prompt)
+            review_data = _extract_json(review_response)
+            if review_data and review_data.get("has_issues"):
+                log_step("REVIEWER", f"Found issues in `{file_path}`.", review_data.get("suggestions"))
+                # In a more advanced version, this would trigger the CODER again (a "fix loop")
+            else:
+                log_step("REVIEWER", f"No major issues found in `{file_path}`.")
+
+        # 5. TESTER
+        log_step("TESTER", "Writing unit tests for all source code...")
+        for file_path in source_files:
+            code_content = read_file(project_dir, file_path)
+            test_file_path = os.path.join("tests", f"test_{os.path.basename(file_path)}")
+            tester_prompt = f"Write a complete pytest test file named `{test_file_path}` for the following code from `{file_path}`:\n\n{code_content}"
+            test_code = generate_with_model("tester", tester_prompt)
+            create_file(project_dir, test_file_path, test_code)
+            log_step("TESTER", f"Generated test file `{test_file_path}`.", test_code)
+
+        # 6. PUBLISHER
+        log_step("PUBLISHER", "Generating final documentation and configuration files...")
+        pub_prompt = f"The project has the following files: {json.dumps(list(os.walk(project_dir)))}. Generate README.md, .gitignore, and Dockerfile."
+        pub_response = generate_with_model("publisher", pub_prompt)
+        pub_data = _extract_json(pub_response)
+        if not pub_data: raise ValueError("Publisher failed to create valid final assets.")
+        for path, content in pub_data.items():
+            create_file(project_dir, path, content)
+        log_step("PUBLISHER", "Final assets created.", json.dumps(pub_data, indent=2))
+        
+        # 7. FINALIZATION
+        log_step("SYSTEM", "Packaging the complete project into a ZIP file...")
+        zip_path = zip_project(project_dir, project_id)
+        update_project_status(project_id, "completed", logs="".join(log_entries), zip_path=zip_path)
+        log_step("SYSTEM", "Project completed successfully!", f"Download available at: {zip_path}")
 
     except Exception as e:
-        error_log = ''.join(log_entries) + f"\n\n❌ FINAL ERROR: {str(e)}"
-        update_project_status(project_id, "failed", error_log)
-        return f"Failed: {str(e)}"
+        tb_str = traceback.format_exc()
+        print(f"--- AGENT CHAIN FAILED for project {project_id} ---\n{tb_str}\n--------------------")
+        error_log = "".join(log_entries) + f"\n\n❌ **CRITICAL ERROR:**\nAn unexpected error occurred and the process had to stop.\n\n**Details:**\n```{str(e)}\n\n{tb_str}```"
+        update_project_status(project_id, "failed", logs=error_log)
+
 
 # ------------------------------ JOB QUEUE ------------------------------
-executor = concurrent.futures.ThreadPoolExecutor(max_workers=2)
+# Set to 1 worker to prevent VRAM OOM errors by running models sequentially.
+executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
 
 def queue_job(project_id, user_id, prompt):
+    print(f"Queuing job for project_id: {project_id}, user_id: {user_id}")
     executor.submit(run_agent_chain, project_id, user_id, prompt)