Update backend.py

backend.py

@@ -1,3 +1,4 @@
+# backend.py — ENHANCED WITH AGENT CONVERSATION LOGS + RAM/VRAM MONITORING
 import sqlite3
 import os
 import json
@@ -10,6 +11,7 @@ import concurrent.futures
 from typing import List, Dict, Any
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import psutil
 
 DB_PATH = "code_agents_pro.db"
 PROJECT_ROOT = "./projects"
@@ -31,8 +33,9 @@ def init_db():
             user_id INTEGER NOT NULL,
             title TEXT,
             description TEXT,
-            status TEXT DEFAULT 'queued', -- queued, running, completed, failed
+            status TEXT DEFAULT 'queued',
             zip_path TEXT,
+            logs TEXT DEFAULT '',
             created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
             FOREIGN KEY (user_id) REFERENCES users(id)
         );
@@ -44,7 +47,6 @@ def init_db():
 init_db()
 
 def hash_password(password):
-    # In production, use bcrypt or argon2. For now, simple hash for demo.
     import hashlib
     return hashlib.sha256(password.encode()).hexdigest()
 
@@ -77,7 +79,7 @@ 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
+        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()
@@ -88,25 +90,25 @@ 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)
-        VALUES (?, ?, ?, 'queued')
+        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
 
-def update_project_status(project_id, status, zip_path=None):
+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 = ? WHERE id = ?
-        """, (status, zip_path, project_id))
+            UPDATE projects SET status = ?, zip_path = ?, logs = ? WHERE id = ?
+        """, (status, zip_path, logs, project_id))
     else:
         cursor.execute("""
-            UPDATE projects SET status = ? WHERE id = ?
-        """, (status, project_id))
+            UPDATE projects SET status = ?, logs = ? WHERE id = ?
+        """, (status, logs, project_id))
     conn.commit()
     conn.close()
 
@@ -114,7 +116,7 @@ def get_project(project_id):
     conn = sqlite3.connect(DB_PATH)
     cursor = conn.cursor()
     cursor.execute("""
-        SELECT id, user_id, title, description, status, zip_path
+        SELECT id, user_id, title, description, status, zip_path, logs
         FROM projects WHERE id = ?
     """, (project_id,))
     row = cursor.fetchone()
@@ -133,9 +135,24 @@ def search_projects(user_id, query):
     conn.close()
     return results
 
-# ------------------------------ MODEL LOADING (PER ROLE) ------------------------------
-# Each role gets the best-suited model — no compromises
+# ------------------------------ 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 ------------------------------
 MODEL_REGISTRY = {
     "understander": "Qwen/Qwen2.5-Coder-0.6B-Instruct",
     "architect": "Qwen/Qwen2.5-Coder-0.6B-Instruct",
@@ -146,9 +163,9 @@ MODEL_REGISTRY = {
     "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": "google/gemma-2-2b-it",
 }
 
-# Load models once at startup
 _MODEL_CACHE = {}
 
 def load_model(model_name):
@@ -298,7 +315,7 @@ If no: list missing items.
 """
 }
 
-# ------------------------------ TOOLS (SAME AS BEFORE) ------------------------------
+# ------------------------------ 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)
@@ -401,81 +418,95 @@ def zip_project(user_id, project_name: str):
                 zf.write(full_path, arcname)
     return zip_path
 
-# ------------------------------ LLM GENERATOR (CORE ENGINE) ------------------------------
+# ------------------------------ LLM GENERATOR (WITH LOGGING) ------------------------------
 def generate_with_model(role: str, prompt: str, context: dict = {}) -> str:
-    model_name = MODEL_REGISTRY[role]
-    tokenizer, model = load_model(model_name)
-
-    # Special formatting per role
-    if role == "understander":
-        messages = [{"role": "system", "content": ROLE_PROMPTS["understander"]},
-                    {"role": "user", "content": prompt}]
-    elif role == "architect":
-        messages = [{"role": "system", "content": ROLE_PROMPTS["architect"]},
-                    {"role": "user", "content": prompt}]
-    elif role == "ceo":
-        messages = [{"role": "system", "content": ROLE_PROMPTS["ceo"]},
-                    {"role": "user", "content": prompt}]
-    elif role == "manager":
-        messages = [{"role": "system", "content": ROLE_PROMPTS["manager"]},
-                    {"role": "user", "content": prompt}]
-    elif role == "worker":
-        messages = [{"role": "system", "content": ROLE_PROMPTS["worker"].format(file=context.get("file"), instructions=context.get("instructions"))},
-                    {"role": "user", "content": ""}]
-    elif role == "reviewer":
-        messages = [{"role": "system", "content": ROLE_PROMPTS["reviewer"]},
-                    {"role": "user", "content": prompt}]
-    elif role == "editor":
-        messages = [{"role": "system", "content": ROLE_PROMPTS["editor"]},
-                    {"role": "user", "content": prompt}]
-    elif role == "tester":
-        messages = [{"role": "system", "content": ROLE_PROMPTS["tester"]},
-                    {"role": "user", "content": prompt}]
-    elif role == "publisher":
-        messages = [{"role": "system", "content": ROLE_PROMPTS["publisher"]},
-                    {"role": "user", "content": prompt}]
-    elif role == "final_ceo":
-        messages = [{"role": "system", "content": ROLE_PROMPTS["final_ceo"]},
-                    {"role": "user", "content": prompt}]
-    else:
-        messages = [{"role": "user", "content": prompt}]
+    try:
+        model_name = MODEL_REGISTRY[role]
+        tokenizer, model = load_model(model_name)
+
+        # Format message based on role
+        if role == "understander":
+            messages = [{"role": "system", "content": ROLE_PROMPTS["understander"]},
+                        {"role": "user", "content": prompt}]
+        elif role == "architect":
+            messages = [{"role": "system", "content": ROLE_PROMPTS["architect"]},
+                        {"role": "user", "content": prompt}]
+        elif role == "ceo":
+            messages = [{"role": "system", "content": ROLE_PROMPTS["ceo"]},
+                        {"role": "user", "content": prompt}]
+        elif role == "manager":
+            messages = [{"role": "system", "content": ROLE_PROMPTS["manager"]},
+                        {"role": "user", "content": prompt}]
+        elif role == "worker":
+            messages = [{"role": "system", "content": ROLE_PROMPTS["worker"].format(file=context.get("file"), instructions=context.get("instructions"))},
+                        {"role": "user", "content": ""}]
+        elif role == "reviewer":
+            messages = [{"role": "system", "content": ROLE_PROMPTS["reviewer"]},
+                        {"role": "user", "content": prompt}]
+        elif role == "editor":
+            messages = [{"role": "system", "content": ROLE_PROMPTS["editor"]},
+                        {"role": "user", "content": prompt}]
+        elif role == "tester":
+            messages = [{"role": "system", "content": ROLE_PROMPTS["tester"]},
+                        {"role": "user", "content": prompt}]
+        elif role == "publisher":
+            messages = [{"role": "system", "content": ROLE_PROMPTS["publisher"]},
+                        {"role": "user", "content": prompt}]
+        elif role == "final_ceo":
+            messages = [{"role": "system", "content": ROLE_PROMPTS["final_ceo"]},
+                        {"role": "user", "content": prompt}]
+        else:
+            messages = [{"role": "user", "content": prompt}]
 
-    text = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
+        text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
 
-    inputs = tokenizer(text, return_tensors="pt").to(model.device)
+        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
-    )
+        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)
+        return response.strip()
 
-    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)}"
 
-# ------------------------------ ASYNC JOB EXECUTOR ------------------------------
+# ------------------------------ ASYNC JOB EXECUTOR (WITH FULL LOGGING) ------------------------------
 def run_agent_chain(project_id, user_id, initial_prompt):
     try:
-        update_project_status(project_id, "running")
+        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)
-        update_project_status(project_id, "running", f"Understood: {understand_result[:100]}...")
+        log_step("Understander", "Analysis complete", understand_result)
 
         # Step 2: Architect
+        log_step("Architect", "Designing structure...", "Started...")
         architect_result = generate_with_model("architect", understand_result)
-        update_project_status(project_id, "running", "Architected structure.")
+        log_step("Architect", "Structure designed", architect_result)
 
         # Parse architect JSON
         try:
@@ -483,16 +514,18 @@ def run_agent_chain(project_id, user_id, initial_prompt):
             arch_data = json.loads(architect_result)
             for file in arch_data.get("structure", []):
                 create_file(user_id, file["path"], file["content"])
-        except:
-            pass
+        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.")
-        update_project_status(project_id, "running", "CEO approved.")
+        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.")
-        update_project_status(project_id, "running", "Manager assigned tasks.")
+        log_step("Manager", "Tasks assigned", manager_response)
 
         try:
             import json
@@ -503,38 +536,44 @@ def run_agent_chain(project_id, user_id, initial_prompt):
                     "instructions": task["instructions"]
                 })
                 create_file(user_id, task["file"], worker_response)
-        except:
-            pass
+                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)
-        update_project_status(project_id, "running", "Reviewer audited code.")
+        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)
-        update_project_status(project_id, "running", "Editor applied fixes.")
+        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)
-        update_project_status(project_id, "running", "Testers wrote unit tests.")
+            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)
-        update_project_status(project_id, "running", "Editor reviewed tests.")
+        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)
-        update_project_status(project_id, "running", "Publisher generated assets.")
+        log_step("Publisher", "Final assets generated", publisher_response)
 
         try:
             import json
@@ -548,25 +587,29 @@ def run_agent_chain(project_id, user_id, initial_prompt):
                     write_requirements(user_id, content)
                 elif fname == ".gitignore":
                     write_gitignore(user_id, content)
-        except:
-            pass
+        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)
-        update_project_status(project_id, "running", "Final CEO verified.")
+        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}")
-        update_project_status(project_id, "completed", zip_path)
+        log_step("System", "ZIP created", f"Saved to {zip_path}")
 
+        update_project_status(project_id, "completed", ''.join(log_entries), zip_path)
         return "Completed successfully"
 
     except Exception as e:
-        update_project_status(project_id, "failed", str(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)}"
 
 # ------------------------------ JOB QUEUE ------------------------------