commit step 7

agent/main.py

@@ -2,6 +2,8 @@ import os
 import time
 import json
 import requests
+import signal
+import sys
 from datetime import datetime
 import openai
 from anthropic import Anthropic
@@ -20,6 +22,17 @@ if "gpt" in LLM_MODEL:
 else:
     llm_client = Anthropic(api_key=LLM_API_KEY)
 
+# Global flag for interrupt handling
+interrupted = False
+
+def signal_handler(sig, frame):
+    global interrupted
+    interrupted = True
+    print(f"\n[{datetime.now()}] Interrupt received, will stop after current instruction")
+
+signal.signal(signal.SIGINT, signal_handler)
+signal.signal(signal.SIGTERM, signal_handler)
+
 def call_mcp(tool, params=None):
     response = requests.post(
         MCP_URL,
@@ -238,10 +251,76 @@ Return ONLY the SQL query, no explanations or markdown."""
             "error": str(e)
         }
 
+def check_workflow_stop(workflow_id):
+    """Check if workflow has been marked to stop"""
+    result = call_mcp("query_graph", {
+        "query": "MATCH (w:Workflow {id: $id}) RETURN w.status as status",
+        "parameters": {"id": workflow_id}
+    })
+    
+    if result['data'] and result['data'][0]['status'] == 'stopped':
+        return True
+    return False
+
+def pause_with_interrupt(duration, instruction_id, workflow_id=None):
+    """Pause for duration seconds with interrupt checking every 10 seconds"""
+    print(f"[{datetime.now()}] Pausing for {duration} seconds for human review")
+    print(f"[{datetime.now()}] You can edit instruction in Neo4j Browser:")
+    print(f"  MATCH (i:Instruction {{id: '{instruction_id}'}}) SET i.parameters = '{{\"key\": \"value\"}}'")
+    
+    # Log pause start
+    call_mcp("write_graph", {
+        "action": "create_node",
+        "label": "Log",
+        "properties": {
+            "type": "pause_started",
+            "instruction_id": instruction_id,
+            "duration": duration,
+            "timestamp": datetime.now().isoformat()
+        }
+    })
+    
+    elapsed = 0
+    while elapsed < duration:
+        # Check every 10 seconds
+        sleep_time = min(10, duration - elapsed)
+        time.sleep(sleep_time)
+        elapsed += sleep_time
+        
+        # Check for workflow stop
+        if workflow_id and check_workflow_stop(workflow_id):
+            print(f"[{datetime.now()}] Workflow stopped during pause")
+            return False
+        
+        # Check for global interrupt
+        if interrupted:
+            print(f"[{datetime.now()}] Interrupted during pause")
+            return False
+        
+        # Show progress
+        remaining = duration - elapsed
+        if remaining > 0 and elapsed % 30 == 0:  # Update every 30 seconds
+            print(f"[{datetime.now()}] Pause remaining: {remaining} seconds")
+    
+    # Log pause end
+    call_mcp("write_graph", {
+        "action": "create_node",
+        "label": "Log",
+        "properties": {
+            "type": "pause_completed",
+            "instruction_id": instruction_id,
+            "timestamp": datetime.now().isoformat()
+        }
+    })
+    
+    print(f"[{datetime.now()}] Pause complete, continuing execution")
+    return True
+
 def main():
+    global interrupted
     print(f"[{datetime.now()}] Agent starting, polling every {POLL_INTERVAL}s")
     
-    while True:
+    while not interrupted:
         try:
             result = call_mcp("get_next_instruction")
             instruction = result.get("instruction")
@@ -249,22 +328,48 @@ def main():
             if instruction:
                 print(f"[{datetime.now()}] Found instruction: {instruction['id']}, type: {instruction['type']}")
                 
+                # Get workflow ID
+                workflow_result = call_mcp("query_graph", {
+                    "query": """
+                        MATCH (w:Workflow)-[:HAS_INSTRUCTION]->(i:Instruction {id: $id})
+                        RETURN w.id as workflow_id
+                    """,
+                    "parameters": {"id": instruction['id']}
+                })
+                workflow_id = workflow_result['data'][0]['workflow_id'] if workflow_result['data'] else None
+                
+                # PAUSE BEFORE EXECUTION
+                pause_duration = instruction.get('pause_duration', 300)
+                if pause_duration > 0:
+                    if not pause_with_interrupt(pause_duration, instruction['id'], workflow_id):
+                        print(f"[{datetime.now()}] Execution cancelled during pause")
+                        continue
+                
+                # Re-fetch instruction to get any edits made during pause
+                refetch_result = call_mcp("query_graph", {
+                    "query": "MATCH (i:Instruction {id: $id}) RETURN i",
+                    "parameters": {"id": instruction['id']}
+                })
+                
+                if refetch_result['data']:
+                    instruction = refetch_result['data'][0]['i']
+                    print(f"[{datetime.now()}] Re-fetched instruction after pause, parameters: {instruction.get('parameters')}")
+                
                 # Update status to executing
                 call_mcp("query_graph", {
                     "query": "MATCH (i:Instruction {id: $id}) SET i.status = 'executing'",
                     "parameters": {"id": instruction['id']}
                 })
                 
-                # Handle different instruction types
+                # Execute based on type
                 if instruction['type'] == 'discover_schema':
                     exec_result = handle_discover_schema(instruction)
                 elif instruction['type'] == 'generate_sql':
                     exec_result = handle_generate_sql(instruction)
                 else:
-                    # Default dummy execution
-                    exec_result = {"status": "success", "result": "Dummy execution"}
+                    exec_result = {"status": "success", "result": "Reviewed"}
                 
-                # Create Execution node
+                # Store execution result
                 exec_node = call_mcp("write_graph", {
                     "action": "create_node",
                     "label": "Execution",
@@ -276,7 +381,7 @@ def main():
                     }
                 })
                 
-                # Link execution to instruction
+                # Link execution
                 call_mcp("query_graph", {
                     "query": """
                         MATCH (i:Instruction {id: $iid}), (e:Execution {id: $eid})
@@ -288,21 +393,25 @@ def main():
                     }
                 })
                 
-                # Update instruction status
+                # Update status
                 final_status = 'complete' if exec_result.get('status') == 'success' else 'failed'
                 call_mcp("query_graph", {
                     "query": "MATCH (i:Instruction {id: $id}) SET i.status = $status",
                     "parameters": {"id": instruction['id'], "status": final_status}
                 })
                 
-                print(f"[{datetime.now()}] Completed instruction: {instruction['id']} with status: {final_status}")
+                print(f"[{datetime.now()}] Completed instruction: {instruction['id']}")
+                
             else:
                 print(f"[{datetime.now()}] No pending instructions")
-                
+            
+            time.sleep(POLL_INTERVAL)
+            
         except Exception as e:
             print(f"[{datetime.now()}] Error: {e}")
-        
-        time.sleep(POLL_INTERVAL)
+            time.sleep(POLL_INTERVAL)
+    
+    print(f"[{datetime.now()}] Agent shutting down")
 
 if __name__ == "__main__":
     main()