Update graph.py

graph.py

@@ -1,4 +1,4 @@
-# graph.py (Enhanced with better context passing)
+# graph.py - Complete enhanced version with proper repo generation
 import json
 import re
 import math
@@ -8,6 +8,7 @@ import shutil
 import zipfile
 import operator
 from typing import TypedDict, List, Dict, Optional, Annotated
+from datetime import datetime
 from langchain_openai import ChatOpenAI
 from langgraph.graph import StateGraph, END
 from memory_manager import memory_manager
@@ -22,7 +23,7 @@ from docx import Document
 from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
 from reportlab.lib.styles import getSampleStyleSheet
 
-# --- Configurable output directory ---
+# --- Configuration ---
 OUT_DIR = os.environ.get("OUT_DIR", "/tmp")
 os.makedirs(OUT_DIR, exist_ok=True)
 EXPORTS_DIR = os.path.join(OUT_DIR, "exports")
@@ -51,7 +52,7 @@ def ensure_int(state, key, default=0):
 def sanitize_path(path: str) -> str:
     return os.path.abspath(path)
 
-# --- Setup & constants ---
+# --- Setup ---
 setup_logging()
 log = get_logger(__name__)
 INITIAL_MAX_REWORK_CYCLES = 3
@@ -59,7 +60,7 @@ GPT4O_INPUT_COST_PER_1K_TOKENS = 0.005
 GPT4O_OUTPUT_COST_PER_1K_TOKENS = 0.015
 AVG_TOKENS_PER_CALL = 2.0
 
-# --- AgentState ---
+# --- State ---
 class AgentState(TypedDict):
     userInput: str
     chatHistory: List[str]
@@ -76,7 +77,7 @@ class AgentState(TypedDict):
     max_loops: int
     status_update: str
 
-# --- LLM & parsing ---
+# --- LLM ---
 llm = ChatOpenAI(model="gpt-4o", temperature=0.1, max_retries=3, request_timeout=60)
 
 def parse_json_from_llm(llm_output: str) -> Optional[dict]:
@@ -94,10 +95,10 @@ def parse_json_from_llm(llm_output: str) -> Optional[dict]:
             json_str = llm_output[start:end+1]
         return json.loads(json_str)
     except Exception as e:
-        log.error(f"JSON parsing failed. Error: {e}. Raw head: {llm_output[:300]}")
+        log.error(f"JSON parsing failed: {e}")
         return None
 
-# --- Artifact detection & normalization ---
+# --- Artifact detection ---
 KNOWN_ARTIFACT_TYPES = {"notebook","excel","word","pdf","image","repo","script"}
 
 def detect_requested_output_types(text: str) -> Dict:
@@ -105,19 +106,17 @@ def detect_requested_output_types(text: str) -> Dict:
         return {"requires_artifact": False, "artifact_type": None, "artifact_hint": None}
     t = text.lower()
     if any(k in t for k in ["jupyter notebook", "jupyter", "notebook", "ipynb"]):
-        return {"requires_artifact": True, "artifact_type": "notebook", "artifact_hint": "jupyter notebook (.ipynb)"}
-    if any(k in t for k in ["excel", ".xlsx", "spreadsheet", "csv", "sheet"]):
-        return {"requires_artifact": True, "artifact_type": "excel", "artifact_hint": "Excel/CSV file"}
-    if any(k in t for k in ["word document", ".docx", "docx", "word file"]):
-        return {"requires_artifact": True, "artifact_type": "word", "artifact_hint": "Word document (.docx)"}
+        return {"requires_artifact": True, "artifact_type": "notebook", "artifact_hint": "jupyter notebook"}
+    if any(k in t for k in ["excel", ".xlsx", "spreadsheet", "csv"]):
+        return {"requires_artifact": True, "artifact_type": "excel", "artifact_hint": "Excel file"}
+    if any(k in t for k in ["word document", ".docx", "docx"]):
+        return {"requires_artifact": True, "artifact_type": "word", "artifact_hint": "Word document"}
     if any(k in t for k in ["pdf", "pdf file"]):
         return {"requires_artifact": True, "artifact_type": "pdf", "artifact_hint": "PDF document"}
-    if any(k in t for k in ["image", "plot", "chart", "png", "jpg", "jpeg"]):
-        return {"requires_artifact": True, "artifact_type": "image", "artifact_hint": "image/plot"}
-    if any(k in t for k in ["repo", "repository", "app repo", "dockerfile", "requirements.txt", "package.json"]):
-        return {"requires_artifact": True, "artifact_type": "repo", "artifact_hint": "application repository (zip)"}
-    if any(k in t for k in [".py", "python script", "r script", ".R", ".r", "java", ".java", "javascript", ".js"]):
-        return {"requires_artifact": True, "artifact_type": "script", "artifact_hint": "language script (py/r/java/js/etc.)"}
+    if any(k in t for k in ["repo", "repository", "app repo", "backend", "codebase"]):
+        return {"requires_artifact": True, "artifact_type": "repo", "artifact_hint": "application repository"}
+    if any(k in t for k in [".py", "python script", "script"]):
+        return {"requires_artifact": True, "artifact_type": "script", "artifact_hint": "Python script"}
     return {"requires_artifact": False, "artifact_type": None, "artifact_hint": None}
 
 def normalize_experiment_type(exp_type: Optional[str], goal_text: str) -> str:
@@ -129,20 +128,20 @@ def normalize_experiment_type(exp_type: Optional[str], goal_text: str) -> str:
         return s
     if "notebook" in s or "ipynb" in s:
         return "notebook"
-    if "excel" in s or "xlsx" in s or "spreadsheet" in s:
+    if "excel" in s or "xlsx" in s:
         return "excel"
     if "word" in s or "docx" in s:
         return "word"
     if "pdf" in s:
         return "pdf"
-    if "repo" in s or "repository" in s or "app" in s:
+    if "repo" in s or "repository" in s or "backend" in s:
         return "repo"
-    if "script" in s or "python" in s or ".py" in s:
+    if "script" in s or "python" in s:
         return "script"
     detection = detect_requested_output_types(goal_text or "")
     return detection.get("artifact_type") or "word"
 
-# --- Notebook & artifact builders ---
+# --- Artifact builders ---
 def write_notebook_from_text(llm_text: str, out_dir: Optional[str]=None) -> str:
     out_dir = out_dir or OUT_DIR
     os.makedirs(out_dir, exist_ok=True)
@@ -159,7 +158,7 @@ def write_notebook_from_text(llm_text: str, out_dir: Optional[str]=None) -> str:
         if i < len(code_blocks) and code_blocks[i].strip():
             cells.append(new_code_cell(code_blocks[i].strip()))
     if not cells:
-        cells = [new_markdown_cell("# Notebook\n\nNo content parsed from LLM output.")]
+        cells = [new_markdown_cell("# Notebook\n\nNo content generated.")]
     nb['cells'] = cells
     uid = uuid.uuid4().hex[:10]
     filename = os.path.join(out_dir, f"generated_notebook_{uid}.ipynb")
@@ -172,16 +171,14 @@ def write_script(code_text: str, language_hint: Optional[str]=None, out_dir: Opt
     ext = ".txt"
     if language_hint:
         l = language_hint.lower()
-        if "python" in l or ".py" in l:
+        if "python" in l:
             ext = ".py"
-        elif l == "r" or l == ".r":
+        elif "r" in l:
             ext = ".R"
-        elif "java" in l or ".java" in l:
+        elif "java" in l:
             ext = ".java"
-        elif "javascript" in l or ".js" in l:
+        elif "javascript" in l:
             ext = ".js"
-        elif "bash" in l or "sh" in l:
-            ext = ".sh"
     uid = uuid.uuid4().hex[:10]
     filename = os.path.join(out_dir, f"generated_script_{uid}{ext}")
     with open(filename, "w", encoding="utf-8") as f:
@@ -214,7 +211,7 @@ def write_excel_from_tables(maybe_table_text: str, out_dir: Optional[str]=None)
             else:
                 df = pd.DataFrame({"content":[str(maybe_table_text)]})
         except Exception:
-            if "," in maybe_table_text or "\t" in maybe_table_text:
+            if "," in maybe_table_text:
                 from io import StringIO
                 df = pd.read_csv(StringIO(maybe_table_text))
             else:
@@ -223,7 +220,7 @@ def write_excel_from_tables(maybe_table_text: str, out_dir: Optional[str]=None)
         return filename
     except Exception as e:
         log.error(f"Excel creation failed: {e}")
-        return write_docx_from_text(f"Failed to create excel. Error: {e}\n\nOriginal:\n{maybe_table_text}", out_dir=out_dir)
+        return write_docx_from_text(f"Excel error: {e}\n\n{maybe_table_text}", out_dir=out_dir)
 
 def write_pdf_from_text(text: str, out_dir: Optional[str]=None) -> str:
     out_dir = out_dir or OUT_DIR
@@ -241,7 +238,7 @@ def write_pdf_from_text(text: str, out_dir: Optional[str]=None) -> str:
         return filename
     except Exception as e:
         log.error(f"PDF creation failed: {e}")
-        return write_docx_from_text(f"Failed to create PDF. Error: {e}\n\nOriginal:\n{text}", out_dir=out_dir)
+        return write_docx_from_text(f"PDF error: {e}\n\n{text}", out_dir=out_dir)
 
 def build_repo_zip(files_map: Dict[str,str], repo_name: str="generated_app", out_dir: Optional[str]=None) -> str:
     out_dir = out_dir or OUT_DIR
@@ -249,14 +246,17 @@ def build_repo_zip(files_map: Dict[str,str], repo_name: str="generated_app", out
     uid = uuid.uuid4().hex[:8]
     repo_dir = os.path.join(out_dir, f"{repo_name}_{uid}")
     os.makedirs(repo_dir, exist_ok=True)
+    
     for rel_path, content in files_map.items():
         dest = os.path.join(repo_dir, rel_path)
         os.makedirs(os.path.dirname(dest), exist_ok=True)
+        
         if isinstance(content, str) and os.path.exists(content):
             shutil.copyfile(content, dest)
         else:
             with open(dest, "w", encoding="utf-8") as fh:
                 fh.write(str(content))
+    
     zip_path = os.path.join(out_dir, f"{repo_name}_{uid}.zip")
     with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
         for root, _, files in os.walk(repo_dir):
@@ -264,283 +264,347 @@ def build_repo_zip(files_map: Dict[str,str], repo_name: str="generated_app", out
                 full = os.path.join(root, f)
                 arc = os.path.relpath(full, repo_dir)
                 zf.write(full, arc)
+    
     return zip_path
 
-# --- Node functions ---
+# --- Nodes ---
 def run_triage_agent(state: AgentState):
     log.info("--- TRIAGE ---")
-    prompt = f"Analyze the user input. Is it a simple conversational greeting or a task? Respond with 'greeting' or 'task'.\n\nUser Input: \"{state.get('userInput','')}\""
+    prompt = f"Is this a greeting or a task? '{state.get('userInput','')}' Reply: 'greeting' or 'task'"
     response = llm.invoke(prompt)
     content = getattr(response, "content", "") or ""
     if 'greeting' in content.lower():
-        log.info("Triage result: Simple Greeting.")
-        return {"draftResponse": "Hello! How can I help you today?", "execution_path": ["Triage Agent"], "status_update": "Responding to greeting."}
-    else:
-        log.info("Triage result: Complex Task.")
-        return {"execution_path": ["Triage Agent"], "status_update": "Request requires a plan. Proceeding..."}
+        return {"draftResponse": "Hello! How can I help?", "execution_path": ["Triage"], "status_update": "Greeting"}
+    return {"execution_path": ["Triage"], "status_update": "Task detected"}
 
 def run_planner_agent(state: AgentState):
-    log.info("--- PLANNER AGENT ---")
-    path = ensure_list(state, 'execution_path') + ["Planner Agent"]
-    prompt = (
-        f"Analyze the user's request. Provide a high-level plan and estimate the number of LLM calls for one loop. "
-        f"User Request: \"{state.get('userInput','')}\". Respond in JSON with keys: 'plan' (list of strings), 'estimated_llm_calls_per_loop' (integer)."
-    )
+    log.info("--- PLANNER ---")
+    path = ensure_list(state, 'execution_path') + ["Planner"]
+    prompt = f"Create a plan for: '{state.get('userInput','')}'. JSON with 'plan' (list), 'estimated_llm_calls_per_loop' (int)"
     response = llm.invoke(prompt)
     plan_data = parse_json_from_llm(getattr(response, "content", "") or "")
     if not plan_data:
-        return {"pmPlan": {"error": "Failed to create a valid plan."}, "execution_path": path, "status_update": "Error: Could not create a plan."}
-    calls_per_loop = plan_data.get('estimated_llm_calls_per_loop', 3)
-    cost_per_loop = (calls_per_loop * AVG_TOKENS_PER_CALL) * ((GPT4O_INPUT_COST_PER_1K_TOKENS + GPT4O_OUTPUT_COST_PER_1K_TOKENS) / 2)
-    estimated_cost = cost_per_loop * (INITIAL_MAX_REWORK_CYCLES + 1)
+        return {"pmPlan": {"error": "Planning failed"}, "execution_path": path, "status_update": "Error"}
+    
+    calls = plan_data.get('estimated_llm_calls_per_loop', 3)
+    cost_per_loop = (calls * AVG_TOKENS_PER_CALL) * ((GPT4O_INPUT_COST_PER_1K_TOKENS + GPT4O_OUTPUT_COST_PER_1K_TOKENS) / 2)
     plan_data['max_loops_initial'] = INITIAL_MAX_REWORK_CYCLES
-    plan_data['estimated_cost_usd'] = round(estimated_cost, 2)
+    plan_data['estimated_cost_usd'] = round(cost_per_loop * (INITIAL_MAX_REWORK_CYCLES + 1), 2)
     plan_data['cost_per_loop_usd'] = max(0.01, round(cost_per_loop, 3))
-    detection = detect_requested_output_types(state.get('userInput','') or state.get('coreObjectivePrompt','') or '')
+    
+    detection = detect_requested_output_types(state.get('userInput',''))
     if detection.get('requires_artifact'):
         plan_data.setdefault('experiment_needed', True)
         plan_data.setdefault('experiment_type', detection.get('artifact_type'))
-        plan_data.setdefault('experiment_goal', f"Produce an artifact: {detection.get('artifact_hint')}. {state.get('userInput','')}")
-    log.info(f"Pre-flight Estimate: {plan_data}")
-    return {"pmPlan": plan_data, "execution_path": path, "status_update": "Plan and cost estimate created. Awaiting approval."}
+        plan_data.setdefault('experiment_goal', state.get('userInput',''))
+    
+    return {"pmPlan": plan_data, "execution_path": path, "status_update": "Plan created"}
 
 def run_memory_retrieval(state: AgentState):
-    log.info("--- MEMORY RETRIEVAL ---")
-    path = ensure_list(state, 'execution_path') + ["Memory Retriever"]
-    relevant_mems = memory_manager.retrieve_relevant_memories(state.get('userInput',''))
-    if relevant_mems:
-        context = "\n".join([f"Memory: {mem.page_content}" for mem in relevant_mems])
-        log.info(f"Found {len(relevant_mems)} relevant memories.")
-    else:
-        context = "No relevant memories found."
-        log.info(context)
-    return {"retrievedMemory": context, "execution_path": path, "status_update": "Searching for relevant past information..."}
+    log.info("--- MEMORY ---")
+    path = ensure_list(state, 'execution_path') + ["Memory"]
+    mems = memory_manager.retrieve_relevant_memories(state.get('userInput',''))
+    context = "\n".join([f"Memory: {m.page_content}" for m in mems]) if mems else "No memories"
+    return {"retrievedMemory": context, "execution_path": path, "status_update": "Memory retrieved"}
 
 def run_intent_agent(state: AgentState):
-    log.info("--- INTENT AGENT ---")
-    path = ensure_list(state, 'execution_path') + ["Intent Agent"]
-    prompt = (f"Refine the user's request into a clear, actionable 'core objective prompt'.\n\nRelevant Memory:\n{state.get('retrievedMemory')}\n\nUser Request: \"{state.get('userInput','')}\"\n\nCore Objective:")
+    log.info("--- INTENT ---")
+    path = ensure_list(state, 'execution_path') + ["Intent"]
+    prompt = f"Refine into clear objective.\n\nMemory: {state.get('retrievedMemory')}\n\nRequest: {state.get('userInput','')}\n\nCore Objective:"
     response = llm.invoke(prompt)
     core_obj = getattr(response, "content", "") or ""
-    detection = detect_requested_output_types(core_obj or state.get('userInput',''))
-    extras = {}
-    if detection.get('requires_artifact'):
-        extras['artifact_detection'] = detection
-    return {"coreObjectivePrompt": core_obj, **extras, "execution_path": path, "status_update": "Clarifying the main objective..."}
+    return {"coreObjectivePrompt": core_obj, "execution_path": path, "status_update": "Objective clarified"}
 
 def run_pm_agent(state: AgentState):
-    log.info("--- PM AGENT ---")
+    log.info("--- PM ---")
     current_cycles = ensure_int(state, 'rework_cycles', 0) + 1
     max_loops_val = ensure_int(state, 'max_loops', 0)
-    log.info(f"Starting work cycle {current_cycles}/{max_loops_val + 1}")
-    path = ensure_list(state, 'execution_path') + ["PM Agent"]
+    path = ensure_list(state, 'execution_path') + ["PM"]
     
-    # BUILD COMPREHENSIVE CONTEXT
     context_parts = [
-        f"=== USER'S ORIGINAL REQUEST ===",
-        f"{state.get('userInput', '')}",
-        f"\n=== CORE OBJECTIVE ===",
-        f"{state.get('coreObjectivePrompt', '')}",
-        f"\n=== RELEVANT MEMORY ===",
-        f"{state.get('retrievedMemory', 'None')}",
+        f"=== USER REQUEST ===\n{state.get('userInput', '')}",
+        f"\n=== OBJECTIVE ===\n{state.get('coreObjectivePrompt', '')}",
+        f"\n=== MEMORY ===\n{state.get('retrievedMemory', 'None')}",
     ]
     
     if state.get('qaFeedback'):
-        context_parts.append(f"\n=== QA FEEDBACK (MUST ADDRESS) ===")
-        context_parts.append(f"{state.get('qaFeedback')}")
-        context_parts.append(f"\n=== PREVIOUS PLAN ===")
-        context_parts.append(f"{json.dumps(state.get('pmPlan', {}).get('plan_steps', []), indent=2)}")
+        context_parts.append(f"\n=== QA FEEDBACK (MUST FIX) ===\n{state.get('qaFeedback')}")
+        context_parts.append(f"\n=== PREVIOUS PLAN ===\n{json.dumps(state.get('pmPlan', {}).get('plan_steps', []), indent=2)}")
     
     full_context = "\n".join(context_parts)
     
-    # ENHANCED PM PROMPT
-    prompt = f"""You are a Project Manager creating a DETAILED, EXECUTABLE plan.
+    prompt = f"""Create DETAILED, EXECUTABLE plan.
 
 {full_context}
 
-Your task is to create a plan where each step is SPECIFIC and ACTIONABLE:
+Each step must be SPECIFIC and ACTIONABLE:
 - State EXACTLY what will be created/analyzed
-- Specify WHAT information/data will be used
-- Define WHAT approach/method will be applied
+- Specify WHAT data/information will be used
+- Define WHAT methods will be applied
 
-Respond in JSON format:
+JSON format:
 {{
-    "plan_steps": [
-        "Specific executable step 1 with clear deliverable...",
-        "Specific executable step 2 with clear action...",
-        "..."
-    ],
+    "plan_steps": ["Specific step 1...", "Specific step 2..."],
     "experiment_needed": true/false,
     "experiment_type": "notebook|script|excel|word|pdf|repo",
-    "experiment_goal": "Detailed description of artifact content and purpose",
-    "experiment_language": "python|r|java|javascript" (if script),
-    "key_requirements": ["Critical requirements that MUST be met"]
+    "experiment_goal": "Detailed artifact description",
+    "key_requirements": ["Critical requirements"]
 }}
 
-CRITICAL: Be specific about:
-- Analysis tasks: WHAT to analyze and HOW
-- Code tasks: WHAT functionality to implement
-- Document tasks: WHAT sections/content to include
-- Using any uploaded files or user-provided data
-"""
+Be specific about using uploaded files, implementing algorithms, creating schemas."""
     
     response = llm.invoke(prompt)
     plan = parse_json_from_llm(getattr(response, "content", "") or "")
     
     if not plan:
-        log.warning("PM Agent did not produce JSON – applying fallback.")
         detection = detect_requested_output_types(state.get('userInput', ''))
         plan = {
-            "plan_steps": [
-                f"Analyze request: {state.get('userInput', '')[:100]}...",
-                "Process relevant information",
-                "Create deliverable with specific details",
-                "Review output quality"
-            ],
+            "plan_steps": ["Analyze request", "Process information", "Create deliverable", "Review"],
             "experiment_needed": detection.get('requires_artifact', False),
             "experiment_type": detection.get('artifact_type', 'word'),
             "experiment_goal": state.get('coreObjectivePrompt', state.get('userInput', ''))
         }
     
-    exp_type_raw = plan.get('experiment_type') or ""
-    plan_goal = plan.get('experiment_goal') or state.get('userInput','') or state.get('coreObjectivePrompt','')
-    normalized = normalize_experiment_type(exp_type_raw, plan_goal)
-    plan['experiment_type'] = normalized
+    exp_type = normalize_experiment_type(plan.get('experiment_type'), plan.get('experiment_goal',''))
+    plan['experiment_type'] = exp_type
     
     if plan.get('experiment_needed') and not plan.get('experiment_goal'):
-        plan['experiment_goal'] = plan_goal
-    
-    log.info(f"Plan: Steps={len(plan.get('plan_steps', []))}, Experiment={plan.get('experiment_needed')}, Type={plan.get('experiment_type')}")
+        plan['experiment_goal'] = state.get('userInput','')
     
-    return {
-        "pmPlan": plan, 
-        "execution_path": path, 
-        "rework_cycles": current_cycles, 
-        "status_update": f"Detailed plan created ({len(plan.get('plan_steps', []))} steps)"
-    }
+    return {"pmPlan": plan, "execution_path": path, "rework_cycles": current_cycles, "status_update": f"Plan created ({len(plan.get('plan_steps', []))} steps)"}
 
 def _extract_code_blocks(text: str, lang_hint: Optional[str]=None) -> List[str]:
     if lang_hint and "python" in (lang_hint or "").lower():
         blocks = re.findall(r"```python\s*(.*?)\s*```", text, re.DOTALL)
         if blocks:
             return blocks
-    blocks = re.findall(r"```(?:\w+)?\s*(.*?)\s*```", text, re.DOTALL)
-    return blocks
+    return re.findall(r"```(?:\w+)?\s*(.*?)\s*```", text, re.DOTALL)
 
 def run_experimenter_agent(state: AgentState):
-    log.info("--- EXPERIMENTER AGENT ---")
-    path = ensure_list(state, 'execution_path') + ["Experimenter Agent"]
+    log.info("--- EXPERIMENTER ---")
+    path = ensure_list(state, 'execution_path') + ["Experimenter"]
     pm = state.get('pmPlan', {}) or {}
     
     if not pm.get('experiment_needed'):
-        return {
-            "experimentCode": None, 
-            "experimentResults": None, 
-            "execution_path": path, 
-            "status_update": "No experiment needed."
-        }
+        return {"experimentCode": None, "experimentResults": None, "execution_path": path, "status_update": "No experiment needed"}
     
     exp_type = normalize_experiment_type(pm.get('experiment_type'), pm.get('experiment_goal',''))
-    goal = pm.get('experiment_goal', 'No goal specified.')
+    goal = pm.get('experiment_goal', 'No goal')
     
-    # BUILD COMPREHENSIVE CONTEXT FOR EXPERIMENTER
+    # BUILD RICH CONTEXT
     context_parts = [
-        f"=== USER'S ORIGINAL REQUEST ===",
-        f"{state.get('userInput', '')}",
-        f"\n=== CORE OBJECTIVE ===",
-        f"{state.get('coreObjectivePrompt', '')}",
-        f"\n=== EXECUTION PLAN ===",
-        f"{json.dumps(pm.get('plan_steps', []), indent=2)}",
-        f"\n=== KEY REQUIREMENTS ===",
-        f"{json.dumps(pm.get('key_requirements', []), indent=2)}",
+        f"=== USER REQUEST ===\n{state.get('userInput', '')}",
+        f"\n=== OBJECTIVE ===\n{state.get('coreObjectivePrompt', '')}",
+        f"\n=== PLAN ===\n{json.dumps(pm.get('plan_steps', []), indent=2)}",
+        f"\n=== REQUIREMENTS ===\n{json.dumps(pm.get('key_requirements', []), indent=2)}",
     ]
     
     if state.get('retrievedMemory'):
-        context_parts.append(f"\n=== RELEVANT PAST CONTEXT ===")
-        context_parts.append(f"{state.get('retrievedMemory', '')}")
+        context_parts.append(f"\n=== CONTEXT ===\n{state.get('retrievedMemory', '')}")
     
     if state.get('qaFeedback'):
-        context_parts.append(f"\n=== FEEDBACK TO ADDRESS ===")
-        context_parts.append(f"{state.get('qaFeedback', '')}")
+        context_parts.append(f"\n=== FEEDBACK TO ADDRESS ===\n{state.get('qaFeedback', '')}")
     
     full_context = "\n".join(context_parts)
     
-    # ENHANCED EXPERIMENTER PROMPT
-    enhanced_prompt = f"""You are creating a HIGH-QUALITY {exp_type} artifact.
+    # REPO REQUIRES SPECIAL HANDLING
+    if exp_type == 'repo':
+        repo_prompt = f"""Create COMPLETE, PRODUCTION-READY application repository.
 
 {full_context}
 
-ARTIFACT GOAL: {goal}
+GOAL: {goal}
+
+CRITICAL REQUIREMENTS:
+
+1. ACTUAL WORKING CODE - Not templates, not documentation, not examples. REAL production code.
+
+2. FILE STRUCTURE - Indicate each file clearly:
+   ### path/to/file.py
+   ```python
+   [Complete working code]
+   ```
+
+3. MUST INCLUDE:
+   - Complete API clients with error handling, retries, rate limiting
+   - Database schema with CREATE TABLE statements
+   - Data processing with real transformation logic
+   - Config management (.env handling)
+   - requirements.txt with ALL dependencies
+   - main.py entry point
+   - Comprehensive README
+
+4. CODE QUALITY:
+   - Environment variables for secrets
+   - Error handling and logging
+   - Docstrings and comments
+   - Real business logic based on request
+   - RUNNABLE out of the box
+
+5. SPECIFIC TO REQUEST:
+   - Use EXACT APIs mentioned (e.g., CricAPI, SportsRadar)
+   - Implement SPECIFIC algorithms (e.g., batting avg, strike rate)
+   - Create EXACT database tables needed
+   - Process SPECIFIC data formats
+
+NO placeholders like "# TODO"
+NO dummy data - implement REAL logic
+NO documentation-style code - PRODUCTION code only
+
+Format each file:
+### path/to/file.py
+```python
+# Complete code here
+```
+
+Generate complete repository:"""
+
+        response = llm.invoke(repo_prompt)
+        llm_text = getattr(response, "content", "") or ""
+        
+        # Parse files from response
+        repo_files = {}
+        
+        # Extract with ### headers
+        file_pattern = r"###\s+([\w\/_\-\.]+)\s*\n```(?:\w+)?\s*\n(.*?)\n```"
+        matches = re.finditer(file_pattern, llm_text, re.DOTALL)
+        
+        for match in matches:
+            filepath = match.group(1).strip()
+            content = match.group(2).strip()
+            repo_files[filepath] = content
+        
+        # Fallback: extract code blocks
+        if not repo_files:
+            code_blocks = re.findall(r"```(?:python|sql)?\s*\n(.*?)\n```", llm_text, re.DOTALL)
+            if code_blocks:
+                for i, block in enumerate(code_blocks):
+                    if len(block) > 50:  # Skip tiny blocks
+                        repo_files[f"module_{i}.py"] = block
+        
+        # Add README if missing
+        if not any('README' in f.upper() for f in repo_files):
+            repo_files["README.md"] = f"""# Generated Application
+
+## Overview
+{goal}
 
-QUALITY REQUIREMENTS:
-1. Use ALL specific details from the user's request
-2. Create PRODUCTION-READY, COMPLETE content (NO templates or placeholders)
-3. Include ACTUAL data, REALISTIC examples, and WORKING implementations
-4. For notebooks: Include markdown explanations, executable code, and visualizations
-5. For scripts: Include error handling, documentation, and real logic
-6. For documents: Provide substantive, detailed content based on context
-7. For analysis: Use specific methodologies and provide concrete insights
+## Files
+{chr(10).join(f'- {f}' for f in sorted(repo_files.keys()))}
 
-Generate complete, high-quality content for '{exp_type}'.
-Use fenced code blocks with language identifiers where appropriate.
+## Setup
+1. `pip install -r requirements.txt`
+2. Copy `.env.example` to `.env` and configure
+3. Run: `python main.py`
+"""
+        
+        # Add requirements.txt
+        if "requirements.txt" not in repo_files:
+            all_code = " ".join(repo_files.values()).lower()
+            deps = []
+            if 'requests' in all_code: deps.append('requests')
+            if 'pandas' in all_code: deps.append('pandas')
+            if 'numpy' in all_code: deps.append('numpy')
+            if 'sqlalchemy' in all_code: deps.append('sqlalchemy')
+            if 'postgresql' in all_code or 'psycopg2' in all_code: deps.append('psycopg2-binary')
+            if 'flask' in all_code: deps.append('flask')
+            if 'fastapi' in all_code: 
+                deps.append('fastapi')
+                deps.append('uvicorn')
+            if 'dotenv' in all_code: deps.append('python-dotenv')
+            
+            repo_files["requirements.txt"] = "\n".join(deps) if deps else "# Dependencies"
+        
+        # Add .env.example
+        if ".env.example" not in repo_files:
+            repo_files[".env.example"] = """# Configuration
+API_KEY=your_key_here
+DATABASE_URL=postgresql://user:pass@localhost/db
+DEBUG=False
 """
+        
+        # Add main.py if missing
+        if not any('main.py' in f for f in repo_files):
+            repo_files["main.py"] = """#!/usr/bin/env python3
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+def main():
+    print("Application starting...")
+    # Add your logic here
+    pass
+
+if __name__ == "__main__":
+    main()
+"""
+        
+        # Build zip
+        zip_path = build_repo_zip(repo_files, repo_name="generated_app", out_dir=OUT_DIR)
+        
+        results = {
+            "success": True,
+            "paths": {"repo_zip": sanitize_path(zip_path)},
+            "files_created": len(repo_files),
+            "context_used": len(full_context)
+        }
+        
+        return {
+            "experimentCode": None,
+            "experimentResults": results,
+            "execution_path": path,
+            "status_update": f"Repository created ({len(repo_files)} files)"
+        }
     
+    # OTHER ARTIFACT TYPES
+    enhanced_prompt = f"""Create HIGH-QUALITY {exp_type} artifact.
+
+{full_context}
+
+GOAL: {goal}
+
+REQUIREMENTS:
+1. Use ALL specific details from request
+2. PRODUCTION-READY, COMPLETE content (NO placeholders)
+3. ACTUAL data, REALISTIC examples, WORKING code
+4. For notebooks: markdown + executable code + visualizations
+5. For scripts: error handling + docs + real logic
+6. For documents: substantive detailed content
+
+Generate complete content for '{exp_type}' with proper code fences."""
+
     response = llm.invoke(enhanced_prompt)
     llm_text = getattr(response, "content", "") or ""
-    out_dir = OUT_DIR
     results = {"success": False, "paths": {}, "stderr": "", "stdout": "", "context_used": len(full_context)}
     
     try:
         if exp_type == 'notebook':
-            nb_path = write_notebook_from_text(llm_text, out_dir=out_dir)
+            nb_path = write_notebook_from_text(llm_text, out_dir=OUT_DIR)
             results.update({"success": True, "paths": {"notebook": sanitize_path(nb_path)}})
-            return {
-                "experimentCode": None, 
-                "experimentResults": results, 
-                "execution_path": path, 
-                "status_update": f"Notebook generated ({len(full_context)} chars context)"
-            }
+            return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": "Notebook created"}
             
         elif exp_type == 'excel':
-            excel_path = write_excel_from_tables(llm_text, out_dir=out_dir)
+            excel_path = write_excel_from_tables(llm_text, out_dir=OUT_DIR)
             results.update({"success": True, "paths": {"excel": sanitize_path(excel_path)}})
-            return {
-                "experimentCode": None, 
-                "experimentResults": results, 
-                "execution_path": path, 
-                "status_update": f"Excel generated"
-            }
+            return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": "Excel created"}
             
         elif exp_type == 'word':
-            docx_path = write_docx_from_text(llm_text, out_dir=out_dir)
+            docx_path = write_docx_from_text(llm_text, out_dir=OUT_DIR)
             results.update({"success": True, "paths": {"docx": sanitize_path(docx_path)}})
-            return {
-                "experimentCode": None, 
-                "experimentResults": results, 
-                "execution_path": path, 
-                "status_update": f"DOCX generated"
-            }
+            return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": "DOCX created"}
             
         elif exp_type == 'pdf':
-            pdf_path = write_pdf_from_text(llm_text, out_dir=out_dir)
+            pdf_path = write_pdf_from_text(llm_text, out_dir=OUT_DIR)
             results.update({"success": True, "paths": {"pdf": sanitize_path(pdf_path)}})
-            return {
-                "experimentCode": None, 
-                "experimentResults": results, 
-                "execution_path": path, 
-                "status_update": f"PDF generated"
-            }
+            return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": "PDF created"}
             
         elif exp_type == 'script':
-            lang_hint = pm.get('experiment_language') or ("python" if ".py" in goal.lower() else None)
+            lang_hint = pm.get('experiment_language') or "python"
             code_blocks = _extract_code_blocks(llm_text, lang_hint)
+            code_text = "\n\n# === BLOCK ===\n\n".join(code_blocks) if code_blocks else llm_text
             
-            if not code_blocks:
-                code_text = llm_text
-            else:
-                code_text = "\n\n# === BLOCK ===\n\n".join(code_blocks)
-            
-            script_path = write_script(code_text, language_hint=lang_hint, out_dir=out_dir)
+            script_path = write_script(code_text, language_hint=lang_hint, out_dir=OUT_DIR)
             exec_results = {}
             
             if script_path.endswith(".py"):
@@ -550,74 +614,37 @@ Use fenced code blocks with language identifiers where appropriate.
                     exec_results = {"stdout":"","stderr":str(e),"success":False}
             
             results.update({
-                "success": True, 
-                "paths": {"script": sanitize_path(script_path)}, 
-                "stdout": exec_results.get("stdout",""), 
+                "success": True,
+                "paths": {"script": sanitize_path(script_path)},
+                "stdout": exec_results.get("stdout",""),
                 "stderr": exec_results.get("stderr","")
             })
-            return {
-                "experimentCode": code_text, 
-                "experimentResults": results, 
-                "execution_path": path, 
-                "status_update": f"Script generated"
-            }
-            
-        elif exp_type == 'repo':
-            repo_files = {}
-            readme = (llm_text[:1000] + "\n\n") if llm_text else "Generated repo"
-            repo_files["README.md"] = readme
-            nb_path = write_notebook_from_text(llm_text, out_dir=out_dir)
-            repo_files["analysis.ipynb"] = nb_path
-            reqs = "nbformat\npandas\nopenpyxl\npython-docx\nreportlab"
-            repo_files["requirements.txt"] = reqs
-            zip_path = build_repo_zip(repo_files, repo_name="generated_app", out_dir=out_dir)
-            results.update({"success": True, "paths": {"repo_zip": sanitize_path(zip_path)}})
-            return {
-                "experimentCode": None, 
-                "experimentResults": results, 
-                "execution_path": path, 
-                "status_update": f"Repository ZIP created"
-            }
-            
+            return {"experimentCode": code_text, "experimentResults": results, "execution_path": path, "status_update": "Script created"}
+        
         else:
-            fallback = write_docx_from_text(llm_text, out_dir=out_dir)
+            fallback = write_docx_from_text(llm_text, out_dir=OUT_DIR)
             results.update({"success": True, "paths": {"docx": sanitize_path(fallback)}})
-            return {
-                "experimentCode": None, 
-                "experimentResults": results, 
-                "execution_path": path, 
-                "status_update": f"Fallback DOCX generated"
-            }
+            return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": "Document created"}
             
     except Exception as e:
         log.error(f"Experimenter failed: {e}")
         results.update({"success": False, "stderr": str(e)})
-        return {
-            "experimentCode": None, 
-            "experimentResults": results, 
-            "execution_path": path, 
-            "status_update": "Error: Experimenter failed."
-        }
+        return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": "Error"}
 
 def run_synthesis_agent(state: AgentState):
-    log.info("--- SYNTHESIS AGENT ---")
+    log.info("--- SYNTHESIS ---")
     _state = state or {}
-    path = ensure_list(_state, 'execution_path') + ["Synthesis Agent"]
+    path = ensure_list(_state, 'execution_path') + ["Synthesis"]
     
     exp_results = _state.get('experimentResults')
     pm_plan = _state.get('pmPlan', {}) or {}
     
-    # BUILD COMPREHENSIVE SYNTHESIS CONTEXT
     synthesis_context = [
-        f"=== ORIGINAL USER REQUEST ===",
-        f"{_state.get('userInput', '')}",
-        f"\n=== CORE OBJECTIVE ===",
-        f"{_state.get('coreObjectivePrompt', '')}",
-        f"\n=== EXECUTION PLAN ===",
-        f"{json.dumps(pm_plan.get('plan_steps', []), indent=2)}",
+        f"=== USER REQUEST ===\n{_state.get('userInput', '')}",
+        f"\n=== OBJECTIVE ===\n{_state.get('coreObjectivePrompt', '')}",
+        f"\n=== PLAN ===\n{json.dumps(pm_plan.get('plan_steps', []), indent=2)}",
     ]
     
-    # Extract artifact details
     artifact_details = []
     artifact_message = ""
     
@@ -630,150 +657,94 @@ def run_synthesis_agent(state: AgentState):
                 artifact_lines.append(f"- **{artifact_type.title()}**: `{os.path.basename(artifact_path)}`")
                 artifact_details.append(f"{artifact_type}: {artifact_path}")
             
-            artifact_message = "\n\n**📦 Artifacts Generated:**\n" + "\n".join(artifact_lines)
-            synthesis_context.append(f"\n=== ARTIFACTS CREATED ===")
-            synthesis_context.append("\n".join(artifact_details))
+            artifact_message = "\n\n**Artifacts Generated:**\n" + "\n".join(artifact_lines)
+            synthesis_context.append(f"\n=== ARTIFACTS ===\n" + "\n".join(artifact_details))
         
         if exp_results.get('stdout'):
-            synthesis_context.append(f"\n=== EXECUTION OUTPUT ===")
-            synthesis_context.append(exp_results.get('stdout', ''))
+            synthesis_context.append(f"\n=== OUTPUT ===\n{exp_results.get('stdout', '')}")
         
-        if exp_results.get('stderr') and not exp_results.get('success'):
-            synthesis_context.append(f"\n=== EXECUTION ERRORS ===")
-            synthesis_context.append(exp_results.get('stderr', ''))
+        if exp_results.get('stderr'):
+            synthesis_context.append(f"\n=== ERRORS ===\n{exp_results.get('stderr', '')}")
     
     full_context = "\n".join(synthesis_context)
     
-    # ENHANCED SYNTHESIS PROMPT
-    synthesis_prompt = f"""You are creating the FINAL RESPONSE after executing a user's request.
+    synthesis_prompt = f"""Create FINAL RESPONSE after executing user's request.
 
 {full_context}
 
-Create a comprehensive, professional response that:
-1. Directly addresses the user's original request
-2. Explains what was accomplished and HOW it was done
-3. References specific artifacts created and explains their PURPOSE
-4. Provides context on how to USE/ACCESS the deliverables
-5. Highlights KEY INSIGHTS or findings
+Create comprehensive response that:
+1. Directly addresses original request
+2. Explains what was accomplished and HOW
+3. References specific artifacts and explains PURPOSE
+4. Provides context on how to USE deliverables
+5. Highlights KEY INSIGHTS
 6. Suggests NEXT STEPS if relevant
 
-TONE: Professional, clear, helpful. Be SPECIFIC about what was created.
+Be SPECIFIC about what was created."""
 
-Write the final response:"""
-    
     response = llm.invoke(synthesis_prompt)
     final_text = getattr(response, "content", "") or ""
     
     if artifact_message:
         final_text = final_text + "\n\n---\n" + artifact_message
     
-    return {
-        "draftResponse": final_text, 
-        "execution_path": path, 
-        "status_update": "Final response synthesized."
-    }
+    return {"draftResponse": final_text, "execution_path": path, "status_update": "Response synthesized"}
 
 def run_qa_agent(state: AgentState):
-    log.info("--- QA AGENT ---")
-    path = ensure_list(state, 'execution_path') + ["QA Agent"]
+    log.info("--- QA ---")
+    path = ensure_list(state, 'execution_path') + ["QA"]
     
-    # Enhanced QA with context awareness
     qa_context = [
-        f"=== ORIGINAL REQUEST ===",
-        f"{state.get('userInput', '')}",
-        f"\n=== CORE OBJECTIVE ===",
-        f"{state.get('coreObjectivePrompt', '')}",
-        f"\n=== DRAFT RESPONSE ===",
-        f"{state.get('draftResponse', '')}",
+        f"=== REQUEST ===\n{state.get('userInput', '')}",
+        f"\n=== OBJECTIVE ===\n{state.get('coreObjectivePrompt', '')}",
+        f"\n=== DRAFT ===\n{state.get('draftResponse', '')}",
     ]
     
     if state.get('experimentResults'):
-        qa_context.append(f"\n=== ARTIFACTS ===")
-        qa_context.append(json.dumps(state.get('experimentResults', {}).get('paths', {}), indent=2))
+        qa_context.append(f"\n=== ARTIFACTS ===\n{json.dumps(state.get('experimentResults', {}).get('paths', {}), indent=2)}")
     
-    full_qa_context = "\n".join(qa_context)
-    
-    prompt = f"""Review the draft response against the core objective.
+    prompt = f"""Review draft against objective.
 
-{full_qa_context}
+{chr(10).join(qa_context)}
 
 Respond ONLY with:
-- 'APPROVED' if the response fully addresses the request with quality and specificity
-- OR provide SPECIFIC, ACTIONABLE feedback for improvement
+- 'APPROVED' if fully addresses request with quality
+- OR provide SPECIFIC, ACTIONABLE feedback"""
 
-Your response:"""
-    
     response = llm.invoke(prompt)
     content = getattr(response, "content", "") or ""
     
     if "APPROVED" in content.upper():
-        return {
-            "approved": True, 
-            "qaFeedback": None, 
-            "execution_path": path, 
-            "status_update": "Response approved!"
-        }
+        return {"approved": True, "qaFeedback": None, "execution_path": path, "status_update": "Approved"}
     else:
-        return {
-            "approved": False, 
-            "qaFeedback": content or "No specific feedback.", 
-            "execution_path": path, 
-            "status_update": "Response needs improvement. Reworking..."
-        }
+        return {"approved": False, "qaFeedback": content or "No feedback", "execution_path": path, "status_update": "Needs improvement"}
 
 def run_archivist_agent(state: AgentState):
-    log.info("--- ARCHIVIST AGENT ---")
-    path = ensure_list(state, 'execution_path') + ["Archivist Agent"]
-    
-    summary_prompt = (
-        f"Create a concise summary of this successful task for long-term memory.\n\n"
-        f"Core Objective: {state.get('coreObjectivePrompt')}\n\n"
-        f"Final Response: {state.get('draftResponse')}\n\n"
-        f"Memory Summary:"
-    )
+    log.info("--- ARCHIVIST ---")
+    path = ensure_list(state, 'execution_path') + ["Archivist"]
     
+    summary_prompt = f"Summarize for memory.\n\nObjective: {state.get('coreObjectivePrompt')}\n\nResponse: {state.get('draftResponse')}\n\nSummary:"
     response = llm.invoke(summary_prompt)
-    memory_manager.add_to_memory(
-        getattr(response,"content",""), 
-        {"objective": state.get('coreObjectivePrompt')}
-    )
+    memory_manager.add_to_memory(getattr(response,"content",""), {"objective": state.get('coreObjectivePrompt')})
     
-    return {
-        "execution_path": path, 
-        "status_update": "Saving key learnings..."
-    }
+    return {"execution_path": path, "status_update": "Saved to memory"}
 
 def run_disclaimer_agent(state: AgentState):
-    log.warning("--- DISCLAIMER AGENT ---")
-    path = ensure_list(state, 'execution_path') + ["Disclaimer Agent"]
-    
-    disclaimer = (
-        "**DISCLAIMER: Budget limit reached. The following is the best available draft and may be incomplete.**\n\n"
-        "---\n\n"
-    )
+    log.warning("--- DISCLAIMER ---")
+    path = ensure_list(state, 'execution_path') + ["Disclaimer"]
     
-    final_response = disclaimer + state.get('draftResponse', "No response was generated.")
+    disclaimer = "**DISCLAIMER: Budget limit reached. Draft may be incomplete.**\n\n---\n\n"
+    final_response = disclaimer + state.get('draftResponse', "No response")
     
-    return {
-        "draftResponse": final_response, 
-        "execution_path": path, 
-        "status_update": "Budget limit reached. Preparing final draft..."
-    }
+    return {"draftResponse": final_response, "execution_path": path, "status_update": "Budget limit"}
 
-# --- Decision functions ---
+# --- Decisions ---
 def should_continue(state: AgentState):
-    log.info("--- DECISION: QA approved? ---")
-    
     if state.get("approved"):
-        log.info("Routing to: Archivist (Success Path)")
         return "archivist_agent"
-    
     if ensure_int(state, "rework_cycles", 0) > ensure_int(state, "max_loops", 0):
-        log.error(f"BUDGET LIMIT REACHED after {ensure_int(state, 'rework_cycles', 0)-1} cycles.")
         return "disclaimer_agent"
-    else:
-        log.info("Routing to: PM Agent for rework")
-        return "pm_agent"
+    return "pm_agent"
 
 def should_run_experiment(state: AgentState):
     pm = state.get('pmPlan', {}) or {}