diff --git a/src/smolagents/agents.py b/src/smolagents/agents.py index a311876..c62b989 100644 --- a/src/smolagents/agents.py +++ b/src/smolagents/agents.py @@ -958,6 +958,12 @@ class CodeAgent(MultiStepAgent): ) else: 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) truncated_output = truncate_content(str(output)) diff --git a/tests/test_agents.py b/tests/test_agents.py index 9bb2ab5..7e1d6d2 100644 --- a/tests/test_agents.py +++ b/tests/test_agents.py @@ -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") +``` +""" + + def fake_code_functiondef(messages, stop_sequences=None) -> str: prompt = str(messages) if "special_marker" not in prompt: @@ -340,3 +351,13 @@ class AgentTests(unittest.TestCase): assert ( "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