src/query_router.py

import re


class QueryRouter:
    NUMERIC_KEYWORDS = [
        "price", "cost", "fee", "rate", "charge", "discount",
        "how much", "how many", "total", "amount", "per",
        "each", "minimum", "maximum", "limit",
    ]
    COMPARISON_KEYWORDS = [
        "compare", "vs", "versus", "better", "worse", "cheaper",
        "expensive", "difference", "cheapest", "best",
    ]
    RANGE_KEYWORDS = [
        "range", "between", "from.*to", "up to", "over",
        "at least", "at most", "minimum", "maximum",
    ]
    CALCULATION_KEYWORDS = [
        "calculate", "compute", "what is the total",
        "sum", "multiply", "divide", "add", "subtract",
    ]
    POLICY_KEYWORDS = [
        "policy", "rule", "regulation", "terms", "condition",
        "eligibility", "requirement", "allowed",
    ]
    SUMMARY_KEYWORDS = [
        "summarize", "summary", "overview", "tell me about",
        "what is", "describe",
    ]

    def route(self, query: str) -> str:
        query_lower = query.lower().strip()

        if any(kw in query_lower for kw in self.CALCULATION_KEYWORDS):
            return "calculation"
        if any(kw in query_lower for kw in self.COMPARISON_KEYWORDS):
            return "comparison"
        if any(kw in query_lower for kw in self.RANGE_KEYWORDS):
            return "range"
        if any(re.search(r"\d+", query_lower) for _ in [1]) and any(kw in query_lower for kw in self.NUMERIC_KEYWORDS):
            return "numeric"
        if any(kw in query_lower for kw in self.POLICY_KEYWORDS):
            return "policy"
        if any(kw in query_lower for kw in self.SUMMARY_KEYWORDS):
            return "summary"

        return "general"