feat: add events_compaction flag to enable event compaction in auto search

MaxKernel/auto_search/algorithms/beam_search.py

@@ -168,6 +168,7 @@ async def run_task(task_idx: int, parent_node: Node, strategy: str) -> Node:
             reference_code=self.reference_code,
             strategy=strategy,
             agent_config=self.agent_config,
+            events_compaction=self.events_compaction,
           )
           if (
             node.execution_status == "SUCCESS"

MaxKernel/auto_search/algorithms/parallel_search.py

@@ -116,6 +116,7 @@ async def run_task(task_idx: int, parent_node: Node, strategy: str) -> Node:
             session_dir=session_dir,
             reference_code=self.reference_code,
             agent_config=self.agent_config,
+            events_compaction=self.events_compaction,
           )
           if (
             node.execution_status == "SUCCESS"

MaxKernel/auto_search/orchestrator.py

@@ -18,12 +18,15 @@ def __init__(
     graph_db_path: Optional[str] = None,
     max_concurrency: int = 2,
     max_worker_retries: int = 1,
+    events_compaction: bool = False,
   ):
     if max_worker_retries < 1:
       raise ValueError(
         f"max_worker_retries must be at least 1, got {max_worker_retries}."
       )
     self.max_worker_retries = max_worker_retries
+    self.events_compaction = events_compaction
+
     # Resolve graph db path and run directory
     if not graph_db_path:
       workdir = os.environ.get("WORKDIR", os.getcwd())

MaxKernel/auto_search/run_batch_search.py

@@ -166,6 +166,11 @@ def parse_args() -> argparse.Namespace:
     default=None,
     help="JSON string of agent config parameters (e.g. '{\"max_iterations\": 5}')",
   )
+  orch_group.add_argument(
+    "--events_compaction",
+    action="store_true",
+    help="Enable event compaction",
+  )
   # Parallel Search Arguments
   parallel_group = parser.add_argument_group(
     "Parallel Search Arguments",
@@ -226,6 +231,7 @@ def main():
     "max_worker_retries": args.max_worker_retries,
     "strategies": args.strategies,
     "agent_config": parsed_agent_config,
+    "events_compaction": args.events_compaction,
     "num_parallel_runs": args.num_parallel_runs,
     "beam_size": args.beam_size,
     "branches_per_node": args.branches_per_node,

MaxKernel/auto_search/run_search.py

@@ -43,6 +43,7 @@ def get_orchestrator(
       strategies=kwargs.get("strategies"),
       max_worker_retries=kwargs.get("max_worker_retries", 1),
       agent_config=kwargs.get("agent_config"),
+      events_compaction=kwargs.get("events_compaction", False),
     )
   elif algorithm == "beam":
     strategies_kwargs = {}
@@ -60,6 +61,7 @@ def get_orchestrator(
       max_depth=kwargs.get("max_depth", 2),
       keep_factor=kwargs.get("keep_factor", 1.0),
       agent_config=kwargs.get("agent_config"),
+      events_compaction=kwargs.get("events_compaction", False),
       **strategies_kwargs,
     )
   elif algorithm == "agentic":
@@ -255,6 +257,11 @@ def parse_args() -> argparse.Namespace:
     default=None,
     help="JSON string of agent config parameters (e.g. '{\"max_iterations\": 5}')",
   )
+  orch_group.add_argument(
+    "--events_compaction",
+    action="store_true",
+    help="Enable event compaction",
+  )
   # Parallel Search Arguments
   parallel_group = parser.add_argument_group(
     "Parallel Search Arguments",
@@ -357,6 +364,7 @@ def main():
     "max_worker_retries": args.max_worker_retries,
     "strategies": args.strategies,
     "agent_config": agent_config,
+    "events_compaction": args.events_compaction,
     # Parallel Search Arguments
     "num_parallel_runs": args.num_parallel_runs,
     # Beam Search Arguments

MaxKernel/auto_search/worker.py

@@ -20,6 +20,7 @@ async def expand_node(
     reference_code: str,
     strategy: Optional[str] = None,
     agent_config: Optional[Dict[str, Any]] = None,
+    events_compaction: bool = False,
   ) -> Node:
     """Expand an ADK session to get a new optimized kernel."""
     os.makedirs(session_dir, exist_ok=True)
@@ -37,6 +38,7 @@ async def expand_node(
         strategy=strategy,
         agent_config=agent_config,
         initial_state=initial_state,
+        events_compaction=events_compaction,
       )
       return self._process_results(
         node_id=node_id,
@@ -102,6 +104,7 @@ async def _run_agent(
     strategy: Optional[str],
     agent_config: Optional[Dict[str, Any]] = None,
     initial_state: Optional[Dict[str, Any]] = None,
+    events_compaction: bool = False,
   ) -> Dict[str, Any]:
     """Sets up a custom AutonomousPipelineAgent and runs the client."""
     agent_config = agent_config or {}
@@ -130,6 +133,7 @@ async def _run_agent(
       session_id=node_id,
       query=query,
       agent=custom_agent,
+      events_compaction=events_compaction,
     )
 
     await client.create_session(initial_state)