diff --git a/src/smolagents/tools.py b/src/smolagents/tools.py index b7f08fd2e..3cd726ee8 100644 --- a/src/smolagents/tools.py +++ b/src/smolagents/tools.py @@ -270,9 +270,87 @@ def to_code_prompt(self) -> str: # Add arguments documentation if self.inputs: - args_descriptions = "\n".join( - f"{arg_name}: {arg_schema['description']}" for arg_name, arg_schema in self.inputs.items() - ) + + def format_nested_args(args_dict, indent_level=0, depth=1, max_depth=3): + """Recursively format nested argument descriptions + + Note: max_depth is set to 3 as deeper nesting levels make it harder for LLMs + to understand and accurately configure parameters. This covers most use cases. + If deeper nesting is needed, this can be discussed and improved in the future. + + Args: + args_dict: The dictionary of arguments + indent_level: The level of indentation + depth: The depth of the nested dictionary + max_depth: Maximum nesting depth to format (fixed at 3) + """ + if depth > max_depth: + return "" + descriptions = [] + indent = " " * indent_level + + for arg_name, arg_schema in args_dict.items(): + if isinstance(arg_schema, dict): + # Check if this is a standard smolagents format (has type and description) + if "type" in arg_schema and "description" in arg_schema: + # Standard smolagents format + if "properties" in arg_schema: + descriptions.append(f"{indent}{arg_name}: {arg_schema['description']}") + nested_descriptions = format_nested_args( + arg_schema["properties"], indent_level + 1, depth + 1, max_depth + ) + if nested_descriptions: + descriptions.append(nested_descriptions) + else: + # Handle array types and other simple types + arg_type = arg_schema.get("type", "string") + if arg_type == "array": + descriptions.append(f"{indent}{arg_name} (array): {arg_schema['description']}") + else: + descriptions.append(f"{indent}{arg_name}: {arg_schema['description']}") + elif "description" in arg_schema: + # MCP format - has description but may not have type in the same way + if "properties" in arg_schema: + descriptions.append(f"{indent}{arg_name}: {arg_schema['description']}") + nested_descriptions = format_nested_args( + arg_schema["properties"], indent_level + 1, depth + 1, max_depth + ) + if nested_descriptions: + descriptions.append(nested_descriptions) + else: + # Handle array types and other simple types in MCP format + arg_type = arg_schema.get("type", "string") + if arg_type == "array": + descriptions.append(f"{indent}{arg_name} (array): {arg_schema['description']}") + else: + descriptions.append(f"{indent}{arg_name}: {arg_schema['description']}") + elif "title" in arg_schema: + # MCP format with title instead of description + if "properties" in arg_schema: + descriptions.append(f"{indent}{arg_name}: {arg_schema['title']}") + nested_descriptions = format_nested_args( + arg_schema["properties"], indent_level + 1, depth + 1, max_depth + ) + if nested_descriptions: + descriptions.append(nested_descriptions) + else: + # Handle array types and other simple types in MCP format + arg_type = arg_schema.get("type", "string") + if arg_type == "array": + descriptions.append(f"{indent}{arg_name} (array): {arg_schema['title']}") + else: + descriptions.append(f"{indent}{arg_name}: {arg_schema['title']}") + else: + # Handle direct nested objects + nested_descriptions = format_nested_args(arg_schema, indent_level, depth + 1, max_depth) + if nested_descriptions: + descriptions.append(nested_descriptions) + else: + descriptions.append(f"{indent}{arg_name}: {arg_schema}") + + return "\n".join(descriptions) + + args_descriptions = format_nested_args(self.inputs) args_doc = f"Args:\n{textwrap.indent(args_descriptions, ' ')}" tool_doc += f"\n\n{args_doc}" diff --git a/tests/test_tools.py b/tests/test_tools.py index a6ab06f27..9365632df 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -83,6 +83,112 @@ def _create_inputs(tool_inputs: dict[str, dict[str | type, str]]) -> dict[str, A class TestTool: + def test_tool_to_code_prompt_with_nested_args(self): + """Test that to_code_prompt correctly formats nested argument descriptions.""" + + class NestedArgsTool(Tool): + name = "nested_args_tool" + description = "Tool with nested arguments" + inputs = { + "config": { + "type": "object", + "description": "Configuration object", + "properties": { + "host": {"type": "string", "description": "Server host"}, + "port": {"type": "integer", "description": "Server port"}, + }, + }, + "simple_arg": {"type": "string", "description": "Simple argument"}, + } + output_type = "string" + + def forward(self, config, simple_arg): + return "test" + + tool = NestedArgsTool() + code_prompt = tool.to_code_prompt() + + # Check that the nested structure is properly formatted + assert "config: Configuration object" in code_prompt + assert " host: Server host" in code_prompt + assert " port: Server port" in code_prompt + assert "simple_arg: Simple argument" in code_prompt + + def test_tool_to_code_prompt_with_deeply_nested_args(self): + """Test that to_code_prompt handles deeply nested arguments up to max_depth=3.""" + + class DeeplyNestedTool(Tool): + name = "deeply_nested_tool" + description = "Tool with deeply nested arguments" + inputs = { + "level1": { + "type": "object", + "description": "Level 1", + "properties": { + "level2": { + "type": "object", + "description": "Level 2", + "properties": { + "level3": { + "type": "object", + "description": "Level 3", + "properties": { + "level4": {"type": "string", "description": "Level 4"}, + }, + }, + }, + }, + }, + }, + } + output_type = "string" + + def forward(self, level1): + return "test" + + tool = DeeplyNestedTool() + code_prompt = tool.to_code_prompt() + + # Check that nested levels are formatted with proper indentation + assert "level1: Level 1" in code_prompt + assert " level2: Level 2" in code_prompt + assert " level3: Level 3" in code_prompt + # Level 4 should NOT be included as it exceeds max_depth=3 + assert "level4" not in code_prompt + + def test_tool_to_code_prompt_with_mixed_nested_and_simple_args(self): + """Test that to_code_prompt handles a mix of nested and simple arguments.""" + + class MixedArgsTool(Tool): + name = "mixed_args_tool" + description = "Tool with mixed argument types" + inputs = { + "simple1": {"type": "string", "description": "First simple arg"}, + "nested": { + "type": "object", + "description": "Nested object", + "properties": { + "field1": {"type": "string", "description": "Nested field 1"}, + "field2": {"type": "integer", "description": "Nested field 2"}, + }, + }, + "simple2": {"type": "integer", "description": "Second simple arg"}, + } + output_type = "string" + + def forward(self, simple1, nested, simple2): + return "test" + + tool = MixedArgsTool() + code_prompt = tool.to_code_prompt() + + # Check that all arguments are properly formatted + assert "simple1: First simple arg" in code_prompt + assert "nested: Nested object" in code_prompt + assert " field1: Nested field 1" in code_prompt + assert " field2: Nested field 2" in code_prompt + assert "simple2: Second simple arg" in code_prompt + @pytest.mark.parametrize( "type_value, should_raise_error, error_contains", [