Merge pull request #64 from AI-Hypercomputer/shangkun-context-management

Implement context management technique to prevent context overflow

Author: shangkunwang01

MaxKernel/auto_agent/agent.py

@@ -4,6 +4,10 @@
 for the human-in-the-loop kernel generation process.
 """
 
+from google.adk.apps.app import App
+
+from auto_agent.config import get_compaction_config
+from auto_agent.constants import EVENTS_COMPACTION
 from auto_agent.subagents.autotuning.agent import autotune_agent
 from auto_agent.subagents.kernel_writing import (
   implement_kernel_agent,
@@ -29,4 +33,16 @@
   max_iterations=5,
 )
 
-__all__ = ["root_agent"]
+if EVENTS_COMPACTION:
+  compaction_config = get_compaction_config()
+else:
+  compaction_config = None
+
+
+app = App(
+  name="auto_agent",
+  root_agent=root_agent,
+  events_compaction_config=compaction_config,
+)
+
+__all__ = ["root_agent", "app"]

MaxKernel/auto_agent/config.py

@@ -2,10 +2,11 @@
 
 import os
 
+from google.adk.apps.app import EventsCompactionConfig
 from google.adk.planners import BuiltInPlanner
 from google.genai import types
 
-from auto_agent.constants import TOP_K, TOP_P
+from auto_agent.constants import TEMPERATURE, TOP_K, TOP_P
 
 # Environment variables
 WORKDIR = os.environ.get("WORKDIR", os.path.dirname(os.path.abspath(__file__)))
@@ -14,17 +15,32 @@
 INCLUDE_THOUGHTS = os.environ.get("INCLUDE_THOUGHTS", "true").lower() == "true"
 MAX_COMPILATION_RETRIES = int(os.environ.get("MAX_COMPILATION_RETRIES", "6"))
 
+# Set events compaction policy to avoid memory overflow
+def get_compaction_config():
+  return EventsCompactionConfig(
+    token_threshold=200000,
+    event_retention_size=100,
+    compaction_interval=1,
+    overlap_size=0,
+  )
+
 # Model configuration
 model_config = types.GenerateContentConfig(
-  temperature=0.5,
+  temperature=TEMPERATURE,
   top_p=TOP_P,
   top_k=TOP_K,
 )
 
-# Planner configuration with thinking/reasoning traces
-thinking_planner = BuiltInPlanner(
-  thinking_config=types.ThinkingConfig(
-    include_thoughts=INCLUDE_THOUGHTS,
-    thinking_level="high",
+
+def get_thinking_planner(level: str = "high") -> BuiltInPlanner:
+  """Returns a BuiltInPlanner configured with the specified thinking level.
+
+  Args:
+    level: The thinking level to use. Can be 'high', 'medium', or 'low'.
+  """
+  return BuiltInPlanner(
+    thinking_config=types.ThinkingConfig(
+      include_thoughts=INCLUDE_THOUGHTS,
+      thinking_level=level,
+    )
   )
-)

MaxKernel/auto_agent/constants.py

@@ -1,10 +1,11 @@
 MODEL_NAME = "gemini-3.1-pro-preview"
 MAX_ITERATIONS = 5
 LLM_GEN_RETRY_COUNT = 3
-TEMPERATURE = 0.1
+TEMPERATURE = 0.5
 TOP_P = 0.9
 TOP_K = 5
 REQUEST_TIMEOUT = 3600 * 3
 TPU_SERVER_PORT = 5463
 CPU_SERVER_PORT = 5464
 EVAL_SERVER_PORT = 1245
+EVENTS_COMPACTION = True

MaxKernel/auto_agent/subagents/__init__.py

@@ -4,14 +4,12 @@
 - kernel_writing: Planning, implementation, and compilation validation
 - testing: Test generation, validation, and execution
 - profiling: Performance profiling and analysis
-- explanation: Explanations for the kernel generation process
 """
 
-from . import explanation, kernel_writing, profiling, testing
+from . import kernel_writing, profiling, testing
 
 __all__ = [
   "kernel_writing",
   "testing",
   "profiling",
-  "explanation",
 ]

MaxKernel/auto_agent/subagents/autotuning/agent.py

@@ -9,7 +9,7 @@
 from google.adk.agents.invocation_context import InvocationContext
 from google.adk.events import Event, EventActions
 
-from auto_agent.config import model_config, thinking_planner
+from auto_agent.config import get_thinking_planner, model_config
 from auto_agent.constants import MODEL_NAME
 from auto_agent.custom_types import CustomLlmAgent
 from auto_agent.subagents.autotuning.autotune_tool import autotune_kernel
@@ -32,7 +32,7 @@
   name="AutotunePlannerAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
+  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],
@@ -165,7 +165,7 @@ async def _run_async_impl(
   name="ApplyBestConfigAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
+  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],
@@ -178,7 +178,6 @@ async def _run_async_impl(
   name="AutotuneSummaryAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
   instruction=summary_prompt.PROMPT,
   description="Summarizes autotuning results.",
   tools=[filesystem_tool_r],

MaxKernel/auto_agent/subagents/explanation/__init__.py

@@ -9,13 +9,11 @@
 from google.adk.events import Event, EventActions
 
 from auto_agent.callbacks import (
-  add_pallas_docs,
   extract_fix_summary,
-  get_tpu_version_callback,
   load_kernel_and_plan_to_state,
   load_single_kernel_to_state,
 )
-from auto_agent.config import model_config, thinking_planner, MAX_COMPILATION_RETRIES
+from auto_agent.config import get_thinking_planner, model_config, 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 (
@@ -267,18 +265,14 @@ async def _run_async_impl(
   name="PlanKernelAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
+  planner=get_thinking_planner("high"),
   instruction=kernel_planning_prompt.PROMPT,
   description="Creates or revises a detailed optimization plan for a Pallas kernel.",
   tools=(
     [search_api_tool, filesystem_tool_rw, vertex_ai_rag_tool]
     if vertex_ai_rag_tool
     else [search_api_tool, filesystem_tool_rw]
   ),
-  before_agent_callback=[
-    add_pallas_docs,
-    get_tpu_version_callback,
-  ],
 )
 
 # Kernel compilation validation agents
@@ -287,7 +281,6 @@ async def _run_async_impl(
   name="ReadFileForValidationAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
   instruction=read_file_prompt.PROMPT,
   description="Reads the kernel file mentioned by the user or from state for validation.",
   tools=[filesystem_tool_r],
@@ -297,7 +290,7 @@ async def _run_async_impl(
   name="FixKernelCompilationAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
+  planner=get_thinking_planner("high"),
   instruction=fix_kernel_compilation.PROMPT,
   description="Fixes compilation errors in the generated kernel while preserving optimization strategy.",
   tools=(
@@ -319,7 +312,7 @@ async def _run_async_impl(
   name="AddDebugStatementsAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
+  planner=get_thinking_planner("high"),
   instruction=add_debug_statements.PROMPT,
   description="Adds strategic debugging statements to diagnose persistent compilation issues.",
   tools=[filesystem_tool_r, write_optimized_kernel_tool],
@@ -331,7 +324,7 @@ async def _run_async_impl(
   name="CleanupDebugStatementsAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
+  planner=get_thinking_planner("high"),
   instruction=cleanup_debug_statements.PROMPT,
   description="Removes debugging statements from successfully compiled kernel.",
   tools=[filesystem_tool_r, write_optimized_kernel_tool],
@@ -358,7 +351,6 @@ async def _run_async_impl(
   name="KernelCompilationSummaryAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
   instruction=kernel_compilation_summary.PROMPT,
   description="Summarizes kernel compilation validation results with full trace on failure.",
   include_contents="none",
@@ -379,7 +371,7 @@ async def _run_async_impl(
   name="ImplementKernelAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
+  planner=get_thinking_planner("high"),
   instruction=kernel_implementation_prompt.PROMPT,
   description="Implements the optimized Pallas kernel following the plan.",
   tools=(