For TooCallingAction fix extraction of function arguments from response

Author: sysradium

src/smolagents/models.py

@@ -126,10 +126,27 @@ def parse_json_if_needed(arguments: Union[str, dict]) -> Union[str, dict]:
             return arguments
 
 
-def parse_tool_args_if_needed(message: ChatMessage) -> ChatMessage:
+def parse_tool_args_if_needed(
+    message: ChatMessage,
+    tool_name_key: str = "name",
+    tool_arguments_key: str = "arguments",
+) -> ChatMessage:
     if message.tool_calls is not None:
         for tool_call in message.tool_calls:
             tool_call.function.arguments = parse_json_if_needed(tool_call.function.arguments)
+    elif isinstance(message.content, str) and message.content.startswith("Action:"):
+        output = message.content.split("Action:", 1)[1].strip()
+        start_index = output.index("{")
+        end_index = output.rindex("}")
+        output = output[start_index : end_index + 1]
+
+        return get_tool_call_chat_message_from_text(
+            output,
+            tool_name_key,
+            tool_arguments_key,
+        )
+
+    __import__("pdb").set_trace()
     return message
 
 
@@ -490,7 +507,7 @@ def __call__(
         self.last_output_token_count = response.usage.completion_tokens
         message = ChatMessage.from_hf_api(response.choices[0].message, raw=response)
         if tools_to_call_from is not None:
-            return parse_tool_args_if_needed(message)
+            return parse_tool_args_if_needed(message, self.tool_name_key, self.tool_arguments_key)
         return message
 
 
@@ -685,6 +702,8 @@ def __call__(
         for _ in self.stream_generate(self.model, self.tokenizer, prompt=prompt_ids, **completion_kwargs):
             self.last_output_token_count += 1
             text += _.text
+
+            found_stop_sequence = False
             for stop_sequence in prepared_stop_sequences:
                 stop_sequence_start = text.rfind(stop_sequence)
                 if stop_sequence_start != -1:
@@ -990,7 +1009,7 @@ def __call__(
         message.raw = response
 
         if tools_to_call_from is not None:
-            return parse_tool_args_if_needed(message)
+            return parse_tool_args_if_needed(message, self.tool_name_key, self.tool_arguments_key)
         return message
 
 
@@ -1076,7 +1095,7 @@ def __call__(
         )
         message.raw = response
         if tools_to_call_from is not None:
-            return parse_tool_args_if_needed(message)
+            return parse_tool_args_if_needed(message, self.tool_name_key, self.tool_arguments_key)
         return message
 
 

tests/test_models.py

@@ -109,7 +109,7 @@ def test_transformers_message_vl_no_tool(self):
         assert output == "Hello! How can"
 
     def test_parse_tool_args_if_needed(self):
-        original_message = ChatMessage(role="user", content=[{"type": "text", "text": "Hello!"}])
+        original_message = ChatMessage(role="user", content={"type": "text", "text": "Hello!"})
         parsed_message = parse_tool_args_if_needed(original_message)
         assert parsed_message == original_message
 
@@ -365,3 +365,16 @@ def test_flatten_messages_as_text_for_all_models(
 
         model = model_class(**{"model_id": "test-model", **model_kwargs})
     assert model.flatten_messages_as_text is expected_flatten_messages_as_text, f"{model_class.__name__} failed"
+
+
+@pytest.mark.parametrize(
+    "content",
+    [
+        'Action:\n{"name": "tool", "arguments": {}}',
+        'Action: {"name": "tool", "arguments": {}}',
+    ],
+)
+def test_parse_tool_args_if_needed_returns_message_with_tools_set_for_tool_calling_agent(content):
+    original_message = ChatMessage(role="user", content=content)
+    parsed_message = parse_tool_args_if_needed(original_message)
+    assert parsed_message.tool_calls