Update backend.py

backend.py

@@ -1,4 +1,4 @@
-# backend.py — FINAL, STABLE & OPTIMIZED VERSION v1.6
+# backend.py — FINAL, TEMPLATE-DRIVEN VERSION v2.0
 import sqlite3
 import os
 import json
@@ -58,12 +58,9 @@ def get_project(project_id): return _db_execute("SELECT * FROM projects WHERE id
 
 # ------------------------------ MODEL LOADING & CACHING ------------------------------
 MODEL_REGISTRY = {
-    "planner": "microsoft/Phi-3-mini-4k-instruct",
-    # --- FINAL, CORRECTED MODEL NAME AS PER YOUR INSTRUCTION ---
-    "architect": "Qwen/Qwen2.5-Coder-0.5B-Instruct",
-    "coder": "Qwen/Qwen2.5-Coder-0.5B-Instruct",
+    "coder": "Qwen/Qwen2-0.5B-Instruct",
     "reviewer": "microsoft/Phi-3-mini-4k-instruct",
-    "tester": "Qwen/Qwen2.5-Coder-0.5B-Instruct",
+    "tester": "Qwen/Qwen2-0.5B-Instruct",
     "publisher": "microsoft/Phi-3-mini-4k-instruct",
 }
 _MODEL_CACHE = {}
@@ -72,25 +69,14 @@ def load_model(model_name):
     if model_name in _MODEL_CACHE:
         return _MODEL_CACHE[model_name]
     
-    # --- THIS IS THE FIX: Conditional Quantization for CPU/GPU compatibility ---
-    model_kwargs = {
-        "device_map": "auto",
-        "trust_remote_code": True,
-        "attn_implementation": "eager",
-    }
+    model_kwargs = {"device_map": "auto", "trust_remote_code": True, "attn_implementation": "eager"}
 
     if torch.cuda.is_available():
         print(f"CUDA is available. Loading model '{model_name}' in 4-bit for GPU acceleration.")
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_use_double_quant=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch.bfloat16
-        )
+        bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16)
         model_kwargs["quantization_config"] = bnb_config
     else:
         print(f"CUDA not available. Loading model '{model_name}' on CPU in default precision.")
-        # No quantization on CPU
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
@@ -101,12 +87,14 @@ def load_model(model_name):
 
 # ------------------------------ AGENT PROMPTS ------------------------------
 ROLE_PROMPTS = {
-    "planner": """You are an expert file planner. Based on the user's request, determine the necessary file structure. Output ONLY a single JSON object with a single key: "files". The "files" key MUST be an array of strings representing complete file paths (e.g., ["src/main.py", "tests/test_main.py", "requirements.txt"]).""",
-    "architect": """You are a software architect. Create initial placeholder content for a list of files. Output ONLY a single JSON object where keys are file paths and values are the initial content (e.g., a comment like '# Main application logic here').""",
     "coder": "You are a professional programmer. Your ONLY job is to write the complete, clean, and functional code for the single file requested. Do NOT add any explanations, introductions, or markdown formatting. Output ONLY the raw source code.",
     "reviewer": """You are a meticulous code reviewer. Analyze the given code for bugs, style issues, and security vulnerabilities. Output ONLY a single JSON object with two keys: "has_issues" (boolean) and "suggestions" (a string containing a bulleted list of required changes).""",
     "tester": "You are a QA engineer. Write a complete pytest test file for the given source code. Cover main functionality and edge cases. Output ONLY the raw source code for the test file.",
-    "publisher": """You are a release manager. Create final documentation and configuration files. Output ONLY a single JSON object where keys are the filenames ("README.md", ".gitignore", "Dockerfile") and values are their complete string content."""
+    "publisher": """You are a release manager. Your job is to create the final project documentation and dependencies based on all the generated code.
+    - Analyze the code to determine the correct Python libraries for requirements.txt.
+    - Write a helpful README.md that explains the project's purpose and how to run it.
+    - Create a standard .gitignore file.
+    Output ONLY a single JSON object where keys are the filenames ("requirements.txt", "README.md", ".gitignore") and values are their complete string content."""
 }
 
 # ------------------------------ FILE SYSTEM & AI TOOLS ------------------------------
@@ -156,9 +144,7 @@ def generate_with_model(role: str, prompt: str) -> str:
         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)
-        
-        outputs = model.generate(**inputs, max_new_tokens=2048, pad_token_id=tokenizer.eos_token_id, use_cache=True)
-        
+        outputs = model.generate(**inputs, max_new_tokens=3072, pad_token_id=tokenizer.eos_token_id, use_cache=True)
         return tokenizer.decode(outputs[0][len(inputs.input_ids[0]):], skip_special_tokens=True).strip()
     except Exception as e:
         print(f"Error during model generation for role {role}: {e}")
@@ -171,59 +157,63 @@ def run_agent_chain(project_id, user_id, initial_prompt):
 
     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"
+        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:
-        log_step("SYSTEM", "Initializing project...")
-
-        log_step("PLANNER", "Determining file structure...")
-        plan_response = generate_with_model("planner", initial_prompt)
-        plan_data = _extract_json(plan_response)
-        if not plan_data or "files" not in plan_data: raise ValueError("Planner failed to create a valid JSON plan.")
-        log_step("PLANNER", "File structure plan created.", json.dumps(plan_data, indent=2))
-
-        log_step("ARCHITECT", "Creating initial file skeletons...")
-        arch_prompt = f"Create initial content for these 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 skeletons.")
-        for path, content in arch_data.items(): create_file(project_dir, path, content)
-        log_step("ARCHITECT", "File skeletons created.", "\n".join(arch_data.keys()))
-
-        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 user's overall request: '{initial_prompt}'.\n\nWrite the full Python code for the file: `{file_path}`."
+        log_step("SYSTEM", "Initializing project with pre-built architecture...")
+
+        # 1. CREATE SKELETON FILES
+        ARCHITECTURE_TEMPLATE = {
+            "app.py": "# Gradio UI will be generated here.",
+            "backend.py": "# Backend logic and database functions will be generated here.",
+            "requirements.txt": "# Dependencies will be generated by the Publisher agent."
+        }
+        for file_path, skeleton_content in ARCHITECTURE_TEMPLATE.items():
+            create_file(project_dir, file_path, skeleton_content)
+        log_step("SYSTEM", "Template files created.", "\n".join(ARCHITECTURE_TEMPLATE.keys()))
+
+        # 2. CODER AGENT FILLS IN THE TEMPLATE
+        files_to_code = ["backend.py", "app.py"]
+        for file_path in files_to_code:
+            log_step("CODER", f"Generating full code for `{file_path}`...")
+            coder_prompt = f"Based on the user's high-level request: '{initial_prompt}'\n\nYour task is to write the complete, final Python code for the file named `{file_path}`. The file should be fully functional and ready to run."
             code = generate_with_model("coder", coder_prompt)
             create_file(project_dir, file_path, code)
             log_step("CODER", f"Finished writing `{file_path}`.", code)
 
-        for file_path in source_files:
+        # 3. REVIEWER AGENT CHECKS THE CODE
+        for file_path in files_to_code:
             code_content = read_file(project_dir, file_path)
             if not code_content: continue
             log_step("REVIEWER", f"Reviewing `{file_path}`...")
             review_response = generate_with_model("reviewer", f"Review this code from `{file_path}`:\n\n{code_content}")
             log_step("REVIEWER", f"Review of `{file_path}` complete.", review_response)
 
-        for file_path in source_files:
-            code_content = read_file(project_dir, file_path)
-            if not code_content: continue
-            test_file_path = os.path.join("tests", f"test_{os.path.basename(file_path)}")
-            log_step("TESTER", f"Writing test for `{file_path}`...")
-            test_code = generate_with_model("tester", f"Write a pytest test file (`{test_file_path}`) for this code from `{file_path}`:\n\n{code_content}")
+        # 4. TESTER AGENT TESTS THE BACKEND
+        log_step("TESTER", "Writing unit tests for `backend.py`...")
+        backend_code = read_file(project_dir, "backend.py")
+        if backend_code:
+            test_file_path = "tests/test_backend.py"
+            tester_prompt = f"Write a complete pytest test file (`{test_file_path}`) for the following backend code:\n\n{backend_code}"
+            test_code = generate_with_model("tester", tester_prompt)
             create_file(project_dir, test_file_path, test_code)
-            log_step("TESTER", f"Generated test `{test_file_path}`.", test_code)
+            log_step("TESTER", f"Generated `{test_file_path}`.", test_code)
 
-        all_files = [os.path.join(r, f).replace(project_dir, '', 1) for r, _, fs in os.walk(project_dir) for f in fs]
-        log_step("PUBLISHER", "Generating final documentation and configuration...")
-        pub_response = generate_with_model("publisher", f"Project file structure: {json.dumps(all_files)}. Generate README.md, .gitignore, and Dockerfile.")
+        # 5. PUBLISHER AGENT CREATES FINAL ASSETS
+        log_step("PUBLISHER", "Generating final documentation and dependencies...")
+        all_code = f"--- app.py ---\n{read_file(project_dir, 'app.py')}\n\n--- backend.py ---\n{read_file(project_dir, 'backend.py')}"
+        pub_prompt = f"Based on the following Python code, generate the project's final assets.\n\n{all_code}"
+        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)
+        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))
         
+        # 6. FINALIZATION
         log_step("SYSTEM", "Packaging 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)
@@ -231,13 +221,12 @@ def run_agent_chain(project_id, user_id, initial_prompt):
 
     except Exception as e:
         tb_str = traceback.format_exc()
-        print(f"--- AGENT CHAIN FAILED for project {project_id} ---\n{tb_str}\n--------------------")
+        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.\n\n**Details:**\n```{str(e)}```"
         update_project_status(project_id, "failed", logs=error_log)
 
 # ------------------------------ JOB QUEUE ------------------------------
 executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
-
 def queue_job(project_id, user_id, prompt):
     print(f"Queuing job for project: {project_id}")
     executor.submit(run_agent_chain, project_id, user_id, prompt)