Update graph.py

graph.py

@@ -1,7 +1,11 @@
-# graph.py (patched)
+# graph.py (patched: artifact generation using nbformat, python-docx, pandas/openpyxl, reportlab)
 import json
 import re
 import math
+import os
+import uuid
+import shutil
+import zipfile
 from typing import TypedDict, List, Dict, Optional
 from langchain_openai import ChatOpenAI
 from langgraph.graph import StateGraph, END
@@ -9,21 +13,26 @@ from memory_manager import memory_manager
 from code_executor import execute_python_code
 from logging_config import setup_logging, get_logger
 
+# External artifact libs
+import nbformat
+from nbformat.v4 import new_notebook, new_markdown_cell, new_code_cell
+import pandas as pd
+from docx import Document
+from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
+from reportlab.lib.styles import getSampleStyleSheet
+
+# --- Helpers ---
 def ensure_list(state, key):
-    """Return a list from state[key], default [] if missing/None/not-list."""
     v = state.get(key) if state else None
     if v is None:
         return []
     if isinstance(v, list):
         return v
-    # if it's a tuple, convert to list
     if isinstance(v, tuple):
         return list(v)
-    # fallback: wrap single scalar value
     return [v]
 
 def ensure_int(state, key, default=0):
-    """Return an int-like value from state[key], default if missing/None/not-int."""
     try:
         v = state.get(key) if state else None
         if v is None:
@@ -32,7 +41,11 @@ def ensure_int(state, key, default=0):
     except Exception:
         return default
 
-# --- Setup Logging & Constants ---
+def sanitize_path(path: str) -> str:
+    # On HF Spaces you may want to move to a served directory. Keep as-is here.
+    return path
+
+# --- Logging & constants ---
 setup_logging()
 log = get_logger(__name__)
 INITIAL_MAX_REWORK_CYCLES = 3
@@ -40,7 +53,7 @@ GPT4O_INPUT_COST_PER_1K_TOKENS = 0.005
 GPT4O_OUTPUT_COST_PER_1K_TOKENS = 0.015
 AVG_TOKENS_PER_CALL = 2.0
 
-# --- Agent State Definition ---
+# --- AgentState ---
 class AgentState(TypedDict):
     userInput: str
     chatHistory: List[str]
@@ -57,74 +70,225 @@ class AgentState(TypedDict):
     max_loops: int
     status_update: str
 
-# --- Helper for Robust JSON Parsing ---
+# --- LLM & parsing ---
+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:
         match = re.search(r"```json\n({.*?})\n```", llm_output, re.DOTALL)
         if match:
             json_str = match.group(1)
         else:
-            json_str = llm_output[llm_output.find('{'):llm_output.rfind('}')+1]
+            start = llm_output.find('{')
+            end = llm_output.rfind('}')
+            if start == -1 or end == -1:
+                return None
+            json_str = llm_output[start:end+1]
         return json.loads(json_str)
-    except (json.JSONDecodeError, AttributeError) as e:
-        log.error(f"JSON parsing failed. Raw output: '{llm_output}'. Error: {e}")
+    except Exception as e:
+        log.error(f"JSON parsing failed. Error: {e}. Raw: {llm_output[:300]}")
         return None
 
-# --- LLM Initialization ---
-llm = ChatOpenAI(model="gpt-4o", temperature=0.1, max_retries=3, request_timeout=60)
+# --- Artifact detection ---
+def detect_requested_output_types(text: str) -> Dict:
+    if not text:
+        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)"}
+    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)"}
+    # scripts for languages
+    if any(k in t for k in [".py", "python script", "r script", ".R", ".r", "java", ".java", "javascript", ".js"]):
+        # heuristic: choose 'script' and later infer language
+        return {"requires_artifact": True, "artifact_type": "script", "artifact_hint": "language script (py/r/java/js/etc.)"}
+    return {"requires_artifact": False, "artifact_type": None, "artifact_hint": None}
+
+# --- Notebook & artifact builders ---
+def write_notebook_from_text(llm_text: str, out_dir: str="/tmp") -> str:
+    """
+    Build a notebook via nbformat from llm_text using fenced python code blocks as code cells and other text as markdown.
+    """
+    code_blocks = re.findall(r"```python\n(.*?)\n```", llm_text, re.DOTALL)
+    # fallback to any fenced blocks
+    if not code_blocks:
+        code_blocks = re.findall(r"```\n(.*?)\n```", llm_text, re.DOTALL)
+    # split markdown by removing code blocks
+    md_parts = re.split(r"```(?:python)?\n.*?\n```", 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 parsed from LLM output.")]
+    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
+
+def write_script(code_text: str, language_hint: Optional[str]=None, out_dir: str="/tmp") -> str:
+    ext = ".txt"
+    if language_hint:
+        l = language_hint.lower()
+        if "python" in l or ".py" in l:
+            ext = ".py"
+        elif l in ("r", ".r"):
+            ext = ".R"
+        elif "java" in l or ".java" in l:
+            ext = ".java"
+        elif "javascript" in l or "node" in l or ".js" 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:
+        f.write(code_text)
+    return filename
+
+def write_docx_from_text(text: str, out_dir: str="/tmp") -> str:
+    doc = Document()
+    # naive: split into paragraphs on double-newline
+    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 write_excel_from_tables(maybe_table_text: str, out_dir: str="/tmp") -> str:
+    """
+    Heuristic: If LLM returns a JSON-convertible table or CSV snippet, attempt to form a DataFrame.
+    Otherwise write a small DataFrame with the provided text.
+    """
+    uid = uuid.uuid4().hex[:10]
+    filename = os.path.join(out_dir, f"generated_excel_{uid}.xlsx")
+    try:
+        # try JSON parse
+        parsed = None
+        try:
+            parsed = json.loads(maybe_table_text)
+            # if parsed is list of dicts
+            if isinstance(parsed, list):
+                df = pd.DataFrame(parsed)
+            elif isinstance(parsed, dict):
+                # dict of lists or single mapping
+                df = pd.DataFrame([parsed])
+            else:
+                df = pd.DataFrame({"content":[str(maybe_table_text)]})
+        except Exception:
+            # fallback: look for CSV text
+            if "," in maybe_table_text or "\t" 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:
+        log.error(f"Excel creation failed: {e}")
+        # write fallback docx with text
+        return write_docx_from_text(f"Failed to create excel. Error: {e}\n\nOriginal:\n{maybe_table_text}", out_dir=out_dir)
+
+def write_pdf_from_text(text: str, out_dir: str="/tmp") -> str:
+    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}")
+        # fallback to docx
+        return write_docx_from_text(f"Failed to create PDF. Error: {e}\n\nOriginal:\n{text}", out_dir=out_dir)
+
+def build_repo_zip(files_map: Dict[str,str], repo_name: str="generated_app", out_dir: str="/tmp") -> str:
+    """
+    files_map: dict of relative path -> absolute local file path/content.
+    If the value is a string and exists as a path, include file. If not a path, create a file with that content.
+    """
+    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:
+            # treat content as file content
+            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):
+            for f in files:
+                full = os.path.join(root, f)
+                arc = os.path.relpath(full, repo_dir)
+                zf.write(full, arc)
+    return zip_path
+
+# --- Node functions (triage/planner/memory/intent/pm/experimenter/synthesis/qa/archivist/disclaimer) ---
+# For brevity reuse earlier implementations but with artifact creation in experimenter
 
-# --- Agent Node Functions ---
 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['userInput']}\""
+    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','')}\""
     response = llm.invoke(prompt)
-
     if 'greeting' in response.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."
-        }
+        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 {"execution_path": ["Triage Agent"], "status_update": "Request requires a plan. Proceeding..."}
 
 def run_planner_agent(state: AgentState):
     log.info("--- ✈️ Running 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['userInput']}\". Respond in JSON with keys: 'plan' (list of strings), "
-        f"'estimated_llm_calls_per_loop' (integer)."
+        f"User Request: \"{state.get('userInput','')}\". Respond in JSON with keys: 'plan' (list of strings), 'estimated_llm_calls_per_loop' (integer)."
     )
-
     response = llm.invoke(prompt)
     plan_data = parse_json_from_llm(response.content)
-
     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
-    )
+    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)
-
     plan_data['max_loops_initial'] = INITIAL_MAX_REWORK_CYCLES
     plan_data['estimated_cost_usd'] = round(estimated_cost, 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 '')
+    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."}
 
 def run_memory_retrieval(state: AgentState):
     log.info("--- 🧠 Accessing Long-Term Memory ---")
     path = ensure_list(state, 'execution_path') + ["Memory Retriever"]
-    relevant_mems = memory_manager.retrieve_relevant_memories(state['userInput'])
+    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.")
@@ -136,79 +300,154 @@ def run_memory_retrieval(state: AgentState):
 def run_intent_agent(state: AgentState):
     log.info("--- 🎯 Running 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\n"
-        f"Relevant Memory:\n{state.get('retrievedMemory')}\n\nUser Request: \"{state.get('userInput')}\"\n\nCore Objective:"
-    )
+    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:")
     response = llm.invoke(prompt)
-    return {"coreObjectivePrompt": response.content, "execution_path": path, "status_update": "Clarifying the main objective..."}
+    core_obj = response.content
+    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..."}
 
 def run_pm_agent(state: AgentState):
     log.info("--- 👷 Running PM Agent ---")
-    # coerce rework_cycles/max_loops to integers if they are None or falsy
     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"]
     feedback = f"QA Feedback (must be addressed): {state.get('qaFeedback')}" if state.get('qaFeedback') else ""
     prompt = (
-        f"Decompose the core objective into a plan. Determine if code execution is needed and define the goal.\n\n"
+        f"Decompose the core objective into a plan. Determine if code execution or artifact generation is needed and define the goal.\n\n"
         f"Core Objective: {state.get('coreObjectivePrompt')}\n\n{feedback}\n\n"
-        f"Respond in JSON with keys: 'plan_steps' (list), 'experiment_needed' (bool), and 'experiment_goal' (str if needed)."
+        f"Respond in JSON with keys: 'plan_steps' (list), 'experiment_needed' (bool), 'experiment_type' (optional string), and 'experiment_goal' (str if needed)."
     )
-
     response = llm.invoke(prompt)
     plan = parse_json_from_llm(response.content)
-
     if not plan:
-        log.error("PM Agent failed to produce a valid JSON plan.")
-        plan = {"plan_steps": ["Error: The Project Manager failed to create a valid plan."], "experiment_needed": False}
-
-    log.info(f"Generated Plan: Experiment Needed = {plan.get('experiment_needed', False)}")
+        log.warning("PM Agent did not produce JSON — applying heuristic fallback.")
+        plan = {"plan_steps": ["Analyze files", "Create notebook if requested", "Synthesize answers"], "experiment_needed": False}
+    intent_detector = state.get('artifact_detection') or {}
+    if intent_detector.get('requires_artifact'):
+        plan['experiment_needed'] = True
+        plan['experiment_type'] = intent_detector.get('artifact_type')
+        plan['experiment_goal'] = f"Produce an artifact: {intent_detector.get('artifact_hint')}. Use document reading and test edge cases for messy files in the folder. {state.get('userInput','')}"
+    if plan.get('experiment_needed') and not plan.get('experiment_type'):
+        detection = detect_requested_output_types(state.get('coreObjectivePrompt','') or state.get('userInput',''))
+        if detection.get('requires_artifact'):
+            plan['experiment_type'] = detection.get('artifact_type')
+            plan['experiment_goal'] = plan.get('experiment_goal') or f"Produce an artifact: {detection.get('artifact_hint')}."
+    log.info(f"Generated Plan: Experiment Needed = {plan.get('experiment_needed', False)}, Type = {plan.get('experiment_type')}")
     return {"pmPlan": plan, "execution_path": path, "rework_cycles": current_cycles, "status_update": "Breaking down the objective into a detailed plan..."}
 
+def _extract_python_blocks(text: str) -> List[str]:
+    return re.findall(r"```python\n(.*?)\n```", text, re.DOTALL) or re.findall(r"```\n(.*?)\n```", text, re.DOTALL)
+
 def run_experimenter_agent(state: AgentState):
     log.info("--- 🔬 Running Experimenter Agent ---")
     path = ensure_list(state, 'execution_path') + ["Experimenter Agent"]
-    if not state.get('pmPlan', {}).get('experiment_needed'):
+    pm = state.get('pmPlan', {}) or {}
+    if not pm.get('experiment_needed'):
         return {"experimentCode": None, "experimentResults": None, "execution_path": path, "status_update": "Proceeding without a code experiment."}
-
-    goal = state.get('pmPlan', {}).get('experiment_goal', 'No goal specified.')
-    prompt = f"Write a Python script to achieve this goal: {goal}. Print results to standard output."
-    response = llm.invoke(prompt)
-    code_match = re.search(r"```python\n(.*?)\n```", response.content, re.DOTALL)
-    if not code_match:
-        log.error("Experimenter failed to generate valid Python code.")
-        results = {"success": False, "stderr": "No valid Python code block was generated."}
-        return {"experimentCode": "# ERROR: No code generated", "experimentResults": results, "execution_path": path, "status_update": "Error: Failed to write necessary code."}
-
-    code = code_match.group(1).strip()
-    results = execute_python_code(code)
-    return {"experimentCode": code, "experimentResults": results, "execution_path": path, "status_update": "Running code to find a solution..."}
+    exp_type = pm.get('experiment_type') or 'notebook'
+    goal = pm.get('experiment_goal', 'No goal specified.')
+    response = llm.invoke(
+        f"Produce content for artifact type '{exp_type}' to achieve: {goal}\n"
+        "Return runnable code in fenced code blocks where appropriate, and explanatory text in plaintext."
+    )
+    llm_text = response.content or ""
+    out_dir = "/tmp"
+    results = {"success": False, "paths": {}, "stderr": "", "stdout": ""}
+    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, "experiment_llm_text": llm_text, "execution_path": path, "status_update": f"Notebook generated at {nb_path}"}
+        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": f"Excel generated at {excel_path}"}
+        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": f"DOCX generated at {docx_path}"}
+        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": f"PDF generated at {pdf_path}"}
+        elif exp_type == 'script':
+            # pick a language hint from plan or goal
+            lang_hint = pm.get('experiment_language') or ("python" if ".py" in goal.lower() else None)
+            # extract code blocks
+            code_blocks = _extract_python_blocks(llm_text)
+            if not code_blocks:
+                # fallback: entire content
+                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)
+            # optionally execute python scripts
+            exec_results = {}
+            if script_path.endswith(".py"):
+                exec_results = execute_python_code(open(script_path,"r",encoding="utf-8").read())
+            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": f"Script generated at {script_path}"}
+        elif exp_type == 'repo':
+            # build a minimal repo by calling LLM for file suggestions or using code blocks
+            # Heuristic: create a simple app repo containing a notebook and README and requirements.txt
+            repo_files = {}
+            # README from first 400 chars as text
+            readme = (llm_text[:1000] + "\n\n") if llm_text else "Generated repo"
+            repo_files["README.md"] = readme
+            # include generated notebook
+            nb_path = write_notebook_from_text(llm_text, out_dir=out_dir)
+            repo_files["analysis.ipynb"] = nb_path
+            # requirements: keep minimal
+            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 at {zip_path}"}
+        else:
+            # fallback: create docx with llm_text
+            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 at {fallback}"}
+    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."}
 
 def run_synthesis_agent(state: AgentState):
     log.info("--- ✍️ Running Synthesis Agent ---")
     path = ensure_list(state, 'execution_path') + ["Synthesis Agent"]
     exp_results = state.get('experimentResults')
     results_summary = "No experiment was conducted."
-    if exp_results:
-        results_summary = f"Experiment Output:\nSTDOUT:\n{exp_results.get('stdout', '')}\nSTDERR:\n{exp_results.get('stderr', '')}"
-
+    artifact_message = ""
+    if exp_results and isinstance(exp_results, dict):
+        paths = exp_results.get("paths") or {}
+        if paths:
+            artifact_lines = []
+            for k,v in paths.items():
+                artifact_lines.append(f"- {k}: `{v}`")
+            artifact_message = "\n\n**Artifacts produced:**\n" + "\n".join(artifact_lines)
+            results_summary = f"Artifacts produced: {list(paths.keys())}"
+        else:
+            results_summary = f"Experiment Output Stdout: {exp_results.get('stdout','')}\nStderr: {exp_results.get('stderr','')}"
     prompt = (
         f"Synthesize all information into a final response.\n\nCore Objective: {state.get('coreObjectivePrompt')}\n\n"
         f"Plan: {state.get('pmPlan', {}).get('plan_steps')}\n\n{results_summary}\n\nFinal Response:"
     )
     response = llm.invoke(prompt)
-    return {"draftResponse": response.content, "execution_path": path, "status_update": "Putting together the final response..."}
+    final_text = response.content or ""
+    if artifact_message:
+        final_text = final_text + "\n\n" + artifact_message
+    return {"draftResponse": final_text, "execution_path": path, "status_update": "Putting together the final response..."}
 
 def run_qa_agent(state: AgentState):
     log.info("--- ✅ Running QA Agent ---")
     path = ensure_list(state, 'execution_path') + ["QA Agent"]
-    prompt = (
-        f"Review the draft response based on the core objective. Respond ONLY with 'APPROVED' or provide concise feedback for rework.\n\n"
-        f"Core Objective: {state.get('coreObjectivePrompt')}\n\nDraft: {state.get('draftResponse')}"
-    )
+    prompt = (f"Review the draft response based on the core objective. Respond ONLY with 'APPROVED' or provide concise feedback for rework.\n\n"
+              f"Core Objective: {state.get('coreObjectivePrompt')}\n\nDraft: {state.get('draftResponse')}")
     response = llm.invoke(prompt)
     if "APPROVED" in response.content.upper():
         return {"approved": True, "qaFeedback": None, "execution_path": path, "status_update": "Response approved!"}
@@ -218,10 +457,8 @@ def run_qa_agent(state: AgentState):
 def run_archivist_agent(state: AgentState):
     log.info("--- 💾 Running 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\nFinal Response: {state.get('draftResponse')}\n\nMemory Summary:"
-    )
+    summary_prompt = (f"Create a concise summary of this successful task for long-term memory.\n\n"
+                      f"Core Objective: {state.get('coreObjectivePrompt')}\n\nFinal Response: {state.get('draftResponse')}\n\nMemory Summary:")
     response = llm.invoke(summary_prompt)
     memory_manager.add_to_memory(response.content, {"objective": state.get('coreObjectivePrompt')})
     return {"execution_path": path, "status_update": "Saving key learnings for future reference..."}
@@ -229,13 +466,11 @@ def run_archivist_agent(state: AgentState):
 def run_disclaimer_agent(state: AgentState):
     log.warning("--- ⚠️ Running Disclaimer Agent ---")
     path = ensure_list(state, 'execution_path') + ["Disclaimer Agent"]
-    disclaimer = (
-        "**DISCLAIMER: The process was stopped after exhausting the budget. The following response is the best available draft and may be incomplete.**\n\n---\n\n"
-    )
+    disclaimer = ("**DISCLAIMER: The process was stopped after exhausting the budget. The following response is the best available draft and may be incomplete.**\n\n---\n\n")
     final_response = disclaimer + state.get('draftResponse', "No response was generated.")
     return {"draftResponse": final_response, "execution_path": path, "status_update": "Budget limit reached. Preparing final draft..."}
 
-# --- Conditional Edges & Graph Definition ---
+# --- Decision & Graph ---
 def should_continue(state: AgentState):
     log.info("--- 🤔 Decision: Is the response QA approved? ---")
     if state.get("approved"):
@@ -249,25 +484,22 @@ def should_continue(state: AgentState):
         return "pm_agent"
 
 def should_run_experiment(state: AgentState):
-    decision = "experimenter_agent" if state.get('pmPlan', {}).get('experiment_needed') else "synthesis_agent"
-    return decision
+    pm = state.get('pmPlan', {}) or {}
+    return "experimenter_agent" if pm.get('experiment_needed') else "synthesis_agent"
 
-# --- Build the Graphs ---
-# 1. Triage Graph
+# --- Build graphs (same as before) ---
 triage_workflow = StateGraph(AgentState)
 triage_workflow.add_node("triage", run_triage_agent)
 triage_workflow.set_entry_point("triage")
 triage_workflow.add_edge("triage", END)
 triage_app = triage_workflow.compile()
 
-# 2. Planner-only Graph
 planner_workflow = StateGraph(AgentState)
 planner_workflow.add_node("planner", run_planner_agent)
 planner_workflow.set_entry_point("planner")
 planner_workflow.add_edge("planner", END)
 planner_app = planner_workflow.compile()
 
-# 3. Full Execution Graph
 main_workflow = StateGraph(AgentState)
 main_workflow.add_node("memory_retriever", run_memory_retrieval)
 main_workflow.add_node("intent_agent", run_intent_agent)
@@ -292,4 +524,4 @@ main_workflow.add_conditional_edges("qa_agent", should_continue, {
     "pm_agent": "pm_agent",
     "disclaimer_agent": "disclaimer_agent"
 })
-main_app = main_workflow.compile()
+main_app = main_workflow.compile()