Add support to declare imports for code agent (#31355)

* Support import declaration in Code Agent

tests/agents/test_agents.py

@@ -141,15 +141,21 @@ Action:
     def test_init_agent_with_different_toolsets(self):
         toolset_1 = []
         agent = ReactCodeAgent(tools=toolset_1, llm_engine=fake_react_code_llm)
-        assert len(agent.toolbox.tools) == 1  # contains only final_answer tool
+        assert (
+            len(agent.toolbox.tools) == 1
+        )  # when no tools are provided, only the final_answer tool is added by default
 
         toolset_2 = [PythonInterpreterTool(), PythonInterpreterTool()]
         agent = ReactCodeAgent(tools=toolset_2, llm_engine=fake_react_code_llm)
-        assert len(agent.toolbox.tools) == 2  # added final_answer tool
+        assert (
+            len(agent.toolbox.tools) == 2
+        )  # deduplication of tools, so only one python_interpreter tool is added in addition to final_answer
 
         toolset_3 = Toolbox(toolset_2)
         agent = ReactCodeAgent(tools=toolset_3, llm_engine=fake_react_code_llm)
-        assert len(agent.toolbox.tools) == 2  # added final_answer tool
+        assert (
+            len(agent.toolbox.tools) == 2
+        )  # same as previous one, where toolset_3 is an instantiation of previous one
 
         # check that add_base_tools will not interfere with existing tools
         with pytest.raises(KeyError) as e:

tests/agents/test_python_interpreter.py

@@ -32,9 +32,19 @@ def add_two(x):
 
 class PythonInterpreterToolTester(unittest.TestCase, ToolTesterMixin):
     def setUp(self):
-        self.tool = load_tool("python_interpreter")
+        self.tool = load_tool("python_interpreter", authorized_imports=["sqlite3"])
         self.tool.setup()
 
+    def test_exact_match_input_spec(self):
+        inputs_spec = self.tool.inputs
+        expected_description = (
+            "The code snippet to evaluate. All variables used in this snippet must be defined in this same snippet, "
+            "else you will get an error. This code can only import the following python libraries: "
+            "['math', 'statistics', 'time', 'itertools', 'stat', 'unicodedata', 'sqlite3', 'queue', 'collections', "
+            "'random', 're']."
+        )
+        self.assertEqual(inputs_spec["code"]["description"], expected_description)
+
     def test_exact_match_arg(self):
         result = self.tool("(2 / 2) * 4")
         self.assertEqual(result, "4.0")