RULE/query_engine/query_executor.py

from __future__ import annotations
"""
query_engine/query_executor.py - QueryWeaver: NL → Cypher → FalkorDB execution.

Responsibility:
  1. Build an LLM prompt using the graph schema description.
  2. Generate a Cypher query.
  3. Validate it is read-only (block DELETE / DROP / SET / REMOVE / MERGE writes).
  4. Execute against FalkorDB and return structured results.
"""


import logging
import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from RULE.query_engine.falkordb_client import FalkorDBClient
from RULE.query_engine.schema_manager import SchemaManager

logger = logging.getLogger(__name__)

__all__ = ["ExecutionResult", "QueryExecutor"]

# ---------------------------------------------------------------------------
# Safety: patterns that must NOT appear in generated Cypher
# ---------------------------------------------------------------------------

_BLOCKED_PATTERNS = re.compile(
    r"\b(DELETE|DROP|SET\s+\w|REMOVE|CREATE|MERGE|DETACH\s+DELETE)\b",
    re.IGNORECASE,
)


@dataclass
class ExecutionResult:
    """Result of a QueryExecutor.execute() call."""
    success: bool
    cypher: str
    rows: List[Dict[str, Any]] = field(default_factory=list)
    error: Optional[str] = None
    graph_name: str = ""

    def row_count(self) -> int:
        return len(self.rows)


class QueryExecutor:
    """
    Translates NL → Cypher and executes against FalkorDB.

    Args:
        client:         Connected (or gracefully degraded) FalkorDBClient.
        schema_manager: SchemaManager instance for schema-aware prompts.
    """

    def __init__(self, client: Optional[FalkorDBClient], schema_manager: SchemaManager) -> None:
        self._client = client
        self._schema_manager = schema_manager

    def execute(self, nl_query: str, graph_name: str) -> ExecutionResult:
        """
        Execute nl_query as a Cypher query against graph_name.

        Pipeline:
          1. Build schema-aware LLM prompt.
          2. Generate Cypher.
          3. Validate safety.
          4. Run against FalkorDB.

        Returns:
            ExecutionResult — always returned; check success flag.
        """
        schema_desc = self._schema_manager.describe_schema(graph_name)

        cypher = self._generate_cypher(nl_query, schema_desc, graph_name)
        if not cypher:
            return ExecutionResult(
                success=False,
                cypher="",
                error="LLM did not generate a valid Cypher query.",
                graph_name=graph_name,
            )

        logger.debug("[QueryExecutor] Generated Cypher:\n%s", cypher)

        if not self._validate_cypher(cypher):
            logger.warning("[QueryExecutor] Cypher blocked by safety validator: %s", cypher)
            return ExecutionResult(
                success=False,
                cypher=cypher,
                error=(
                    "Generated Cypher contains write operations "
                    "(DELETE/DROP/SET/CREATE/MERGE). Only read queries are permitted."
                ),
                graph_name=graph_name,
            )

        if self._client is None or not self._client.is_connected():
            return ExecutionResult(
                success=False,
                cypher=cypher,
                error="FalkorDB is not connected. Start FalkorDB and retry.",
                graph_name=graph_name,
            )

        rows = self._client.execute_cypher(cypher, graph_name=graph_name)
        return ExecutionResult(
            success=True,
            cypher=cypher,
            rows=rows,
            graph_name=graph_name,
        )

    @staticmethod
    def _call_llm(prompt: str) -> str:
        """Invoke LLM via RULE's call_gemini."""
        try:
            from pandas_rule import call_gemini
            return call_gemini(prompt)
        except Exception as exc:
            logger.debug("[QueryExecutor] LLM call failed: %s", exc)
            return ""

    def _generate_cypher(self, nl_query: str, schema_desc: str, graph_name: str) -> str:
        """Ask the LLM to generate a Cypher query for nl_query."""
        prompt = f"""You are a Cypher query generator for FalkorDB.

Graph Schema:
{schema_desc}

User request: "{nl_query}"

Rules:
- Generate ONLY a single valid Cypher query.
- Use only MATCH and RETURN clauses (read-only).
- Do NOT use DELETE, DROP, SET, CREATE, MERGE, or REMOVE.
- Do NOT wrap the query in markdown code fences.
- Do NOT add any explanation. Output ONLY the Cypher query.

Example:
  User: "Which rules depend on AML_RULE?"
  Cypher: MATCH (r:Rule)-[:DEPENDS_ON]->(p:Rule {{name:'AML_RULE'}}) RETURN r.name, r.severity

Cypher:"""

        raw = self._call_llm(prompt).strip()
        raw = re.sub(r"```(?:cypher)?", "", raw, flags=re.IGNORECASE).strip()
        raw = raw.strip("`").strip()
        raw = raw.split(";")[0].strip()

        return raw

    @staticmethod
    def _validate_cypher(cypher: str) -> bool:
        """Return True if the Cypher query is safe (read-only)."""
        if not cypher or not cypher.strip():
            return False
        return _BLOCKED_PATTERNS.search(cypher) is None