import yaml
import os

class Router:
    def __init__(self, routing_rules_path='routing_rules.yaml'):
        self.routing_rules_path = routing_rules_path
        self.routing_rules = self._load_routing_rules()

    def _load_routing_rules(self):
        try:
            with open(self.routing_rules_path, 'r') as f:
                return yaml.safe_load(f)
        except FileNotFoundError:
            print(f"Routing rules file not found at {self.routing_rules_path}. Using default routing.")
            return {}
        except yaml.YAMLError as e:
            print(f"Error parsing routing rules file: {e}. Using default routing.")
            return {}

    def get_model_route(self, task_description):
        """Routes the request to the optimal model based on the task description."""
        for rule in self.routing_rules.get('rules', []):
            if 'contains' in rule and rule['contains'].lower() in task_description.lower():
                return rule['model']

        # Default model if no rule matches
        return self.routing_rules.get('default_model', 'gpt-3.5-turbo') # Default model


if __name__ == '__main__':
    # Example Usage
    router = Router()

    task1 = "Summarize this article on quantum physics."
    model1 = router.get_model_route(task1)
    print(f"Task: {task1}, Model: {model1}")

    task2 = "Write python code to calculate the factorial of a number."
    model2 = router.get_model_route(task2)
    print(f"Task: {task2}, Model: {model2}")

    task3 = "Translate this sentence into Spanish."
    model3 = router.get_model_route(task3)
    print(f"Task: {task3}, Model: {model3}")
