from pydantic import BaseModel, ValidationError
from typing import Any

class EntitySearchResult(BaseModel):
    entity_id: str
    name: str
    entity_type: str
    confidence: float

def call_tool_with_validation(tool_name: str, args: dict) -> Any:
    raw_response = dispatch_tool(tool_name, args)

    # Validate response matches what the model expects
    if tool_name == "search_entities":
        try:
            validated = [EntitySearchResult(**item) for item in raw_response]
            return validated
        except ValidationError as e:
            # Response shape changed — surface this immediately
            raise ToolResponseSchemaError(
                f"Tool '{tool_name}' returned unexpected shape. "
                f"Registration is out of sync with implementation.\n{e}"
            )

    return raw_response