Refactor pipeline agent planning, dynamic model selection, and test generation

MaxKernel/auto_agent/agent_client/auto_agent_client.py

@@ -14,34 +14,35 @@
     "dotenv not installed, skipping loading environment variables"
   )
 
+from google.adk.apps.app import App
 from google.adk.runners import Runner
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.genai.types import Content, Part
 
 from auto_agent.agent import root_agent
+from auto_agent.config import get_compaction_config
 
 logger = logging.getLogger(__name__)
 
 
 class AutoAgentClient:
-  user_id: str
-  session_id: str
-  query: str
-  app_name: str = "auto_agent"
-
   def __init__(
     self,
     user_id: str,
     session_id: str,
     query: str,
     agent: Optional[Any] = None,
+    app_name: str = "auto_agent",
+    events_compaction: bool = False,
   ):
     self.user_id = user_id
     self.session_id = session_id
     self.query = query
     self.agent = agent or root_agent
     self.session_service = InMemorySessionService()
     self.session = None
+    self.app_name = app_name
+    self.events_compaction = events_compaction
 
   async def create_session(
     self, initial_state: Optional[dict[str, Any]] = None
@@ -77,9 +78,19 @@ async def run_async(self) -> None:
     if not self.session:
       await self.create_session()
 
+    # Compaction configuration must be passed via an App object
+    if self.events_compaction:
+      compaction_config = get_compaction_config()
+    else:
+      compaction_config = None
+    app = App(
+      name=self.app_name,
+      root_agent=self.agent,
+      events_compaction_config=compaction_config,
+    )
+
     runner = Runner(
-      app_name=self.app_name,
-      agent=self.agent,
+      app=app,
       session_service=self.session_service,
     )
 
@@ -124,6 +135,11 @@ def main():
     default="client_query.txt",
     help="File containing the query to send",
   )
+  parser.add_argument(
+    "--events_compaction",
+    action="store_true",
+    help="Enable event compaction",
+  )
   args = parser.parse_args()
 
   user_id = args.user_id
@@ -132,7 +148,12 @@ def main():
   query = read_query_from_file(args.query_file)
 
   # Create client instance
-  client = AutoAgentClient(user_id, session_id, query)
+  client = AutoAgentClient(
+    user_id=user_id,
+    session_id=session_id,
+    query=query,
+    events_compaction=args.events_compaction,
+  )
 
   logger.info(
     f"Generating script for user {user_id} in session {session_id} with query: {query}"

MaxKernel/auto_agent/agent_client/run_batch_agent_call.py

@@ -27,6 +27,11 @@ def parse_args():
     default=2,
     help="Maximum number of retries for failed tasks",
   )
+  parser.add_argument(
+    "--events_compaction",
+    action="store_true",
+    help="Enable event compaction",
+  )
   parser.add_argument(
     "--log_file", type=str, default=None, help="File to save logs to"
   )
@@ -61,6 +66,7 @@ async def process_problem(
   data_dir: str,
   max_retries: int,
   sem: asyncio.Semaphore,
+  events_compaction: bool = False,
 ):
   async with sem:
     problem_dir = os.path.join(data_dir, problem_id)
@@ -88,6 +94,7 @@ async def process_problem(
         user_id=user_id,
         session_id=session_id,
         query=query,
+        events_compaction=events_compaction,
       )
 
       session_file = os.path.join(
@@ -159,7 +166,9 @@ async def main_async(args):
   sem = asyncio.Semaphore(args.num_concurrent)
 
   tasks = [
-    process_problem(problem, args.data_dir, args.max_retries, sem)
+    process_problem(
+      problem, args.data_dir, args.max_retries, sem, args.events_compaction
+    )
     for problem in problems
   ]
 

MaxKernel/auto_agent/config.py

@@ -19,9 +19,9 @@
 # Set events compaction policy to avoid memory overflow
 def get_compaction_config():
   return EventsCompactionConfig(
-    token_threshold=200000,
-    event_retention_size=100,
-    compaction_interval=1,
+    token_threshold=300000,
+    event_retention_size=5,
+    compaction_interval=0,
     overlap_size=0,
   )
 

MaxKernel/auto_agent/custom_types.py

@@ -8,10 +8,6 @@
 from google.adk.models.google_llm import Gemini
 from google.genai import Client, types
 
-from auto_agent.constants import (
-  MODEL_NAME,
-)
-
 
 class TimeoutGemini(Gemini):
   @cached_property
@@ -46,7 +42,7 @@ def __init__(self, *args, **kwargs):
     # If model is a string, use the pre-configured gemini_model with retry support
     if "model" in kwargs and isinstance(kwargs["model"], str):
       gemini_model = TimeoutGemini(
-        model=MODEL_NAME,
+        model=kwargs["model"],
         retry_options=types.HttpRetryOptions(
           initial_delay=1,
           attempts=10,

MaxKernel/auto_agent/subagents/autotuning/agent.py

@@ -25,18 +25,25 @@
 )
 from auto_agent.tools.search_api_tool import search_api_tool
 
+
 # 1. Planner Agent
 # This agent identifies parameters, creates the template, and defines the search space.
 # It saves them to session state instead of calling the tool directly.
-autotune_planner_agent = CustomLlmAgent(
-  name="AutotunePlannerAgent",
-  model=MODEL_NAME,
-  generate_content_config=model_config,
-  planner=get_thinking_planner("high"),
-  instruction=autotune_prompt.PROMPT,
-  description="Prepares code template and search space for auto-tuning Pallas kernels.",
-  tools=[filesystem_tool_r, write_autotune_specs_tool, search_api_tool],
-)
+def create_autotune_planner_agent(
+  model_name: str = MODEL_NAME,
+) -> CustomLlmAgent:
+  return CustomLlmAgent(
+    name="AutotunePlannerAgent",
+    model=model_name,
+    generate_content_config=model_config,
+    planner=get_thinking_planner("high"),
+    instruction=autotune_prompt.PROMPT,
+    description="Prepares code template and search space for auto-tuning Pallas kernels.",
+    tools=[filesystem_tool_r, write_autotune_specs_tool, search_api_tool],
+  )
+
+
+autotune_planner_agent = create_autotune_planner_agent()
 
 
 # 2. Runner Agent
@@ -154,52 +161,86 @@ async def _run_async_impl(
       )
 
 
-autotune_runner = AutotuneRunner(
-  name="AutotuneRunner",
-  output_key="autotune_results",
-)
+def create_autotune_runner() -> AutotuneRunner:
+  return AutotuneRunner(
+    name="AutotuneRunner",
+    output_key="autotune_results",
+  )
+
+
+autotune_runner = create_autotune_runner()
 
 
 # 3. Apply Best Config Agent
-apply_best_config_agent = CustomLlmAgent(
-  name="ApplyBestConfigAgent",
-  model=MODEL_NAME,
-  generate_content_config=model_config,
-  planner=get_thinking_planner("high"),
-  instruction=apply_best_config_prompt.PROMPT,
-  description="Applies autotuning results to the optimized kernel file.",
-  tools=[filesystem_tool_r, write_optimized_kernel_tool],
-)
+def create_apply_best_config_agent(
+  model_name: str = MODEL_NAME,
+) -> CustomLlmAgent:
+  return CustomLlmAgent(
+    name="ApplyBestConfigAgent",
+    model=model_name,
+    generate_content_config=model_config,
+    planner=get_thinking_planner("high"),
+    instruction=apply_best_config_prompt.PROMPT,
+    description="Applies autotuning results to the optimized kernel file.",
+    tools=[filesystem_tool_r, write_optimized_kernel_tool],
+  )
+
+
+apply_best_config_agent = create_apply_best_config_agent()
 
 
 # 4. Summarizer Agent
 # This agent reads results from state and talks to the user.
-autotune_summary_agent = CustomLlmAgent(
-  name="AutotuneSummaryAgent",
-  model=MODEL_NAME,
-  generate_content_config=model_config,
-  instruction=summary_prompt.PROMPT,
-  description="Summarizes autotuning results.",
-  tools=[filesystem_tool_r],
-  output_key="autotuning_summary",
-)
+def create_autotune_summary_agent(
+  model_name: str = MODEL_NAME,
+) -> CustomLlmAgent:
+  return CustomLlmAgent(
+    name="AutotuneSummaryAgent",
+    model=model_name,
+    generate_content_config=model_config,
+    instruction=summary_prompt.PROMPT,
+    description="Summarizes autotuning results.",
+    tools=[filesystem_tool_r],
+    output_key="autotuning_summary",
+  )
+
+
+autotune_summary_agent = create_autotune_summary_agent()
 
 
 class CombinedAutotuneAgent(BaseAgent):
   """Chains autotuning steps and conditionally applies best config."""
 
-  def __init__(self, name: str):
-    super().__init__(name=name)
+  planner_agent: Optional[BaseAgent] = None
+  runner_agent: Optional[BaseAgent] = None
+  apply_config_agent: Optional[BaseAgent] = None
+  summary_agent: Optional[BaseAgent] = None
+
+  def __init__(
+    self,
+    name: str,
+    planner_agent: BaseAgent,
+    runner_agent: BaseAgent,
+    apply_config_agent: BaseAgent,
+    summary_agent: BaseAgent,
+  ):
+    super().__init__(
+      name=name,
+      planner_agent=planner_agent,
+      runner_agent=runner_agent,
+      apply_config_agent=apply_config_agent,
+      summary_agent=summary_agent,
+    )
 
   async def _run_async_impl(
     self, ctx: InvocationContext
   ) -> AsyncGenerator[Event, None]:
     logging.info(f"[{self.name}] Running AutotunePlannerAgent...")
-    async for event in autotune_planner_agent.run_async(ctx):
+    async for event in self.planner_agent.run_async(ctx):
       yield event
 
     logging.info(f"[{self.name}] Running AutotuneRunner...")
-    async for event in autotune_runner.run_async(ctx):
+    async for event in self.runner_agent.run_async(ctx):
       yield event
 
     autotune_results = ctx.session.state.get("autotune_results", {})
@@ -209,7 +250,7 @@ async def _run_async_impl(
       and autotune_results.get("best_time_ms") is not None
     ):
       logging.info(f"[{self.name}] Running ApplyBestConfigAgent...")
-      async for event in apply_best_config_agent.run_async(ctx):
+      async for event in self.apply_config_agent.run_async(ctx):
         yield event
     else:
       logging.warning(
@@ -218,10 +259,22 @@ async def _run_async_impl(
       )
 
     logging.info(f"[{self.name}] Running AutotuneSummaryAgent...")
-    async for event in autotune_summary_agent.run_async(ctx):
+    async for event in self.summary_agent.run_async(ctx):
       yield event
 
 
-autotune_agent = CombinedAutotuneAgent(name="AutotuneAgent")
+def create_autotune_agent(
+  model_name: str = MODEL_NAME,
+) -> CombinedAutotuneAgent:
+  return CombinedAutotuneAgent(
+    name="AutotuneAgent",
+    planner_agent=create_autotune_planner_agent(model_name),
+    runner_agent=create_autotune_runner(),
+    apply_config_agent=create_apply_best_config_agent(model_name),
+    summary_agent=create_autotune_summary_agent(model_name),
+  )
+
+
+autotune_agent = create_autotune_agent()
 
-__all__ = ["autotune_agent"]
+__all__ = ["autotune_agent", "create_autotune_agent"]