Merge pull request #60 from AI-Hypercomputer/luke-maxkernel-updates-after-testing

Author: LukeA12

MaxKernel/auto_agent/config.py

@@ -12,6 +12,7 @@
 TPU_VERSION = os.environ.get("TPU_VERSION", "")
 RAG_CORPUS = os.environ.get("RAG_CORPUS", "")
 INCLUDE_THOUGHTS = os.environ.get("INCLUDE_THOUGHTS", "true").lower() == "true"
+MAX_COMPILATION_RETRIES = int(os.environ.get("MAX_COMPILATION_RETRIES", "6"))
 
 # Model configuration
 model_config = types.GenerateContentConfig(

MaxKernel/auto_agent/subagents/kernel_writing/agent.py

@@ -15,7 +15,7 @@
   load_kernel_and_plan_to_state,
   load_single_kernel_to_state,
 )
-from auto_agent.config import model_config, thinking_planner
+from auto_agent.config import model_config, thinking_planner, MAX_COMPILATION_RETRIES
 from auto_agent.constants import MODEL_NAME
 from auto_agent.custom_types import CustomLlmAgent
 from auto_agent.subagents.kernel_writing.kernel_compilation import (
@@ -45,15 +45,15 @@ class KernelCompilationValidationLoop(BaseAgent):
   compilation_checker: Optional[BaseAgent] = None
   fix_agent: Optional[BaseAgent] = None
   debug_agent: Optional[BaseAgent] = None
-  max_retries: int = 4
+  max_retries: int = 6
 
   def __init__(
     self,
     name: str,
     compilation_checker: BaseAgent,
     fix_agent: BaseAgent,
     debug_agent: Optional[BaseAgent] = None,
-    max_retries: int = 4,
+    max_retries: int = 6,
   ):
     super().__init__(
       name=name,
@@ -351,7 +351,7 @@ async def _run_async_impl(
   compilation_checker=kernel_compilation_checker_for_validation,
   fix_agent=fix_kernel_compilation_agent,
   debug_agent=add_debug_statements_agent,
-  max_retries=6,
+  max_retries=MAX_COMPILATION_RETRIES,
 )
 
 kernel_compilation_summary_agent = CustomLlmAgent(

MaxKernel/auto_agent/subagents/kernel_writing/prompts/ask_validation_prompt.py

@@ -8,7 +8,7 @@
 YOUR TASK:
 Present the user with clear options for what to do next:
 
-1. **Validate & Auto-Fix Compilation**: Run automatic compilation validation with error fixing (up to 4 attempts)
+1. **Validate & Auto-Fix Compilation**: Run automatic compilation validation with error fixing (includes up to 6 automatic error-fixing attempts; you can ask me to change this limit at any time)
 2. **Generate Tests**: Skip validation and proceed directly to test generation
 3. **Something Else**: Ask a question or perform another action
 

MaxKernel/hitl_agent/agent.py

@@ -31,7 +31,11 @@
   unified_test_agent,
   validated_test_generation_agent,
 )
-from hitl_agent.tools.tools import filesystem_tool_r
+from hitl_agent.tools.tools import (
+  filesystem_tool_r,
+  set_working_directory,
+  set_max_compilation_retries,
+)
 
 # Root orchestration agent
 root_agent = CustomLlmAgent(
@@ -56,7 +60,9 @@
     autotune_agent,  # Step 7: Auto-tune kernel
   ],
   tools=[
-    filesystem_tool_r
+    filesystem_tool_r,
+    set_working_directory,
+    set_max_compilation_retries,
   ],  # Read-only access - orchestrator delegates writes to sub-agents
   instruction=interactive_prompt.PROMPT,
   description="Orchestrates the human-in-the-loop kernel generation process with GPU to JAX conversion capability.",

MaxKernel/hitl_agent/config.py

@@ -12,6 +12,7 @@
 TPU_VERSION = os.environ.get("TPU_VERSION", "")
 RAG_CORPUS = os.environ.get("RAG_CORPUS", "")
 INCLUDE_THOUGHTS = os.environ.get("INCLUDE_THOUGHTS", "true").lower() == "true"
+MAX_COMPILATION_RETRIES = int(os.environ.get("MAX_COMPILATION_RETRIES", "6"))
 
 # Model configuration
 model_config = types.GenerateContentConfig(

MaxKernel/hitl_agent/prompts/interactive_prompt.py

@@ -10,10 +10,12 @@
 ### Your Capabilities
 
 1.  **Filesystem Tool (Read-Only)**: You have a `filesystem_tool` that can **only read and list files** - you cannot write files. All file writing must be delegated to the appropriate sub-agents. Assume that all files you need to read are located in {workdir}. For example, if the user asks to read file `script.py`, you should call the tool with the path `{workdir}/script.py`.
-2.  **Specialist Sub-Agents**: You have a team of sub-agents you can delegate tasks to:
+2.  **Set Working Directory Tool**: You have a `set_working_directory` tool that can update the active workspace/working directory path for the session. Use this tool when the user requests to switch, set, or change the active working directory or workspace folder.
+3.  **Set Max Compilation Retries Tool**: You have a `set_max_compilation_retries` tool that can update the maximum number of compilation validation and auto-fixing attempts. Use this tool when the user requests to set or change the maximum compilation retries or fixing attempts count.
+4.  **Specialist Sub-Agents**: You have a team of sub-agents you can delegate tasks to:
     * **PlanKernelAgent**: Creates or revises optimization plans for Pallas kernels. Automatically presents plans to the user with approval options. Use for both new plans ("optimize X") and revisions ("change Y in the plan").
     * **ImplementKernelAgent**: Implements a Pallas kernel following an approved plan. Use ONLY after user approves the plan. Automatically asks the user about validation options after completion.
-    * **ValidateKernelCompilationAgent**: Validates kernel compilation with automatic error fixing and debugging (up to 4 attempts). Use when user requests validation or wants to check compilation.
+    * **ValidateKernelCompilationAgent**: Validates kernel compilation with automatic error fixing and debugging (up to 6 attempts by default; the user can ask you to change this limit). Use when user requests validation or wants to check compilation.
     * **ExplanationAgent**: Explains TPU or Pallas concepts.
     * **GenerateTestFileAgent**: Generates a comprehensive pytest test file with compilation, correctness, and performance tests for kernel files.
     * **UnifiedTestAgent**: Executes the generated pytest test file on TPU and provides comprehensive results including full tracebacks. Automatically manages server lifecycle (starts/stops TPU and eval servers as needed).
@@ -33,9 +35,19 @@
     * Is it test execution? (e.g., "Run the tests", "Test the kernels")
     * Is it profiling? (e.g., "Profile the kernel", "What are the bottlenecks?")
     * Is it GPU conversion? (e.g., "Convert CUDA to JAX", "Write JAX from PyTorch")
+    * Is it setting/changing the workspace or working directory? (e.g., "Change working directory to /path/to/folder", "Set workspace folder to Y")
+    * Is it changing the maximum compilation retries or attempts? (e.g., "Set max retries to 8", "Change max fixing attempts to 10")
 
 2.  **Execute the Plan**:
 
+    * **If the request is to CHANGE/SET the workspace or working directory**:
+        * **Action**: Call `set_working_directory` tool with the absolute path.
+        * **Note**: Set `persist=True` ONLY if the user explicitly asks to persist, save, or make the change permanent. Otherwise, leave it as `False` (default).
+
+    * **If the request is to CHANGE/SET the maximum compilation retries**:
+        * **Action**: Call `set_max_compilation_retries` tool with the integer value.
+        * **Note**: Set `persist=True` ONLY if the user explicitly asks to persist, save, or make the change permanent. Otherwise, leave it as `False` (default).
+
     * **If the request is simple explanation** (like "What is a TPU?"):
         * **Action**: Delegate to `ExplanationAgent`.
 

MaxKernel/hitl_agent/subagents/kernel_writing/agent.py

@@ -15,7 +15,7 @@
   load_single_kernel_to_state,
   save_kernel_and_plan_paths,
 )
-from hitl_agent.config import model_config, thinking_planner
+from hitl_agent.config import model_config, thinking_planner, MAX_COMPILATION_RETRIES
 from hitl_agent.constants import MODEL_NAME
 from hitl_agent.custom_types import CustomLlmAgent
 from hitl_agent.subagents.kernel_writing.kernel_compilation import (
@@ -41,15 +41,15 @@ class KernelCompilationValidationLoop(BaseAgent):
   compilation_checker: Optional[BaseAgent] = None
   fix_agent: Optional[BaseAgent] = None
   debug_agent: Optional[BaseAgent] = None
-  max_retries: int = 4
+  max_retries: int = 6
 
   def __init__(
     self,
     name: str,
     compilation_checker: BaseAgent,
     fix_agent: BaseAgent,
     debug_agent: Optional[BaseAgent] = None,
-    max_retries: int = 4,
+    max_retries: int = 6,
   ):
     super().__init__(
       name=name,
@@ -372,7 +372,7 @@ async def _run_async_impl(
   compilation_checker=kernel_compilation_checker_for_validation,
   fix_agent=fix_kernel_compilation_agent,
   debug_agent=add_debug_statements_agent,
-  max_retries=4,
+  max_retries=MAX_COMPILATION_RETRIES,
 )
 
 kernel_compilation_summary_agent = CustomLlmAgent(

MaxKernel/hitl_agent/subagents/kernel_writing/prompts/ask_validation_prompt.py

@@ -8,7 +8,7 @@
 YOUR TASK:
 Present the user with clear options for what to do next:
 
-1. **Validate & Auto-Fix Compilation**: Run automatic compilation validation with error fixing (up to 4 attempts)
+1. **Validate & Auto-Fix Compilation**: Run automatic compilation validation with error fixing (includes up to 6 automatic error-fixing attempts; you can ask me to change this limit at any time)
 2. **Generate Tests**: Skip validation and proceed directly to test generation
 3. **Something Else**: Ask a question or perform another action
 

MaxKernel/hitl_agent/tools/filesystem_tools.py

@@ -0,0 +1,40 @@
+"""Filesystem tools configuration for the HITL agent."""
+
+import os
+from google.adk.tools.mcp_tool.mcp_session_manager import StdioConnectionParams
+from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset
+from mcp import StdioServerParameters
+
+from hitl_agent.config import WORKDIR
+
+# Read-only filesystem tool for orchestration agent (no write access)
+filesystem_tool_r = MCPToolset(
+  connection_params=StdioConnectionParams(
+    server_params=StdioServerParameters(
+      command="npx",
+      args=[
+        "-y",  # Argument for npx to auto-confirm install
+        "@modelcontextprotocol/server-filesystem@0.5.1",
+        os.path.abspath(WORKDIR),
+      ],
+    ),
+  ),
+  # Optional: Filter which tools from the MCP server are exposed
+  tool_filter=["list_directory", "read_file"],
+)
+
+# Read-write filesystem tool for sub-agents
+filesystem_tool_rw = MCPToolset(
+  connection_params=StdioConnectionParams(
+    server_params=StdioServerParameters(
+      command="npx",
+      args=[
+        "-y",  # Argument for npx to auto-confirm install
+        "@modelcontextprotocol/server-filesystem@0.5.1",
+        os.path.abspath(WORKDIR),
+      ],
+    ),
+  ),
+  # Optional: Filter which tools from the MCP server are exposed
+  tool_filter=["list_directory", "read_file", "write_file"],
+)

MaxKernel/hitl_agent/tools/retries_tool.py

@@ -0,0 +1,60 @@
+"""Tool for dynamically updating the max compilation retries setting."""
+
+import logging
+import os
+from google.adk.tools import FunctionTool, ToolContext
+
+
+async def set_max_compilation_retries_fn(
+    retries: int, tool_context: ToolContext, persist: bool = False
+) -> str:
+  """Set the maximum number of compilation validation and auto-fixing attempts.
+
+  Args:
+      retries: The maximum number of attempts (must be positive).
+      persist: Whether to permanently save this change to the configuration file (default is False).
+
+  Returns:
+      A string containing the operation result description:
+      - If successful: "Successfully updated maximum compilation retries to: <retries> (persisted: <persist>)"
+      - If failed: "Error: The number of retries must be a positive integer."
+  """
+  if retries <= 0:
+    return "Error: The number of retries must be a positive integer."
+
+  # Update in-memory config variables
+  import hitl_agent.config as hitl_cfg
+  hitl_cfg.MAX_COMPILATION_RETRIES = retries
+  os.environ["MAX_COMPILATION_RETRIES"] = str(retries)
+
+  # Update .env file if persist is requested
+  if persist:
+    try:
+      env_path = ".env"
+      if os.path.exists(env_path):
+        with open(env_path, "r") as f:
+          lines = f.readlines()
+        with open(env_path, "w") as f:
+          updated = False
+          for line in lines:
+            if line.strip().startswith("MAX_COMPILATION_RETRIES="):
+              f.write(f"MAX_COMPILATION_RETRIES={retries}\n")
+              updated = True
+            else:
+              f.write(line)
+          if not updated:
+            f.write(f"MAX_COMPILATION_RETRIES={retries}\n")
+    except Exception as e:
+      logging.error(f"Failed to update .env: {e}")
+
+  # Update the active validation loops in-memory
+  try:
+    from hitl_agent.subagents.kernel_writing.agent import kernel_compilation_validation_loop as hitl_loop
+    hitl_loop.max_retries = retries
+  except Exception as e:
+    logging.warning(f"Could not update hitl_loop max_retries: {e}")
+
+  return f"Successfully updated maximum compilation retries to: {retries} (persisted: {persist})"
+
+
+set_max_compilation_retries = FunctionTool(set_max_compilation_retries_fn)