feat: enhanced support for results directory and more metadata

Signed-off-by: vsoch <[email protected]>

Author: vsoch

fractale/agent/base.py

@@ -1,9 +1,9 @@
-import json
 import os
 import sys
 
 import google.generativeai as genai
 
+from fractale.agent.decorators import callback, save_logs
 import fractale.agent.defaults as defaults
 import fractale.agent.logger as logger
 import fractale.utils as utils
@@ -23,20 +23,29 @@ class Agent:
 
     # name and description should be on the class
 
-    def __init__(self, use_cache=False):
+    def __init__(self, use_cache=False, results_dir=None, incremental=False):
 
         # Max attempts defaults to unlimited
         # We start counting at 1 for the user to see.
         # Eat your heart out, Matlab.
         self.attempts = 1
         self.max_attempts = None
 
+        # For now, assume this is for the manager.
+        # We can optionally save incremental result objects
+        self.results_dir = results_dir or os.getcwd()
+        self.save_incremental = incremental
+
         # The user can save if desired - caching the context to skip steps that already run.
         self.setup_cache(use_cache)
 
+        # This supports saving custom logs and step (attempt) metadata
+        self.metadata = {}
+
         # Custom initialization functions
         self.init()
 
+    @callback(save_logs)
     def run(self, context):
         """
         Run the agent - a wrapper around internal function _run that prepares it.

fractale/agent/decorators.py

@@ -0,0 +1,32 @@
+import time
+
+def callback(callback_func):
+    """
+    A decorator that executes a callback function after the decorated function.
+
+    We need this so that specific functions for the agent can return objects that save
+    one or more metadata items (automatically).
+    """
+    def decorator(func):
+        def wrapper(self, *args, **kwargs):
+            # This is the original function
+            start = time.time()
+            result = func(self, *args, **kwargs)
+            # Get the result and pass to the callback!
+            end = time.time()
+            callback_func(self, result, end-start)
+            return result
+        return wrapper
+    return decorator
+
+
+def save_logs(instance, context, elapsed_time):
+    """
+    If defined (requested by the user) save the stage result.
+    """
+    result = context.get('result')
+    if not instance.save_incremental or not result:
+        return
+    if "logs" not in instance.metadata:
+        instance.metadata['logs'] = []  
+    instance.metadata['logs'].append({"log": result, "elapsed_time_seconds": elapsed_time})

fractale/agent/kubernetes_job/agent.py

@@ -151,6 +151,9 @@ def get_diagnostics(self, job_name, namespace):
         Helper to collect rich error data for a failed job.
         """
         print("[yellow]Gathering diagnostics for failed job...[/yellow]")
+        print('TODO streamline')
+        import IPython 
+        IPython.embed()
 
         describe_job_cmd = ["kubectl", "describe", "job", job_name, "-n", namespace]
         job_description = subprocess.run(

fractale/agent/manager/agent.py

@@ -78,9 +78,11 @@ def save_results(self, tracker):
 
         Just ploop into pwd for now, we can eventually take a path.
         """
+        if not os.path.exists(self.results_dir):
+            os.makedirs(self.results_dir)
         now = datetime.now()
         timestamp = now.strftime("%Y-%m-%d_%H-%M-%S")
-        results_file = os.path.join(os.getcwd(), f"results-{timestamp}.json")
+        results_file = os.path.join(self.results_dir, f"results-{timestamp}.json")
         utils.write_json(tracker, results_file)
 
     def run(self, context):
@@ -167,7 +169,16 @@ def run_tasks(self, context, plan):
 
             # Keep track of running the agent and the time it took
             # Also keep result of each build step (we assume there is one)
-            tracker.append([step.agent, timer.elapsed_time, context.get("result")])
+            # We will eventually want to also save the log.
+            tracker.append(
+                {
+                    "agent": step.agent,
+                    "elapsed_time_seconds": timer.elapsed_time,
+                    "result": context.get("result"),
+                    "attempts": step.attempts,
+                    "logs": step.logs,
+                }
+            )
 
             # If we are successful, we go to the next step.
             # Not setting a return code indicates success.

fractale/cli/__init__.py

@@ -97,6 +97,16 @@ def get_parser():
         action="store_true",
         default=False,
     )
+    agent.add_argument(
+        "--results",
+        help="Save to a custom results directory.",
+    )
+    agent.add_argument(
+        "--incremental",
+        help="Save incremental results for later inspection",
+        action="store_true",
+        default=False,
+    )
 
     # Add agent parsers
     parsers.register(agents)

fractale/cli/agent.py

@@ -19,7 +19,9 @@ def main(args, extra, **kwargs):
 
     # Get the agent and run!
     # Save determines if we want to save state to an output directory
-    agent = agents[args.agent_name](use_cache=args.use_cache)
+    agent = agents[args.agent_name](
+        use_cache=args.use_cache, results_dir=args.results, incremental=args.incremental
+    )
 
     # This is the context - we can remove variables not needed
     context = vars(args)