diff --git a/mellea/backends/litellm.py b/mellea/backends/litellm.py index 226251ce2..6fab9d6b4 100644 --- a/mellea/backends/litellm.py +++ b/mellea/backends/litellm.py @@ -41,6 +41,7 @@ add_tools_from_context_actions, add_tools_from_model_options, convert_tools_to_json, + validate_tool_arguments, ) format: None = None # typing this variable in order to shadow the global format function and ensure mypy checks for errors @@ -603,7 +604,12 @@ def _extract_model_tool_requests( # Returns the args as a string. Parse it here. args = json.loads(tool_args) - model_tool_calls[tool_name] = ModelToolCall(tool_name, func, args) + + # Validate and coerce argument types + validated_args = validate_tool_arguments(func, args, strict=False) + model_tool_calls[tool_name] = ModelToolCall( + tool_name, func, validated_args + ) if len(model_tool_calls) > 0: return model_tool_calls diff --git a/mellea/backends/ollama.py b/mellea/backends/ollama.py index d20e2aa17..19efc1e33 100644 --- a/mellea/backends/ollama.py +++ b/mellea/backends/ollama.py @@ -502,6 +502,8 @@ async def generate_from_raw( def _extract_model_tool_requests( self, tools: dict[str, AbstractMelleaTool], chat_response: ollama.ChatResponse ) -> dict[str, ModelToolCall] | None: + from .tools import validate_tool_arguments + model_tool_calls: dict[str, ModelToolCall] = {} if chat_response.message.tool_calls: @@ -514,8 +516,11 @@ def _extract_model_tool_requests( continue # skip this function if we can't find it. args = tool.function.arguments + + # Validate and coerce argument types + validated_args = validate_tool_arguments(func, args, strict=False) model_tool_calls[tool.function.name] = ModelToolCall( - tool.function.name, func, args + tool.function.name, func, validated_args ) if len(model_tool_calls) > 0: diff --git a/mellea/backends/tools.py b/mellea/backends/tools.py index 09249c66f..ec24b96b7 100644 --- a/mellea/backends/tools.py +++ b/mellea/backends/tools.py @@ -211,10 +211,200 @@ def parse_tools(llm_response: str) -> list[tuple[str, Mapping]]: tool_name, tool_arguments = find_func(possible_tool) if tool_name is not None and tool_arguments is not None: tools.append((tool_name, tool_arguments)) - return tools +def validate_tool_arguments( + tool: AbstractMelleaTool, + args: Mapping[str, Any], + *, + coerce_types: bool = True, + strict: bool = False, +) -> dict[str, Any]: + """Validate and optionally coerce tool arguments against tool's JSON schema. + + This function validates tool call arguments extracted from LLM responses against + the tool's JSON schema from as_json_tool. It can automatically coerce common type + mismatches (e.g., string "30" to int 30) and provides detailed error messages. + + Args: + tool: The MelleaTool instance to validate against + args: Raw arguments from model (post-JSON parsing) + coerce_types: If True, attempt type coercion for common cases (default: True) + strict: If True, raise ValidationError on failures; if False, log warnings + and return original args (default: False) + + Returns: + Validated and optionally coerced arguments dict + + Raises: + ValidationError: If strict=True and validation fails + + Examples: + >>> def get_weather(location: str, days: int = 1) -> dict: + ... return {"location": location, "days": days} + >>> tool = MelleaTool.from_callable(get_weather) + + >>> # LLM returns days as string + >>> args = {"location": "Boston", "days": "3"} + >>> validated = validate_tool_arguments(tool, args) + >>> validated + {'location': 'Boston', 'days': 3} + + >>> # Strict mode raises on validation errors + >>> bad_args = {"location": "Boston", "days": "not_a_number"} + >>> validate_tool_arguments(tool, bad_args, strict=True) + Traceback (most recent call last): + ... + pydantic.ValidationError: ... + """ + from pydantic import ValidationError, create_model + + from ..core import FancyLogger + + # Extract JSON schema from tool + tool_schema = tool.as_json_tool.get("function", {}) + tool_name = tool_schema.get("name", "unknown_tool") + parameters = tool_schema.get("parameters", {}) + properties = parameters.get("properties", {}) + required_fields = parameters.get("required", []) + + # Map JSON schema types to Python types + JSON_TYPE_TO_PYTHON = { + "string": str, + "integer": int, + "number": float, + "boolean": bool, + "array": list, + "object": dict, + } + + # Build Pydantic model from JSON schema + field_definitions: dict[str, Any] = {} + + for param_name, param_schema in properties.items(): + # Get type from JSON schema + json_type = param_schema.get("type", "string") + + # Handle comma-separated types (e.g., "integer, string" for Union types) + if isinstance(json_type, str) and "," in json_type: + # Create Union type for multiple types + type_list = [t.strip() for t in json_type.split(",")] + python_types = [JSON_TYPE_TO_PYTHON.get(t, Any) for t in type_list] + # Remove duplicates while preserving order + seen = set() + unique_types = [] + for t in python_types: + if t not in seen: + seen.add(t) + unique_types.append(t) + + if len(unique_types) == 1: + param_type = unique_types[0] + else: + # Use modern union syntax (Python 3.10+) + from functools import reduce + from operator import or_ + + param_type = reduce(or_, unique_types) + else: + # Map to Python type + param_type = JSON_TYPE_TO_PYTHON.get(json_type, Any) + + # Determine if parameter is required + if param_name in required_fields: + # Required parameter + field_definitions[param_name] = (param_type, ...) + else: + # Optional parameter (default to None) + field_definitions[param_name] = (param_type, None) + + # Configure model for type coercion if requested + if coerce_types: + model_config = ConfigDict( + str_strip_whitespace=True, + strict=False, # Allow type coercion + extra="forbid" if strict else "allow", # Handle extra fields + # Enable coercion modes for common LLM output issues + coerce_numbers_to_str=True, # Allow int/float -> str + ) + else: + model_config = ConfigDict( + strict=True, # No coercion + extra="forbid" if strict else "allow", + ) + + # Create dynamic Pydantic model for validation + ValidatorModel = create_model( + f"{tool_name}_Validator", __config__=model_config, **field_definitions + ) + + try: + # Validate using Pydantic + validated_model = ValidatorModel(**args) + validated_args = validated_model.model_dump() + + # In lenient mode with extra="allow", Pydantic includes extra fields + # but we need to preserve them from the original args + if not strict: + # Add back any extra fields that weren't in the model + for key, value in args.items(): + if key not in field_definitions: + validated_args[key] = value + + # Log successful validation with coercion details + coerced_fields = [] + for key, original_value in args.items(): + validated_value = validated_args.get(key) + if type(original_value) is not type(validated_value): + coerced_fields.append( + f"{key}: {type(original_value).__name__} → {type(validated_value).__name__}" + ) + + if coerced_fields and coerce_types: + FancyLogger.get_logger().debug( + f"Tool '{tool_name}' arguments coerced: {', '.join(coerced_fields)}" + ) + + return validated_args + + except ValidationError as e: + # Format error message + error_details = [] + for error in e.errors(): + field = ".".join(str(loc) for loc in error["loc"]) + msg = error["msg"] + error_details.append(f" - {field}: {msg}") + + error_msg = f"Tool argument validation failed for '{tool_name}':\n" + "\n".join( + error_details + ) + + if strict: + # Re-raise with enhanced message + FancyLogger.get_logger().error(error_msg) + raise + else: + # Log warning and return original args + FancyLogger.get_logger().warning( + error_msg + "\nReturning original arguments without validation." + ) + return dict(args) + + except Exception as e: + # Catch any other errors during validation + error_msg = f"Unexpected error validating tool '{tool_name}' arguments: {e}" + + if strict: + FancyLogger.get_logger().error(error_msg) + raise + else: + FancyLogger.get_logger().warning( + error_msg + "\nReturning original arguments without validation." + ) + return dict(args) + + # Below functions and classes extracted from Ollama Python SDK (v0.6.1) # so that all backends don't need it installed. # https://github.com/ollama/ollama-python/blob/60e7b2f9ce710eeb57ef2986c46ea612ae7516af/ollama/_types.py#L19-L101 diff --git a/mellea/backends/utils.py b/mellea/backends/utils.py index 2c3c00f6d..3a5c5cf71 100644 --- a/mellea/backends/utils.py +++ b/mellea/backends/utils.py @@ -9,7 +9,7 @@ from ..core.base import AbstractMelleaTool from ..formatters import ChatFormatter from ..stdlib.components import Message -from .tools import parse_tools +from .tools import parse_tools, validate_tool_arguments # Chat = dict[Literal["role", "content"], str] # external apply_chat_template type hint is weaker # Chat = dict[str, str | list[dict[str, Any]] ] # for multi-modal models @@ -75,7 +75,9 @@ def to_tool_calls( if len(param_map) == 0: tool_args = {} - model_tool_calls[tool_name] = ModelToolCall(tool_name, func, tool_args) + # Validate and coerce argument types + validated_args = validate_tool_arguments(func, tool_args, strict=False) + model_tool_calls[tool_name] = ModelToolCall(tool_name, func, validated_args) if len(model_tool_calls) > 0: return model_tool_calls diff --git a/mellea/backends/watsonx.py b/mellea/backends/watsonx.py index 0fe1c2bdc..27401e537 100644 --- a/mellea/backends/watsonx.py +++ b/mellea/backends/watsonx.py @@ -44,6 +44,7 @@ add_tools_from_context_actions, add_tools_from_model_options, convert_tools_to_json, + validate_tool_arguments, ) format: None = None # typing this variable in order to shadow the global format function and ensure mypy checks for errors @@ -590,7 +591,10 @@ def _extract_model_tool_requests( # Watsonx returns the args as a string. Parse it here. args = json.loads(tool_args) - model_tool_calls[tool_name] = ModelToolCall(tool_name, func, args) + + # Validate and coerce argument types + validated_args = validate_tool_arguments(func, args, strict=False) + model_tool_calls[tool_name] = ModelToolCall(tool_name, func, validated_args) if len(model_tool_calls) > 0: return model_tool_calls diff --git a/mellea/helpers/openai_compatible_helpers.py b/mellea/helpers/openai_compatible_helpers.py index 7374b157f..5afec1232 100644 --- a/mellea/helpers/openai_compatible_helpers.py +++ b/mellea/helpers/openai_compatible_helpers.py @@ -4,6 +4,7 @@ from collections.abc import Callable from typing import Any +from ..backends.tools import validate_tool_arguments from ..core import FancyLogger, ModelToolCall from ..core.base import AbstractMelleaTool from ..stdlib.components import Document, Message @@ -31,7 +32,10 @@ def extract_model_tool_requests( if tool_args is not None: # Returns the args as a string. Parse it here. args = json.loads(tool_args) - model_tool_calls[tool_name] = ModelToolCall(tool_name, func, args) + + # Validate and coerce argument types + validated_args = validate_tool_arguments(func, args, strict=False) + model_tool_calls[tool_name] = ModelToolCall(tool_name, func, validated_args) if len(model_tool_calls) > 0: return model_tool_calls diff --git a/test/backends/test_tool_validation_integration.py b/test/backends/test_tool_validation_integration.py new file mode 100644 index 000000000..2ffd488fa --- /dev/null +++ b/test/backends/test_tool_validation_integration.py @@ -0,0 +1,459 @@ +"""Integration tests for the validate_tool_arguments function. + +These tests verify that the validation function works correctly with +the actual tool call flow. +""" + +from typing import Any, Optional, Union + +import pytest +from pydantic import ValidationError + +from mellea.backends.tools import MelleaTool, validate_tool_arguments +from mellea.core import ModelToolCall + +# ============================================================================ +# Test Fixtures - Tool Functions +# ============================================================================ + + +def simple_tool(message: str) -> str: + """A simple tool that takes a string. + + Args: + message: The message to process + """ + return f"Processed: {message}" + + +def typed_tool(name: str, age: int, score: float, active: bool) -> dict: + """Tool with multiple primitive types. + + Args: + name: Person's name + age: Person's age in years + score: Performance score + active: Whether person is active + """ + return {"name": name, "age": age, "score": score, "active": active} + + +def optional_tool(required: str, optional: str | None = None) -> str: + """Tool with optional parameters. + + Args: + required: A required parameter + optional: An optional parameter + """ + return f"{required}:{optional or 'none'}" + + +def union_tool(value: str | int) -> str: + """Tool with union type parameter. + + Args: + value: Can be string or integer + """ + return f"Value: {value} (type: {type(value).__name__})" + + +def list_tool(items: list[str]) -> int: + """Tool with list parameter. + + Args: + items: List of string items + """ + return len(items) + + +def dict_tool(config: dict[str, Any]) -> str: + """Tool with dict parameter. + + Args: + config: Configuration dictionary + """ + import json + + return json.dumps(config) + + +def no_params_tool() -> str: + """Tool with no parameters.""" + return "No params needed" + + +def untyped_param(message) -> str: + """A tool with an untyped parameter. + + Args: + message: The message to process (no type hint) + """ + return f"Processed: {message}" + + +# ============================================================================ +# Test Cases: Type Coercion +# ============================================================================ + + +class TestTypeCoercion: + """Test automatic type coercion with validation.""" + + def test_string_to_int_coercion(self): + """Test that string "30" is coerced to int 30.""" + args = {"name": "Test", "age": "30", "score": 95.5, "active": True} + tool = MelleaTool.from_callable(typed_tool) + validated = validate_tool_arguments(tool, args, coerce_types=True) + + assert validated["age"] == 30 + assert isinstance(validated["age"], int) + + def test_string_to_float_coercion(self): + """Test that string "95.5" is coerced to float 95.5.""" + args = {"name": "Test", "age": 30, "score": "95.5", "active": True} + tool = MelleaTool.from_callable(typed_tool) + validated = validate_tool_arguments(tool, args, coerce_types=True) + + assert validated["score"] == 95.5 + assert isinstance(validated["score"], float) + + def test_int_to_float_coercion(self): + """Test that int 95 is coerced to float 95.0.""" + args = {"name": "Test", "age": 30, "score": 95, "active": True} + tool = MelleaTool.from_callable(typed_tool) + validated = validate_tool_arguments(tool, args, coerce_types=True) + + assert validated["score"] == 95.0 + assert isinstance(validated["score"], float) + + def test_int_to_string_coercion(self): + """Test that int 123 is coerced to string "123".""" + args = {"message": 123} + tool = MelleaTool.from_callable(simple_tool) + validated = validate_tool_arguments(tool, args, coerce_types=True) + + assert validated["message"] == "123" + assert isinstance(validated["message"], str) + + def test_bool_coercion_from_int(self): + """Test that int 1/0 is coerced to bool True/False.""" + args = {"name": "Test", "age": 30, "score": 95.5, "active": 1} + tool = MelleaTool.from_callable(typed_tool) + validated = validate_tool_arguments(tool, args, coerce_types=True) + + assert validated["active"] is True + assert isinstance(validated["active"], bool) + + args["active"] = 0 + validated = validate_tool_arguments(tool, args, coerce_types=True) + assert validated["active"] is False + + +class TestValidationModes: + """Test strict vs. lenient validation modes.""" + + def test_lenient_mode_with_invalid_type(self): + """Test that lenient mode returns original args on validation failure.""" + args = {"name": "Test", "age": "not_a_number", "score": 95.5, "active": True} + tool = MelleaTool.from_callable(typed_tool) + validated = validate_tool_arguments(tool, args, strict=False) + + # Should return original args + assert validated == args + assert validated["age"] == "not_a_number" + + def test_strict_mode_with_invalid_type(self): + """Test that strict mode raises ValidationError on failure.""" + args = {"name": "Test", "age": "not_a_number", "score": 95.5, "active": True} + tool = MelleaTool.from_callable(typed_tool) + + with pytest.raises(ValidationError): + validate_tool_arguments(tool, args, strict=True) + + def test_lenient_mode_with_missing_required(self): + """Test lenient mode with missing required parameter.""" + args = {"optional": "value"} # Missing 'required' + tool = MelleaTool.from_callable(optional_tool) + validated = validate_tool_arguments(tool, args, strict=False) + + # Should return original args + assert validated == args + + def test_strict_mode_with_missing_required(self): + """Test strict mode with missing required parameter.""" + args = {"optional": "value"} # Missing 'required' + tool = MelleaTool.from_callable(optional_tool) + + with pytest.raises(ValidationError): + validate_tool_arguments(tool, args, strict=True) + + +class TestWithModelToolCall: + """Test validation integrated with ModelToolCall.""" + + def test_validated_tool_call_with_coercion(self): + """Test that validated args work correctly with ModelToolCall.""" + # LLM returns age as string + args = {"name": "Alice", "age": "30", "score": "95.5", "active": True} + tool = MelleaTool.from_callable(typed_tool) + + # Validate and coerce + validated_args = validate_tool_arguments(tool, args, coerce_types=True) + + # Create tool call with validated args + tool_call = ModelToolCall("typed_tool", tool, validated_args) + result = tool_call.call_func() + + # Verify result has correct types + assert result["age"] == 30 + assert isinstance(result["age"], int) + assert result["score"] == 95.5 + assert isinstance(result["score"], float) + + def test_unvalidated_vs_validated_comparison(self): + """Compare behavior with and without validation.""" + args = {"name": "Bob", "age": "25", "score": "88.7", "active": True} + tool = MelleaTool.from_callable(typed_tool) + + # Without validation - types stay as strings + unvalidated_call = ModelToolCall("typed_tool", tool, args) + unvalidated_result = unvalidated_call.call_func() + assert isinstance(unvalidated_result["age"], str) # Still string! + + # With validation - types are coerced + validated_args = validate_tool_arguments(tool, args, coerce_types=True) + validated_call = ModelToolCall("typed_tool", tool, validated_args) + validated_result = validated_call.call_func() + assert isinstance(validated_result["age"], int) # Correctly coerced! + + +class TestOptionalParameters: + """Test validation with optional parameters.""" + + def test_optional_param_provided(self): + """Test validation when optional parameter is provided.""" + args = {"required": "value1", "optional": "value2"} + tool = MelleaTool.from_callable(optional_tool) + validated = validate_tool_arguments(tool, args) + + assert validated == args + + def test_optional_param_omitted(self): + """Test validation when optional parameter is omitted.""" + args = {"required": "value1"} + tool = MelleaTool.from_callable(optional_tool) + validated = validate_tool_arguments(tool, args) + + assert validated["required"] == "value1" + assert "optional" not in validated or validated.get("optional") is None + + def test_optional_param_none(self): + """Test validation when optional parameter is explicitly None.""" + args = {"required": "value1", "optional": None} + tool = MelleaTool.from_callable(optional_tool) + validated = validate_tool_arguments(tool, args) + + assert validated["required"] == "value1" + assert validated["optional"] is None + + +class TestComplexTypes: + """Test validation with complex types.""" + + def test_list_parameter(self): + """Test validation with list parameter.""" + args = {"items": ["apple", "banana", "cherry"]} + tool = MelleaTool.from_callable(list_tool) + validated = validate_tool_arguments(tool, args) + + assert validated["items"] == ["apple", "banana", "cherry"] + assert isinstance(validated["items"], list) + + def test_dict_parameter(self): + """Test validation with dict parameter.""" + args = {"config": {"key1": "value1", "key2": 42, "key3": True}} + tool = MelleaTool.from_callable(dict_tool) + validated = validate_tool_arguments(tool, args) + + assert validated["config"] == args["config"] + assert isinstance(validated["config"], dict) + + def test_empty_list(self): + """Test validation with empty list.""" + args = {"items": []} + tool = MelleaTool.from_callable(list_tool) + validated = validate_tool_arguments(tool, args) + + assert validated["items"] == [] + + +class TestUnionTypes: + """Test validation with union types.""" + + def test_union_with_string(self): + """Test union type with string value.""" + args = {"value": "hello"} + tool = MelleaTool.from_callable(union_tool) + validated = validate_tool_arguments(tool, args) + + assert validated["value"] == "hello" + assert isinstance(validated["value"], str) + + def test_union_with_int(self): + """Test union type with integer value.""" + args = {"value": 42} + tool = MelleaTool.from_callable(union_tool) + validated = validate_tool_arguments(tool, args) + + assert validated["value"] == 42 + assert isinstance(validated["value"], int) + + def test_union_with_string_number(self): + """Test union type with string that looks like number.""" + args = {"value": "42"} + tool = MelleaTool.from_callable(union_tool) + validated = validate_tool_arguments(tool, args, coerce_types=True) + + # Pydantic will try to coerce to the first matching type + # Result depends on Union order and Pydantic's coercion rules + assert validated["value"] in ["42", 42] + + +class TestEdgeCases: + """Test edge cases.""" + + def test_no_parameters_tool(self): + """Test validation with no-parameter tool.""" + args = {} + tool = MelleaTool.from_callable(no_params_tool) + validated = validate_tool_arguments(tool, args) + + assert validated == {} + + def test_no_parameters_with_extra_args(self): + """Test that extra args for no-param tool are handled.""" + args = {"fake_param": "should_be_ignored"} + tool = MelleaTool.from_callable(no_params_tool) + + # In lenient mode, returns original args + validated = validate_tool_arguments(tool, args, strict=False) + assert validated == args + + # In strict mode, should raise + with pytest.raises(ValidationError): + validate_tool_arguments(tool, args, strict=True) + + def test_whitespace_stripping(self): + """Test that whitespace is stripped from strings.""" + args = {"message": " hello world "} + tool = MelleaTool.from_callable(simple_tool) + validated = validate_tool_arguments(tool, args, coerce_types=True) + + assert validated["message"] == "hello world" + + def test_empty_string(self): + """Test validation with empty string.""" + args = {"message": ""} + tool = MelleaTool.from_callable(simple_tool) + validated = validate_tool_arguments(tool, args) + + assert validated["message"] == "" + + +class TestErrorMessages: + """Test that error messages are helpful.""" + + def test_missing_required_error_message(self): + """Test error message for missing required parameter.""" + args = {} + tool = MelleaTool.from_callable(simple_tool) + + try: + validate_tool_arguments(tool, args, strict=True) + pytest.fail("Should have raised ValidationError") + except ValidationError as e: + error_str = str(e) + assert "message" in error_str.lower() + assert "required" in error_str.lower() or "missing" in error_str.lower() + + def test_type_mismatch_error_message(self): + """Test error message for type mismatch.""" + args = {"name": "Test", "age": "not_a_number", "score": 95.5, "active": True} + tool = MelleaTool.from_callable(typed_tool) + + try: + validate_tool_arguments(tool, args, strict=True) + pytest.fail("Should have raised ValidationError") + except ValidationError as e: + error_str = str(e) + assert "age" in error_str.lower() + + +class TestUntypedParameters: + """Test validation with untyped parameters.""" + + def test_untyped_parameter_accepts_string(self): + """Test that untyped parameters accept string values.""" + args = {"message": "test"} + tool = MelleaTool.from_callable(untyped_param) + validated = validate_tool_arguments(tool, args) + + assert validated["message"] == "test" + + def test_untyped_parameter_accepts_int(self): + """Test that untyped parameters accept integer values. + + Note: Without type hints, validation may coerce to string for safety. + """ + args = {"message": 123} + tool = MelleaTool.from_callable(untyped_param) + validated = validate_tool_arguments(tool, args) + + # Validation may coerce to string when no type hint is present + assert validated["message"] in [123, "123"] + + def test_untyped_parameter_accepts_dict(self): + """Test untyped parameter with complex type (dict).""" + args = {"message": {"key": "value", "number": 42}} + tool = MelleaTool.from_callable(untyped_param) + validated = validate_tool_arguments(tool, args) + + assert validated["message"] == {"key": "value", "number": 42} + + def test_untyped_parameter_accepts_list(self): + """Test untyped parameter with list.""" + args = {"message": ["item1", "item2", "item3"]} + tool = MelleaTool.from_callable(untyped_param) + validated = validate_tool_arguments(tool, args) + + assert validated["message"] == ["item1", "item2", "item3"] + + def test_untyped_parameter_accepts_bool(self): + """Test untyped parameter with boolean.""" + args = {"message": True} + tool = MelleaTool.from_callable(untyped_param) + validated = validate_tool_arguments(tool, args) + + assert validated["message"] is True + + def test_untyped_parameter_accepts_none(self): + """Test untyped parameter with None.""" + args = {"message": None} + tool = MelleaTool.from_callable(untyped_param) + validated = validate_tool_arguments(tool, args) + + assert validated["message"] is None + + def test_untyped_parameter_no_coercion(self): + """Test that untyped parameters don't get coerced.""" + args = {"message": "123"} + tool = MelleaTool.from_callable(untyped_param) + validated = validate_tool_arguments(tool, args, coerce_types=True) + + # Should remain as string since there's no type hint to coerce to + assert validated["message"] == "123" + assert isinstance(validated["message"], str) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])