Update graph.py

graph.py

@@ -1,4 +1,5 @@
-# graph.py - Domain-Agnostic Production Implementation
+# graph.py - Enhanced with better loop control and cost tracking
+
 import json
 import re
 import math
@@ -8,6 +9,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
@@ -18,13 +20,40 @@ from logging_config import setup_logging, get_logger
 import nbformat
 from nbformat.v4 import new_notebook, new_markdown_cell, new_code_cell
 import pandas as pd
-from docx import Document as DocxDocument
+from docx import Document
 from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
 from reportlab.lib.styles import getSampleStyleSheet
 
+# --- Configuration ---
 OUT_DIR = os.environ.get("OUT_DIR", "/tmp")
 os.makedirs(OUT_DIR, exist_ok=True)
+EXPORTS_DIR = os.path.join(OUT_DIR, "exports")
+os.makedirs(EXPORTS_DIR, exist_ok=True)
+
+# --- Helpers ---
+def ensure_list(state, key):
+    v = state.get(key) if state else None
+    if v is None:
+        return []
+    if isinstance(v, list):
+        return v
+    if isinstance(v, tuple):
+        return list(v)
+    return [v]
+
+def ensure_int(state, key, default=0):
+    try:
+        v = state.get(key) if state else None
+        if v is None:
+            return default
+        return int(v)
+    except Exception:
+        return default
 
+def sanitize_path(path: str) -> str:
+    return os.path.abspath(path)
+
+# --- Setup ---
 setup_logging()
 log = get_logger(__name__)
 INITIAL_MAX_REWORK_CYCLES = 3
@@ -32,6 +61,7 @@ GPT4O_INPUT_COST_PER_1K_TOKENS = 0.005
 GPT4O_OUTPUT_COST_PER_1K_TOKENS = 0.015
 AVG_TOKENS_PER_CALL = 2.0
 
+# --- State ---
 class AgentState(TypedDict):
     userInput: str
     chatHistory: List[str]
@@ -47,22 +77,13 @@ class AgentState(TypedDict):
     rework_cycles: int
     max_loops: int
     status_update: str
+    # NEW: For real-time cost tracking
+    current_cost: float
+    budget_exceeded: bool
 
-llm = ChatOpenAI(model="gpt-4o", temperature=0.1, max_retries=3, request_timeout=60)
 
-def ensure_list(state, key):
-    v = state.get(key) if state else None
-    return [] if v is None else (list(v) if isinstance(v, (list, tuple)) else [v])
-
-def ensure_int(state, key, default=0):
-    try:
-        v = state.get(key) if state else None
-        return int(v) if v is not None else default
-    except Exception:
-        return default
-
-def sanitize_path(path: str) -> str:
-    return os.path.abspath(path)
+# --- 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]:
     try:
@@ -82,375 +103,149 @@ def parse_json_from_llm(llm_output: str) -> Optional[dict]:
         log.error(f"JSON parsing failed: {e}")
         return None
 
-KNOWN_ARTIFACT_TYPES = {"notebook", "excel", "word", "pdf", "repo", "script"}
+# --- Artifact detection ---
+KNOWN_ARTIFACT_TYPES = {"notebook","excel","word","pdf","image","repo","script"}
 
 def detect_requested_output_types(text: str) -> Dict:
-    """Detect artifact type from user request - domain agnostic."""
     if not text:
-        return {"requires_artifact": False, "artifact_type": None}
+        return {"requires_artifact": False, "artifact_type": None, "artifact_hint": None}
     t = text.lower()
-    
-    # Repository indicators
-    if any(k in t for k in ["repo", "repository", "application", "app", "codebase", "project", "package"]):
-        return {"requires_artifact": True, "artifact_type": "repo"}
-    
-    # Notebook indicators
-    if any(k in t for k in ["jupyter", "notebook", "ipynb", "analysis", "exploration"]):
-        return {"requires_artifact": True, "artifact_type": "notebook"}
-    
-    # Document indicators
-    if any(k in t for k in ["document", "report", "documentation", ".docx", "word"]):
-        return {"requires_artifact": True, "artifact_type": "word"}
-    
-    # Spreadsheet indicators
-    if any(k in t for k in ["spreadsheet", "excel", ".xlsx", "table", "csv"]):
-        return {"requires_artifact": True, "artifact_type": "excel"}
-    
-    # Script indicators
-    if any(k in t for k in ["script", ".py", "automation", "tool"]):
-        return {"requires_artifact": True, "artifact_type": "script"}
-    
-    return {"requires_artifact": False, "artifact_type": None}
+    if any(k in t for k in ["jupyter notebook", "jupyter", "notebook", "ipynb"]):
+        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 ["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:
     if not exp_type:
         detection = detect_requested_output_types(goal_text or "")
         return detection.get("artifact_type") or "word"
     s = exp_type.strip().lower()
-    return s if s in KNOWN_ARTIFACT_TYPES else "word"
-
-def generate_generic_sample_data(format_type: str, num_records: int = 5) -> str:
-    """Generate domain-neutral sample data."""
-    if format_type == 'csv':
-        rows = [f"{i},Item {i},{100 * i},Category {(i % 3) + 1}" for i in range(1, num_records + 1)]
-        return "id,name,value,category\n" + "\n".join(rows)
-    else:  # json
-        items = [{"id": i, "name": f"Item {i}", "value": 100 * i, "category": f"Category {(i % 3) + 1}"} for i in range(1, num_records + 1)]
-        return json.dumps({"data": items, "metadata": {"total": num_records, "generated": True}}, indent=2)
-
-def generate_default_readme(goal: str) -> str:
-    return f"""# Generated Application
-
-## Overview
-{goal}
-
-## Quick Start
-
-```bash
-# Install dependencies
-pip install -r requirements.txt
-
-# Run demonstration
-python main.py --demo
-
-# Run tests
-pytest tests/ -v
-```
-
-## Project Structure
-
-```
-.
-├── main.py              # Entry point
-├── src/                 # Source modules
-├── tests/               # Test suite
-├── data/                # Sample data
-├── docs/                # Documentation
-└── requirements.txt     # Dependencies
-```
-
-## Testing
-
-```bash
-pytest tests/ -v --cov=src
-```
-
-## Documentation
-
-See docs/ directory for detailed documentation.
-"""
-
-def generate_default_requirements() -> str:
-    return """python-dotenv>=1.0.0
-pytest>=7.0.0
-pytest-cov>=4.0.0
-"""
-
-def generate_default_main() -> str:
-    return """#!/usr/bin/env python3
-\"\"\"Main entry point.\"\"\"
-import argparse
-import logging
+    if s in KNOWN_ARTIFACT_TYPES:
+        return s
+    if "notebook" in s or "ipynb" in s:
+        return "notebook"
+    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 "backend" in s:
+        return "repo"
+    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"
+
+# --- 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)
+    code_blocks = re.findall(r"```python\s*(.*?)\s*```", llm_text, re.DOTALL)
+    if not code_blocks:
+        code_blocks = re.findall(r"```\s*(.*?)\s*```", llm_text, re.DOTALL)
+    md_parts = re.split(r"```(?:python)?\s*.*?\s*```", llm_text, flags=re.DOTALL)
+    nb = new_notebook()
+    cells = []
+    max_len = max(len(md_parts), len(code_blocks))
+    for i in range(max_len):
+        if i < len(md_parts) and md_parts[i].strip():
+            cells.append(new_markdown_cell(md_parts[i].strip()))
+        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 generated.")]
+    nb['cells'] = cells
+    uid = uuid.uuid4().hex[:10]
+    filename = os.path.join(out_dir, f"generated_notebook_{uid}.ipynb")
+    nbformat.write(nb, filename)
+    return filename
 
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+def write_script(code_text: str, language_hint: Optional[str]=None, out_dir: Optional[str]=None) -> str:
+    out_dir = out_dir or OUT_DIR
+    os.makedirs(out_dir, exist_ok=True)
+    ext = ".txt"
+    if language_hint:
+        l = language_hint.lower()
+        if "python" in l:
+            ext = ".py"
+        elif "r" in l:
+            ext = ".R"
+        elif "java" in l:
+            ext = ".java"
+        elif "javascript" in l:
+            ext = ".js"
+    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:
+        f.write(code_text)
+    return filename
 
-def run_demo():
-    \"\"\"Run demonstration.\"\"\"
-    logger.info("Running demo...")
-    logger.info("Demo completed successfully")
-    return 0
+def write_docx_from_text(text: str, out_dir: Optional[str]=None) -> str:
+    out_dir = out_dir or OUT_DIR
+    os.makedirs(out_dir, exist_ok=True)
+    doc = Document()
+    for para in [p.strip() for p in text.split("\n\n") if p.strip()]:
+        doc.add_paragraph(para)
+    uid = uuid.uuid4().hex[:10]
+    filename = os.path.join(out_dir, f"generated_doc_{uid}.docx")
+    doc.save(filename)
+    return filename
 
-def main():
-    parser = argparse.ArgumentParser(description="Application")
-    parser.add_argument('--demo', action='store_true', help='Run demo mode')
-    parser.add_argument('--input', type=str, help='Input file')
-    
-    args = parser.parse_args()
-    
+def write_excel_from_tables(maybe_table_text: str, out_dir: Optional[str]=None) -> str:
+    out_dir = out_dir or OUT_DIR
+    os.makedirs(out_dir, exist_ok=True)
+    uid = uuid.uuid4().hex[:10]
+    filename = os.path.join(out_dir, f"generated_excel_{uid}.xlsx")
     try:
-        if args.demo:
-            return run_demo()
-        elif args.input:
-            logger.info(f"Processing: {args.input}")
-            return 0
-        else:
-            parser.print_help()
-            return 1
+        try:
+            parsed = json.loads(maybe_table_text)
+            if isinstance(parsed, list):
+                df = pd.DataFrame(parsed)
+            elif isinstance(parsed, dict):
+                df = pd.DataFrame([parsed])
+            else:
+                df = pd.DataFrame({"content":[str(maybe_table_text)]})
+        except Exception:
+            if "," in maybe_table_text:
+                from io import StringIO
+                df = pd.read_csv(StringIO(maybe_table_text))
+            else:
+                df = pd.DataFrame({"content":[maybe_table_text]})
+        df.to_excel(filename, index=False, engine="openpyxl")
+        return filename
     except Exception as e:
-        logger.error(f"Error: {e}", exc_info=True)
-        return 1
-
-if __name__ == "__main__":
-    exit(main())
-"""
-
-def generate_default_test() -> str:
-    return """\"\"\"Test suite.\"\"\"
-import pytest
-from pathlib import Path
-
-def test_main_imports():
-    \"\"\"Verify main module imports.\"\"\"
-    import main
-    assert hasattr(main, 'main')
-
-def test_sample_data_exists():
-    \"\"\"Verify sample data present.\"\"\"
-    data_path = Path("data/sample_data.csv")
-    assert data_path.exists(), "Sample data file missing"
-
-def test_requirements_file():
-    \"\"\"Verify requirements.txt exists.\"\"\"
-    assert Path("requirements.txt").exists()
-"""
-
-def generate_architecture_doc(goal: str) -> str:
-    return f"""# Architecture Documentation
-
-## System Overview
-
-**Purpose:** {goal}
-
-## Design Principles
-
-1. **Modularity**: Clear separation of concerns
-2. **Testability**: Comprehensive test coverage
-3. **Local-First**: No external dependencies for demo
-4. **Documentation**: Self-documenting code
-
-## Technology Stack
-
-- Python 3.8+
-- pytest for testing
-- Standard library focused
-
-## Component Structure
-
-Components are organized by responsibility with clear interfaces between modules.
-
-## Testing Strategy
-
-- Unit tests for individual functions
-- Integration tests for workflows
-- Sample data for validation
-"""
-
-def generate_evaluation_guide() -> str:
-    return """# Evaluation Guide
-
-## Quick Evaluation (5 minutes)
-
-### Step 1: Install
-```bash
-pip install -r requirements.txt
-```
-
-### Step 2: Run Demo
-```bash
-python main.py --demo
-```
-
-### Step 3: Run Tests
-```bash
-pytest tests/ -v
-```
-
-## Quality Checklist
-
-- [ ] All dependencies install successfully
-- [ ] Demo runs without errors
-- [ ] All tests pass
-- [ ] Sample data is present
-- [ ] Documentation is complete
-
-## Expected Results
-
-- Demo should complete successfully
-- Tests should show 100% pass rate
-- Output files should be generated
-"""
-
-def generate_caveats_doc() -> str:
-    return """# Caveats and Limitations
-
-## General Limitations
-
-- Sample data for demonstration only
-- Not optimized for production scale
-- Manual review recommended for critical use
-
-## Scalability Considerations
-
-Current implementation suitable for:
-- Development and testing
-- Small to medium datasets
-- Single-machine execution
-
-For production scale:
-- Consider distributed processing
-- Implement proper monitoring
-- Add comprehensive error handling
+        log.error(f"Excel creation failed: {e}")
+        return write_docx_from_text(f"Excel error: {e}\n\n{maybe_table_text}", out_dir=out_dir)
 
-## Best Practices
-
-1. Test on small dataset first
-2. Implement proper logging
-3. Regular quality audits
-4. Maintain documentation
-"""
-
-def validate_repository_quality(repo_files: Dict[str, str]) -> Dict:
-    """Domain-agnostic quality validation."""
-    score = 0
-    issues = []
-    warnings = []
-    successes = []
-    
-    # Critical files (30 points)
-    critical = {'README.md': 10, 'requirements.txt': 8, 'main.py': 12}
-    for filename, points in critical.items():
-        if any(filename in path for path in repo_files.keys()):
-            score += points
-            successes.append(f"✓ {filename}")
-        else:
-            issues.append(f"Missing {filename} (-{points})")
-    
-    # Documentation (20 points)
-    doc_files = {'ARCHITECTURE.md': 7, 'EVALUATION_GUIDE.md': 7, 'CAVEATS.md': 6}
-    for filename, points in doc_files.items():
-        if any(filename in path for path in repo_files.keys()):
-            score += points
-            successes.append(f"✓ {filename}")
-        else:
-            warnings.append(f"Missing {filename}")
-    
-    # Sample data (15 points)
-    has_data = any('data/' in f or 'sample' in f.lower() for f in repo_files.keys())
-    if has_data:
-        score += 15
-        successes.append("✓ Sample data")
-    else:
-        issues.append("No sample data (-15)")
-    
-    # Tests (15 points)
-    test_files = [f for f in repo_files.keys() if 'test' in f.lower()]
-    if test_files:
-        score += 15
-        successes.append(f"✓ {len(test_files)} test(s)")
-    else:
-        issues.append("No tests (-15)")
-    
-    # Code quality (20 points)
-    all_code = "\n".join(str(v) for v in repo_files.values())
-    
-    if 'def ' in all_code or 'class ' in all_code:
-        score += 5
-    if '"""' in all_code:
-        score += 3
-    if 'logging' in all_code.lower():
-        score += 3
-    if 'try:' in all_code and 'except' in all_code:
-        score += 4
-    if '->' in all_code:
-        score += 3
-    if '.env' in all_code.lower() or 'dotenv' in all_code.lower():
-        score += 2
-    
-    # Anti-patterns (penalties)
-    if '# TODO' in all_code or '#TODO' in all_code:
-        penalty = min(10, all_code.count('# TODO') * 2)
-        score -= penalty
-        issues.append(f"TODO placeholders (-{penalty})")
-    
-    if 'example.com' in all_code.lower():
-        score -= 5
-        issues.append("Dummy URLs (-5)")
-    
-    if 'pass  # implement' in all_code:
-        score -= 5
-        issues.append("Placeholder code (-5)")
-    
-    # Wrong file types
-    doc_in_py = [f for f in repo_files.keys() if f.endswith('.py') and any(
-        kw in f.lower() for kw in ['readme', 'doc', 'guide', 'architecture']
-    )]
-    if doc_in_py:
-        score -= 5
-        issues.append(f"Docs in .py files (-5)")
-    
-    final_score = max(0, min(100, score))
-    
-    if final_score >= 90:
-        grade = "A"
-    elif final_score >= 80:
-        grade = "B"
-    elif final_score >= 70:
-        grade = "C"
-    elif final_score >= 60:
-        grade = "D"
-    else:
-        grade = "F"
-    
-    return {
-        "score": final_score,
-        "grade": grade,
-        "issues": issues,
-        "warnings": warnings,
-        "successes": successes,
-        "file_count": len(repo_files),
-        "has_tests": bool(test_files),
-        "has_data": has_data
-    }
-
-def extract_files_from_llm(llm_text: str) -> Dict[str, str]:
-    """Extract files from LLM response."""
-    files = {}
-    
-    # Pattern: ### filename
-    for match in re.finditer(r"###\s+([\w\/_\-\.]+)\s*\n```(?:\w+)?\s*\n(.*?)\n```", llm_text, re.DOTALL):
-        files[match.group(1).strip()] = match.group(2).strip()
-    
-    # Pattern: **filename**
-    for match in re.finditer(r"\*\*([\w\/_\-\.]+)\*\*\s*\n```(?:\w+)?\s*\n(.*?)\n```", llm_text, re.DOTALL):
-        fp = match.group(1).strip()
-        if fp not in files:
-            files[fp] = match.group(2).strip()
-    
-    return files
+def write_pdf_from_text(text: str, out_dir: Optional[str]=None) -> str:
+    out_dir = out_dir or OUT_DIR
+    os.makedirs(out_dir, exist_ok=True)
+    uid = uuid.uuid4().hex[:10]
+    filename = os.path.join(out_dir, f"generated_doc_{uid}.pdf")
+    try:
+        doc = SimpleDocTemplate(filename)
+        styles = getSampleStyleSheet()
+        flowables = []
+        for para in [p.strip() for p in text.split("\n\n") if p.strip()]:
+            flowables.append(Paragraph(para.replace("\n","<br/>"), styles["Normal"]))
+            flowables.append(Spacer(1, 8))
+        doc.build(flowables)
+        return filename
+    except Exception as e:
+        log.error(f"PDF creation failed: {e}")
+        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:
-    """Build repository ZIP file."""
     out_dir = out_dir or OUT_DIR
     os.makedirs(out_dir, exist_ok=True)
     uid = uuid.uuid4().hex[:8]
@@ -460,8 +255,12 @@ def build_repo_zip(files_map: Dict[str,str], repo_name: str="generated_app", out
     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)
-        with open(dest, "w", encoding="utf-8") as fh:
-            fh.write(str(content))
+        
+        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:
@@ -473,93 +272,91 @@ def build_repo_zip(files_map: Dict[str,str], repo_name: str="generated_app", out
     
     return zip_path
 
-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)
-    nb = new_notebook()
-    nb['cells'] = [new_markdown_cell("# Generated Notebook\n\nContent generated by AI system.")]
-    uid = uuid.uuid4().hex[:10]
-    filename = os.path.join(out_dir, f"notebook_{uid}.ipynb")
-    nbformat.write(nb, filename)
-    return filename
-
-def write_docx_from_text(text: str, out_dir: Optional[str]=None) -> str:
-    out_dir = out_dir or OUT_DIR
-    os.makedirs(out_dir, exist_ok=True)
-    doc = DocxDocument()
-    for para in [p.strip() for p in text.split("\n\n") if p.strip()]:
-        doc.add_paragraph(para)
-    uid = uuid.uuid4().hex[:10]
-    filename = os.path.join(out_dir, f"document_{uid}.docx")
-    doc.save(filename)
-    return filename
-
-# ===================================================================
-# AGENTS
-# ===================================================================
-
+# --- Nodes ---
 def run_triage_agent(state: AgentState):
     log.info("--- TRIAGE ---")
-    prompt = f"Is this a greeting or task request? '{state.get('userInput','')}' Respond: 'greeting' or 'task'"
+    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():
-        return {"draftResponse": "Hello! How can I assist you today?", "execution_path": ["Triage"], "status_update": "Greeting"}
-    return {"execution_path": ["Triage"], "status_update": "Task"}
+        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 ---")
     path = ensure_list(state, 'execution_path') + ["Planner"]
-    prompt = f"Create execution plan for: '{state.get('userInput','')}'. JSON with 'plan' (list), 'estimated_llm_calls_per_loop' (int)"
+    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": "Planning failed"}, "execution_path": path}
+        return {"pmPlan": {"error": "Planning failed"}, "execution_path": path, "status_update": "Error"}
     
     calls = plan_data.get('estimated_llm_calls_per_loop', 3)
-    cost = calls * AVG_TOKENS_PER_CALL * (GPT4O_INPUT_COST_PER_1K_TOKENS + GPT4O_OUTPUT_COST_PER_1K_TOKENS) / 2
-    plan_data['estimated_cost_usd'] = round(cost * (INITIAL_MAX_REWORK_CYCLES + 1), 2)
+    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(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',''))
     if detection.get('requires_artifact'):
-        plan_data['experiment_needed'] = True
-        plan_data['experiment_type'] = detection.get('artifact_type')
+        plan_data.setdefault('experiment_needed', True)
+        plan_data.setdefault('experiment_type', detection.get('artifact_type'))
+        plan_data.setdefault('experiment_goal', state.get('userInput',''))
     
-    return {"pmPlan": plan_data, "execution_path": path, "status_update": "Planned"}
+    return {"pmPlan": plan_data, "execution_path": path, "status_update": "Plan created"}
 
 def run_memory_retrieval(state: AgentState):
     log.info("--- MEMORY ---")
     path = ensure_list(state, 'execution_path') + ["Memory"]
     mems = memory_manager.retrieve_relevant_memories(state.get('userInput',''))
-    context = "\n".join([m.page_content for m in mems]) if mems else "No prior context"
-    return {"retrievedMemory": context, "execution_path": path, "status_update": "Retrieved"}
+    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 ---")
     path = ensure_list(state, 'execution_path') + ["Intent"]
-    prompt = f"Clarify core objective.\n\nContext: {state.get('retrievedMemory', '')}\n\nRequest: {state.get('userInput','')}\n\nObjective:"
+    prompt = f"Refine into clear objective.\n\nMemory: {state.get('retrievedMemory')}\n\nRequest: {state.get('userInput','')}\n\nCore Objective:"
     response = llm.invoke(prompt)
-    return {"coreObjectivePrompt": getattr(response, "content", ""), "execution_path": path, "status_update": "Clarified"}
+    core_obj = getattr(response, "content", "") or ""
+    return {"coreObjectivePrompt": core_obj, "execution_path": path, "status_update": "Objective clarified"}
 
 def run_pm_agent(state: AgentState):
     log.info("--- PM ---")
     current_cycles = ensure_int(state, 'rework_cycles', 0) + 1
+    max_loops_val = ensure_int(state, 'max_loops', 0)
     path = ensure_list(state, 'execution_path') + ["PM"]
     
-    context = f"User Request: {state.get('userInput', '')}\n\nObjective: {state.get('coreObjectivePrompt', '')}"
+    context_parts = [
+        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 += f"\n\nQA Feedback: {state.get('qaFeedback')}"
+        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)
     
-    prompt = f"""{context}
+    prompt = f"""Create DETAILED, EXECUTABLE plan.
+
+{full_context}
 
-Create detailed execution plan. JSON format:
+Each step must be SPECIFIC and ACTIONABLE:
+- State EXACTLY what will be created/analyzed
+- Specify WHAT data/information will be used
+- Define WHAT methods will be applied
+
+JSON format:
 {{
-    "plan_steps": ["Specific step 1", "Specific step 2"],
+    "plan_steps": ["Specific step 1...", "Specific step 2..."],
     "experiment_needed": true/false,
-    "experiment_type": "repo|notebook|word|excel|script",
-    "experiment_goal": "What to create"
-}}"""
+    "experiment_type": "notebook|script|excel|word|pdf|repo",
+    "experiment_goal": "Detailed artifact description",
+    "key_requirements": ["Critical requirements"]
+}}
+
+Be specific about using uploaded files, implementing algorithms, creating schemas."""
     
     response = llm.invoke(prompt)
     plan = parse_json_from_llm(getattr(response, "content", "") or "")
@@ -567,14 +364,26 @@ Create detailed execution plan. JSON format:
     if not plan:
         detection = detect_requested_output_types(state.get('userInput', ''))
         plan = {
-            "plan_steps": ["Analyze requirements", "Implement solution", "Validate output"],
+            "plan_steps": ["Analyze request", "Process information", "Create deliverable", "Review"],
             "experiment_needed": detection.get('requires_artifact', False),
-            "experiment_type": detection.get('artifact_type', 'word')
+            "experiment_type": detection.get('artifact_type', 'word'),
+            "experiment_goal": state.get('coreObjectivePrompt', state.get('userInput', ''))
         }
     
-    plan['experiment_type'] = normalize_experiment_type(plan.get('experiment_type'), plan.get('experiment_goal', ''))
+    exp_type = normalize_experiment_type(plan.get('experiment_type'), plan.get('experiment_goal',''))
+    plan['experiment_type'] = exp_type
     
-    return {"pmPlan": plan, "execution_path": path, "rework_cycles": current_cycles, "status_update": "Planned"}
+    if plan.get('experiment_needed') and not plan.get('experiment_goal'):
+        plan['experiment_goal'] = state.get('userInput','')
+    
+    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
+    return re.findall(r"```(?:\w+)?\s*(.*?)\s*```", text, re.DOTALL)
 
 def run_experimenter_agent(state: AgentState):
     log.info("--- EXPERIMENTER ---")
@@ -582,140 +391,400 @@ def run_experimenter_agent(state: AgentState):
     pm = state.get('pmPlan', {}) or {}
     
     if not pm.get('experiment_needed'):
-        return {"experimentCode": None, "experimentResults": None, "execution_path": path, "status_update": "Skipped"}
+        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', state.get('userInput', ''))
+    goal = pm.get('experiment_goal', 'No goal')
+    
+    # BUILD RICH CONTEXT
+    context_parts = [
+        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=== CONTEXT ===\n{state.get('retrievedMemory', '')}")
     
+    if state.get('qaFeedback'):
+        context_parts.append(f"\n=== FEEDBACK TO ADDRESS ===\n{state.get('qaFeedback', '')}")
+    
+    full_context = "\n".join(context_parts)
+    
+    # REPO REQUIRES SPECIAL HANDLING
     if exp_type == 'repo':
-        repo_prompt = f"""Create complete, functional repository.
-
-Goal: {goal}
-
-Requirements:
-- README.md with usage instructions
-- requirements.txt with dependencies
-- main.py entry point with --demo flag
-- tests/ directory with pytest tests
-- data/ directory with sample data
-- docs/ with ARCHITECTURE.md, EVALUATION_GUIDE.md, CAVEATS.md
-- .env.example for configuration
-- NO TODO comments or placeholders
-- NO example.com URLs
-- Documentation in .md files (not .py)
-
-Format each file:
-### path/to/file.ext
-```language
-[complete code]
-```"""
-
-        response = llm.invoke(repo_prompt)
-        llm_text = getattr(response, "content", "") or ""
+        repo_prompt = f"""Create COMPLETE, PRODUCTION-READY application repository.
+
+{full_context}
+
+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]
+   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
+
+    CODE QUALITY:
+
+    Environment variables for secrets
+
+    Error handling and logging
+
+    Docstrings and comments
+
+    Real business logic based on request
+
+    RUNNABLE out of the box
+
+    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
+    # 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}
+
+                                        Files
+                                        {chr(10).join(f'- {f}' for f in sorted(repo_files.keys()))}
+
+                                        Setup
+                                        pip install -r requirements.txt
+
+                                        Copy .env.example to .env and configure
+
+                                        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:
+
+    Use ALL specific details from request
+
+    PRODUCTION-READY, COMPLETE content (NO placeholders)
+
+    ACTUAL data, REALISTIC examples, WORKING code
+
+    For notebooks: markdown + executable code + visualizations
+
+    For scripts: error handling + docs + real logic
+
+    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 ""
+    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)
+            results.update({"success": True, "paths": {"notebook": sanitize_path(nb_path)}})
+            return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": "Notebook created"}
         
-        repo_files = extract_files_from_llm(llm_text)
+        elif exp_type == 'excel':
+            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": "Excel created"}
         
-        # Ensure critical files
-        if not any('README' in f.upper() for f in repo_files.keys()):
-            repo_files['README.md'] = generate_default_readme(goal)
-        if 'requirements.txt' not in repo_files:
-            repo_files['requirements.txt'] = generate_default_requirements()
-        if not any('main.py' in f for f in repo_files.keys()):
-            repo_files['main.py'] = generate_default_main()
-        if not any('test' in f.lower() for f in repo_files.keys()):
-            repo_files['tests/test_main.py'] = generate_default_test()
-        if not any('data/' in f or 'sample' in f.lower() for f in repo_files.keys()):
-            repo_files['data/sample_data.csv'] = generate_generic_sample_data('csv')
-        if not any('ARCHITECTURE' in f.upper() for f in repo_files.keys()):
-            repo_files['docs/ARCHITECTURE.md'] = generate_architecture_doc(goal)
-        if not any('EVALUATION' in f.upper() for f in repo_files.keys()):
-            repo_files['docs/EVALUATION_GUIDE.md'] = generate_evaluation_guide()
-        if not any('CAVEAT' in f.upper() for f in repo_files.keys()):
-            repo_files['docs/CAVEATS.md'] = generate_caveats_doc()
+        elif exp_type == 'word':
+            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": "DOCX created"}
         
-        zip_path = build_repo_zip(repo_files)
-        validation = validate_repository_quality(repo_files)
+        elif exp_type == 'pdf':
+            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": "PDF created"}
         
-        results = {
-            "success": True,
-            "paths": {"repo_zip": sanitize_path(zip_path)},
-            "files_created": len(repo_files),
-            "validation": validation
-        }
+        elif exp_type == 'script':
+            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
         
-        return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": f"Repository ({validation['score']}/100)"}
-    
-    elif exp_type == 'notebook':
-        nb_path = write_notebook_from_text("", out_dir=OUT_DIR)
-        return {"experimentCode": None, "experimentResults": {"success": True, "paths": {"notebook": sanitize_path(nb_path)}}, "execution_path": path, "status_update": "Notebook created"}
+            script_path = write_script(code_text, language_hint=lang_hint, out_dir=OUT_DIR)
+            exec_results = {}
+        
+            if script_path.endswith(".py"):
+                try:
+                    exec_results = execute_python_code(code_text)
+                except Exception as e:
+                    exec_results = {"stdout":"","stderr":str(e),"success":False}
+        
+            results.update({
+                "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": "Script created"}
     
-    else:
-        doc_path = write_docx_from_text("Generated content", out_dir=OUT_DIR)
-        return {"experimentCode": None, "experimentResults": {"success": True, "paths": {"docx": sanitize_path(doc_path)}}, "execution_path": path, "status_update": "Document created"}
-
+        else:
+            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": "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"}
 def run_synthesis_agent(state: AgentState):
     log.info("--- SYNTHESIS ---")
-    path = ensure_list(state, 'execution_path') + ["Synthesis"]
+    _state = state or {}
+    path = ensure_list(_state, 'execution_path') + ["Synthesis"]
+
+    exp_results = _state.get('experimentResults')
+    pm_plan = _state.get('pmPlan', {}) or {}
+
+    synthesis_context = [
+        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)}",
+    ]
+
+    artifact_details = []
+    artifact_message = ""
+
+    if exp_results and isinstance(exp_results, dict):
+        paths = exp_results.get("paths") or {}
     
-    exp_results = state.get('experimentResults', {})
-    validation = exp_results.get('validation', {})
+        if paths:
+            artifact_lines = []
+            for artifact_type, artifact_path in paths.items():
+                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 ===\n" + "\n".join(artifact_details))
     
-    context = f"Request: {state.get('userInput', '')}\n\nObjective: {state.get('coreObjectivePrompt', '')}"
-    if validation:
-        context += f"\n\nQuality: {validation.get('score', 0)}/100 (Grade {validation.get('grade', 'N/A')})"
+        if exp_results.get('stdout'):
+            synthesis_context.append(f"\n=== OUTPUT ===\n{exp_results.get('stdout', '')}")
     
-    prompt = f"""{context}
+        if exp_results.get('stderr'):
+            synthesis_context.append(f"\n=== ERRORS ===\n{exp_results.get('stderr', '')}")
 
-Create comprehensive final response explaining what was delivered and how to use it."""
-    
-    response = llm.invoke(prompt)
+    full_context = "\n".join(synthesis_context)
+
+    synthesis_prompt = f"""Create FINAL RESPONSE after executing user's request.
+    {full_context}
+
+    Create comprehensive response that:
+
+    Directly addresses original request
+
+    Explains what was accomplished and HOW
+
+    References specific artifacts and explains PURPOSE
+
+    Provides context on how to USE deliverables
+
+    Highlights KEY INSIGHTS
+
+    Suggests NEXT STEPS if relevant
+
+    Be SPECIFIC about what was created."""
+
+    response = llm.invoke(synthesis_prompt)
     final_text = getattr(response, "content", "") or ""
-    
-    # Add artifacts
-    if exp_results and exp_results.get("paths"):
-        artifacts = [f"- {k.title()}: {os.path.basename(v)}" for k, v in exp_results["paths"].items()]
-        final_text += "\n\n**Generated Artifacts:**\n" + "\n".join(artifacts)
-    
-    # Add quality note
-    if validation.get('score', 100) < 70:
-        final_text += f"\n\n⚠️ Quality: {validation['score']}/100 - Review recommended"
-    
-    return {"draftResponse": final_text, "execution_path": path, "status_update": "Synthesized"}
+
+    if artifact_message:
+        final_text = final_text + "\n\n---\n" + artifact_message
+
+    return {"draftResponse": final_text, "execution_path": path, "status_update": "Response synthesized"}
 
 def run_qa_agent(state: AgentState):
     log.info("--- QA ---")
     path = ensure_list(state, 'execution_path') + ["QA"]
-    
-    exp_results = state.get('experimentResults', {})
-    validation = exp_results.get('validation', {})
-    
-    if validation.get('score', 100) < 70:
-        issues_str = ', '.join(validation.get('issues', [])[:3])
-        feedback = f"Quality insufficient ({validation['score']}/100). Issues: {issues_str}"
-        return {"approved": False, "qaFeedback": feedback, "execution_path": path, "status_update": "Failed QA"}
-    
-    prompt = f"Review response: {state.get('draftResponse', '')[:300]}... Respond 'APPROVED' or provide feedback"
+
+    qa_context = [
+        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 ===\n{json.dumps(state.get('experimentResults', {}).get('paths', {}), indent=2)}")
+
+    # MODIFIED PROMPT: Encourage convergence
+    prompt = f"""You are a QA reviewer. Review the draft response against the user's objective.
+    {chr(10).join(qa_context)}
+
+    Review Instructions:
+
+    Does the draft and its artifacts COMPLETELY satisfy ALL parts of the user's request?
+
+    Is the quality of the work high?
+
+    If this is a re-submission (rework cycle > 1), has the previous feedback been successfully addressed?
+
+    Response Format:
+
+    If the work is complete and high-quality, respond ONLY with the word 'APPROVED'.
+
+    Otherwise, provide SPECIFIC, ACTIONABLE, and NOVEL feedback on what must be changed. Do not repeat previous feedback if it has already been actioned.
+    """
+
     response = llm.invoke(prompt)
     content = getattr(response, "content", "") or ""
-    
+
     if "APPROVED" in content.upper():
         return {"approved": True, "qaFeedback": None, "execution_path": path, "status_update": "Approved"}
-    return {"approved": False, "qaFeedback": content, "execution_path": path, "status_update": "Revision needed"}
+    else:
+    # Sanitize feedback to ensure it's a useful string
+        feedback = content.replace("APPROVED", "").strip()
+        if not feedback:
+            feedback = "General quality improvements required."
+            return {"approved": False, "qaFeedback": feedback, "execution_path": path, "status_update": "Needs improvement"}
 
 def run_archivist_agent(state: AgentState):
     log.info("--- ARCHIVIST ---")
     path = ensure_list(state, 'execution_path') + ["Archivist"]
-    summary = state.get('coreObjectivePrompt', '')[:200]
-    memory_manager.add_to_memory(summary, {"objective": state.get('coreObjectivePrompt')})
-    return {"execution_path": path, "status_update": "Archived"}
+
+    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')})
+
+    return {"execution_path": path, "status_update": "Saved to memory"}
 
 def run_disclaimer_agent(state: AgentState):
     log.warning("--- DISCLAIMER ---")
     path = ensure_list(state, 'execution_path') + ["Disclaimer"]
-    disclaimer = "**Budget limit reached. Response may be incomplete.**\n\n"
-    return {"draftResponse": disclaimer + state.get('draftResponse', ""), "execution_path": path, "status_update": "Limited"}
+
+    reason = "Budget limit reached." if state.get('budget_exceeded') else "Rework limit reached."
+    disclaimer = f"**DISCLAIMER: {reason} 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": reason}
 
 def should_continue(state: AgentState):
+    # NEW: Check for budget excess first
+    if state.get("budget_exceeded"):
+        return "disclaimer_agent"
     if state.get("approved"):
         return "archivist_agent"
     if ensure_int(state, "rework_cycles", 0) > ensure_int(state, "max_loops", 0):
@@ -726,10 +795,7 @@ def should_run_experiment(state: AgentState):
     pm = state.get('pmPlan', {}) or {}
     return "experimenter_agent" if pm.get('experiment_needed') else "synthesis_agent"
 
-# ===================================================================
-# WORKFLOW GRAPHS
-# ===================================================================
-
+#--- Build graphs ---
 triage_workflow = StateGraph(AgentState)
 triage_workflow.add_node("triage", run_triage_agent)
 triage_workflow.set_entry_point("triage")
@@ -762,9 +828,11 @@ main_workflow.add_edge("disclaimer_agent", END)
 
 main_workflow.add_conditional_edges("pm_agent", should_run_experiment)
 main_workflow.add_conditional_edges("qa_agent", should_continue, {
-    "archivist_agent": "archivist_agent",
-    "pm_agent": "pm_agent",
-    "disclaimer_agent": "disclaimer_agent"
+"archivist_agent": "archivist_agent",
+"pm_agent": "pm_agent",
+"disclaimer_agent": "disclaimer_agent"
 })
 
-main_app = main_workflow.compile()
+main_app = main_workflow.compile()
+
+