Update graph.py

graph.py

@@ -87,43 +87,47 @@ llm = ChatOpenAI(model="gpt-4o", temperature=0.1, max_retries=3, request_timeout
 
 def parse_json_from_llm(llm_output: str) -> Optional[dict]:
     """
-    Robust JSON extraction from LLM text.
-    Tries these strategies in order:
-      1. extract explicit ```json ... ``` fenced block
-      2. extract best-balanced {...} substring
-      3. json.loads()
-      4. ast.literal_eval() -> convert to JSON
-      5. conservative single-quote to double-quote + remove trailing commas -> json.loads()
-
-    Returns parsed dict or None on failure. Logs the original content for debugging.
+    More robust JSON extraction:
+      - Looks for explicit ```json {} ``` blocks
+      - Falls back to the last balanced {...} substring in the output
+      - Tries ast.literal_eval for Python-like dicts
+      - Performs conservative cleanup (remove trailing commas, comments, safe single->double quote) and retries
+    Returns dict or None. Logs failures for debugging.
     """
-    import json
     import re
+    import json
     import ast
     from logging import getLogger
 
     logger = getLogger(__name__)
 
-    if not llm_output or not llm_output.strip():
+    if not llm_output or not isinstance(llm_output, str) or not llm_output.strip():
         return None
 
     text = llm_output.strip()
 
-    # 1) Try explicit ```json ... ``` codeblock first (best)
-    m = re.search(r"```json\s*(\{.*?\})\s*```", text, re.DOTALL | re.IGNORECASE)
-    if m:
-        candidate = m.group(1).strip()
+    # 1) explicit fenced JSON block
+    match = re.search(r"```json\s*({.*?})\s*```", text, re.DOTALL | re.IGNORECASE)
+    if match:
+        candidate = match.group(1).strip()
         try:
             return json.loads(candidate)
         except Exception as e:
             logger.debug(f"json.loads failed on triple-backtick json block: {e}")
 
-    # 2) Try to extract a balanced JSON-like substring (first balanced {...})
-    #    This is safer than taking first '{' to last '}'.
+    # 2) any code-fence containing a JSON-like object
+    match2 = re.search(r"```(?:json|python|text)?\s*({.*?})\s*```", text, re.DOTALL | re.IGNORECASE)
+    if match2:
+        candidate = match2.group(1).strip()
+        try:
+            return json.loads(candidate)
+        except Exception as e:
+            logger.debug(f"json.loads failed on fenced candidate: {e}")
+
+    # 3) find first balanced {...} substring
     def find_balanced_brace_substring(s: str):
         start_idx = None
         depth = 0
-        best = None
         for i, ch in enumerate(s):
             if ch == '{':
                 if start_idx is None:
@@ -133,19 +137,12 @@ def parse_json_from_llm(llm_output: str) -> Optional[dict]:
                 if depth > 0:
                     depth -= 1
                 if depth == 0 and start_idx is not None:
-                    best = s[start_idx:i+1]
-                    return best
+                    return s[start_idx:i+1]
         return None
 
-    candidate = None
-    # Try extracting from ```...``` blocks as fallback (any code fences)
-    m2 = re.search(r"```(?:json|python|text)?\s*(\{.*?\})\s*```", text, re.DOTALL | re.IGNORECASE)
-    if m2:
-        candidate = m2.group(1).strip()
-    if not candidate:
-        candidate = find_balanced_brace_substring(text)
+    candidate = find_balanced_brace_substring(text)
 
-    # If still nothing, as last resort take substring between first { and last }
+    # 4) fallback: last { ... } block heuristically
     if not candidate:
         first = text.find('{')
         last = text.rfind('}')
@@ -159,45 +156,39 @@ def parse_json_from_llm(llm_output: str) -> Optional[dict]:
         except Exception as e:
             logger.debug(f"json.loads failed on candidate substring: {e}")
 
-        # try ast.literal_eval (handles single quotes, Python tuples, etc.)
+        # try ast.literal_eval (handles single quotes & Python literals)
         try:
             parsed = ast.literal_eval(candidate)
-            # ast.literal_eval may return non-dict (list etc.). Coerce to dict if needed.
             if isinstance(parsed, (dict, list)):
+                # convert to a strict JSON-compatible dict/list
                 return json.loads(json.dumps(parsed))
-            else:
-                logger.debug("ast.literal_eval produced non-dict result; returning None.")
         except Exception as e:
             logger.debug(f"ast.literal_eval failed: {e}")
 
-        # final attempt: conservative cleanup for common LLM deviations:
-        #  - replace single quotes with double quotes only where safe-ish
-        #  - remove trailing commas before } or ]
+        # conservative cleanup: remove comments, trailing commas, and handle simple single-quote strings
         cleaned = candidate
-
-        # Remove JavaScript/JSON-like line comments (// ...) and block comments (/* ... */)
-        cleaned = re.sub(r"//.*?$", "", cleaned, flags=re.MULTILINE)
-        cleaned = re.sub(r"/\*.*?\*/", "", cleaned, flags=re.DOTALL)
-
-        # Remove trailing commas (",\s*}" or ",\s*]")
-        cleaned = re.sub(r",\s*([}\]])", r"\1", cleaned)
-
-        # Replace single-quoted strings with double-quoted strings where it looks like a string token
-        # This is conservative: replace only when a single quote is preceded by : or [, or { or start-of-string.
-        def single_to_double(match):
-            inner = match.group(1)
-            inner_escaped = inner.replace('"', '\\"')
-            return f'"{inner_escaped}"'
-        cleaned = re.sub(r"(?<=[:\{\[,]\s*)'([^']*?)'", single_to_double, cleaned)
-
-        # Try json.loads again
         try:
+            # remove line comments //...
+            cleaned = re.sub(r"//.*?$", "", cleaned, flags=re.MULTILINE)
+            # remove block comments /* ... */
+            cleaned = re.sub(r"/\*.*?\*/", "", cleaned, flags=re.DOTALL)
+            # remove trailing commas before } or ]
+            cleaned = re.sub(r",\s*([}\]])", r"\1", cleaned)
+
+            # replace single-quoted strings with double quotes when likely safe
+            def _single_to_double(m):
+                inner = m.group(1)
+                inner_escaped = inner.replace('"', '\\"')
+                return f'"{inner_escaped}"'
+            cleaned = re.sub(r"(?<=[:\{\[,]\s*)'([^']*?)'", _single_to_double, cleaned)
+
+            # final attempt
             return json.loads(cleaned)
         except Exception as e:
             logger.debug(f"json.loads still failed after cleanup: {e}")
 
-    # As last safety, log the original LLM output for inspection (do not raise)
-    logger.error("parse_json_from_llm failed to parse LLM output. LLM output preview (200 chars): %s", text[:200])
+    # nothing parsed — log preview and return None
+    logger.error("parse_json_from_llm failed to parse LLM output. LLM output preview (200 chars): %s", text[:200].replace("\n","\\n"))
     return None
 
 
@@ -830,7 +821,6 @@ def run_qa_agent(state: AgentState):
     if state.get('experimentResults'):
         qa_context.append(f"\n=== ARTIFACTS ===\n{json.dumps(state.get('experimentResults', {}).get('paths', {}), indent=2)}")
 
-    # MODIFIED PROMPT: Encourage convergence
     prompt = f"""You are a QA reviewer. Review the draft response against the user's objective.
     {chr(10).join(qa_context)}
 
@@ -842,24 +832,45 @@ def run_qa_agent(state: AgentState):
 
     If this is a re-submission (rework cycle > 1), has the previous feedback been successfully addressed?
 
-    Response Format:
+    Response Format (required JSON or a single word 'APPROVED'):
 
-    If the work is complete and high-quality, respond ONLY with the word 'APPROVED'.
+    Either return EXACTLY the single word:
+      APPROVED
 
-    Otherwise, provide SPECIFIC, ACTIONABLE, and NOVEL feedback on what must be changed. Do not repeat previous feedback if it has already been actioned.
+    Or return JSON like:
+    {{
+      "approved": false,
+      "feedback": "Specific, actionable items to fix (bullet list or numbered).",
+      "required_changes": ["..."]
+    }}
     """
 
-    response = llm.invoke(prompt)
-    content = getattr(response, "content", "") or ""
+    try:
+        response = llm.invoke(prompt)
+        content = getattr(response, "content", "") or ""
+    except Exception as e:
+        log.exception("QA LLM call failed: %s", e)
+        # Fail-safe: mark as not approved with conservative feedback
+        return {"approved": False, "qaFeedback": "QA LLM failed; manual review required.", "execution_path": path, "status_update": "QA failed"}
 
-    if "APPROVED" in content.upper():
+    # If LLM returned APPROVED word, treat as approved
+    if "APPROVED" in content.strip().upper() and len(content.strip()) <= 20:
         return {"approved": True, "qaFeedback": None, "execution_path": path, "status_update": "Approved"}
-    else:
-    # Sanitize feedback to ensure it's a useful string
-        feedback = content.replace("APPROVED", "").strip()
-        if not feedback:
-            feedback = "General quality improvements required."
-            return {"approved": False, "qaFeedback": feedback, "execution_path": path, "status_update": "Needs improvement"}
+
+    # Else try JSON parse
+    parsed = parse_json_from_llm(content)
+    if isinstance(parsed, dict):
+        approved = bool(parsed.get("approved", False))
+        feedback = parsed.get("feedback") or parsed.get("qaFeedback") or parsed.get("required_changes") or ""
+        # Normalize feedback to string
+        if isinstance(feedback, list):
+            feedback = "\n".join([str(x) for x in feedback])
+        elif not isinstance(feedback, str):
+            feedback = str(feedback)
+        return {"approved": approved, "qaFeedback": feedback if not approved else None, "execution_path": path, "status_update": "QA completed"}
+    # Fallback: return raw text as feedback (not approved)
+    safe_feedback = content.strip()[:2000] or "QA produced no actionable output."
+    return {"approved": False, "qaFeedback": safe_feedback, "execution_path": path, "status_update": "QA needs rework"}
 
 def run_archivist_agent(state: AgentState):
     log.info("--- ARCHIVIST ---")