test: refactor pod/job into abstractions

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

Author: vsoch

fractale/agent/base.py

@@ -3,11 +3,11 @@
 
 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
 from fractale.agent.context import get_context
+from fractale.agent.decorators import callback, save_logs
 
 
 class Agent:
@@ -23,18 +23,17 @@ class Agent:
 
     # name and description should be on the class
 
-    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
+    def __init__(
+        self, use_cache=False, results_dir=None, save_incremental=False, max_attempts=None
+    ):
+        self.attempts = 0
+        self.max_attempts = max_attempts
 
-        # For now, assume this is for the manager.
+        # For now, assume these are for the manager.
+        # They get added to other agents via the step creation
         # We can optionally save incremental result objects
         self.results_dir = results_dir or os.getcwd()
-        self.save_incremental = incremental
+        self.save_incremental = save_incremental
 
         # The user can save if desired - caching the context to skip steps that already run.
         self.setup_cache(use_cache)
@@ -50,9 +49,6 @@ def run(self, context):
         """
         Run the agent - a wrapper around internal function _run that prepares it.
         """
-        # Init attempts. Each agent has an internal counter for total attempts
-        self.attempts = self.attempts or 1
-
         # Load cached context. This is assumed to override user provided args
         # If we have a saved context, we assume we want to use it, return early
         cached_context = self.load_cache()
@@ -66,6 +62,7 @@ def run(self, context):
         context = get_context(context)
 
         # Run, wrapping with a load and save of cache
+        # This will return here when the internal loop is done
         context = self._run(context)
         self.save_cache(context)
         return context
@@ -79,6 +76,17 @@ def print_result(self, result):
         """
         pass
 
+    def reset_context(self, context):
+        """
+        Remove output and any stateful variables. This is assuming we
+        are starting again.
+        """
+        for key in ["result", "error_message"]:
+            if key in context.data:
+                del context.data[key]
+        # We don't need a return here, but let's be explicit
+        return context
+
     def setup_cache(self, use_cache=False):
         """
         Setup (or load) a cache.
@@ -132,10 +140,7 @@ def reached_max_attempts(self):
         # Unset (None) or 1.
         if not self.max_attempts:
             return False
-        return self.attempts >= self.max_attempts
-
-    def set_max_attempts(self, max_attempts):
-        self.max_attempts = max_attempts
+        return self.attempts > self.max_attempts
 
     def add_shared_arguments(self, agent):
         """
@@ -214,14 +219,6 @@ def _run(self, context):
         assert context
         raise NotImplementedError(f"The {self.name} agent is missing internal 'run' function")
 
-    def get_initial_prompt(self, context):
-        """
-        Get the initial prompt (with details) to provide context to the manager.
-
-        If we don't do this, the manager can provide a bad instruction for how to fix the error.
-        """
-        return self.get_prompt(context)
-
     def get_prompt(self, context):
         """
         This function should take the same context as run and return the parsed prompt that

fractale/agent/build/agent.py

@@ -57,6 +57,17 @@ def _add_arguments(self, subparser):
         )
         return build
 
+    def reset_context(self, context):
+        """
+        Remove output and any stateful variables. This is assuming we
+        are starting again.
+        """
+        for key in ["result", "dockerfile", "error_message"]:
+            if key in context.data:
+                del context.data[key]
+        # We don't need a return here, but let's be explicit
+        return context
+
     def get_prompt(self, context):
         """
         Get the prompt for the LLM. We expose this so the manager can take it
@@ -125,6 +136,7 @@ def _run(self, context):
             agent = DebugAgent()
             # This updates the error message to be the output
             context = agent.run(context, requires=prompts.requires)
+            print("\n[bold cyan] Requesting Correction from Build Agent[/bold cyan]")
 
             # If we have reached the max attempts...
             if self.reached_max_attempts():
@@ -139,10 +151,9 @@ def _run(self, context):
                 logger.exit(f"Max attempts {self.max_attempts} reached.", title="Agent Failure")
 
             self.attempts += 1
-            print("\n[bold cyan] Requesting Correction from Build Agent[/bold cyan]")
 
             # Update the context with error message
-            return self.run(context)
+            return self._run(context)
 
         # Add generation line
         self.write_file(context, context.result)
@@ -219,6 +230,17 @@ def build(self, context):
         shutil.rmtree(build_dir, ignore_errors=True)
         return (p.returncode, p.stdout + p.stderr)
 
+    def save_dockerfile(self, dockerfile):
+        """
+        Save logs to metadata
+        """
+        if self.save_incremental:
+            if "steps" not in self.metadata:
+                self.metadata["steps"] = []
+            self.metadata["steps"].append(
+                {"item": dockerfile, "type": "dockerfile", "attempt": self.attempts}
+            )
+
     def generate_dockerfile(self, context):
         """
         Generates or refines a Dockerfile using the Gemini API.
@@ -241,6 +263,7 @@ def generate_dockerfile(self, context):
                 dockerfile = match.group(1).strip()
             else:
                 dockerfile = content.strip()
+            self.save_dockerfile(dockerfile)
 
             # The result is saved as a build step
             # The dockerfile is the argument used internally

fractale/agent/decorators.py

@@ -1,32 +1,45 @@
 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)
+            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')
+    return save_general(instance, context, elapsed_time, "final-result")
+
+
+def save_general(instance, context, elapsed_time, result_type):
+    """
+    Shared saving function.
+    """
+    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})
+        instance.metadata["logs"] = []
+    instance.metadata["logs"].append(
+        {"item": result, "type": result_type, "total_seconds": elapsed_time}
+    )