refactor: renaming init_params and run_params to reflect their schema. (#332) (#333)

We’re currently using the JSON Schema standard for these fields

Author: lucasgomide

generate_tool_specs.py

@@ -5,9 +5,17 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Type
 
+from pydantic import BaseModel
+
 from crewai_tools import tools
-from crewai.tools.base_tool import EnvVar
+from crewai.tools.base_tool import BaseTool, EnvVar
+
+from pydantic.json_schema import GenerateJsonSchema
+from pydantic_core import PydanticOmit
 
+class SchemaGenerator(GenerateJsonSchema):
+    def handle_invalid_for_json_schema(self, schema, error_info):
+        raise PydanticOmit
 
 class ToolSpecExtractor:
     def __init__(self) -> None:
@@ -22,20 +30,21 @@ def extract_all_tools(self) -> List[Dict[str, Any]]:
                     self.extract_tool_info(obj)
                     self.processed_tools.add(name)
         return self.tools_spec
-
-    def extract_tool_info(self, tool_class: Type) -> None:
+    def extract_tool_info(self, tool_class: BaseTool) -> None:
         try:
             core_schema = tool_class.__pydantic_core_schema__
             if not core_schema:
                 return
 
             schema = self._unwrap_schema(core_schema)
             fields = schema.get("schema", {}).get("fields", {})
+
             tool_info = {
                 "name": tool_class.__name__,
                 "humanized_name": self._extract_field_default(fields.get("name"), fallback=tool_class.__name__),
                 "description": self._extract_field_default(fields.get("description")).strip(),
-                "run_params": self._extract_params(fields.get("args_schema")),
+                "run_params_schema": self._extract_params(fields.get("args_schema")),
+                "init_params_schema": self._extract_init_params(tool_class),
                 "env_vars": self._extract_env_vars(fields.get("env_vars")),
                 "package_dependencies": self._extract_field_default(fields.get("package_dependencies"), fallback=[]),
             }
@@ -60,35 +69,17 @@ def _extract_field_default(self, field: Optional[Dict], fallback: str = "") -> s
 
     def _extract_params(self, args_schema_field: Optional[Dict]) -> List[Dict[str, str]]:
         if not args_schema_field:
-            return []
+            return {}
 
         args_schema_class = args_schema_field.get("schema", {}).get("default")
         if not (inspect.isclass(args_schema_class) and hasattr(args_schema_class, "__pydantic_core_schema__")):
-            return []
+            return {}
 
         try:
-            core_schema = args_schema_class.__pydantic_core_schema__
-            schema = self._unwrap_schema(core_schema)
-            fields = schema.get("schema", {}).get("fields", {})
-
-            params = []
-            for name, info in fields.items():
-                _type = self._extract_param_type(info)
-                if _type == "union":
-                    breakpoint()
-                param = {
-                    "name": name,
-                    "description": self._extract_field_description_from_metadata(info),
-                    "type": _type,
-                    "default": self._extract_field_default(info),
-                }
-                params.append(param)
-
-            return params
-
+            return args_schema_class.model_json_schema(schema_generator=SchemaGenerator, mode='validation')
         except Exception as e:
             print(f"Error extracting params from {args_schema_class}: {e}")
-            return []
+            return {}
 
     def _extract_env_vars(self, env_vars_field: Optional[Dict]) -> List[Dict[str, str]]:
         if not env_vars_field:
@@ -105,47 +96,18 @@ def _extract_env_vars(self, env_vars_field: Optional[Dict]) -> List[Dict[str, st
                 })
         return env_vars
 
-    def _extract_field_description_from_metadata(self, field: Dict) -> str:
-        if metadata := field.get("metadata"):
-            return metadata.get("pydantic_js_updates", {}).get("description", "")
-        return ""
-
-    def _extract_param_type(self, info: Dict) -> Optional[str]:
-        schema = info.get("schema", {})
-        schema = self._unwrap_schema(schema)
-
-        if schema.get("type") == "nullable":
-            inner = schema.get("schema", {})
-            return self._schema_type_to_str(inner)
-
-        return self._schema_type_to_str(schema)
-
-    def _schema_type_to_str(self, schema: Dict) -> str:
-        schema_type = schema.get("type", "")
-
-        if schema_type == "list" and "items_schema" in schema:
-            item_type = self._schema_type_to_str(schema["items_schema"])
-            return f"list[{item_type}]"
-
-        if schema_type == "union" and "choices" in schema:
-            choices = schema["choices"]
-            item_types = [self._schema_type_to_str(choice) for choice in choices]
-            return f"union[{', '.join(item_types)}]"
-
-        if schema_type == "dict" and "keys_schema" in schema and "values_schema" in schema:
-            key_type = self._schema_type_to_str(schema["keys_schema"])
-            value_type = self._schema_type_to_str(schema["values_schema"])
-            return f"dict[{key_type}, {value_type}]"
-
-        return {
-            "str": "str",
-            "int": "int",
-            "float": "float",
-            "bool": "bool",
-            "list": "list",
-            "dict": "dict",
-            "any": "any",
-        }.get(schema_type, schema_type or "unknown")
+    def _extract_init_params(self, tool_class: BaseTool) -> dict:
+        ignored_init_params = ['name', 'description', 'env_vars', 'args_schema', 'description_updated', 'cache_function', 'result_as_answer', 'max_usage_count', 'current_usage_count', 'package_dependencies']
+
+        json_schema = tool_class.model_json_schema(schema_generator=SchemaGenerator, mode='serialization')
+
+        properties = {}
+        for key, value in json_schema['properties'].items():
+            if key not in ignored_init_params:
+                properties[key] = value
+
+        json_schema['properties'] = properties
+        return json_schema
 
     def save_to_json(self, output_path: str) -> None:
         with open(output_path, "w", encoding="utf-8") as f:

tests/test_generate_tool_specs.py

@@ -44,91 +44,106 @@ def test_unwrap_schema(extractor):
     assert result["value"] == "test"
 
 
-@pytest.mark.parametrize(
-    "schema, expected",
-    [
-        ({"type": "str"}, "str"),
-        ({"type": "list", "items_schema": {"type": "str"}}, "list[str]"),
-        ({"type": "dict", "keys_schema": {"type": "str"}, "values_schema": {"type": "int"}}, "dict[str, int]"),
-        ({"type": "union", "choices": [{"type": "str"}, {"type": "int"}]}, "union[str, int]"),
-        ({"type": "custom_type"}, "custom_type"),
-        ({}, "unknown"),
-    ]
-)
-def test_schema_type_to_str(extractor, schema, expected):
-    assert extractor._schema_type_to_str(schema) == expected
-
-
-@pytest.mark.parametrize(
-    "info, expected_type",
-    [
-        ({"schema": {"type": "str"}}, "str"),
-        ({"schema": {"type": "nullable", "schema": {"type": "int"}}}, "int"),
-        ({"schema": {"type": "default", "schema": {"type": "list", "items_schema": {"type": "str"}}}}, "list[str]"),
-    ]
-)
-def test_extract_param_type(extractor, info, expected_type):
-    assert extractor._extract_param_type(info) == expected_type
-
-
-def test_extract_all_tools(extractor):
+@pytest.fixture
+def mock_tool_extractor(extractor):
     with mock.patch("generate_tool_specs.dir", return_value=["MockTool"]), \
          mock.patch("generate_tool_specs.getattr", return_value=MockTool):
         extractor.extract_all_tools()
-
         assert len(extractor.tools_spec) == 1
-        tool_info = extractor.tools_spec[0]
-
-        assert tool_info.keys() == {
-            "name",
-            "humanized_name",
-            "description",
-            "run_params",
-            "env_vars",
-            "init_params",
-            "package_dependencies",
-        }
-
-        assert tool_info["name"] == "MockTool"
-        assert tool_info["humanized_name"] == "Mock Search Tool"
-        assert tool_info["description"] == "A tool that mocks search functionality"
+        return extractor.tools_spec[0]
+
+def test_extract_basic_tool_info(mock_tool_extractor):
+    tool_info = mock_tool_extractor
+
+    assert tool_info.keys() == {
+        "name",
+        "humanized_name",
+        "description",
+        "run_params_schema",
+        "env_vars",
+        "init_params_schema",
+        "package_dependencies",
+    }
 
-        assert len(tool_info["env_vars"]) == 2
-        api_key_var, rate_limit_var = tool_info["env_vars"]
+    assert tool_info["name"] == "MockTool"
+    assert tool_info["humanized_name"] == "Mock Search Tool"
+    assert tool_info["description"] == "A tool that mocks search functionality"
 
-        assert api_key_var["name"] == "SERPER_API_KEY"
-        assert api_key_var["description"] == "API key for Serper"
-        assert api_key_var["required"] == True
-        assert api_key_var["default"] == None
+def test_extract_init_params_schema(mock_tool_extractor):
+    tool_info = mock_tool_extractor
+    init_params_schema = tool_info["init_params_schema"]
 
-        assert rate_limit_var["name"] == "API_RATE_LIMIT"
-        assert rate_limit_var["description"] == "API rate limit"
-        assert rate_limit_var["required"] == False
-        assert rate_limit_var["default"] == "100"
+    assert init_params_schema.keys() == {
+        "$defs",
+        "properties",
+        "title",
+        "type",
+    }
 
-        assert len(tool_info["run_params"]) == 3
+    another_parameter = init_params_schema['properties']['another_parameter']
+    assert another_parameter["description"] == ""
+    assert another_parameter["default"] == "Another way to define a default value"
+    assert another_parameter["type"] == "string"
+
+    my_parameter = init_params_schema['properties']['my_parameter']
+    assert my_parameter["description"] == "What a description"
+    assert my_parameter["default"] == "This is default value"
+    assert my_parameter["type"] == "string"
+
+    my_parameter_bool = init_params_schema['properties']['my_parameter_bool']
+    assert my_parameter_bool["default"] == False
+    assert my_parameter_bool["type"] == "boolean"
+
+def test_extract_env_vars(mock_tool_extractor):
+    tool_info = mock_tool_extractor
+
+    assert len(tool_info["env_vars"]) == 2
+    api_key_var, rate_limit_var = tool_info["env_vars"]
+    assert api_key_var["name"] == "SERPER_API_KEY"
+    assert api_key_var["description"] == "API key for Serper"
+    assert api_key_var["required"] == True
+    assert api_key_var["default"] == None
+
+    assert rate_limit_var["name"] == "API_RATE_LIMIT"
+    assert rate_limit_var["description"] == "API rate limit"
+    assert rate_limit_var["required"] == False
+    assert rate_limit_var["default"] == "100"
+
+def test_extract_run_params_schema(mock_tool_extractor):
+    tool_info = mock_tool_extractor
+
+    run_params_schema = tool_info["run_params_schema"]
+    assert run_params_schema.keys() == {
+        "properties",
+        "required",
+        "title",
+        "type",
+    }
 
-        params = {p["name"]: p for p in tool_info["run_params"]}
-        assert params["query"]["description"] == "The query parameter"
-        assert params["query"]["type"] == "str"
-        assert params["query"]["default"] == ""
+    query_param = run_params_schema["properties"]["query"]
+    assert query_param["description"] == "The query parameter"
+    assert query_param["type"] == "string"
 
-        assert params["count"]["type"] == "int"
-        assert params["count"]["default"] == 5
+    count_param = run_params_schema["properties"]["count"]
+    assert count_param["type"] == "integer"
+    assert count_param["default"] == 5
 
-        assert params["filters"]["description"] == "Optional filters to apply"
-        assert params["filters"]["type"] == "list[str]"
-        assert params["filters"]["default"] == ""
+    filters_param = run_params_schema["properties"]["filters"]
+    assert filters_param["description"] == "Optional filters to apply"
+    assert filters_param["default"] == None
+    assert filters_param['anyOf'] == [{'items': {'type': 'string'}, 'type': 'array'}, {'type': 'null'}]
 
-        assert tool_info["package_dependencies"] == ["this-is-a-required-package", "another-required-package"]
+def test_extract_package_dependencies(mock_tool_extractor):
+    tool_info = mock_tool_extractor
+    assert tool_info["package_dependencies"] == ["this-is-a-required-package", "another-required-package"]
 
 
 def test_save_to_json(extractor, tmp_path):
     extractor.tools_spec = [{
         "name": "TestTool",
         "humanized_name": "Test Tool",
         "description": "A test tool",
-        "run_params": [
+        "run_params_schema": [
             {"name": "param1", "description": "Test parameter", "type": "str"}
         ]
     }]
@@ -144,20 +159,4 @@ def test_save_to_json(extractor, tmp_path):
     assert "tools" in data
     assert len(data["tools"]) == 1
     assert data["tools"][0]["humanized_name"] == "Test Tool"
-    assert data["tools"][0]["run_params"][0]["name"] == "param1"
-
-
-@pytest.mark.integration
-def test_full_extraction_process():
-    extractor = ToolSpecExtractor()
-    specs = extractor.extract_all_tools()
-
-    assert len(specs) > 0
-
-    for tool in specs:
-        assert "name" in tool
-        assert "humanized_name" in tool and tool["humanized_name"]
-        assert "description" in tool
-        assert isinstance(tool["run_params"], list)
-        for param in tool["run_params"]:
-            assert "name" in param and param["name"]
+    assert data["tools"][0]["run_params_schema"][0]["name"] == "param1"