RULE/query_engine/query_classifier.py

from __future__ import annotations
"""
query_engine/query_classifier.py - NL query type classifier for QueryWeaver.

Determines whether a natural-language query should be handled by:
  - The Pandas/DataFrame engine  ("dataframe_query")
  - The FalkorDB graph engine    ("graph_query")
  - A domain rule engine         ("rule_query")

Strategy:
  1. Fast keyword pre-filter — if clearly a graph query, skip the LLM call.
  2. LLM call (via call_gemini from RULE's pandas_rule module).
  3. JSON parse failure → fallback to "dataframe_query" (safe default).
"""


import json
import logging
import re
from dataclasses import dataclass
from typing import List, Literal, Optional

logger = logging.getLogger(__name__)

__all__ = ["ClassificationResult", "QueryClassifier"]

QueryType = Literal["dataframe_query", "graph_query", "rule_query"]

# Keywords that strongly suggest a graph / relationship query
_GRAPH_KEYWORDS = [
    "depend", "depends on", "dependency", "dependencies",
    "triggered by", "triggers", "which rules", "rule depend",
    "fraud rule", "linked to", "related to", "connected to",
    "relationship", "graph", "path", "neighbor", "neighbours",
    "ancestors", "descendants", "upstream", "downstream",
    "conflicts with", "conflict",
]

# Keywords that strongly suggest a rule / compliance graph query
_RULE_GRAPH_KEYWORDS = [
    "which customers triggered",
    "which accounts triggered",
    "triggered fraud",
    "triggered aml",
    "rule triggered",
    "rule dependency",
    "aml_rule",
    "aml rule",
]


@dataclass
class ClassificationResult:
    """Result of classifying a single natural-language query."""
    query_type: QueryType
    dataset: Optional[str]
    confidence: float
    reasoning: str

    def is_graph(self) -> bool:
        return self.query_type == "graph_query"

    def is_dataframe(self) -> bool:
        return self.query_type == "dataframe_query"


class QueryClassifier:
    """
    Classifies NL queries as dataframe / graph / rule queries.

    Args:
        df_columns: Column names of the currently loaded DataFrame.
        registry_summary: Text summary of registered graph datasets.
    """

    @staticmethod
    def _call_llm(prompt: str) -> str:
        """
        Invoke the LLM. Uses RULE's existing call_gemini from pandas_rule.
        Falls back gracefully if unavailable.
        """
        try:
            from pandas_rule import call_gemini
            return call_gemini(prompt)
        except Exception as exc:
            logger.debug("[Classifier] LLM call failed: %s", exc)
            return ""

    def classify(
        self,
        nl_query: str,
        df_columns: Optional[List[str]] = None,
        registry_summary: str = "",
    ) -> ClassificationResult:
        """
        Classify nl_query and return a ClassificationResult.

        Args:
            nl_query:         Raw natural-language query from the user.
            df_columns:       Column names of the active DataFrame (may be None).
            registry_summary: Text summary of registered graph datasets.
        """
        logger.debug("[Classifier] Classifying query: %r", nl_query)

        # Step 1: keyword pre-filter
        pre_result = self._keyword_classify(nl_query)
        if pre_result is not None:
            logger.debug(
                "[Classifier] Keyword fast-path → %s (dataset: %s)",
                pre_result.query_type, pre_result.dataset,
            )
            return pre_result

        # Step 2: LLM classification
        llm_result = self._llm_classify(nl_query, df_columns or [], registry_summary)
        if llm_result is not None:
            logger.debug(
                "[Classifier] LLM classified → %s (dataset: %s, confidence: %.2f)",
                llm_result.query_type, llm_result.dataset, llm_result.confidence,
            )
            return llm_result

        # Step 3: safe fallback
        logger.debug("[Classifier] Falling back to dataframe_query.")
        return ClassificationResult(
            query_type="dataframe_query",
            dataset=None,
            confidence=0.5,
            reasoning="Classification failed; defaulting to DataFrame engine.",
        )

    def _keyword_classify(self, nl_query: str) -> Optional[ClassificationResult]:
        """Return a result if the query matches strong keyword patterns."""
        q_lower = nl_query.lower()

        for kw in _RULE_GRAPH_KEYWORDS:
            if kw in q_lower:
                dataset = self._guess_dataset(q_lower)
                return ClassificationResult(
                    query_type="graph_query",
                    dataset=dataset or "rules",
                    confidence=0.90,
                    reasoning=f"Matched rule/graph keyword: '{kw}'",
                )

        graph_hits = [kw for kw in _GRAPH_KEYWORDS if kw in q_lower]
        if len(graph_hits) >= 2:
            dataset = self._guess_dataset(q_lower)
            return ClassificationResult(
                query_type="graph_query",
                dataset=dataset,
                confidence=0.85,
                reasoning=f"Matched graph keywords: {graph_hits[:3]}",
            )

        return None

    @staticmethod
    def _guess_dataset(q_lower: str) -> Optional[str]:
        """Guess the most likely dataset from query keywords."""
        if any(w in q_lower for w in ["fraud", "aml", "transaction", "account"]):
            return "fraud"
        if any(w in q_lower for w in ["rule", "policy", "depend", "trigger"]):
            return "rules"
        if any(w in q_lower for w in ["sales", "customer", "product", "purchase"]):
            return "sales"
        return None

    def _llm_classify(
        self,
        nl_query: str,
        df_columns: List[str],
        registry_summary: str,
    ) -> Optional[ClassificationResult]:
        """Use the LLM to classify the query. Returns None on any failure."""
        col_info = ", ".join(df_columns[:30]) if df_columns else "(no DataFrame loaded)"

        prompt = f"""You are a query routing assistant for a rule engine system.

Your task: classify the user's query into ONE of these types:
  - "dataframe_query" — filtering, aggregating, or transforming tabular data
  - "graph_query"     — relationships, dependencies, networks, or graph traversal
  - "rule_query"      — compliance/AML rule execution logic

Context:
  Active DataFrame columns: {col_info}
  {registry_summary}

User query: "{nl_query}"

Respond with ONLY valid JSON in this exact format (no markdown, no extra text):
{{
  "query_type": "dataframe_query" | "graph_query" | "rule_query",
  "dataset": "<dataset name if graph_query, else null>",
  "confidence": <float 0.0-1.0>,
  "reasoning": "<one sentence>"
}}"""

        raw = self._call_llm(prompt)
        if not raw:
            return None

        raw = re.sub(r"```(?:json)?", "", raw).strip()

        try:
            data = json.loads(raw)
            qt = data.get("query_type", "dataframe_query")
            if qt not in ("dataframe_query", "graph_query", "rule_query"):
                qt = "dataframe_query"
            return ClassificationResult(
                query_type=qt,
                dataset=data.get("dataset"),
                confidence=float(data.get("confidence", 0.7)),
                reasoning=str(data.get("reasoning", "")),
            )
        except (json.JSONDecodeError, TypeError, ValueError) as exc:
            logger.debug("[Classifier] JSON parse error: %s — raw: %r", exc, raw[:200])
            return None