Update graph.py

graph.py

@@ -1,4 +1,4 @@
-# graph.py (revised, robust)
+# graph.py (final patched)
 import json
 import re
 import math
@@ -6,53 +6,38 @@ import os
 import uuid
 import shutil
 import zipfile
-from typing import TypedDict, List, Dict, Optional, Any
+import operator
+from typing import TypedDict, List, Dict, Optional, Annotated
 from langchain_openai import ChatOpenAI
 from langgraph.graph import StateGraph, END
 from memory_manager import memory_manager
 from code_executor import execute_python_code
 from logging_config import setup_logging, get_logger
 
-# Artifact libs (optional — if not installed, artifact-writing functions will fallback to writing text)
-try:
-    import nbformat
-    from nbformat.v4 import new_notebook, new_markdown_cell, new_code_cell
-except Exception:
-    nbformat = None
-
-try:
-    import pandas as pd
-except Exception:
-    pd = None
-
-try:
-    from docx import Document
-except Exception:
-    Document = None
-
-try:
-    from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
-    from reportlab.lib.styles import getSampleStyleSheet
-except Exception:
-    SimpleDocTemplate = None
+# 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: Optional[dict], key: str) -> List[Any]:
+def ensure_list(state, key):
     """Return a list from state[key], default [] if missing/None/not-list."""
-    v = (state or {}).get(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)
-    # if e.g. a single string, keep as list
     return [v]
 
-def ensure_int(state: Optional[dict], key: str, default: int = 0) -> int:
+def ensure_int(state, key, default=0):
     """Return an int from state[key], default if missing/invalid."""
     try:
-        v = (state or {}).get(key)
+        v = state.get(key) if state else None
         if v is None:
             return default
         return int(v)
@@ -60,17 +45,8 @@ def ensure_int(state: Optional[dict], key: str, default: int = 0) -> int:
         return default
 
 def sanitize_path(path: str) -> str:
-    """Return an absolute normalized path, safe for UI linking."""
-    try:
-        return os.path.abspath(path)
-    except Exception:
-        return path
-
-def safe_write_text_file(text: str, out_path: str) -> str:
-    os.makedirs(os.path.dirname(out_path), exist_ok=True)
-    with open(out_path, "w", encoding="utf-8") as fh:
-        fh.write(text or "")
-    return out_path
+    """Sanitize/normalize output path for return to UI."""
+    return os.path.abspath(path)
 
 # --- Setup & constants ---
 setup_logging()
@@ -79,11 +55,9 @@ INITIAL_MAX_REWORK_CYCLES = 3
 GPT4O_INPUT_COST_PER_1K_TOKENS = 0.005
 GPT4O_OUTPUT_COST_PER_1K_TOKENS = 0.015
 AVG_TOKENS_PER_CALL = 2.0
-OUT_DIR = "outputs"
-os.makedirs(OUT_DIR, exist_ok=True)
 
-# --- AgentState definition (non-total so keys optional) ---
-class AgentState(TypedDict, total=False):
+# --- AgentState ---
+class AgentState(TypedDict):
     userInput: str
     chatHistory: List[str]
     coreObjectivePrompt: str
@@ -94,7 +68,8 @@ class AgentState(TypedDict, total=False):
     draftResponse: str
     qaFeedback: Optional[str]
     approved: bool
-    execution_path: List[str]
+    # <-- important: Annotated with operator.add to declare accumulation behavior
+    execution_path: Annotated[List[str], operator.add]
     rework_cycles: int
     max_loops: int
     status_update: str
@@ -103,16 +78,14 @@ class AgentState(TypedDict, total=False):
 llm = ChatOpenAI(model="gpt-4o", temperature=0.1, max_retries=3, request_timeout=60)
 
 def parse_json_from_llm(llm_output: str) -> Optional[dict]:
-    """Robustly try to extract a single JSON object from LLM text."""
+    """Robustly try to extract JSON object from LLM text."""
     try:
         if not llm_output:
             return None
-        # Prefer fenced json blocks
         match = re.search(r"```json\s*({.*?})\s*```", llm_output, re.DOTALL)
         if match:
             json_str = match.group(1)
         else:
-            # fallback: first {...}
             start = llm_output.find('{')
             end = llm_output.rfind('}')
             if start == -1 or end == -1:
@@ -120,14 +93,14 @@ 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.debug(f"parse_json_from_llm failed: {e}. snippet: {str(llm_output)[:300]}")
+        log.error(f"JSON parsing failed. Error: {e}. Raw head: {llm_output[:300]}")
         return None
 
-# --- Artifact type detection & normalization ---
-KNOWN_ARTIFACT_TYPES = {"notebook", "excel", "word", "pdf", "image", "repo", "script", "docx"}
+# --- Artifact detection & normalization ---
+KNOWN_ARTIFACT_TYPES = {"notebook","excel","word","pdf","image","repo","script"}
 
-def detect_requested_output_types(text: str) -> Dict[str, Optional[str]]:
-    """Simple heuristics to detect requested artifact types from user text."""
+def detect_requested_output_types(text: str) -> Dict:
+    """Heuristic detect requested artifact type from text."""
     if not text:
         return {"requires_artifact": False, "artifact_type": None, "artifact_hint": None}
     t = text.lower()
@@ -137,65 +110,65 @@ def detect_requested_output_types(text: str) -> Dict[str, Optional[str]]:
         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 "pdf" in t:
+    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", ".js"]):
-        return {"requires_artifact": True, "artifact_type": "script", "artifact_hint": "language script file"}
+    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.)"}
     return {"requires_artifact": False, "artifact_type": None, "artifact_hint": None}
 
 def normalize_experiment_type(exp_type: Optional[str], goal_text: str) -> str:
-    """Map arbitrary LLM experiment_type to a known artifact type; fallback to detection."""
-    if exp_type:
-        s = exp_type.strip().lower()
-        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:
-            return "repo"
-        if "script" in s or ".py" in s or "python" in s:
-            return "script"
+    """Map arbitrary LLM returned experiment_type into known set or infer from goal_text."""
+    if not exp_type:
+        detection = detect_requested_output_types(goal_text or "")
+        return detection.get("artifact_type") or "docx"
+    s = exp_type.strip().lower()
+    # direct mapping heuristics
+    if s in KNOWN_ARTIFACT_TYPES:
+        return s
+    # common synonyms
+    if "notebook" in s or "ipynb" in s:
+        return "notebook"
+    if "excel" in s or "xlsx" in s or "spreadsheet" 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:
+        return "repo"
+    if "script" in s or "python" in s or ".py" in s:
+        return "script"
+    # fallback to detection from goal
     detection = detect_requested_output_types(goal_text or "")
     return detection.get("artifact_type") or "docx"
 
-# --- Artifact writers (safe fallbacks) ---
-def write_notebook_from_text(llm_text: str, out_dir: str = OUT_DIR) -> str:
-    """Create an .ipynb from free text/code blocks using nbformat if available."""
-    os.makedirs(out_dir, exist_ok=True)
+# --- Notebook & artifact builders ---
+def write_notebook_from_text(llm_text: str, out_dir: str="/tmp") -> str:
+    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 parsed from LLM output.")]
+    nb['cells'] = cells
     uid = uuid.uuid4().hex[:10]
     filename = os.path.join(out_dir, f"generated_notebook_{uid}.ipynb")
-    if nbformat:
-        code_blocks = re.findall(r"```python\s*(.*?)\s*```", llm_text, re.DOTALL)
-        generic_code = re.findall(r"```(?:\w+)?\s*(.*?)\s*```", llm_text, re.DOTALL)
-        # if language-specific code found, use that; otherwise use generic fenced blocks
-        code_blocks = code_blocks or generic_code
-        md_parts = re.split(r"```(?:\w+)?\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\nNo content parsed from LLM output.")]
-        nb["cells"] = cells
-        nbformat.write(nb, filename)
-        return filename
-    # fallback: write plain text file with .ipynb extension (so downloads still work)
-    return safe_write_text_file("# nbformat not installed\n\n" + (llm_text or ""), filename)
+    nbformat.write(nb, filename)
+    return filename
 
-def write_script(code_text: str, language_hint: Optional[str] = None, out_dir: str = OUT_DIR) -> str:
-    os.makedirs(out_dir, exist_ok=True)
+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()
@@ -203,95 +176,79 @@ def write_script(code_text: str, language_hint: Optional[str] = None, out_dir: s
             ext = ".py"
         elif l == "r" or l == ".r":
             ext = ".R"
-        elif "java" in l:
+        elif "java" in l or ".java" in l:
             ext = ".java"
-        elif "javascript" in l or "js" in l:
+        elif "javascript" 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}")
-    return safe_write_text_file(code_text or "", filename)
-
-def write_docx_from_text(text: str, out_dir: str = OUT_DIR) -> str:
-    os.makedirs(out_dir, exist_ok=True)
+    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()
+    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")
-    if Document:
-        doc = Document()
-        for para in [p.strip() for p in (text or "").split("\n\n") if p.strip()]:
-            doc.add_paragraph(para)
-        doc.save(filename)
-        return filename
-    # fallback to txt file with .docx extension
-    return safe_write_text_file((text or ""), filename)
+    doc.save(filename)
+    return filename
 
-def write_excel_from_tables(maybe_table_text: str, out_dir: str = OUT_DIR) -> str:
-    os.makedirs(out_dir, exist_ok=True)
+def write_excel_from_tables(maybe_table_text: str, out_dir: str="/tmp") -> str:
     uid = uuid.uuid4().hex[:10]
     filename = os.path.join(out_dir, f"generated_excel_{uid}.xlsx")
-    if pd:
+    try:
         try:
-            # try JSON table first
-            parsed = None
-            try:
-                parsed = json.loads(maybe_table_text)
-            except Exception:
-                parsed = None
+            parsed = json.loads(maybe_table_text)
             if isinstance(parsed, list):
                 df = pd.DataFrame(parsed)
             elif isinstance(parsed, dict):
                 df = pd.DataFrame([parsed])
             else:
-                # try csv-like
-                if "," in (maybe_table_text or "") or "\t" in (maybe_table_text or ""):
-                    from io import StringIO
-                    df = pd.read_csv(StringIO(maybe_table_text))
-                else:
-                    df = pd.DataFrame({"content": [maybe_table_text or ""]})
-            # write with engine openpyxl if available
-            try:
-                df.to_excel(filename, index=False, engine="openpyxl")
-            except Exception:
-                df.to_excel(filename, index=False)
-            return filename
-        except Exception as e:
-            log.debug(f"write_excel_from_tables failed: {e}")
-            return write_docx_from_text(f"Failed to create excel: {e}\n\nOriginal:\n{maybe_table_text}", out_dir=out_dir)
-    return write_docx_from_text(maybe_table_text, out_dir=out_dir)
-
-def write_pdf_from_text(text: str, out_dir: str = OUT_DIR) -> str:
-    os.makedirs(out_dir, exist_ok=True)
+                df = pd.DataFrame({"content":[str(maybe_table_text)]})
+        except Exception:
+            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}")
+        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")
-    if SimpleDocTemplate:
-        try:
-            doc = SimpleDocTemplate(filename)
-            styles = getSampleStyleSheet()
-            flowables = []
-            for para in [p.strip() for p in (text or "").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.debug(f"write_pdf_from_text failed: {e}")
-            return write_docx_from_text(f"Failed to create PDF: {e}\n\nOriginal:\n{text}", out_dir=out_dir)
-    return write_docx_from_text(text, out_dir=out_dir)
-
-def build_repo_zip(files_map: Dict[str, str], repo_name: str = "generated_app", out_dir: str = OUT_DIR) -> str:
-    os.makedirs(out_dir, exist_ok=True)
+    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"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:
     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 or {}).items():
+    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 or ""))
+                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):
@@ -301,91 +258,97 @@ def build_repo_zip(files_map: Dict[str, str], repo_name: str = "generated_app",
                 zf.write(full, arc)
     return zip_path
 
-# --- Node functions (all robust) ---
+# --- Node functions ---
 def run_triage_agent(state: AgentState):
-    path = ensure_list(state, 'execution_path') + ["Triage Agent"]
-    prompt = f"Analyze the user input. Is it a simple conversational greeting or a task? Respond with 'greeting' or 'task'.\n\nUser Input: \"{(state or {}).get('userInput','')}\""
+    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','')}\""
     response = llm.invoke(prompt)
     content = getattr(response, "content", "") or ""
-    if "greeting" in content.lower():
-        return {"draftResponse": "Hello! How can I help you today?", "execution_path": path, "status_update": "Responding to greeting."}
-    return {"execution_path": path, "status_update": "Request requires a plan. Proceeding..."}
+    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..."}
 
 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 or {}).get('userInput','')}\". Respond in JSON with keys: 'plan' (list of strings), '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(getattr(response, "content", "") or "")
     if not plan_data:
-        plan_data = {"plan": ["Could not parse plan — defaulting minimal plan."], "estimated_llm_calls_per_loop": 3}
-    # cost math
-    calls_per_loop = int(plan_data.get('estimated_llm_calls_per_loop', 3))
+        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)
-    plan_data.setdefault('max_loops_initial', INITIAL_MAX_REWORK_CYCLES)
+    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))
-    # If user requested artifact, set experiment flags
-    detection = detect_requested_output_types((state or {}).get('userInput','') or (state or {}).get('coreObjectivePrompt','') or "")
-    if detection.get("requires_artifact"):
+    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"]
-    try:
-        relevant_mems = memory_manager.retrieve_relevant_memories((state or {}).get('userInput',''))
-    except Exception:
-        relevant_mems = []
+    relevant_mems = memory_manager.retrieve_relevant_memories(state.get('userInput',''))
     if relevant_mems:
-        context = "\n".join([f"Memory: {getattr(mem,'page_content',str(mem))}" for mem in 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..."}
 
 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\nRelevant Memory:\n{(state or {}).get('retrievedMemory')}\n\n"
-        f"User Request: \"{(state or {}).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)
-    core_obj = getattr(response, "content", "") or (state or {}).get('coreObjectivePrompt', "")
-    detection = detect_requested_output_types(core_obj or (state or {}).get('userInput',''))
+    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..."}
 
 def run_pm_agent(state: AgentState):
-    # increment cycles safely
+    log.info("--- 👷 Running PM Agent ---")
     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: {state.get('qaFeedback')}" if state.get('qaFeedback') else ""
+    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 or artifact generation is needed and define the goal.\n\n"
         f"Core Objective: {state.get('coreObjectivePrompt')}\n\n{feedback}\n\n"
-        "Respond in JSON with keys: 'plan_steps' (list), 'experiment_needed' (bool), 'experiment_type' (optional), 'experiment_goal' (optional)."
+        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(getattr(response, "content", "") or "")
     if not plan:
-        plan = {"plan_steps": ["Analyze files", "If asked, generate artifact", "Synthesize results"], "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}
+    # normalize experiment type
     exp_type_raw = plan.get('experiment_type') or ""
-    plan_goal = plan.get('experiment_goal') or (state or {}).get('userInput','') or (state or {}).get('coreObjectivePrompt','')
+    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
     if plan.get('experiment_needed') and not plan.get('experiment_goal'):
         plan['experiment_goal'] = plan_goal
+    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_code_blocks(text: str, lang_hint: Optional[str] = None) -> List[str]:
+def _extract_code_blocks(text: str, lang_hint: Optional[str]=None) -> List[str]:
+    # prefer specific language fences, fallback to generic fenced blocks
     if lang_hint and "python" in (lang_hint or "").lower():
         blocks = re.findall(r"```python\s*(.*?)\s*```", text, re.DOTALL)
         if blocks:
@@ -394,130 +357,143 @@ def _extract_code_blocks(text: str, lang_hint: Optional[str] = None) -> List[str
     return blocks
 
 def run_experimenter_agent(state: AgentState):
+    log.info("--- 🔬 Running Experimenter Agent ---")
     path = ensure_list(state, 'execution_path') + ["Experimenter Agent"]
-    pm = (state or {}).get('pmPlan') or {}
+    pm = state.get('pmPlan', {}) or {}
     if not pm.get('experiment_needed'):
-        # return consistent shape: experimentResults.paths is empty dict
-        return {"experimentCode": None, "experimentResults": {"success": False, "paths": {}, "stdout": "", "stderr": ""}, "execution_path": path, "status_update": "No experiment required."}
+        return {"experimentCode": None, "experimentResults": None, "execution_path": path, "status_update": "Proceeding without a code experiment."}
     exp_type = normalize_experiment_type(pm.get('experiment_type'), pm.get('experiment_goal',''))
-    goal = pm.get('experiment_goal') or "Produce artifact"
-    response = llm.invoke(f"Produce content for artifact type '{exp_type}' to achieve: {goal}\nReturn runnable code in fenced code blocks where appropriate, and explanatory text otherwise.")
+    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 otherwise."
+    )
     llm_text = getattr(response, "content", "") or ""
-    results = {"success": False, "paths": {}, "stdout": "", "stderr": ""}
+    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)
+        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": f"Notebook generated."}
-        if exp_type == "excel":
-            excel_path = write_excel_from_tables(llm_text, out_dir=OUT_DIR)
+            return {"experimentCode": None, "experimentResults": results, "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."}
-        if exp_type in ("word", "docx"):
-            docx_path = write_docx_from_text(llm_text, out_dir=OUT_DIR)
+            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."}
-        if exp_type == "pdf":
-            pdf_path = write_pdf_from_text(llm_text, out_dir=OUT_DIR)
+            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."}
-        if exp_type == "script":
-            lang_hint = pm.get('experiment_language') or ("python" if ".py" in (pm.get('experiment_goal','') or "").lower() else None)
+            return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": f"PDF generated at {pdf_path}"}
+        elif exp_type == 'script':
+            lang_hint = pm.get('experiment_language') or ("python" if ".py" in goal.lower() else None)
             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
-            script_path = write_script(code_text, language_hint=lang_hint, out_dir=OUT_DIR)
+            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)
             exec_results = {}
             if script_path.endswith(".py"):
                 try:
                     exec_results = execute_python_code(open(script_path,"r",encoding="utf-8").read())
                 except Exception as e:
-                    exec_results = {"stdout": "", "stderr": str(e), "success": False}
+                    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": f"Script generated."}
-        if exp_type == "repo":
-            files_map = {"README.md": llm_text or "Generated repo"}
-            try:
-                nb_path = write_notebook_from_text(llm_text, out_dir=OUT_DIR)
-                files_map["analysis.ipynb"] = open(nb_path,"r",encoding="utf-8").read() if os.path.exists(nb_path) else llm_text
-            except Exception:
-                files_map["analysis.ipynb"] = llm_text
-            files_map["requirements.txt"] = "nbformat\npandas\nopenpyxl\npython-docx\nreportlab"
-            zip_path = build_repo_zip(files_map, repo_name="generated_app", out_dir=OUT_DIR)
+            return {"experimentCode": code_text, "experimentResults": results, "execution_path": path, "status_update": f"Script generated at {script_path}"}
+        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."}
-        # fallback: write docx
-        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": "Fallback DOCX generated."}
+            return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": f"Repository ZIP created at {zip_path}"}
+        else:
+            # safe fallback: write docx
+            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.exception("Experimenter failed")
+        log.error(f"Experimenter failed: {e}")
         results.update({"success": False, "stderr": str(e)})
-        return {"experimentCode": None, "experimentResults": results, "execution_path": path, "status_update": "Experimenter error."}
+        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 or {}).get('experimentResults') or {}
+    exp_results = state.get('experimentResults')
+    results_summary = "No experiment was conducted."
     artifact_message = ""
     if exp_results and isinstance(exp_results, dict):
         paths = exp_results.get("paths") or {}
         if paths:
-            artifact_lines = [f"- {k}: `{v}`" for k,v in paths.items()]
+            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"
-        f"{artifact_message}\n\nFinal Response:"
+        f"Plan: {state.get('pmPlan', {}).get('plan_steps')}\n\n{results_summary}\n\nFinal Response:"
     )
     response = llm.invoke(prompt)
     final_text = getattr(response, "content", "") or ""
-    if not final_text.strip():
-        final_text = "I prepared the requested artifact(s). Please download them from the provided paths."
     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)
     content = getattr(response, "content", "") or ""
     if "APPROVED" in content.upper():
         return {"approved": True, "qaFeedback": None, "execution_path": path, "status_update": "Response approved!"}
-    return {"approved": False, "qaFeedback": content or "No specific feedback.", "execution_path": path, "status_update": "Response needs rework."}
+    else:
+        return {"approved": False, "qaFeedback": content or "No specific feedback.", "execution_path": path, "status_update": "Response needs improvement. Reworking..."}
 
 def run_archivist_agent(state: AgentState):
+    log.info("--- 💾 Running Archivist Agent ---")
     path = ensure_list(state, 'execution_path') + ["Archivist Agent"]
-    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(prompt)
-    memory_text = getattr(response, "content", "") or state.get('draftResponse','') or ""
-    try:
-        memory_manager.add_to_memory(memory_text, {"objective": state.get('coreObjectivePrompt')})
-    except Exception:
-        log.debug("Failed to write memory; continuing.")
-    return {"execution_path": path, "status_update": "Saved learnings to memory."}
+    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(getattr(response,"content",""), {"objective": state.get('coreObjectivePrompt')})
+    return {"execution_path": path, "status_update": "Saving key learnings for future reference..."}
 
 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 may be incomplete.**\n\n---\n\n")
-    final_response = disclaimer + (state.get('draftResponse') or "No response was generated.")
-    return {"draftResponse": final_response, "execution_path": path, "status_update": "Budget limit reached."}
+    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..."}
 
-# --- Decision functions ---
+# --- Decision & Graph ---
 def should_continue(state: AgentState):
+    log.info("--- 🤔 Decision: Is the response 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"
-    return "pm_agent"
+    else:
+        log.info("Routing to: PM Agent for rework")
+        return "pm_agent"
 
 def should_run_experiment(state: AgentState):
-    pm = (state or {}).get('pmPlan') or {}
+    pm = state.get('pmPlan', {}) or {}
     return "experimenter_agent" if pm.get('experiment_needed') else "synthesis_agent"
 
 # --- Build graphs ---
@@ -546,8 +522,6 @@ main_workflow.add_node("disclaimer_agent", run_disclaimer_agent)
 main_workflow.set_entry_point("memory_retriever")
 main_workflow.add_edge("memory_retriever", "intent_agent")
 main_workflow.add_edge("intent_agent", "pm_agent")
-# add direct edges; conditional edges override routing internally
-main_workflow.add_edge("pm_agent", "experimenter_agent")
 main_workflow.add_edge("experimenter_agent", "synthesis_agent")
 main_workflow.add_edge("synthesis_agent", "qa_agent")
 main_workflow.add_edge("archivist_agent", END)
@@ -559,4 +533,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()