models/tool_routing.py

"""
tool_routing.py
~~~~~~~~~~~~~~~
Modular tool implementation and routing engine.
Provides safe execution of calculation, datetime query, web search,
and context QA lookup tools.

Usage
-----
    from models.tool_routing import ToolRouter

    router = ToolRouter()
    reply, executed = router.route_and_execute("What is the current time?")
    if executed:
        print("Tool Output:", reply)
"""

from __future__ import annotations

import os
import ast
import operator
from datetime import datetime
import re
from typing import Dict, Any, Tuple

from models.logger_config import logger

# ── Safe Calculator Evaluation ────────────────────────────────────────────────
# Supported AST operators to avoid unsafe eval() execution
_SAFE_OPERATORS = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.truediv,
    ast.Pow: operator.pow,
    ast.USub: operator.neg,
}


def _safe_eval(node) -> float:
    """Recursively evaluate AST nodes safely."""
    if isinstance(node, ast.Num):  # python < 3.8 compatibility
        return node.n
    elif isinstance(node, ast.Constant):  # python 3.8+ compatibility
        if isinstance(node.value, (int, float)):
            return node.value
        raise TypeError("Unsupported constant type")
    elif isinstance(node, ast.BinOp):
        op_type = type(node.op)
        if op_type in _SAFE_OPERATORS:
            return _SAFE_OPERATORS[op_type](_safe_eval(node.left), _safe_eval(node.right))
        raise ValueError(f"Unsupported operator: {op_type}")
    elif isinstance(node, ast.UnaryOp):
        op_type = type(node.op)
        if op_type in _SAFE_OPERATORS:
            return _SAFE_OPERATORS[op_type](_safe_eval(node.operand))
        raise ValueError(f"Unsupported operator: {op_type}")
    raise ValueError(f"Unsupported syntax: {type(node)}")


def safe_calculate(expression: str) -> str:
    """Parse and calculate math expressions safely using AST representation."""
    clean_expr = re.sub(r"[^\d+\-*/().\s]", "", expression)
    if not clean_expr.strip():
        return "Error: Empty expression or invalid characters."
    try:
        node = ast.parse(clean_expr, mode="eval").body
        res = _safe_eval(node)
        return str(res)
    except Exception as e:
        return f"Error evaluating expression '{expression}': {e}"


# ── Tool Registry ─────────────────────────────────────────────────────────────

class ToolRouter:
    """
    Modular execution router for AI assistant tools.
    Integrates calculator, datetime, web search, and retrieval QA capabilities.
    """

    def __init__(self) -> None:
        pass

    def tool_calculator(self, expression: str) -> str:
        """Safe math calculation tool."""
        logger.info(f"[TOOL] Executing Calculator: {expression}")
        return f"[Calculator Result] {safe_calculate(expression)}"

    def tool_datetime(self) -> str:
        """Returns the current date and time."""
        logger.info("[TOOL] Executing Datetime")
        now = datetime.now()
        return f"[Datetime Result] Current date and time is {now.strftime('%Y-%m-%d %H:%M:%S')}."

    def tool_web_retrieval(self, query: str) -> str:
        """Mocked web retrieval interface returning contextual details."""
        logger.info(f"[TOOL] Executing Web Retrieval: {query}")
        query_lower = query.lower()
        
        # Simple simulated search summaries
        if "weather" in query_lower:
            return "[Web Result] Weather forecast is currently sunny, 22°C with light wind."
        elif "groq" in query_lower:
            return "[Web Result] Groq is an AI infrastructure company that builds LPU (Language Processing Unit) chips for high-speed inference."
        elif "qwen" in query_lower:
            return "[Web Result] Qwen is a series of open-source large language models developed by Alibaba Group, optimized for chat and reasoning."
        
        return f"[Web Result] Web search matches for '{query}': No highly relevant resources found."

    def tool_retrieval_qa(self, query: str, context: str) -> str:
        """Extract matching sentence or answer from a local text context."""
        logger.info(f"[TOOL] Executing Retrieval QA on query: {query}")
        sentences = re.split(r"(?<=[.!?])\s+", context)
        matches = []
        for word in query.lower().split():
            if len(word) > 3:  # focus on key terms
                for sentence in sentences:
                    if word in sentence.lower() and sentence not in matches:
                        matches.append(sentence)
        if matches:
            return f"[QA Result] Found matching context: " + " ".join(matches[:2])
        return "[QA Result] No direct matches found in current text context."

    def route_and_execute(self, user_input: str, qa_context: str = "") -> Tuple[str, bool]:
        """
        Scan input text for specific tool triggers and execute matched tool.
        Returns:
        --------
        (output_text, executed)
        """
        input_lower = user_input.lower()

        # 1. Calculator Trigger
        calc_match = re.search(r"\b(?:calculate|math|calc)\s+([\d+\-*/().\s]+)", input_lower)
        if calc_match:
            expr = calc_match.group(1).strip()
            return self.tool_calculator(expr), True

        # 2. Datetime Trigger
        if any(term in input_lower for term in ["current time", "what time is it", "today's date", "what is the date"]):
            return self.tool_datetime(), True

        # 3. Web Retrieval Trigger
        search_match = re.search(r"\b(?:search\s+for|web\s+search|lookup|find\s+info\s+on)\s+(.+)", input_lower)
        if search_match:
            q = search_match.group(1).strip()
            return self.tool_web_retrieval(q), True

        # 4. Retrieval QA Trigger (if context is present and question is asked)
        if qa_context and any(term in input_lower for term in ["who", "what", "where", "explain", "find"]):
            return self.tool_retrieval_qa(user_input, qa_context), True

        return "", False