agents/restaurant_agent.py

"""
Restaurant agent for dining and food establishment searches.
"""
from typing import Dict
from .base_agent import BaseAgent


class RestaurantAgent(BaseAgent):
    """Specialized agent for finding restaurants and dining options."""
    
    def _define_capabilities(self) -> Dict[str, float]:
        """Define restaurant search capabilities."""
        return {
            "restaurant": 0.9,
            "restaurants": 0.9,
            "food": 0.8,
            "dining": 0.8,
            "eat": 0.7,
            "dinner": 0.7,
            "lunch": 0.7,
            "breakfast": 0.7,
            "cafe": 0.7,
            "bar": 0.7
        }
    
    def execute(self, task: str) -> str:
        """Execute restaurant search task."""
        try:
            # Clean up the task
            task = task.strip()
            
            # Add "restaurant" prefix if not present for better matching
            if not any(x in task.lower() for x in ["restaurant", "food", "dining", "eat"]):
                task = f"Find restaurants in {task}"
            
            # Use the base agent's execution method as a fallback
            return self._execute_through_agent(task)
            
        except Exception as e:
            return f"❌ Error searching for restaurants: {str(e)}"