refactor: remove explanation subagent and implement dynamic thinking level configuration for agents

Author: shangkunwang01

MaxKernel/auto_agent/config.py

@@ -20,10 +20,16 @@
   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/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

@@ -15,7 +15,7 @@
   load_kernel_and_plan_to_state,
   load_single_kernel_to_state,
 )
-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.kernel_writing.kernel_compilation import (
@@ -267,7 +267,7 @@ 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=(
@@ -287,7 +287,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 +296,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 +318,7 @@ async def _run_async_impl(
   name="AddDebugStatementsAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
+  planner=get_thinking_planner("medium"),
   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 +330,7 @@ async def _run_async_impl(
   name="CleanupDebugStatementsAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
+  planner=get_thinking_planner("low"),
   instruction=cleanup_debug_statements.PROMPT,
   description="Removes debugging statements from successfully compiled kernel.",
   tools=[filesystem_tool_r, write_optimized_kernel_tool],
@@ -358,7 +357,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 +377,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=(

MaxKernel/auto_agent/subagents/explanation/agent.py

@@ -11,7 +11,7 @@
   load_profiling_script_to_state,
   load_single_kernel_to_state,
 )
-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.profiling import offline_tools
@@ -32,7 +32,7 @@
   name="GenerateProfilingScriptAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
+  planner=get_thinking_planner("medium"),
   instruction=gen_profiling_script.PROMPT,
   description="Generates a profiling script to identify performance bottlenecks in the kernel code and writes it to a file.",
   tools=[filesystem_tool_r, write_profiling_script_tool],
@@ -44,7 +44,6 @@
   name="ReadProfilingScriptAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
   instruction=read_profiling_script_prompt.PROMPT,
   description="Loads the generated profiling script file contents from disk into memory for execution.",
   before_agent_callback=load_profiling_script_to_state,
@@ -110,7 +109,7 @@ async def _run_async_impl(
   name="SummarizeProfileAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
+  planner=get_thinking_planner("high"),
   instruction=analyze_profile_prompt.PROMPT,
   description=(
     "Summarizes the profiling results of the kernel and performs deep"

MaxKernel/auto_agent/subagents/explanation/prompts/__init__.py

@@ -16,7 +16,7 @@
 from google.adk.events import Event, EventActions
 
 from auto_agent.client_utils.eval_client import call_eval_server_async
-from auto_agent.config import model_config, thinking_planner
+from auto_agent.config import get_thinking_planner, model_config
 from auto_agent.constants import EVAL_SERVER_PORT, MODEL_NAME, REQUEST_TIMEOUT
 from auto_agent.custom_types import CustomLlmAgent
 from auto_agent.subagents.testing.prompts import (
@@ -860,7 +860,6 @@ async def _run_async_impl(
   name="ValidationSummaryAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
   instruction=validation_summary.PROMPT,
   description="Summarizes validation results and provides next steps to the user.",
 )
@@ -870,7 +869,7 @@ async def _run_async_impl(
   name="GenerateTestFileAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
+  planner=get_thinking_planner("high"),
   instruction=gen_test_file.PROMPT,
   description="Generates a comprehensive pytest test file.",
   tools=(
@@ -914,7 +913,7 @@ async def _run_async_impl(
   name="FixTestScriptAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
+  planner=get_thinking_planner("high"),
   instruction=fix_test_script.PROMPT,
   description="Fixes validation errors in the generated test file.",
   tools=[filesystem_tool_r, write_test_file_tool, search_api_tool],
@@ -954,7 +953,7 @@ async def _run_async_impl(
   name="SummarizeTestResultsAgent",
   model=MODEL_NAME,
   generate_content_config=model_config,
-  planner=thinking_planner,
+  planner=get_thinking_planner("medium"),
   instruction=summarize_test_results_prompt.PROMPT,
   description="Analyzes pytest test results and provides recommendations.",
   tools=(