RULE/query_engine/schema_manager.py

from __future__ import annotations
"""
query_engine/schema_manager.py - Graph schema registry for QueryWeaver.

Maintains schema descriptions for named FalkorDB graphs.
Schema information is used to build LLM prompts for Cypher generation.

Pre-seeded schemas: rules_graph, fraud_graph, sales_graph, rulzai_graph.
"""


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

logger = logging.getLogger(__name__)


@dataclass
class GraphSchema:
    """Schema description for a single FalkorDB graph."""

    graph_name: str
    node_labels: List[str] = field(default_factory=list)
    relationship_types: List[str] = field(default_factory=list)
    properties: Dict[str, List[str]] = field(default_factory=dict)
    description: str = ""

    def to_prompt_str(self) -> str:
        """Format schema as a human-readable string for LLM prompts."""
        lines = [f"Graph: {self.graph_name}"]
        if self.description:
            lines.append(f"Description: {self.description}")
        if self.node_labels:
            lines.append(f"Node Labels: {', '.join(self.node_labels)}")
        if self.relationship_types:
            lines.append(f"Relationship Types: {', '.join(self.relationship_types)}")
        if self.properties:
            lines.append("Properties:")
            for label, props in self.properties.items():
                lines.append(f"  {label}: {', '.join(props)}")
        return "\n".join(lines)


class SchemaManager:
    """
    Manages schema information for all registered FalkorDB graphs.

    Usage:
        manager = SchemaManager()
        manager.register(GraphSchema(graph_name="my_graph", ...))
        schema_str = manager.describe_schema("my_graph")
    """

    def __init__(self):
        self._schemas: Dict[str, GraphSchema] = {}
        self._seed_defaults()

    def _seed_defaults(self) -> None:
        """Seed pre-defined schemas for known graphs."""
        defaults = [
            GraphSchema(
                graph_name="rules_graph",
                description="AML compliance rule dependency graph.",
                node_labels=["Rule", "Condition", "Action"],
                relationship_types=["DEPENDS_ON", "HAS_CONDITION", "TRIGGERS"],
                properties={
                    "Rule": ["rule_id", "name", "severity", "rule_type", "applies_to"],
                    "Condition": ["field", "operator", "value"],
                    "Action": ["decision", "auto_block", "sla_hours"],
                },
            ),
            GraphSchema(
                graph_name="fraud_graph",
                description="Fraud detection entity relationship graph.",
                node_labels=["Account", "Transaction", "Customer", "Alert", "Rule"],
                relationship_types=["OWNS", "SENT", "TRIGGERED", "LINKED_TO"],
                properties={
                    "Account": ["account_id", "risk_score", "country"],
                    "Transaction": [
                        "transaction_id", "amount_usd", "currency",
                        "timestamp", "match_score"
                    ],
                    "Customer": ["customer_id", "name", "pep_flag", "entity_type"],
                    "Alert": ["alert_id", "rule_id", "decision", "severity"],
                    "Rule": ["rule_id", "name", "severity"],
                },
            ),
            GraphSchema(
                graph_name="sales_graph",
                description="Sales product and customer relationship graph.",
                node_labels=["Customer", "Product", "Order", "Category"],
                relationship_types=["PURCHASED", "BELONGS_TO", "PLACED"],
                properties={
                    "Customer": ["customer_id", "name", "region", "segment"],
                    "Product": ["product_id", "name", "price", "category"],
                    "Order": ["order_id", "amount", "date", "status"],
                    "Category": ["name", "parent_category"],
                },
            ),
            GraphSchema(
                graph_name="rulzai_graph",
                description="General RULZAI analysis graph for rule and dataset relationships.",
                node_labels=["Rule", "Dataset", "Column", "Alert"],
                relationship_types=["EVALUATES", "HAS_COLUMN", "GENERATES"],
                properties={
                    "Rule": ["rule_id", "name", "outcome", "condition"],
                    "Dataset": ["name", "rows", "columns_count"],
                    "Column": ["name", "dtype", "distinct_count"],
                    "Alert": ["alert_id", "severity", "timestamp"],
                },
            ),
        ]
        for schema in defaults:
            self._schemas[schema.graph_name] = schema

    def register(self, schema: GraphSchema) -> None:
        """Register or update a graph schema."""
        self._schemas[schema.graph_name] = schema
        logger.debug("[SchemaManager] Registered schema for '%s'.", schema.graph_name)

    def get_schema(self, graph_name: str) -> Optional[GraphSchema]:
        """Retrieve a schema by graph name."""
        return self._schemas.get(graph_name)

    def describe_schema(self, graph_name: str) -> str:
        """
        Return a human-readable schema description for LLM prompt construction.

        Returns a generic description if the schema is not registered.
        """
        schema = self._schemas.get(graph_name)
        if schema:
            return schema.to_prompt_str()
        return (
            f"Graph: {graph_name}\n"
            f"(No schema registered — use generic MATCH/RETURN patterns.)"
        )

    def list_graphs(self) -> List[str]:
        """Return list of all registered graph names."""
        return list(self._schemas.keys())

    def register_from_ingest(
        self,
        graph_name: str,
        node_labels: List[str],
        relationship_types: List[str],
        properties: Dict[str, List[str]],
        description: str = "",
    ) -> None:
        """Register a schema discovered during graph auto-ingestion."""
        schema = GraphSchema(
            graph_name=graph_name,
            node_labels=node_labels,
            relationship_types=relationship_types,
            properties=properties,
            description=description or f"Auto-ingested graph: {graph_name}",
        )
        self.register(schema)
        logger.info("[SchemaManager] Auto-registered schema for '%s'.", graph_name)


# ─── Singleton ─────────────────────────────────────────────────────────────────

_schema_manager_instance: Optional[SchemaManager] = None


def get_schema_manager() -> SchemaManager:
    """Return the singleton SchemaManager instance."""
    global _schema_manager_instance
    if _schema_manager_instance is None:
        _schema_manager_instance = SchemaManager()
    return _schema_manager_instance