Update graph.py

graph.py

@@ -1,4 +1,4 @@
-# graph.py (patched: artifact generation using nbformat, python-docx, pandas/openpyxl, reportlab)
+# graph.py (final patched)
 import json
 import re
 import math
@@ -13,7 +13,7 @@ from memory_manager import memory_manager
 from code_executor import execute_python_code
 from logging_config import setup_logging, get_logger
 
-# External artifact libs
+# Artifact libs
 import nbformat
 from nbformat.v4 import new_notebook, new_markdown_cell, new_code_cell
 import pandas as pd
@@ -23,6 +23,7 @@ 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 []
@@ -33,6 +34,7 @@ def ensure_list(state, key):
     return [v]
 
 def ensure_int(state, key, default=0):
+    """Return an int from state[key], default if missing/invalid."""
     try:
         v = state.get(key) if state else None
         if v is None:
@@ -42,10 +44,10 @@ def ensure_int(state, key, default=0):
         return default
 
 def sanitize_path(path: str) -> str:
-    # On HF Spaces you may want to move to a served directory. Keep as-is here.
-    return path
+    """Sanitize/normalize output path for return to UI."""
+    return os.path.abspath(path)
 
-# --- Logging & constants ---
+# --- Setup & constants ---
 setup_logging()
 log = get_logger(__name__)
 INITIAL_MAX_REWORK_CYCLES = 3
@@ -74,8 +76,11 @@ class AgentState(TypedDict):
 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 JSON object from LLM text."""
     try:
-        match = re.search(r"```json\n({.*?})\n```", llm_output, re.DOTALL)
+        if not llm_output:
+            return None
+        match = re.search(r"```json\s*({.*?})\s*```", llm_output, re.DOTALL)
         if match:
             json_str = match.group(1)
         else:
@@ -86,11 +91,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.error(f"JSON parsing failed. Error: {e}. Raw: {llm_output[:300]}")
+        log.error(f"JSON parsing failed. Error: {e}. Raw head: {llm_output[:300]}")
         return None
 
-# --- Artifact detection ---
+# --- Artifact detection & normalization ---
+KNOWN_ARTIFACT_TYPES = {"notebook","excel","word","pdf","image","repo","script"}
+
 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()
@@ -106,23 +114,42 @@ def detect_requested_output_types(text: str) -> Dict:
         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}
 
+def normalize_experiment_type(exp_type: Optional[str], goal_text: str) -> str:
+    """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"
+
 # --- 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
+    code_blocks = re.findall(r"```python\s*(.*?)\s*```", llm_text, re.DOTALL)
     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)
+        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))
@@ -145,11 +172,11 @@ def write_script(code_text: str, language_hint: Optional[str]=None, out_dir: str
         l = language_hint.lower()
         if "python" in l or ".py" in l:
             ext = ".py"
-        elif l in ("r", ".r"):
+        elif l == "r" or l == ".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:
+        elif "javascript" in l or ".js" in l:
             ext = ".js"
         elif "bash" in l or "sh" in l:
             ext = ".sh"
@@ -161,7 +188,6 @@ def write_script(code_text: str, language_hint: Optional[str]=None, out_dir: str
 
 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]
@@ -170,27 +196,18 @@ def write_docx_from_text(text: str, out_dir: str="/tmp") -> str:
     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))
@@ -200,7 +217,6 @@ def write_excel_from_tables(maybe_table_text: str, out_dir: str="/tmp") -> str:
         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:
@@ -217,14 +233,9 @@ def write_pdf_from_text(text: str, out_dir: str="/tmp") -> str:
         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)
@@ -234,7 +245,6 @@ def build_repo_zip(files_map: Dict[str,str], repo_name: str="generated_app", out
         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")
@@ -246,14 +256,13 @@ def build_repo_zip(files_map: Dict[str,str], repo_name: str="generated_app", out
                 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
-
+# --- 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.get('userInput','')}\""
     response = llm.invoke(prompt)
-    if 'greeting' in response.content.lower():
+    content = getattr(response, "content", "") or ""
+    if 'greeting' in content.lower():
         log.info("Triage result: Simple Greeting.")
         return {"draftResponse": "Hello! How can I help you today?", "execution_path": ["Triage Agent"], "status_update": "Responding to greeting."}
     else:
@@ -268,7 +277,7 @@ def run_planner_agent(state: AgentState):
         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)
+    plan_data = parse_json_from_llm(getattr(response, "content", "") or "")
     if not plan_data:
         return {"pmPlan": {"error": "Failed to create a valid plan."}, "execution_path": path, "status_update": "Error: Could not create a plan."}
     calls_per_loop = plan_data.get('estimated_llm_calls_per_loop', 3)
@@ -302,7 +311,7 @@ def run_intent_agent(state: AgentState):
     path = ensure_list(state, 'execution_path') + ["Intent Agent"]
     prompt = (f"Refine the user's request into a clear, actionable 'core objective prompt'.\n\nRelevant Memory:\n{state.get('retrievedMemory')}\n\nUser Request: \"{state.get('userInput','')}\"\n\nCore Objective:")
     response = llm.invoke(prompt)
-    core_obj = response.content
+    core_obj = getattr(response, "content", "") or ""
     detection = detect_requested_output_types(core_obj or state.get('userInput',''))
     extras = {}
     if detection.get('requires_artifact'):
@@ -322,25 +331,28 @@ def run_pm_agent(state: AgentState):
         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)
+    plan = parse_json_from_llm(getattr(response, "content", "") or "")
     if not plan:
         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')}."
+    # normalize experiment type
+    exp_type_raw = plan.get('experiment_type') or ""
+    plan_goal = plan.get('experiment_goal') or state.get('userInput','') or state.get('coreObjectivePrompt','')
+    normalized = normalize_experiment_type(exp_type_raw, plan_goal)
+    plan['experiment_type'] = normalized
+    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_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 _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:
+            return blocks
+    blocks = re.findall(r"```(?:\w+)?\s*(.*?)\s*```", text, re.DOTALL)
+    return blocks
 
 def run_experimenter_agent(state: AgentState):
     log.info("--- 🔬 Running Experimenter Agent ---")
@@ -348,20 +360,20 @@ 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": "Proceeding without a code experiment."}
-    exp_type = pm.get('experiment_type') or 'notebook'
+    exp_type = normalize_experiment_type(pm.get('experiment_type'), pm.get('experiment_goal',''))
     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."
+        "Return runnable code in fenced code blocks where appropriate, and explanatory text otherwise."
     )
-    llm_text = response.content or ""
+    llm_text = getattr(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}"}
+            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)}})
@@ -375,40 +387,34 @@ def run_experimenter_agent(state: AgentState):
             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)
+            code_blocks = _extract_code_blocks(llm_text, lang_hint)
             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())
+                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}
             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
+            # 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}"}
@@ -438,7 +444,7 @@ def run_synthesis_agent(state: AgentState):
         f"Plan: {state.get('pmPlan', {}).get('plan_steps')}\n\n{results_summary}\n\nFinal Response:"
     )
     response = llm.invoke(prompt)
-    final_text = response.content or ""
+    final_text = getattr(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..."}
@@ -449,10 +455,11 @@ def run_qa_agent(state: AgentState):
     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():
+    content = getattr(response, "content", "") or ""
+    if "APPROVED" in content.upper():
         return {"approved": True, "qaFeedback": None, "execution_path": path, "status_update": "Response approved!"}
     else:
-        return {"approved": False, "qaFeedback": response.content or "No specific feedback.", "execution_path": path, "status_update": "Response needs improvement. Reworking..."}
+        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 ---")
@@ -460,7 +467,7 @@ def run_archivist_agent(state: AgentState):
     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')})
+    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):
@@ -487,7 +494,7 @@ def should_run_experiment(state: AgentState):
     pm = state.get('pmPlan', {}) or {}
     return "experimenter_agent" if pm.get('experiment_needed') else "synthesis_agent"
 
-# --- Build graphs (same as before) ---
+# --- Build graphs ---
 triage_workflow = StateGraph(AgentState)
 triage_workflow.add_node("triage", run_triage_agent)
 triage_workflow.set_entry_point("triage")