tools/restaurant_tool.py

"""
Restaurant search tool for MCP server.
"""
from smolagents import Tool
from typing import Optional, Dict, Any
from services.restaurant_service import restaurant_service


class RestaurantSearchTool(Tool):
    """Tool for searching restaurants using enhanced restaurant service."""
    
    def __init__(self):
        self.name = "restaurant_search"
        self.description = "Search for restaurants and dining options in a location"
        self.input_type = "object"
        self.output_type = "object"
        self.inputs = {
            "location": {
                "type": "string",
                "description": "Location to search for restaurants"
            },
            "cuisine": {
                "type": "string",
                "description": "Type of cuisine (optional)",
                "optional": True,
                "nullable": True
            }
        }
        self.outputs = {
            "location": {
                "type": "string",
                "description": "Search location"
            },
            "total_found": {
                "type": "integer",
                "description": "Total number of restaurants found"
            },
            "restaurants": {
                "type": "array",
                "description": "List of restaurants found"
            }
        }
        self.required_inputs = ["location"]
        self.is_initialized = True

    def forward(self, location: str, cuisine: Optional[str] = None) -> dict:
        """Search for restaurants using enhanced restaurant service."""
        try:
            # The service now correctly handles location and cuisine parsing
            results = restaurant_service.search_restaurants(location, cuisine=cuisine)
            return results
            
        except Exception as e:
            return {
                "location": location,
                "total_found": 0,
                "restaurants": [],
                "error": f"Unexpected error: {str(e)}"
            }