Merge pull request #56 from AI-Hypercomputer/nina_gpu_code_fix

Separate execution tool and routing tool to different responses

Author: NinaCai

MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/prompts/convert_simplified_to_jax_prompt.py

@@ -39,8 +39,8 @@
 4. Do NOT proceed to transfer_to_agent if the file write failed
 
 After SUCCESSFULLY writing the file (check the tool response for success):
-1. IMMEDIATELY call transfer_to_agent('ValidateSyntaxAgent') in the SAME response
-2. Do NOT output any message, confirmation, or status update before calling transfer_to_agent
-3. The transfer must be the ONLY action after successful file write
+1. Transfer to ValidateSyntaxAgent. Do NOT call transfer_to_agent in the same turn/response as write_file_direct. Wait for the tool response first, and then call:
+   transfer_to_agent('ValidateSyntaxAgent')
+2. Do NOT output any message, confirmation, or status update. The transfer must be the only action after successful file write.
 
 If you cannot write the file after multiple attempts, inform the user about the tool failure and ask them to manually create the file before proceeding."""

MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/prompts/generate_test_prompt.py

@@ -54,7 +54,7 @@ def run_validation_test():
    - content=<the complete test script>
    - Note: The file will be written to the same directory as the user-provided GPU code file automatically
    - Example: If the user provided "/home/user/project/kernel.cu", the test will be written to "/home/user/project/test_correctness.py"
-4. After writing successfully, IMMEDIATELY transfer to RunCorrectnessTestAgent in the SAME response:
+4. After writing successfully, transfer to RunCorrectnessTestAgent. Do NOT call transfer_to_agent in the same turn/response as write_file_direct. Wait for the tool response first, and then call:
    transfer_to_agent('RunCorrectnessTestAgent')
 
-CRITICAL: Do NOT output any message, confirmation, or status update before calling transfer_to_agent. The transfer must be the ONLY action after writing the file."""
+CRITICAL: Do NOT output any message, confirmation, or status update. The transfer must be the only action after writing the file."""

MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/prompts/identify_framework_prompt.py

@@ -23,10 +23,10 @@
    
    Where <FRAMEWORK_NAME> is the exact framework you detected (e.g., "CUDA", "Triton", "PyTorch CUDA")
 
-4. After successfully saving the framework, IMMEDIATELY transfer to the next agent:
+4. After successfully saving the framework, transfer to the next agent. Do NOT call transfer_to_agent in the same turn/response as save_framework_detection. Wait for the tool response first, and then call:
    transfer_to_agent('AnalyzePlanAndWriteAgent')
-   
-   Do this in the SAME response - do not wait for user input.
+
+   Do not wait for user input.
 
 IMPORTANT:
 - ALWAYS call save_framework_detection before transferring to ensure state is populated

MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/prompts/run_test_routing_prompt.py

@@ -6,7 +6,7 @@
 1. The correctness test code has been loaded into state (in the 'correctness_test_code' key) via before_agent_callback
 2. Call the run_correctness_test tool to execute the test
    - The tool will read correctness_test_code from state automatically and save results to correctness_test_results
-3. After the tool completes, IMMEDIATELY transfer to GenerateAndWriteSummaryAgent WITHOUT any status messages:
+3. After the tool completes, transfer to GenerateAndWriteSummaryAgent. Do NOT call transfer_to_agent in the same turn/response as run_correctness_test. Wait for the tool response first, and then call:
    transfer_to_agent('GenerateAndWriteSummaryAgent')
 
-CRITICAL: Do NOT output any message, acknowledgment, or status update to the user. Simply call run_correctness_test and then transfer to the next agent immediately."""
+CRITICAL: Do NOT output any message, acknowledgment, or status update. The transfer must be the only action after run_correctness_test completes."""

MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/prompts/validate_compilation_routing_prompt.py

@@ -6,7 +6,7 @@
 1. The JAX code is already loaded in state (in the 'jax_code' key)
 2. Call the check_jax_compilation tool to validate compilation
    - The tool will read jax_code from state automatically and save results to compilation_results
-3. After the tool completes, IMMEDIATELY transfer to ValidateShapesAgent WITHOUT any status messages:
+3. After the tool completes, transfer to ValidateShapesAgent. Do NOT call transfer_to_agent in the same turn/response as check_jax_compilation. Wait for the tool response first, and then call:
    transfer_to_agent('ValidateShapesAgent')
 
-CRITICAL: Do NOT output any message, acknowledgment, or status update to the user. Simply call check_jax_compilation and then transfer to the next agent immediately."""
+CRITICAL: Do NOT output any message, acknowledgment, or status update. The transfer must be the only action after check_jax_compilation completes."""

MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/prompts/validate_shapes_routing_prompt.py

@@ -6,7 +6,7 @@
 1. The JAX code is already loaded in state (in the 'jax_code' key)
 2. Call the validate_shapes tool to validate tensor shapes
    - The tool will read jax_code from state automatically and save results to shape_validation_results
-3. After the tool completes, IMMEDIATELY transfer to GenerateCorrectnessTestAgent WITHOUT any status messages:
+3. After the tool completes, transfer to GenerateCorrectnessTestAgent. Do NOT call transfer_to_agent in the same turn/response as validate_shapes. Wait for the tool response first, and then call:
    transfer_to_agent('GenerateCorrectnessTestAgent')
 
-CRITICAL: Do NOT output any message, acknowledgment, or status update to the user. Simply call validate_shapes and then transfer to the next agent immediately."""
+CRITICAL: Do NOT output any message, acknowledgment, or status update. The transfer must be the only action after validate_shapes completes."""

MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/prompts/validate_syntax_routing_prompt.py

@@ -6,7 +6,7 @@
 1. The JAX code has been loaded into state (in the 'jax_code' key) via before_agent_callback
 2. Call the check_jax_syntax tool to validate the JAX code
    - The tool will read jax_code from state automatically and save results to syntax_validation_results
-3. Based on the validation results returned by the tool, transfer to the appropriate agent
+3. After the tool completes, transfer to the appropriate agent. Do NOT call transfer_to_agent in the same turn/response as check_jax_syntax. Wait for the tool response first, and then transfer based on the results.
 
 ROUTING LOGIC:
 - If check_jax_syntax returns errors (message contains "JAX Syntax Validation Failed"):
@@ -15,4 +15,4 @@
 - If syntax is valid (message contains "JAX Syntax Validation Passed"):
   transfer_to_agent('ValidateCompilationAgent')
 
-CRITICAL: Do NOT output any message, acknowledgment, or status update. Simply call check_jax_syntax and then transfer to the appropriate agent immediately based on the results."""
+CRITICAL: Do NOT output any message, acknowledgment, or status update. The transfer must be the only action after check_jax_syntax completes."""

MaxKernel/hitl_agent/subagents/gpu_to_jax_agent/prompts/write_readme_prompt.py

@@ -32,7 +32,7 @@
 2. Find the original GPU code path and simplified code from conversation history
 3. Create the comprehensive README
 4. Call write_file_direct with path="SIMPLIFICATION_SUMMARY.md" and your README content
-5. After successfully writing, IMMEDIATELY transfer to ConvertToJaxAgent in the SAME response:
+5. After successfully writing the file, transfer to ConvertToJaxAgent. Do NOT call transfer_to_agent in the same turn/response as write_file_direct. Wait for the write_file_direct tool response first, and then call:
    transfer_to_agent('ConvertToJaxAgent')
 
-CRITICAL: Do NOT output any message, confirmation, or status update before calling transfer_to_agent. The transfer must be the ONLY action after writing the file."""
+CRITICAL: Do NOT output any message, confirmation, or status update. The transfer must be the only action after the file has been successfully written."""