Repository: huggingface/smolagents

Issue:
The call.func parameter type of the local_python_executor.evaluate_call function might be ast.Subscript
Hi
While testing Python code execution, I discovered an issue where the call.func parameter type could potentially be ast.Subscript.
Please consider supporting this case.

Discussion:
@nvrxq: I think the PR would be good for this, also I fixed a list of function.  #194@aymeric-roucher: #194 should have fixed this, closing this issue!

Solution patch:
diff --git a/src/smolagents/local_python_executor.py b/src/smolagents/local_python_executor.py
index e54e059..1b9b35d 100644
--- a/src/smolagents/local_python_executor.py
+++ b/src/smolagents/local_python_executor.py
@@ -591,7 +591,11 @@ def evaluate_call(
     custom_tools: Dict[str, Callable],
     authorized_imports: List[str],
 ) -> Any:
-    if not (isinstance(call.func, ast.Attribute) or isinstance(call.func, ast.Name)):
+    if not (
+        isinstance(call.func, ast.Attribute)
+        or isinstance(call.func, ast.Name)
+        or isinstance(call.func, ast.Subscript)
+    ):
         raise InterpreterError(f"This is not a correct function: {call.func}).")
     if isinstance(call.func, ast.Attribute):
         obj = evaluate_ast(
@@ -617,6 +621,23 @@ def evaluate_call(
                 f"It is not permitted to evaluate other functions than the provided tools or functions defined in previous code (tried to execute {call.func.id})."
             )
 
+    elif isinstance(call.func, ast.Subscript):
+        value = evaluate_ast(
+            call.func.value, state, static_tools, custom_tools, authorized_imports
+        )
+        index = evaluate_ast(
+            call.func.slice, state, static_tools, custom_tools, authorized_imports
+        )
+        if isinstance(value, (list, tuple)):
+            func = value[index]
+        else:
+            raise InterpreterError(
+                f"Cannot subscript object of type {type(value).__name__}"
+            )
+
+        if not callable(func):
+            raise InterpreterError(f"This is not a correct function: {call.func}).")
+        func_name = None
     args = []
     for arg in call.args:
         if isinstance(arg, ast.Starred):
@@ -726,6 +747,8 @@ def evaluate_name(
         return state[name.id]
     elif name.id in static_tools:
         return static_tools[name.id]
+    elif name.id in custom_tools:
+        return custom_tools[name.id]
     elif name.id in ERRORS:
         return ERRORS[name.id]
     close_matches = difflib.get_close_matches(name.id, list(state.keys()))
diff --git a/src/smolagents/models.py b/src/smolagents/models.py
index f25ced9..6a1ae77 100644
--- a/src/smolagents/models.py
+++ b/src/smolagents/models.py
@@ -14,7 +14,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from dataclasses import dataclass
+from dataclasses import dataclass, asdict
 import json
 import logging
 import os
@@ -54,12 +54,29 @@ if _is_package_available("litellm"):
     import litellm
 
 
+def get_dict_from_nested_dataclasses(obj):
+    def convert(obj):
+        if hasattr(obj, "__dataclass_fields__"):
+            return {k: convert(v) for k, v in asdict(obj).items()}
+        return obj
+
+    return convert(obj)
+
+
 @dataclass
 class ChatMessageToolCallDefinition:
     arguments: Any
     name: str
     description: Optional[str] = None
 
+    @classmethod
+    def from_hf_api(cls, tool_call_definition) -> "ChatMessageToolCallDefinition":
+        return cls(
+            arguments=tool_call_definition.arguments,
+            name=tool_call_definition.name,
+            description=tool_call_definition.description,
+        )
+
 
 @dataclass
 class ChatMessageToolCall:
@@ -67,6 +84,14 @@ class ChatMessageToolCall:
     id: str
     type: str
 
+    @classmethod
+    def from_hf_api(cls, tool_call) -> "ChatMessageToolCall":
+        return cls(
+            function=ChatMessageToolCallDefinition.from_hf_api(tool_call.function),
+            id=tool_call.id,
+            type=tool_call.type,
+        )
+
 
 @dataclass
 class ChatMessage:
@@ -74,6 +99,19 @@ class ChatMessage:
     content: Optional[str] = None
     tool_calls: Optional[List[ChatMessageToolCall]] = None
 
+    def model_dump_json(self):
+        return json.dumps(get_dict_from_nested_dataclasses(self))
+
+    @classmethod
+    def from_hf_api(cls, message) -> "ChatMessage":
+        tool_calls = None
+        if getattr(message, "tool_calls", None) is not None:
+            tool_calls = [
+                ChatMessageToolCall.from_hf_api(tool_call)
+                for tool_call in message.tool_calls
+            ]
+        return cls(role=message.role, content=message.content, tool_calls=tool_calls)
+
 
 class MessageRole(str, Enum):
     USER = "user"
@@ -283,7 +321,7 @@ class HfApiModel(Model):
             )
         self.last_input_token_count = response.usage.prompt_tokens
         self.last_output_token_count = response.usage.completion_tokens
-        return response.choices[0].message
+        return ChatMessage.from_hf_api(response.choices[0].message)
 
 
 class TransformersModel(Model):
@@ -315,14 +353,18 @@ class TransformersModel(Model):
         logger.info(f"Using device: {self.device}")
         try:
             self.tokenizer = AutoTokenizer.from_pretrained(model_id)
-            self.model = AutoModelForCausalLM.from_pretrained(model_id, device_map=self.device)
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_id, device_map=self.device
+            )
         except Exception as e:
             logger.warning(
                 f"Failed to load tokenizer and model for {model_id=}: {e}. Loading default tokenizer and model instead from {default_model_id=}."
             )
             self.model_id = default_model_id
             self.tokenizer = AutoTokenizer.from_pretrained(default_model_id)
-            self.model = AutoModelForCausalLM.from_pretrained(model_id, device_map=self.device)
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_id, device_map=self.device
+            )
 
     def make_stopping_criteria(self, stop_sequences: List[str]) -> StoppingCriteriaList:
         class StopOnStrings(StoppingCriteria):
@@ -551,4 +593,5 @@ __all__ = [
     "HfApiModel",
     "LiteLLMModel",
     "OpenAIServerModel",
+    "ChatMessage",
 ]


Test patch:
diff --git a/tests/test_models.py b/tests/test_models.py
index dbd93ce..5a3a821 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -13,9 +13,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import unittest
+import json
 from typing import Optional
 
-from smolagents import models, tool
+from smolagents import models, tool, ChatMessage, HfApiModel
 
 
 class ModelTests(unittest.TestCase):
@@ -38,3 +39,13 @@ class ModelTests(unittest.TestCase):
                 "properties"
             ]["celsius"]
         )
+
+    def test_chatmessage_has_model_dumps_json(self):
+        message = ChatMessage("user", "Hello!")
+        data = json.loads(message.model_dump_json())
+        assert data["content"] == "Hello!"
+
+    def test_get_hfapi_message_no_tool(self):
+        model = HfApiModel()
+        messages = [{"role": "user", "content": "Hello!"}]
+        model(messages, stop_sequences=["great"])
diff --git a/tests/test_python_interpreter.py b/tests/test_python_interpreter.py
index 58f250c..6e5b3c2 100644
--- a/tests/test_python_interpreter.py
+++ b/tests/test_python_interpreter.py
@@ -60,6 +60,14 @@ class PythonInterpreterTester(unittest.TestCase):
             in str(e)
         )
 
+    def test_subscript_call(self):
+        code = """def foo(x,y):return x*y\n\ndef boo(y):\n\treturn y**3\nfun = [foo, boo]\nresult_foo = fun[0](4,2)\nresult_boo = fun[1](4)"""
+        state = {}
+        result, _ = evaluate_python_code(code, BASE_PYTHON_TOOLS, state=state)
+        assert result == 64
+        assert state["result_foo"] == 8
+        assert state["result_boo"] == 64
+
     def test_evaluate_call(self):
         code = "y = add_two(x)"
         state = {"x": 3}