src/facts.py

"""
FactExtractor + KnowledgeGraph + GraphQuery.

FactExtractor: lightweight regex NER (persons, orgs, locations, dates,
numbers, emails, urls, keywords) that parses web snippets / image captions
into entity lists. No model dependency.

KnowledgeGraph: in-memory directed graph (entity -> (relation -> entities))
with dedup, so facts can be queried across turns / across search results.

GraphQuery: given a question, extract entities and pull connected facts
so the verifier can cross-examine generations against grounded facts.
"""

import json
import re
from dataclasses import dataclass, field
from typing import Dict, List, Set, Tuple


@dataclass
class Fact:
    subject: str
    relation: str
    object: str
    source: str = ""

    def to_text(self) -> str:
        return f"{self.subject} {self.relation} {self.object}"


_ORG_WORDS = {"inc", "corp", "company", "ltd", "llc", "gmbh", "co", "university",
              "institute", "labs", "foundation", "group", "association", "bank", "airlines"}
_PERSON_RE = re.compile(r"\b([A-Z][a-z]+ [A-Z][a-z]+)\b")
_PROPER_NOUN_RE = re.compile(r"\b([A-Z][A-Za-z]+)\b")
_SENT_START_RE = re.compile(r"(^|[.!?]\s+|\n|\(|\"|')")
_STOP = {"When", "What", "Where", "How", "Who", "Why", "Which", "The", "A", "An",
         "I", "You", "It", "This", "That", "There", "His", "Her", "Its", "We", "They",
         "Do", "Does", "Did", "Is", "Are", "Was", "Were", "Can", "Could", "Would", "Should"}
_COMMON_CAPS = {"In", "On", "At", "For", "From", "With", "By", "As", "Of", "To", "And",
                "Or", "But", "So", "Also", "However", "Meanwhile", "Currently", "Recently",
                "Today", "Yesterday", "Now", "New", "More", "Most", "Many", "Some", "About",
                "After", "Before", "During", "Since", "While", "Within", "Under", "Over",
                "According", "Despite", "Using", "Via", "Per", "Among", "Between", "Both",
                "Each", "Every", "First", "Second", "Third", "Next", "Last", "Other", "Only",
                "Our", "Their", "Your", "His", "Her", "Our", "These", "Those", "Another"}
_ORG_RE = re.compile(r"\b([A-Z][A-Za-z0-9&\.\- ]{1,40}(?:Inc|Corp|Ltd|LLC|GmbH|Company|University|Institute|Labs|Foundation|Group|Association|Bank|Airlines))\b")
_LOC_RE = re.compile(r"\b(?:in|at|from|near)\s+([A-Z][a-zA-Z]+)\b")
_DATE_RE = re.compile(r"\b((?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]* \d{1,2},? \d{4}|\d{4}|\d{1,2}/\d{1,2}/\d{2,4})\b")
_NUM_RE = re.compile(r"\b\d+(?:\.\d+)?(?:%| million| billion| trillion)?\b")
_EMAIL_RE = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}")
_URL_RE = re.compile(r"https?://\S+")
_REL_RE = re.compile(r"\b(is a|is|was founded in|founded in|located in|based in|part of|known for|released in|launched in|created by|developed by|owned by|has a population of|is the capital of|won|beat|discovered|acquired|merged with|reported|announced|raised|produces|manufactures|sells|provides|founded by|led by|headquartered in)\b")

_KEYWORD_DIMS = {
    "code": r"\b(python|function|debug|api|compile|syntax|variable|loop)\b",
    "math": r"\b(equation|derivative|integral|theorem|solve|calculate)\b",
    "data": r"\b(data|dataset|statistics|chart|correlation)\b",
    "time": r"\b(today|now|latest|current|news|price|weather)\b",
    "place": r"\b(where|location|city|country|capital)\b",
}


class FactExtractor:
    def extract(self, text: str) -> Dict[str, Set[str]]:
        """Return {type: {entity,...}}."""
        out: Dict[str, Set[str]] = {}
        for name, pat in _KEYWORD_DIMS.items():
            hits = set(re.findall(pat, text.lower()))
            if hits:
                out[name] = hits
        for label, pat in [
            ("PERSON", _PERSON_RE), ("ORG", _ORG_RE), ("LOCATION", _LOC_RE),
            ("DATE", _DATE_RE), ("NUMBER", _NUM_RE), ("EMAIL", _EMAIL_RE), ("URL", _URL_RE),
        ]:
            hits = set(pat.findall(text))
            # NUMBER regex can over-match; keep only 1-2 digit numbers + units
            if label == "NUMBER":
                hits = set(m for m in hits if len(m) <= 4 or not m.replace(".", "").isdigit())
            if hits:
                out[label] = hits
        # proper-noun candidates (brands like "SpaceX"/"Google" with no ORG
        # suffix) -> ORG; two-capital names are already tagged PERSON above.
        # Sentence-initial capitalized COMMON words (The/In/Also...) are
        # dropped; real proper nouns that start a sentence (Google/SpaceX)
        # are kept.
        filtered = set()
        for m in _PROPER_NOUN_RE.finditer(text):
            word = m.group(0)
            if word in _STOP or word in _COMMON_CAPS:
                continue
            before = text[max(0, m.start() - 6):m.start()]
            at_sentence_start = bool(re.search(r"(^|[.!?]\s|\n|[:;]\s|\(|\"|')", before))
            if at_sentence_start:
                # keep sentence-start word only if it is not a common English cap
                if word in _COMMON_CAPS or word in _STOP:
                    continue
            filtered.add(word)
        if filtered:
            existing = out.get("ORG", set())
            existing |= filtered
            out["ORG"] = existing
        return out

    def triples(self, text: str) -> List[Fact]:
        """Naive subject-relation-object extraction on sentences."""
        facts = []
        for sent in re.split(r"(?<=[.!?])\s+", text):
            m = _REL_RE.search(sent)
            if not m:
                continue
            rel = m.group(0)
            before = sent[: m.start()].strip()
            after = sent[m.end():].strip()
            subj = before.split(",")[0].strip().strip(".:;\"'[]()")
            obj = re.sub(r"\s+", " ", after)
            obj = re.split(r"[.;]", obj)[0].strip().strip("\"()")
            # drop question-phrase subjects and overly long spans (noise)
            if (subj.lower().startswith(("what", "who", "how", "why", "where", "when",
                                         "which", "is", "are", "the", "also", "this"))) or len(subj) > 40:
                continue
            # skip relations that are really "X is Y" generic copulas with no entity
            if rel == "is" and not re.search(r"[A-Z][a-z]+", subj):
                continue
            if subj and obj and len(obj) < 120:
                facts.append(Fact(subject=subj, relation=rel, object=obj))
        return facts


class KnowledgeGraph:
    def __init__(self):
        # entity -> {relation: {entity,...}}  (subject -> object edges)
        self.nodes: Dict[str, Dict[str, Set[str]]] = {}
        # object -> {subject,...}  (reverse index so queries can match objects)
        self.rev: Dict[str, Set[str]] = {}
        self.sources: Dict[str, Set[str]] = {}

    def add_fact(self, fact: Fact):
        subj = fact.subject.lower()
        obj = fact.object.lower()
        self.nodes.setdefault(subj, {}).setdefault(fact.relation.lower(), set()).add(obj)
        self.rev.setdefault(obj, set()).add(subj)
        self.sources.setdefault(subj, set()).add(fact.source)

    def add_many(self, facts: List[Fact]):
        for f in facts:
            self.add_fact(f)

    def neighbors(self, entity: str, depth: int = 1) -> Set[str]:
        seen: Set[str] = {entity.lower()}
        frontier = [entity.lower()]
        for _ in range(depth):
            nxt = []
            for ent in frontier:
                for rel, targets in self.nodes.get(ent, {}).items():
                    for t in targets:
                        if t not in seen:
                            seen.add(t)
                            nxt.append(t)
            frontier = nxt
        return seen - {entity.lower()}

    def to_text(self, entities: List[str], depth: int = 1, max_facts: int = 20) -> str:
        out = []
        seen_facts = set()
        for ent in entities:
            ent = ent.lower()
            # entity as subject
            for rel, targets in self.nodes.get(ent, {}).items():
                for t in list(targets)[:max_facts]:
                    key = (ent, rel, t)
                    if key not in seen_facts:
                        seen_facts.add(key)
                        out.append(f"{ent} {rel} {t}")
            # entity as object (reverse edges)
            for subj in self.rev.get(ent, set()):
                for rel, targets in self.nodes.get(subj, {}).items():
                    for t in list(targets)[:3]:
                        key = (subj, rel, t)
                        if key not in seen_facts:
                            seen_facts.add(key)
                            out.append(f"{subj} {rel} {t}")
            for n in self.neighbors(ent, depth):
                for rel, targets in self.nodes.get(n, {}).items():
                    for t in list(targets)[:3]:
                        key = (n, rel, t)
                        if key not in seen_facts:
                            seen_facts.add(key)
                            out.append(f"{n} {rel} {t}")
            if len(out) >= max_facts:
                break
        return " | ".join(out[:max_facts])

    def save(self, path: str):
        with open(path, "w", encoding="utf-8") as f:
            json.dump({
                "nodes": {k: {r: sorted(v) for r, v in rels.items()}
                          for k, rels in self.nodes.items()},
                "rev": {k: sorted(v) for k, v in self.rev.items()},
                "sources": {k: sorted(v) for k, v in self.sources.items()},
            }, f)

    def load(self, path: str):
        data = json.load(open(path, encoding="utf-8"))
        self.nodes = {k: {r: set(v) for r, v in rels.items()}
                      for k, rels in data.get("nodes", {}).items()}
        self.rev = {k: set(v) for k, v in data.get("rev", {}).items()}
        self.sources = {k: set(v) for k, v in data.get("sources", {}).items()}


class GraphQuery:
    def __init__(self, extractor: Optional[FactExtractor] = None):
        self.extractor = extractor or FactExtractor()

    def facts_for_question(self, graph: KnowledgeGraph, question: str, depth: int = 1) -> str:
        ents = self.extractor.extract(question)
        query_entities = set()
        for kind in ("PERSON", "ORG", "LOCATION", "DATE", "NUMBER", "code", "math", "data"):
            query_entities |= ents.get(kind, set())
        result = graph.to_text(list(query_entities), depth=depth)
        if result:
            return result
        # fallback: content-word overlap against every stored fact so a plain
        # question like "capital of france" still surfaces relevant edges
        query_tokens = {w for w in re.findall(r"[a-z0-9]{3,}", question.lower())}
        hits = []
        for subj, rels in graph.nodes.items():
            for rel, targets in rels.items():
                for t in targets:
                    fact_tokens = set(re.findall(r"[a-z0-9]{3,}", f"{subj} {t}".lower()))
                    if query_tokens & fact_tokens:
                        hits.append(f"{subj} {rel} {t}")
        return " | ".join(hits[:10])