prolog_kb.py

"""
Prolog knowledge base: convert NLP facts into Prolog predicates and query them.

Uses pyswip to interface with SWI-Prolog for efficient unification and backtracking.
"""

import re
from pyswip import Prolog


def _safe_atom(s):
    """Convert string to a safe Prolog atom (lowercase, alphanumeric + underscore)."""
    s = s.lower().strip()
    s = re.sub(r"[^a-z0-9_]", "_", s)
    s = re.sub(r"_+", "_", s)
    s = s.strip("_")
    if not s or not s[0].isalpha():
        s = "x_" + s
    return s[:60]


def _safe_string(s):
    """Escape a string for Prolog string literals."""
    return s.replace("\\", "\\\\").replace("'", "\\'").replace('"', '\\"')


class PrologKB:
    """
    A Prolog-based knowledge base for session facts.

    Facts asserted:
        entity(SessionAtom, EntityAtom, TypeAtom)
        mentions(SessionAtom, WordAtom)
        relation(SessionAtom, SubjAtom, PredAtom, ObjAtom)
        temporal(SessionAtom, TimeAtom)
        preference(SessionAtom, PrefAtom)
        trigram(SessionAtom, W1, W2, W3)
        proper_noun(SessionAtom, NameAtom)
        noun_phrase(SessionAtom, PhraseAtom)
        keyword(SessionAtom, KeywordAtom)
        quoted(SessionAtom, PhraseAtom)

    Query helpers for finding sessions by various criteria.
    """

    def __init__(self):
        self.prolog = Prolog()
        self._session_atoms = {}  # session_id -> atom
        self._fact_count = 0

        # Define helper rules
        self._define_rules()

    def _define_rules(self):
        """Define Prolog rules for complex queries."""
        # Find sessions that share an entity
        self.prolog.assertz(
            "shared_entity(S1, S2, E) :- entity(S1, E, _), entity(S2, E, _), S1 \\= S2"
        )

        # Find sessions connected by a relation subject or object
        self.prolog.assertz(
            "related_session(S1, S2) :- relation(S1, Subj, _, _), "
            "relation(S2, Subj, _, _), S1 \\= S2"
        )
        self.prolog.assertz(
            "related_session(S1, S2) :- relation(S1, _, _, Obj), "
            "relation(S2, _, _, Obj), S1 \\= S2"
        )

        # Find sessions mentioning a specific keyword
        # (already direct: mentions/2)

        # Find sessions with matching trigram patterns
        self.prolog.assertz(
            "trigram_match(S1, S2) :- trigram(S1, A, B, C), trigram(S2, A, B, C), S1 \\= S2"
        )

    def _session_atom(self, session_id):
        """Get or create an atom for a session ID."""
        if session_id not in self._session_atoms:
            self._session_atoms[session_id] = _safe_atom(session_id)
        return self._session_atoms[session_id]

    def add_session_facts(self, facts):
        """
        Add all facts from a session to the Prolog KB.

        Args:
            facts: dict from nlp_extract.extract_session_facts()
        """
        sid = self._session_atom(facts["session_id"])

        # Entities
        for ent_text, ent_type in facts.get("entities", []):
            atom = _safe_atom(ent_text)
            type_atom = _safe_atom(ent_type)
            self._assert(f"entity({sid}, {atom}, {type_atom})")

        # Keywords / mentions
        for word in facts.get("keywords", []):
            atom = _safe_atom(word)
            self._assert(f"mentions({sid}, {atom})")

        # Relations (SVO triples)
        for subj, pred, obj in facts.get("relations", []):
            s = _safe_atom(subj)
            p = _safe_atom(pred)
            o = _safe_atom(obj)
            self._assert(f"relation({sid}, {s}, {p}, {o})")

        # Temporal
        for t in facts.get("temporal", []):
            atom = _safe_atom(t)
            self._assert(f"temporal({sid}, {atom})")

        # Preferences
        for pref in facts.get("preferences", []):
            atom = _safe_atom(pref)
            self._assert(f"preference({sid}, {atom})")

        # Trigrams (limit to avoid explosion — keep first 50 most unique)
        trigrams = facts.get("trigrams", [])[:50]
        for w1, w2, w3 in trigrams:
            a1, a2, a3 = _safe_atom(w1), _safe_atom(w2), _safe_atom(w3)
            self._assert(f"trigram({sid}, {a1}, {a2}, {a3})")

        # Proper nouns
        for name in facts.get("proper_nouns", []):
            atom = _safe_atom(name)
            self._assert(f"proper_noun({sid}, {atom})")

        # Noun phrases (first 30)
        for phrase in facts.get("noun_phrases", [])[:30]:
            atom = _safe_atom(phrase)
            self._assert(f"noun_phrase({sid}, {atom})")

        # Quoted phrases
        for phrase in facts.get("quoted_phrases", []):
            atom = _safe_atom(phrase)
            self._assert(f"quoted({sid}, {atom})")

    def _assert(self, clause):
        """Assert a clause, silently skip duplicates."""
        try:
            self.prolog.assertz(clause)
            self._fact_count += 1
        except Exception:
            pass  # Skip malformed assertions

    def query_by_entities(self, entity_atoms):
        """Find sessions containing any of the given entities."""
        scores = {}
        for ent in entity_atoms:
            atom = _safe_atom(ent)
            try:
                results = list(self.prolog.query(f"entity(S, {atom}, _)"))
                for r in results:
                    s = str(r["S"])
                    scores[s] = scores.get(s, 0) + 2.0  # Entity match is strong signal
            except Exception:
                pass
        return scores

    def query_by_keywords(self, keyword_atoms):
        """Find sessions mentioning any of the given keywords."""
        scores = {}
        for kw in keyword_atoms:
            atom = _safe_atom(kw)
            try:
                results = list(self.prolog.query(f"mentions(S, {atom})"))
                for r in results:
                    s = str(r["S"])
                    scores[s] = scores.get(s, 0) + 1.0
            except Exception:
                pass
        return scores

    def query_by_proper_nouns(self, names):
        """Find sessions mentioning proper nouns (person names)."""
        scores = {}
        for name in names:
            atom = _safe_atom(name)
            try:
                results = list(self.prolog.query(f"proper_noun(S, {atom})"))
                for r in results:
                    s = str(r["S"])
                    scores[s] = scores.get(s, 0) + 3.0  # Names are very strong signal
            except Exception:
                pass
        return scores

    def query_by_quoted_phrases(self, phrases):
        """Find sessions containing quoted phrases."""
        scores = {}
        for phrase in phrases:
            atom = _safe_atom(phrase)
            try:
                results = list(self.prolog.query(f"quoted(S, {atom})"))
                for r in results:
                    s = str(r["S"])
                    scores[s] = scores.get(s, 0) + 4.0  # Exact quotes very strong
            except Exception:
                pass
        return scores

    def query_by_preferences(self, pref_keywords):
        """Find sessions with preferences matching keywords."""
        scores = {}
        for kw in pref_keywords:
            atom = _safe_atom(kw)
            try:
                results = list(self.prolog.query(f"preference(S, P), sub_atom(P, _, _, _, {atom})"))
            except Exception:
                results = []

            # Fallback: direct keyword match in preference atoms
            if not results:
                try:
                    all_prefs = list(self.prolog.query("preference(S, P)"))
                    for r in all_prefs:
                        if atom in str(r["P"]):
                            scores[str(r["S"])] = scores.get(str(r["S"]), 0) + 1.5
                except Exception:
                    pass
            else:
                for r in results:
                    s = str(r["S"])
                    scores[s] = scores.get(s, 0) + 1.5
        return scores

    def query_by_temporal(self, temporal_atoms):
        """Find sessions with matching temporal markers."""
        scores = {}
        for t in temporal_atoms:
            atom = _safe_atom(t)
            try:
                results = list(self.prolog.query(f"temporal(S, {atom})"))
                for r in results:
                    s = str(r["S"])
                    scores[s] = scores.get(s, 0) + 1.0
            except Exception:
                pass
        return scores

    def query_by_trigrams(self, query_trigrams):
        """Find sessions sharing trigrams with the query."""
        scores = {}
        for w1, w2, w3 in query_trigrams[:20]:  # Limit for speed
            a1, a2, a3 = _safe_atom(w1), _safe_atom(w2), _safe_atom(w3)
            try:
                results = list(self.prolog.query(f"trigram(S, {a1}, {a2}, {a3})"))
                for r in results:
                    s = str(r["S"])
                    scores[s] = scores.get(s, 0) + 2.5  # Trigram match is strong
            except Exception:
                pass
        return scores

    def query_by_noun_phrases(self, phrases):
        """Find sessions containing specific noun phrases."""
        scores = {}
        for phrase in phrases:
            atom = _safe_atom(phrase)
            try:
                results = list(self.prolog.query(f"noun_phrase(S, {atom})"))
                for r in results:
                    s = str(r["S"])
                    scores[s] = scores.get(s, 0) + 1.5
            except Exception:
                pass
        return scores

    def query_by_relations(self, subj=None, pred=None, obj=None):
        """
        Find sessions with matching relation triples.
        Any of subj/pred/obj can be None (wildcard).
        """
        scores = {}
        s_part = _safe_atom(subj) if subj else "Subj"
        p_part = _safe_atom(pred) if pred else "Pred"
        o_part = _safe_atom(obj) if obj else "Obj"

        try:
            query_str = f"relation(S, {s_part}, {p_part}, {o_part})"
            results = list(self.prolog.query(query_str))
            for r in results:
                s = str(r["S"])
                scores[s] = scores.get(s, 0) + 2.0
        except Exception:
            pass
        return scores

    def comprehensive_query(self, query_facts):
        """
        Run all applicable queries and merge scores.

        Args:
            query_facts: dict from nlp_extract.extract_facts() on the question text

        Returns:
            dict of session_atom -> score
        """
        all_scores = {}

        def merge(new_scores):
            for s, sc in new_scores.items():
                all_scores[s] = all_scores.get(s, 0) + sc

        # Entity matching
        if query_facts.get("entities"):
            ent_names = [e[0] for e in query_facts["entities"]]
            merge(self.query_by_entities(ent_names))

        # Proper noun matching (strong signal for person-name queries)
        if query_facts.get("proper_nouns"):
            merge(self.query_by_proper_nouns(query_facts["proper_nouns"]))

        # Keyword matching
        if query_facts.get("keywords"):
            merge(self.query_by_keywords(query_facts["keywords"]))

        # Quoted phrase matching
        if query_facts.get("quoted_phrases"):
            merge(self.query_by_quoted_phrases(query_facts["quoted_phrases"]))

        # Temporal matching
        if query_facts.get("temporal"):
            merge(self.query_by_temporal(query_facts["temporal"]))

        # Noun phrase matching
        if query_facts.get("noun_phrases"):
            merge(self.query_by_noun_phrases(query_facts["noun_phrases"]))

        # Trigram matching
        if query_facts.get("trigrams"):
            merge(self.query_by_trigrams(query_facts["trigrams"]))

        # Relation matching (try each relation from query)
        for subj, pred, obj in query_facts.get("relations", []):
            merge(self.query_by_relations(subj=subj, pred=pred, obj=obj))
            # Also try partial matches
            merge(self.query_by_relations(subj=subj))
            merge(self.query_by_relations(obj=obj))

        return all_scores

    def query_sessions_with_entity_type(self, entity_type):
        """Find sessions containing entities of a specific type."""
        scores = {}
        type_atom = _safe_atom(entity_type)
        try:
            results = list(self.prolog.query(f"entity(S, _, {type_atom})"))
            for r in results:
                s = str(r["S"])
                scores[s] = scores.get(s, 0) + 0.5
        except Exception:
            pass
        return scores

    @property
    def fact_count(self):
        return self._fact_count

    def get_session_id_from_atom(self, atom):
        """Reverse lookup: atom -> original session_id."""
        for sid, a in self._session_atoms.items():
            if a == atom:
                return sid
        return None