from __future__ import annotations """ query_engine/query_interface.py - Safe query interface for business users. Provides a structured query interface that validates inputs against a whitelist of allowed patterns and executes them via the rule engine. """ import logging import re from dataclasses import dataclass, field from typing import Any, Dict, List, Optional import pandas as pd logger = logging.getLogger(__name__) __all__ = ["QueryInterface", "QueryResult"] # ─── Allowed pattern whitelist ───────────────────────────────────────────────── ALLOWED_PATTERNS = [ # Aggregation queries re.compile(r"(count|sum|average|avg|mean|max|min)\s+(of\s+)?[\w\s]+", re.I), re.compile(r"how many\s+[\w\s]+", re.I), re.compile(r"total\s+[\w\s]+", re.I), # Filtering re.compile(r"show\s+(me\s+)?[\w\s]+(where|with|having)\s+[\w\s]+", re.I), re.compile(r"filter\s+[\w\s]+", re.I), re.compile(r"list\s+(all\s+)?[\w\s]+", re.I), # Comparison / ranking re.compile(r"top\s+\d+\s+[\w\s]+", re.I), re.compile(r"[\w\s]+(greater|less|above|below|equal)\s+[\w\s]+", re.I), # Graph patterns re.compile(r"which\s+[\w\s]+(depend|trigger|link|connect)\w*", re.I), re.compile(r"[\w\s]+(depend|trigger|link|connect)\w+\s+on\s+[\w\s]+", re.I), ] @dataclass class QueryResult: """Result of a QueryInterface.execute() call.""" success: bool query: str answer: str = "" data: List[Dict[str, Any]] = field(default_factory=list) row_count: int = 0 error: Optional[str] = None engine: str = "dataframe" def to_dict(self) -> Dict[str, Any]: return { "success": self.success, "query": self.query, "answer": self.answer, "data": self.data, "row_count": self.row_count, "error": self.error, "engine": self.engine, } class QueryInterface: """ Safe business query interface with whitelist validation. Validates queries against allowed patterns before execution. Does NOT allow arbitrary AI-generated answers — only structured results. Usage: interface = QueryInterface(df=df) result = interface.execute("How many transactions have match_score > 90?") """ def __init__( self, df: Optional[pd.DataFrame] = None, strict_whitelist: bool = False, ): """ Args: df: DataFrame to query against. strict_whitelist: If True, reject queries not matching whitelist. """ self._df = df self._strict = strict_whitelist def execute(self, query: str) -> QueryResult: """ Execute a safe business query. Args: query: Natural language query string. Returns: QueryResult with structured data. """ query = query.strip() if not query: return QueryResult(success=False, query=query, error="Empty query.") # Validate against whitelist if self._strict and not self._matches_whitelist(query): return QueryResult( success=False, query=query, error=( "Query does not match allowed patterns. " "Use structured queries like 'count of X', 'show Y where Z', etc." ), ) if self._df is None or self._df.empty: return QueryResult( success=False, query=query, error="No dataset loaded. Upload a dataset first.", ) # Parse and execute result = self._dispatch_query(query) return result def _matches_whitelist(self, query: str) -> bool: """Check if query matches any allowed pattern.""" for pattern in ALLOWED_PATTERNS: if pattern.search(query): return True return False def _dispatch_query(self, query: str) -> QueryResult: """Simple pattern-based query dispatch for common business queries.""" q_lower = query.lower() df = self._df try: # Count queries if any(kw in q_lower for kw in ["how many", "count"]): return self._execute_count(query, df, q_lower) # Aggregation queries if any(kw in q_lower for kw in ["sum", "total", "average", "avg", "mean", "max", "min"]): return self._execute_aggregation(query, df, q_lower) # Top N queries top_match = re.search(r"top\s+(\d+)", q_lower) if top_match: return self._execute_top_n(query, df, int(top_match.group(1))) # List / show all if any(kw in q_lower for kw in ["list", "show", "filter"]): return self._execute_list(query, df, q_lower) # Default: return preview return QueryResult( success=True, query=query, answer=f"Dataset has {len(df)} rows and {len(df.columns)} columns.", data=df.head(20).to_dict(orient="records"), row_count=len(df), engine="dataframe", ) except Exception as exc: logger.error("[QueryInterface] Query dispatch failed: %s", exc) return QueryResult( success=False, query=query, error=str(exc), ) def _execute_count(self, query: str, df: pd.DataFrame, q_lower: str) -> QueryResult: """Handle count queries.""" total = len(df) return QueryResult( success=True, query=query, answer=f"Total count: {total} rows.", data=[{"count": total}], row_count=1, engine="dataframe", ) def _execute_aggregation(self, query: str, df: pd.DataFrame, q_lower: str) -> QueryResult: """Handle aggregation queries on numeric columns.""" numeric_cols = df.select_dtypes(include="number").columns.tolist() if not numeric_cols: return QueryResult( success=False, query=query, error="No numeric columns found." ) stats = {} for col in numeric_cols[:5]: # limit to first 5 numeric cols stats[col] = { "sum": float(df[col].sum()), "mean": float(df[col].mean()), "min": float(df[col].min()), "max": float(df[col].max()), } return QueryResult( success=True, query=query, answer=f"Aggregation statistics for {len(numeric_cols)} numeric columns.", data=[stats], row_count=len(numeric_cols), engine="dataframe", ) def _execute_top_n(self, query: str, df: pd.DataFrame, n: int) -> QueryResult: """Handle top-N queries.""" numeric_cols = df.select_dtypes(include="number").columns.tolist() if numeric_cols: sort_col = numeric_cols[0] result_df = df.nlargest(n, sort_col) else: result_df = df.head(n) return QueryResult( success=True, query=query, answer=f"Top {n} rows returned.", data=result_df.to_dict(orient="records"), row_count=len(result_df), engine="dataframe", ) def _execute_list(self, query: str, df: pd.DataFrame, q_lower: str) -> QueryResult: """Handle list/show/filter queries.""" preview = df.head(20).to_dict(orient="records") return QueryResult( success=True, query=query, answer=f"Showing first {len(preview)} of {len(df)} rows.", data=preview, row_count=len(df), engine="dataframe", )