core/query_processor.py

import re
import logging
from typing import List

logger = logging.getLogger("axiom.query")


class QueryProcessor:
    """
    Cleans, validates, and expands the user query
    before it hits the retrieval pipeline.
    """

    BANNED_PATTERNS = [
        r"ignore previous instructions",
        r"you are now",
        r"forget everything",
        r"act as",
    ]

    def __init__(self):
        logger.info("Ready.")

    def clean(self, query: str) -> str:
        """Basic cleaning — strip, normalize whitespace."""
        query = query.strip()
        query = re.sub(r"\s+", " ", query)
        return query

    def is_safe(self, query: str) -> bool:
        """Guardrail — reject prompt injection attempts."""
        lowered = query.lower()
        for pattern in self.BANNED_PATTERNS:
            if re.search(pattern, lowered):
                return False
        return True

    def expand(self, query: str) -> str:
        """
        Simple query expansion:
        appends 'explain', 'definition', 'overview'
        to help retrieval find introductory chunks.
        """
        expansion_suffix = "definition overview explanation"
        return f"{query} {expansion_suffix}"

    def process(self, raw_query: str) -> dict:
        """
        Full query processing pipeline.
        Returns dict with cleaned query, expanded query, and safety flag.
        """
        cleaned  = self.clean(raw_query)
        safe     = self.is_safe(cleaned)
        expanded = self.expand(cleaned) if safe else ""

        return {
            "original":  raw_query,
            "cleaned":   cleaned,
            "expanded":  expanded,
            "is_safe":   safe
        }