Add warning about missing imports in CodeAgent error logs
This commit is contained in:
		
							parent
							
								
									87bbc3f1cd
								
							
						
					
					
						commit
						3f79baee71
					
				|  | @ -958,6 +958,12 @@ class CodeAgent(MultiStepAgent): | ||||||
|                 ) |                 ) | ||||||
|             else: |             else: | ||||||
|                 error_msg = f"Code execution failed: {str(e)}" |                 error_msg = f"Code execution failed: {str(e)}" | ||||||
|  |             if "Import of " in str(e) and " is not allowed" in str(e): | ||||||
|  |                 console.print( | ||||||
|  |                     "[bold red]Code execution failed due to an unauthorized import. Consider passing said import under additional_authorized_imports when initializing your CodeAgent." | ||||||
|  |                 ) | ||||||
|  |             else: | ||||||
|  |                 console.print("PLACEHOLDER" + str(e)) | ||||||
|             raise AgentExecutionError(error_msg) |             raise AgentExecutionError(error_msg) | ||||||
| 
 | 
 | ||||||
|         truncated_output = truncate_content(str(output)) |         truncated_output = truncate_content(str(output)) | ||||||
|  |  | ||||||
|  | @ -130,6 +130,17 @@ final_answer("got an error") | ||||||
| """ | """ | ||||||
| 
 | 
 | ||||||
| 
 | 
 | ||||||
|  | def fake_code_model_import(messages, stop_sequences=None) -> str: | ||||||
|  |     return """ | ||||||
|  | Thought: I can answer the question | ||||||
|  | Code: | ||||||
|  | ```py | ||||||
|  | import numpy as np | ||||||
|  | final_answer("got an error") | ||||||
|  | ```<end_code> | ||||||
|  | """ | ||||||
|  | 
 | ||||||
|  | 
 | ||||||
| def fake_code_functiondef(messages, stop_sequences=None) -> str: | def fake_code_functiondef(messages, stop_sequences=None) -> str: | ||||||
|     prompt = str(messages) |     prompt = str(messages) | ||||||
|     if "special_marker" not in prompt: |     if "special_marker" not in prompt: | ||||||
|  | @ -340,3 +351,13 @@ class AgentTests(unittest.TestCase): | ||||||
|         assert ( |         assert ( | ||||||
|             "You can also give requests to team members." in manager_agent.system_prompt |             "You can also give requests to team members." in manager_agent.system_prompt | ||||||
|         ) |         ) | ||||||
|  | 
 | ||||||
|  |     def test_code_agent_missing_import_triggers_advice_in_error_log(self): | ||||||
|  |         agent = CodeAgent(tools=[], model=fake_code_model_import) | ||||||
|  | 
 | ||||||
|  |         from smolagents.agents import console | ||||||
|  | 
 | ||||||
|  |         with console.capture() as capture: | ||||||
|  |             agent.run("Count to 3") | ||||||
|  |         str_output = capture.get() | ||||||
|  |         assert "import under additional_authorized_imports" in str_output | ||||||
|  |  | ||||||
		Loading…
	
		Reference in New Issue