Repository: ethz-spylab/agentdojo

Issue:
Problem parsing list argument in LLaMA model
The problem boils down to the following. If LLaMA outputs a list as argument to a function

```
<function-thoughts>I can use the `send_email` function to send an email to the specified recipient with the given subject and body.</function-thoughts>
<function-call>[send_email(recipients=["mark.black-2134@gmail.com"], subject="Important message!", body="Hey, how is it going?")]</function-call>
```

this goes through the function [parse_tool_calls_from_python_function](https://github.com/ethz-spylab/agentdojo/blob/34e9a686ac5ec83657ad0dad285d6d32c4d256da/src/agentdojo/ast_utils.py#L53) and throws error:

```
agentdojo.ast_utils.ASTParsingError: The keyword values should be literals. Got: <ast.List object at 0x7f4e78077df0>
```

because it fails to parse the list correctly [here](https://github.com/ethz-spylab/agentdojo/blob/34e9a686ac5ec83657ad0dad285d6d32c4d256da/src/agentdojo/ast_utils.py#L44).

A minimal reproducible example follows:

```
output = """[send_email(recipients=["mark.black-2134@gmail.com"], subject="Important message!", body="Hey, how is it going?")]"""
parsed_model_output = ast.parse(output)
if not isinstance(parsed_model_output.body[0], ast.Expr):
        raise ASTParsingError("Error parsing the model output.")
if not isinstance(parsed_model_output.body[0].value, ast.List):
    if not isinstance(parsed_model_output.body[0].value, ast.Call):
        raise ASTParsingError(f"The function calls should be given a list. Got {parsed_model_output.body[0].value}")
    calls_list = [parsed_model_output.body[0].value]
else:
    calls_list = parsed_model_output.body[0].value.elts

call = calls_list[0]

class ASTParsingError(Exception): ...

args = {}
for keyword in call.keywords:
    if not isinstance(keyword.arg, str):
        raise ASTParsingError(f"The keyword arguments should be strings. Got {keyword.arg}")
    if not isinstance(keyword.value, ast.Constant):
        raise ASTParsingError(f"The keyword values should be literals. Got: {keyword.value}")
    args[keyword.arg] = keyword.value.value
```

I am not sure if this is supposed to be the case but lists as arguments seem reasonable in this case. I implemented a workaround using:

```
def parse_tool_call_from_ast(call: ast.expr) -> FunctionCall:
    if not isinstance(call, ast.Call):
        raise ASTParsingError("The function calls should be in the form of a call.")
    if not isinstance(call.func, ast.Name):
        raise ASTParsingError("The function calls should be in the form of a function name.")
    function_name = call.func.id
    args = {}
    for keyword in call.keywords:
        if not isinstance(keyword.arg, str):
            raise ASTParsingError(f"The keyword arguments should be strings. Got {keyword.arg}")
        if not isinstance(keyword.value, ast.Constant) and not isinstance(keyword.value, ast.List):
            raise ASTParsingError(f"The keyword values should be literals. Got: {keyword.value}")
        if isinstance(keyword.value, ast.List):
            # TODO: workaround for lists. Get the value of each element in the list
            args[keyword.arg] = [element.value for element in keyword.value.elts]
        else:
            args[keyword.arg] = keyword.value.value

    random_string = "".join(random.choices(string.ascii_lowercase + string.digits, k=10))

    return FunctionCall(function=function_name, args=args, id=f"tool_call_{random_string}")
```

Not sure if this was intended or if there is a better way to handle this. I leave up to you to find the best solution! :)

Discussion:

Patch:
diff --git a/src/agentdojo/ast_utils.py b/src/agentdojo/ast_utils.py
index 54c793a2..b8f0c8e3 100644
--- a/src/agentdojo/ast_utils.py
+++ b/src/agentdojo/ast_utils.py
@@ -2,7 +2,7 @@ import ast
 import random
 import string
 
-from agentdojo.functions_runtime import FunctionCall
+from agentdojo.functions_runtime import FunctionCall, FunctionCallArgTypes
 
 
 class ASTParsingError(Exception): ...
@@ -31,6 +31,20 @@ def create_python_function_from_tool_call(func_call_dict: FunctionCall) -> str:
         raise ValueError(f"Error creating AST from dictionary: {e}")
 
 
+def parse_arg_value(arg_value: ast.expr | None) -> FunctionCallArgTypes:
+    match arg_value:
+        case ast.Constant(value=value):
+            return value
+        case ast.List(elts=elements):
+            return [parse_arg_value(element) for element in elements]
+        case ast.Dict(keys=keys, values=values):
+            return {parse_arg_value(key): parse_arg_value(value) for key, value in zip(keys, values)}
+        case None:
+            return None
+        case _:
+            raise ASTParsingError(f"The keyword values should be literals. Got: {arg_value}")
+
+
 def parse_tool_call_from_ast(call: ast.expr) -> FunctionCall:
     if not isinstance(call, ast.Call):
         raise ASTParsingError("The function calls should be in the form of a call.")
@@ -41,9 +55,7 @@ def parse_tool_call_from_ast(call: ast.expr) -> FunctionCall:
     for keyword in call.keywords:
         if not isinstance(keyword.arg, str):
             raise ASTParsingError(f"The keyword arguments should be strings. Got {keyword.arg}")
-        if not isinstance(keyword.value, ast.Constant):
-            raise ASTParsingError(f"The keyword values should be literals. Got: {keyword.value}")
-        args[keyword.arg] = keyword.value.value
+        args[keyword.arg] = parse_arg_value(keyword.value)
 
     random_string = "".join(random.choices(string.ascii_lowercase + string.digits, k=10))
 


Test patch:
diff --git a/tests/test_ast_utils.py b/tests/test_ast_utils.py
index 8dd5e3f8..84870e40 100644
--- a/tests/test_ast_utils.py
+++ b/tests/test_ast_utils.py
@@ -27,3 +27,16 @@ def test_parse_tool_calls_from_python_function():
     assert tool_calls[1]["args"] == {"a": 1, "b": 2}
     assert tool_calls[2]["function"] == "f"
     assert tool_calls[2]["args"] == {}
+
+
+def test_parse_tool_calls_from_python_function_weird_args():
+    model_output = "[send_email(recipients=['example@example.com'], subject='Hello', body='Hello', attachments={'type': 'document', 'id': '12'})]"
+    tool_calls = parse_tool_calls_from_python_function(model_output)
+    assert len(tool_calls) == 1
+    assert tool_calls[0]["function"] == "send_email"
+    assert tool_calls[0]["args"] == {
+        "recipients": ["example@example.com"],
+        "subject": "Hello",
+        "body": "Hello",
+        "attachments": {"type": "document", "id": "12"},
+    }