From 0f426445427f948a1ada4affe9d7e035a840fe6b Mon Sep 17 00:00:00 2001 From: shangkunwang Date: Thu, 30 Jul 2026 21:13:49 +0000 Subject: [PATCH] feat: add events_compaction flag to enable event compaction in auto search --- MaxKernel/auto_search/algorithms/beam_search.py | 1 + MaxKernel/auto_search/algorithms/parallel_search.py | 1 + MaxKernel/auto_search/orchestrator.py | 3 +++ MaxKernel/auto_search/run_batch_search.py | 6 ++++++ MaxKernel/auto_search/run_search.py | 8 ++++++++ MaxKernel/auto_search/worker.py | 4 ++++ 6 files changed, 23 insertions(+) diff --git a/MaxKernel/auto_search/algorithms/beam_search.py b/MaxKernel/auto_search/algorithms/beam_search.py index f3ed48b..75a0a91 100644 --- a/MaxKernel/auto_search/algorithms/beam_search.py +++ b/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" diff --git a/MaxKernel/auto_search/algorithms/parallel_search.py b/MaxKernel/auto_search/algorithms/parallel_search.py index 368f0ab..e9e17c0 100644 --- a/MaxKernel/auto_search/algorithms/parallel_search.py +++ b/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" diff --git a/MaxKernel/auto_search/orchestrator.py b/MaxKernel/auto_search/orchestrator.py index a5c2a36..1ab4d82 100644 --- a/MaxKernel/auto_search/orchestrator.py +++ b/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()) diff --git a/MaxKernel/auto_search/run_batch_search.py b/MaxKernel/auto_search/run_batch_search.py index 9cdf491..8321d8b 100644 --- a/MaxKernel/auto_search/run_batch_search.py +++ b/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, diff --git a/MaxKernel/auto_search/run_search.py b/MaxKernel/auto_search/run_search.py index 1e71597..12c3983 100644 --- a/MaxKernel/auto_search/run_search.py +++ b/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 diff --git a/MaxKernel/auto_search/worker.py b/MaxKernel/auto_search/worker.py index ddbb6fb..76d0d3a 100644 --- a/MaxKernel/auto_search/worker.py +++ b/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)