Configure app for Hugging Face deployment

graph.py

@@ -1,4 +1,4 @@
-# graph.py (Updated with Triage Agent)
+# graph.py (Updated with User-Facing Status Updates)
 
 import json
 import re
@@ -13,10 +13,10 @@ from logging_config import setup_logging, get_logger
 # --- Setup Logging & Constants ---
 setup_logging()
 log = get_logger(__name__)
-INITIAL_MAX_REWORK_CYCLES = 3 
+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 
+AVG_TOKENS_PER_CALL = 2.0
 
 # --- Agent State Definition ---
 class AgentState(TypedDict):
@@ -33,6 +33,8 @@ class AgentState(TypedDict):
     execution_path: List[str]
     rework_cycles: int
     max_loops: int
+    # NEW: Add a field for user-facing status updates
+    status_update: str
 
 # --- Helper for Robust JSON Parsing ---
 def parse_json_from_llm(llm_output: str) -> Optional[dict]:
@@ -51,57 +53,48 @@ def parse_json_from_llm(llm_output: str) -> Optional[dict]:
 # --- LLM Initialization ---
 llm = ChatOpenAI(model="gpt-4o", temperature=0.1, max_retries=3, request_timeout=60)
 
-# --- NEW: Triage Agent ---
+# --- Agent Node Functions (Updated with Status Messages) ---
 def run_triage_agent(state: AgentState):
     log.info("---  triage ---")
-    prompt = f"Analyze the user input. Is it a simple conversational greeting (e.g., 'hello', 'hi', 'how are you?') that doesn't require a complex plan, or is it a task/question that needs to be decomposed and executed? 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['userInput']}\""
     response = llm.invoke(prompt)
-    
+
     if 'greeting' in response.content.lower():
-        log.info("Triage result: Simple Greeting. Bypassing main workflow.")
-        # Set a direct, simple response
+        log.info("Triage result: Simple Greeting.")
         return {
             "draftResponse": "Hello! How can I help you today?",
-            "execution_path": ["Triage Agent"]
+            "execution_path": ["Triage Agent"],
+            "status_update": "Responding to greeting."
         }
     else:
-        log.info("Triage result: Complex Task. Proceeding to planner.")
-        # Let the main workflow handle it
-        return { "execution_path": ["Triage Agent"] }
-
-def should_start_full_workflow(state: AgentState):
-    """Conditional edge to decide if we run the full graph or end early."""
-    if state.get("draftResponse"):
-        # If the triage agent already set a response, we are done.
-        return "end_early"
-    else:
-        return "planner"
+        log.info("Triage result: Complex Task.")
+        return {
+            "execution_path": ["Triage Agent"],
+            "status_update": "Request requires a plan. Proceeding..."
+        }
 
-# --- Agent Node Functions (Planner, PM, etc. are unchanged) ---
 def run_planner_agent(state: AgentState):
-    log.info("--- ✈️ Running Planner Agent for Pre-flight Check ---")
+    log.info("--- ✈️ Running Planner Agent ---")
     path = state.get('execution_path', []) + ["Planner Agent"]
-    prompt = f"Analyze the user's request and provide a high-level plan and cost estimate. 1. Decompose the request into high-level steps. 2. Count the likely number of specialist LLM calls needed for ONE successful loop. User Request: \"{state['userInput']}\". Provide the output ONLY as a valid JSON object within ```json code fences with keys: 'plan' (list of strings), 'estimated_llm_calls_per_loop' (integer)."
-    
+    prompt = f"Analyze the user's request. Provide a high-level plan and estimate the number of LLM calls for one loop. User Request: \"{state['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}
+        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['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))
-    
+
     log.info(f"Pre-flight Estimate: {plan_data}")
-    return {"pmPlan": plan_data, "execution_path": path}
+    return {"pmPlan": plan_data, "execution_path": path, "status_update": "Plan and cost estimate created. Awaiting approval."}
 
-# --- All other agent functions remain the same ---
-# ... (run_pm_agent, run_memory_retrieval, etc.)
 def run_memory_retrieval(state: AgentState):
     log.info("--- 🧠 Accessing Long-Term Memory ---")
     path = state.get('execution_path', []) + ["Memory Retriever"]
@@ -112,14 +105,14 @@ def run_memory_retrieval(state: AgentState):
     else:
         context = "No relevant memories found."
         log.info(context)
-    return {"retrievedMemory": context, "execution_path": path}
+    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 = state.get('execution_path', []) + ["Intent Agent"]
-    prompt = f"As an expert prompt engineer, analyze the user's request, conversation history, and retrieved memories. Refine it into a clear, actionable 'core objective prompt'.\n\nRelevant 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}
+    return {"coreObjectivePrompt": response.content, "execution_path": path, "status_update": "Clarifying the main objective..."}
 
 def run_pm_agent(state: AgentState):
     log.info("--- 👷 Running PM Agent ---")
@@ -128,35 +121,36 @@ def run_pm_agent(state: AgentState):
 
     path = state.get('execution_path', []) + ["PM Agent"]
     feedback = f"QA Feedback (must be addressed): {state.get('qaFeedback')}" if state.get('qaFeedback') else ""
-    prompt = f"As the Project Manager, decompose the core objective into a plan. Identify if code execution is needed. If so, define the goal of the experiment.\n\nCore Objective: {state.get('coreObjectivePrompt')}\n\n{feedback}\n\nProvide your plan ONLY as a valid JSON object within ```json code fences with keys: 'plan_steps' (list of strings) and 'experiment_needed' (boolean). If 'experiment_needed' is true, also include 'experiment_goal' (string)."
-    
+    prompt = f"Decompose the core objective into a plan. Determine if code execution is needed and define the goal.\n\nCore Objective: {state.get('coreObjectivePrompt')}\n\n{feedback}\n\nRespond in JSON with keys: 'plan_steps' (list), 'experiment_needed' (bool), 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)}")
-    return {"pmPlan": plan, "execution_path": path, "rework_cycles": current_cycles}
+    return {"pmPlan": plan, "execution_path": path, "rework_cycles": current_cycles, "status_update": "Breaking down the objective into a detailed plan..."}
 
 def run_experimenter_agent(state: AgentState):
     log.info("--- 🔬 Running Experimenter Agent ---")
     path = state.get('execution_path', []) + ["Experimenter Agent"]
     if not state.get('pmPlan', {}).get('experiment_needed'):
-        return {"experimentCode": None, "experimentResults": None, "execution_path": path}
-    
+        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. Do not require user input. Print results to standard output.\n\nGoal: {goal}\n\nPython Code:\n```python\n# Your code here\n```"
+    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 within fences.")
-        return {"experimentCode": "# ERROR: No code generated", "experimentResults": {"success": False, "stderr": "No valid Python code block was generated."}, "execution_path": path}
-    
+        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}
+    return {"experimentCode": code, "experimentResults": results, "execution_path": path, "status_update": "Running code to find a solution..."}
 
 def run_synthesis_agent(state: AgentState):
     log.info("--- ���️ Running Synthesis Agent ---")
@@ -164,51 +158,45 @@ def run_synthesis_agent(state: AgentState):
     exp_results = state.get('experimentResults')
     results_summary = "No experiment was conducted."
     if exp_results:
-        results_summary = f"An experiment was conducted. Output:\nSTDOUT:\n{exp_results.get('stdout', '')}\nSTDERR:\n{exp_results.get('stderr', '')}"
-    
+        results_summary = f"Experiment Output:\nSTDOUT:\n{exp_results.get('stdout', '')}\nSTDERR:\n{exp_results.get('stderr', '')}"
+
     prompt = f"Synthesize all information into a final response.\n\nCore Objective: {state.get('coreObjectivePrompt')}\n\nPlan: {state.get('pmPlan', {}).get('plan_steps')}\n\n{results_summary}\n\nFinal Response:"
     response = llm.invoke(prompt)
-    return {"draftResponse": response.content, "execution_path": path}
+    return {"draftResponse": response.content, "execution_path": path, "status_update": "Putting together the final response..."}
 
 def run_qa_agent(state: AgentState):
     log.info("--- ✅ Running QA Agent ---")
     path = state.get('execution_path', []) + ["QA Agent"]
-    prompt = f"Review the draft response based on the core objective. Respond ONLY with 'APPROVED' if it is complete and accurate. Otherwise, provide concise feedback for rework.\n\nCore Objective: {state.get('coreObjectivePrompt')}\n\nDraft Response: {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\nCore 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}
+        return {"approved": True, "qaFeedback": None, "execution_path": path, "status_update": "Response approved!"}
     else:
-        return {"approved": False, "qaFeedback": response.content or "No specific feedback provided.", "execution_path": path}
+        return {"approved": False, "qaFeedback": response.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 = state.get('execution_path', []) + ["Archivist Agent"]
-    summary_prompt = f"Create a concise summary of this successful task for long-term memory. Focus on key insights and results.\n\nCore Objective: {state.get('coreObjectivePrompt')}\n\nFinal Response: {state.get('draftResponse')}\n\nMemory Summary:"
-    
-    # FIX: Use the correct variable name 'summary_prompt' instead of 'prompt'
+    summary_prompt = f"Create a concise summary of this successful task for long-term memory.\n\nCore 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}
+    return {"execution_path": path, "status_update": "Saving key learnings for future reference..."}
 
 def run_disclaimer_agent(state: AgentState):
     log.warning("--- ⚠️ Running Disclaimer Agent ---")
     path = state.get('execution_path', []) + ["Disclaimer Agent"]
-    disclaimer = "**DISCLAIMER: The process was stopped after exhausting the user-defined budget. The following response is the best available draft but may be incomplete, unverified, or contain errors. Please review it carefully.**\n\n---\n\n"
-    final_response = disclaimer + state.get('draftResponse', "No response was generated before the budget was exhausted.")
-    return {"draftResponse": final_response, "execution_path": path}
-
+    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 ---
-# ... (should_continue and should_run_experiment are unchanged)
 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 state.get("rework_cycles", 0) > state.get("max_loops", 0):
-        log.error(f"BUDGET LIMIT REACHED. Aborting task after {state.get('rework_cycles', 0)-1} rework cycles.")
+        log.error(f"BUDGET LIMIT REACHED after {state.get('rework_cycles', 0)-1} cycles.")
         return "disclaimer_agent"
     else:
         log.info("Routing to: PM Agent for rework")
@@ -218,36 +206,22 @@ def should_run_experiment(state: AgentState):
     decision = "experimenter_agent" if state.get('pmPlan', {}).get('experiment_needed') else "synthesis_agent"
     return decision
 
-# --- Build the Graph (Updated with Triage) ---
-workflow = StateGraph(AgentState)
-workflow.add_node("triage", run_triage_agent)
-workflow.add_node("memory_retriever", run_memory_retrieval)
-# ... add all other nodes
-workflow.add_node("intent_agent", run_intent_agent)
-workflow.add_node("pm_agent", run_pm_agent)
-workflow.add_node("experimenter_agent", run_experimenter_agent)
-workflow.add_node("synthesis_agent", run_synthesis_agent)
-workflow.add_node("qa_agent", run_qa_agent)
-workflow.add_node("archivist_agent", run_archivist_agent)
-workflow.add_node("disclaimer_agent", run_disclaimer_agent)
-
-# NEW Triage Logic
-# We now have two separate graphs that are conditionally called by the app
-# 1. A simple Triage graph
+# --- Build the Graphs ---
+# 1. Triage Graph
 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
+# 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
+# 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)