feat: add prompt saving functionality and improve error handling in DebugSaver

Author: provos

examples/deepsearch/deepsearch/debug.py

@@ -53,6 +53,14 @@ def __init__(
         self._replay_threads: Dict[str, threading.Thread] = {}
         self._lock = threading.RLock()
 
+    def save_prompt(self, session_id: str, prompt: str):
+        """Save prompt for replay."""
+        if self.mode == "replay":
+            return
+
+        with open(self.output_dir / f"{session_id}.prompt", "w") as fp:
+            fp.write(prompt)
+
     def capture(self, func_name: str = None):
         # In replay mode, return a no-op decorator
         if self.mode == "replay":
@@ -88,9 +96,13 @@ def wrapper(*args, **kwargs):
                     "timestamp": datetime.now().isoformat(),
                 }
 
-                # Write the pickled call data
-                pickle.dump(call_data, self.fp)
-                self.fp.flush()
+                try:
+                    # Write the pickled call data
+                    pickle.dump(call_data, self.fp)
+                    self.fp.flush()
+                    logging.debug(f"Captured call to {func_name}: {call_data}")
+                except Exception as e:
+                    logging.error(f"Failed to capture call data: {e}")
 
                 return func(*args, **kwargs)
 
@@ -246,6 +258,14 @@ def load_replay_session(self, session_id: str):
 
     def load_replays(self):
         """Load all replay data into memory."""
+
+        if not any(self.output_dir.glob("*.prompt")):
+            logging.warning(f"No replay files found in {self.output_dir}")
+            return
+
         for session_file in self.output_dir.glob("*.prompt"):
             session_id = session_file.stem
-            self.load_replay_session(session_id)
+            try:
+                self.load_replay_session(session_id)
+            except Exception as e:
+                logging.error(f"Failed to load replay session {session_id}: {e}")

examples/deepsearch/deepsearch/deepsearch.py

@@ -115,7 +115,6 @@ def worker():
     return worker_thread
 
 
-@debug_saver.capture("notify") if debug_saver else lambda x: x
 def notify(metadata, message: Response):
     """Callback to receive notifications from the graph."""
     global task_queue, session_manager
@@ -272,14 +271,16 @@ def handle_message(data):
 
     # If in replay mode, trigger replay with current session
     global debug_saver
-    if debug_saver and debug_saver.mode == "replay":
-        debug_saver.start_replay_session(message.strip(), current_metadata)
-        return
+    if debug_saver:
+        if debug_saver.mode == "replay":
+            debug_saver.start_replay_session(message.strip(), current_metadata)
+            return
+        else:
+            debug_saver.save_prompt(session_id, message.strip())
 
     # Update session timestamp on activity
     session_manager.update_session_timestamp(session_id)
 
-    @debug_saver.capture("notify_planai") if debug_saver else lambda x: x
     def wrapped_notify_planai(*args, **kwargs):
         return notify_planai(*args, **kwargs)
 
@@ -295,7 +296,6 @@ def wrapped_notify_planai(*args, **kwargs):
     session_metadata["provenance"] = provenance
 
 
-@debug_saver.capture("notify_planai") if debug_saver else lambda x: x
 def notify_planai(
     metadata: Dict[str, Any],
     prefix: ProvenanceChain,
@@ -312,7 +312,7 @@ def notify_planai(
         global session_manager
         # get the metadata for this session
         session_metadata = session_manager.metadata(session_id)
-        if session_metadata.get("started"):
+        if session_metadata.get("started") and False:  # ignore for now
             # this indicates that we failed the task
             task_queue.put(
                 (
@@ -360,6 +360,14 @@ def notify_planai(
     )
 
 
+def patch_notify_functions():
+    """Patch the notify functions with debug_saver decorators."""
+    global notify, notify_planai, debug_saver
+    if debug_saver:
+        notify = debug_saver.capture("notify")(notify)
+        notify_planai = debug_saver.capture("notify_planai")(notify_planai)
+
+
 def main():
     import argparse
 
@@ -386,6 +394,9 @@ def main():
             replay_delay=args.replay_delay,
         )
 
+        # Patch the notify functions after debug_saver is initialized
+        patch_notify_functions()
+
         if args.replay:
             # Register the original functions for replay
             debug_saver.register_replay_handler("notify", notify)