feat: enable events compaction in auto agent client

Author: shangkunwang01

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,
   )