feat: add helper agents

a helper agent is intended to run (and not retry)
and be one level lower than a step. This will allow
for getting back large error outputs and streamlining
to better get to an issue.

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

Author: vsoch

README.md

@@ -12,12 +12,21 @@ This library is primarily being used for development for the descriptive thrust
 ### Agents
 
 The `fractale agent` command provides means to run build, job generation, and deployment agents.
-This part of the library is under development.
+This part of the library is under development. There are three kinds of agents:
+
+ - `step` agents are experts on doing specific tasks (do hold state)
+ - `manager` agents know how to orchestrate step agents and choose between them (don't hold state, but could)
+ - `helper` agents are used by step agents to do small tasks (e.g., suggest a fix for an error)
+
+The design is simple in that each agent is responding to state of error vs. success. In the case of a step agent, the return code determines to continue or try again. In the case of a helper, the input is typically an erroneous response (or something that needs changing) with respect to a goal.
+For a manager, we are making a choice based on a previous erroneous step.
 
 See [examples/agent](examples/agent) for an example.
 
 #### To do items
 
+- refactor manager to not handle prompt, just get step when retries come back.
+- then need to decide how to handle kubernetes job creating additional structures.
 - Get basic runner working
 - Add in ability to get log and optimize - the manager will need to use goal
 - We likely want the manager to be able to edit the prompt.
@@ -28,21 +37,29 @@ See [examples/agent](examples/agent) for an example.
 
 **And experiment ideas**
 
-- How do we define stability? 
+- How do we define stability?
 - What are the increments of change (e.g., "adding a library")? We should be able to keep track of times for each stage and what changed, and an analyzer LLM can look at result and understand (categorize) most salient contributions to change.
   - We also can time the time it takes to do subsequent changes, when relevant. For example, if we are building, we should be able to use cached layers (and the build times speed up) if the LLM is changing content later in the Dockerfile.
 - We can also save the successful results (Dockerfile builds, for example) and compare for similarity. How consistent is the LLM?
 - How does specificity of the prompt influence the result?
 - For an experiment, we would want to do a build -> deploy and successful run for a series of apps and get distributions of attempts, reasons for failure, and a general sense of similarity / differences.
 - For the optimization experiment, we'd want to do the same, but understand gradients of change that led to improvement.
 
-## Observations
+#### Observations
 
 - Specifying cpu seems important - if you don't it wants to do GPU
 - If you ask for a specific example, it sometimes tries to download data (tell it where data is)
 - Always include common issues in the initial prompt
 - If you are too specific about instance types, it adds node selectors/affinity, and that often doesn't work.
 
+#### Ideas
+
+- The manager agent is currently generated an updated prompt AND choosing the step.
+ - Arguably we should have a separation of responsibility so a step can ask to fix an error without a manager.
+- I think we need one more level of agent - a step agent should have helper agents that can:
+ - take an error message and analyze to get a fix.
+
+
 ### Job Specifications
 
 #### Simple

examples/agent/README.md

@@ -10,7 +10,7 @@ The build agent will use the Gemini API to generate a Dockerfile and then build
 Here is how to first ask the build agent to generate a lammps container for Google cloud.
 
 ```bash
-fractale agent build lammps --environment "google cloud" --outfile dockerfile
+fractale agent build lammps --environment "google cloud CPU" --outfile Dockerfile.lammps
 ```
 
 That might generate the [Dockerfile](Dockerfile) here, and a container that defaults to the application name "lammps"
@@ -27,7 +27,7 @@ kind load docker-image lammps
 To start, we will assume a kind cluster running and tell the agent the image is loaded into it (and so the pull policy will be never). 
 
 ```bash
-fractale agent kubernetes-job lammps --environment "google cloud CPU" --context-file ./Dockerfile --no-pull 
+fractale agent kubernetes-job lammps --environment "google cloud CPU" --context-file ./Dockerfile --no-pull
 ```
 
 ## Manager

examples/agent/plans/run-lammps.yaml

@@ -22,7 +22,6 @@ plan:
     environment: "google cloud CPU instance in Kubernetes" 
     max_attempts: 1
     details: |
-      Please execute the reaxff HNS example, and assume the data in the PWD,
+      Please execute the in.reaxff.hns example, and assume the data in the PWD,
       Run lammpss with params -v x 2 -v y 2 -v z 2 -in ./in.reaxff.hns
-      and with -nocite flag for CPU. Do not try to generate configmap data.
-      Do not add any nodeSelector or affinity rules since we are testing.
+      and with -nocite flag for CPU.

fractale/agent/base.py

@@ -1,5 +1,12 @@
+import os
+import sys
+
+import google.generativeai as genai
+
+import fractale.agent.defaults as defaults
 import fractale.utils as utils
 
+
 class Agent:
     """
     A base for an agent. Each agent should:
@@ -64,7 +71,7 @@ def write_file(self, context, content, add_comment=True):
             content += f"\n# Generated by fractale {self.name} agent"
         utils.write_file(content, outfile)
 
-    def get_code_block(self, content, code_type):             
+    def get_code_block(self, content, code_type):
         """
         Parse a code block from the response
         """
@@ -76,22 +83,13 @@ def get_code_block(self, content, code_type):
             content = content[: -len("```")]
         return content
 
-    def ask_gemini(self, prompt, with_history=True):
+    def get_result(self, context):
         """
-        Ask gemini adds a wrapper with some error handling.
+        Return either the entire context or single result.
         """
-        try:
-            if with_history:
-                response = self.chat.send_message(prompt)
-            else:
-                response = self.model.generate_content(prompt)
-
-            # This line can fail. If it succeeds, return entire response
-            return response.text.strip()
-
-        except ValueError as e:
-            print(f"[Error] The API response was blocked and contained no text: {str(e)}")
-            return "GEMINI ERROR: The API returned an error (or stop) and we need to try again."
+        if context.is_managed:
+            return context
+        return context.result
 
     def run(self, context):
         """
@@ -115,3 +113,34 @@ def get_prompt(self, context):
         """
         assert context
         raise NotImplementedError(f"The {self.name} agent is missing a 'get_prompt' function")
+
+
+class GeminiAgent(Agent):
+    """
+    A base for an agent that uses the Gemini API.
+    """
+
+    def init(self):
+        self.model = genai.GenerativeModel(defaults.gemini_model)
+        self.chat = self.model.start_chat()
+        try:
+            genai.configure(api_key=os.environ["GEMINI_API_KEY"])
+        except KeyError:
+            sys.exit("ERROR: GEMINI_API_KEY environment variable not set.")
+
+    def ask_gemini(self, prompt, with_history=True):
+        """
+        Ask gemini adds a wrapper with some error handling.
+        """
+        try:
+            if with_history:
+                response = self.chat.send_message(prompt)
+            else:
+                response = self.model.generate_content(prompt)
+
+            # This line can fail. If it succeeds, return entire response
+            return response.text.strip()
+
+        except ValueError as e:
+            print(f"[Error] The API response was blocked and contained no text: {str(e)}")
+            return "GEMINI ERROR: The API returned an error (or stop) and we need to try again."

fractale/agent/build/agent.py

@@ -1,7 +1,7 @@
-from fractale.agent.base import Agent
+from fractale.agent.base import GeminiAgent
 import fractale.agent.build.prompts as prompts
 from fractale.agent.context import get_context
-import fractale.agent.defaults as defaults
+from fractale.agent.errors import DebugAgent
 
 import fractale.utils as utils
 import argparse
@@ -18,29 +18,23 @@
 import subprocess
 import textwrap
 
-import google.generativeai as genai
 
 # regular expression in case LLM does not follow my instructions!
 dockerfile_pattern = r"```(?:dockerfile)?\n(.*?)```"
 
 
-class BuildAgent(Agent):
+class BuildAgent(GeminiAgent):
     """
     Builder agent.
+
+    Observations from v:
+    1. Holding the context (chat) seems to take longer.
+    2. Don't forget to ask for CPU - GPU will take a lot longer.
     """
 
     name = "build"
     description = "builder agent"
 
-    def init(self):
-        """
-        Custom initialization. I want to try using the same model
-        agent across requests. I'm not sure if that means it's the same
-        context (I don't think so).
-        """
-        model = genai.GenerativeModel("gemini-2.5-pro")
-        self.chat = model.start_chat()
-
     def add_arguments(self, subparser):
         """
         Add arguments for the plugin to show up in argparse
@@ -102,6 +96,8 @@ def run(self, context):
         1. Populate a context.
         2. Call supporting functions with the context.
         3. Parse the result and update context, taking appropriate action.
+        4. The current object to generate should be put into result.
+        5. The current issue or error goes into error_message.
         """
         # Create or get global context
         context = get_context(context)
@@ -110,15 +106,14 @@ def run(self, context):
         # Start at 1 since we are showing to a user.
         self.attempts = self.attempts or 1
 
-        try:
-            genai.configure(api_key=os.environ["GEMINI_API_KEY"])
-        except KeyError:
-            sys.exit("ERROR: GEMINI_API_KEY environment variable not set.")
-
         # This will either generate fresh or rebuild erroneous Dockerfile
         # We don't return the dockerfile because it is updated in the context
         self.generate_dockerfile(context)
-        print(Panel(context.dockerfile, title="[green]Dockerfile or Response[/green]", border_style="green"))
+        print(
+            Panel(
+                context.result, title="[green]Dockerfile or Response[/green]", border_style="green"
+            )
+        )
 
         # Set the container on the context for a next step to use it...
         container = context.get("container") or self.generate_name(context.application)
@@ -127,50 +122,53 @@ def run(self, context):
         # Build it! We might want to only allow a certain number of retries or incremental changes.
         return_code, output = self.build(context)
         if return_code == 0:
+            self.print_dockerfile(context.result)
             print(
                 Panel(
                     f"[bold green]✅ Build complete in {self.attempts} attempts[/bold green]",
                     title="Success",
                     border_style="green",
                 )
             )
+
         else:
             print(
                 Panel(
                     "[bold red]❌ Build failed[/bold red]", title="Build Status", border_style="red"
                 )
             )
+            # Ask the debug agent to better instruct the error message
+            # This becomes a more guided output
+            context.error_message = output
+            agent = DebugAgent()
+            # This updates the error message to be the output
+            context = agent.run(context, requires=prompts.requires)
+
+            # TODO: test this idea extending to manager
+            # manager should not be deciding what to do on failure,
+            # but decidin what to do (step) AFTER reach limit
             # If we are returning a failure:
             # 1. Set context.return_code
             # 2. error message is the result
-            if self.return_on_failure():
-                context.return_code = -1
-                context.result = output
-                return self.get_result(context)
+            # if self.return_on_failure():
+            #    context.return_code = -1
+            #    # TODO we should not have the manager parse error...
+            #    context.result = context.error_message
+            #    return self.get_result(context)
 
             self.attempts += 1
             print("\n[bold cyan] Requesting Correction from Build Agent[/bold cyan]")
 
             # Update the context with error message
-            context.error_message = output
             return self.run(context)
 
         # Add generation line
-        self.write_file(context, context.dockerfile)
-        self.print_dockerfile(context.dockerfile)
+        self.write_file(context, context.result)
 
         # Assume being called by a human that wants Dockerfile back,
         # unless we are being managed
         return self.get_result(context)
 
-    def get_result(self, context):
-        """
-        Return either the entire context or single result.
-        """
-        if context.is_managed:
-            return context
-        return context.dockerfile
-
     def print_dockerfile(self, dockerfile):
         """
         Print Dockerfile with highlighted Syntax
@@ -223,7 +221,7 @@ def build(self, context):
 
         # Write the Dockerfile to the temporary directory
         utils.write_file(dockerfile, os.path.join(build_dir, "Dockerfile"))
-        
+
         # If only one max attempt, don't print here, not important to show.
         if self.max_attempts is not None and self.max_attempts > 1:
             print(
@@ -252,15 +250,15 @@ def generate_dockerfile(self, context):
         """
         prompt = self.get_prompt(context)
         print("Sending build prompt to Gemini...")
-        print(textwrap.indent(prompt, "> ", predicate=lambda _: True))
+        print(textwrap.indent(prompt[0:1000], "> ", predicate=lambda _: True))
 
         # The API can error and not return a response.text.
         content = self.ask_gemini(prompt)
         print("Received Dockerfile response from Gemini...")
 
         # Try to remove Dockerfile from code block
         try:
-            content = self.get_code_block(content, 'dockerfile')
+            content = self.get_code_block(content, "dockerfile")
 
             # If we are getting commentary...
             match = re.search(dockerfile_pattern, content, re.DOTALL)
@@ -270,7 +268,8 @@ def generate_dockerfile(self, context):
                 dockerfile = content.strip()
 
             # The result is saved as a build step
-            context.dockerfile = dockerfile
+            # The dockerfile is the argument used internally
             context.result = dockerfile
+            context.dockerfile = dockerfile
         except Exception as e:
             sys.exit(f"Error parsing response from Gemini: {e}\n{content}")