feat: propagate parent evaluation metrics into child node ADK sessions

Author: shangkunwang01

MaxKernel/auto_agent/agent_client/auto_agent_client.py

@@ -43,11 +43,14 @@ def __init__(
     self.session_service = InMemorySessionService()
     self.session = None
 
-  async def create_session(self) -> None:
+  async def create_session(
+    self, initial_state: Optional[dict[str, Any]] = None
+  ) -> None:
     self.session = await self.session_service.create_session(
       app_name=self.app_name,
       user_id=self.user_id,
       session_id=self.session_id,
+      state=initial_state,
     )
 
   def get_session_data(self) -> dict:

MaxKernel/auto_search/utils/visualize_graph.py

@@ -61,7 +61,7 @@ def visualize_graph(
 
 if __name__ == "__main__":
   logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
-  
+
   if len(sys.argv) < 2:
     print(
       "Usage: python -m auto_search.utils.visualize_graph <path_to_graph.json> [output_graph_path]"

MaxKernel/auto_search/worker.py

@@ -26,6 +26,8 @@ async def expand_node(
 
     # Prepare inputs files for the agent based on reference_code and parent_node
     self._prepare_inputs(session_dir, parent_node, reference_code)
+    # Prepare initial state
+    initial_state = self._prepare_initial_state(parent_node)
 
     # Run the agent and process results
     try:
@@ -34,6 +36,7 @@ async def expand_node(
         session_dir=session_dir,
         strategy=strategy,
         agent_config=agent_config,
+        initial_state=initial_state,
       )
       return self._process_results(
         node_id=node_id,
@@ -76,12 +79,29 @@ def _prepare_inputs(
       with open(kernel_plan_path, "w") as f:
         f.write(parent_node.plan)
 
+  def _prepare_initial_state(self, parent_node: Node) -> Dict[str, Any]:
+    initial_state = {}
+    if parent_node.evaluation:
+      initial_state = {
+        "kernel_compilation_status": {
+          "success": parent_node.evaluation.compiled,
+          "message": parent_node.evaluation.compilation_error,
+        },
+        "test_results": {
+          "success": parent_node.evaluation.correct,
+          "output": parent_node.evaluation.test_error,
+        },
+        "profiling_summary": parent_node.evaluation.profiling_summary,
+      }
+    return initial_state
+
   async def _run_agent(
     self,
     node_id: str,
     session_dir: str,
     strategy: Optional[str],
     agent_config: Optional[Dict[str, Any]] = None,
+    initial_state: Optional[Dict[str, Any]] = None,
   ) -> Dict[str, Any]:
     """Sets up a custom AutonomousPipelineAgent and runs the client."""
     agent_config = agent_config or {}
@@ -111,7 +131,8 @@ async def _run_agent(
       query=query,
       agent=custom_agent,
     )
-    await client.create_session()
+
+    await client.create_session(initial_state)
     await client.run_async()
     try:
       session_json_path = os.path.join(session_dir, "session.json")
@@ -144,7 +165,13 @@ def _process_results(
 
     comp_status = best_run.get("compilation_status", {})
     compiled = comp_status.get("success", False)
-    compilation_error = comp_status.get("message") if not compiled else None
+    compilation_error = None
+    if not compiled:
+      compilation_error = comp_status.get("message", "")
+      if comp_status.get("final_errors"):
+        compilation_error += "\n\nFinal Errors:\n" + comp_status.get(
+          "final_errors"
+        )
 
     test_status = best_run.get("test_status", {})
     correct = test_status.get("success", False)