falsify/seed.py

"""
FALSIFY demo corpus — the "Company X recall" investigation.

This module builds the Session-1 belief graph that the demo revises in Session 2.
It is intentionally small and hand-authored so the belief-revision mechanics are
legible on screen in under 30 seconds, and so the cascade result is deterministic
(the winning demo must not depend on a flaky LLM).

The investigation
-----------------
Question Q:  "Did Company X know about the defect before the recall?"

Competing hypotheses:
    A — "X knew via the QA report dated March 2021"      (supported by E_qa)
    B — "X knew via a supplier email in January 2021"    (supported by E_email)
    C — "X did not know before the recall"               (unsupported)

Evidence:
    E_qa    — the March-2021 QA report       (supports A)
    E_email — the January-2021 supplier email (supports B)

Conclusion:
    K — "Company X knew about the defect by March 2021"
        depends_on E_qa  (critical=True)   <-- the propagation rail

The Session-2 fact (NEW_FACT) is a forensic finding that the March QA report was
back-dated. It contradicts E_qa. Refuting E_qa must cascade:

    E_qa -> refuted
      -> K (depends_on E_qa, critical) -> invalidated
      -> A (only supporter E_qa now dead) -> superseded ; B promoted (new frontier)
      -> K orphaned (no alive consumer) -> forgotten (hard-deleted, graph + vector)
      -> E_qa kept as refuted provenance (it is NEW_FACT's supersedes anchor)

The scoreboard then shows FALSIFY answering via B (Jan 2021) while a plain vector
(RAG) baseline still cites the refuted March QA report.
"""

from __future__ import annotations

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

# NOTE: importing falsify.models runs falsify/__init__, which sets the Cognee env
# defaults (access-control off, session cache on) before Cognee is imported.
from falsify import graph_ops
from falsify.edges import DEPENDS_ON, SUPPORTS
from falsify.models import (
    Conclusion,
    Evidence,
    Hypothesis,
    InvestigationQuestion,
)

logger = logging.getLogger("falsify.seed")

# The research question this whole graph hangs off of.
QUESTION_TEXT = "Did Company X know about the defect before the recall?"

# The Session-2 fact that triggers belief revision. It contradicts E_qa.
NEW_FACT = (
    "A forensic audit found that the March 2021 QA report was back-dated: it was "
    "actually created shortly after the product recall, not before it."
)

# The SECOND contradiction, used only by the diamond scenario. It contradicts
# E_email. Dropped after NEW_FACT to demonstrate iterative revision: once BOTH
# legs of a multi-source conclusion are refuted, the conclusion finally collapses.
NEW_FACT_2 = (
    "A forensic analysis of email headers shows the January 2021 supplier email was "
    "fabricated: the sender domain was not registered until April 2021, and the DKIM "
    "signature is invalid."
)

# Human-readable stable keys -> used by --demo mode to pin the refutation target
# deterministically (so the cascade runs on real graph APIs even if the LLM judge
# is unavailable). These map to the ``key`` field on the seeded nodes below.
REFUTED_EVIDENCE_KEY = "E_qa"
REFUTED_EVIDENCE_KEY_2 = "E_email"  # second-phase target for the diamond scenario


@dataclass
class SeededGraph:
    """Handles to the nodes created by :func:`build_investigation`.

    ``ids`` maps a stable human key (e.g. ``"E_qa"``, ``"A"``, ``"K"``) to the
    string node id in the graph, so the demo/tests can reference specific nodes
    without re-querying. ``labels`` maps the same keys to display strings.
    """

    question_id: str
    ids: Dict[str, str] = field(default_factory=dict)
    labels: Dict[str, str] = field(default_factory=dict)

    @property
    def refuted_target_id(self) -> str:
        """Node id of the evidence the Session-2 fact contradicts (E_qa)."""
        return self.ids[REFUTED_EVIDENCE_KEY]


async def build_investigation() -> SeededGraph:
    """Create and persist the Session-1 belief graph. Returns a :class:`SeededGraph`.

    The caller is responsible for starting from a clean state (e.g. via
    ``cognee.forget(everything=True)`` and ``cognee.low_level.setup()``). This
    function only builds; it does not prune.

    Nodes are persisted with Cognee's ``add_data_points`` (writing both the graph
    node and the vector row for each ``Embeddable`` field). Typed edges with
    properties are then added explicitly through the graph engine so we control the
    exact relationship names and edge properties (``critical`` / ``weight``).
    """
    # Import here so the module import stays cheap and env defaults are already set.
    from cognee.tasks.storage import add_data_points

    # ------------------------------------------------------------------ nodes
    question = InvestigationQuestion(question=QUESTION_TEXT, source_id="investigation")

    hyp_a = Hypothesis(
        statement="Company X knew via the QA report dated March 2021.",
        question_id=str(question.id),
        prior=0.5,
        confidence=0.6,
        source_id="analyst",
    )
    hyp_b = Hypothesis(
        statement="Company X knew via a supplier email in January 2021.",
        question_id=str(question.id),
        prior=0.5,
        confidence=0.55,
        source_id="analyst",
    )
    hyp_c = Hypothesis(
        statement="Company X did not know about the defect before the recall.",
        question_id=str(question.id),
        prior=0.5,
        confidence=0.4,
        source_id="analyst",
    )

    ev_qa = Evidence(
        claim="A QA report dated March 2021 documented the defect internally.",
        source_id="qa_report_2021_03",
        quote="Internal QA report, dated 2021-03-15, flags the defect.",
        stance="supports",
        asserted_at="2021-03-15",
        confidence=0.8,
    )
    ev_email = Evidence(
        claim="A supplier email in January 2021 warned Company X about the defect.",
        source_id="supplier_email_2021_01",
        quote="Supplier email, 2021-01-20: 'we have observed the defect in test units.'",
        stance="supports",
        asserted_at="2021-01-20",
        confidence=0.7,
    )

    conclusion_k = Conclusion(
        statement="Company X knew about the defect by March 2021.",
        confidence=0.8,
        depends_on_ids=[str(ev_qa.id)],
        source_id="analyst",
    )

    nodes: List = [question, hyp_a, hyp_b, hyp_c, ev_qa, ev_email, conclusion_k]

    logger.info("Persisting %d belief nodes via add_data_points", len(nodes))
    await add_data_points(nodes)

    # ------------------------------------------------------------------ edges
    # Evidence -> Hypothesis (supports, weighted)
    await graph_ops.add_edge(str(ev_qa.id), str(hyp_a.id), SUPPORTS, {"weight": 0.8})
    await graph_ops.add_edge(str(ev_email.id), str(hyp_b.id), SUPPORTS, {"weight": 0.7})

    # Conclusion -> Evidence (depends_on, critical) — THE propagation rail
    await graph_ops.add_edge(
        str(conclusion_k.id), str(ev_qa.id), DEPENDS_ON, {"critical": True}
    )

    # Hypothesis -> Question (answers) — keeps the graph connected for visualization
    for hyp in (hyp_a, hyp_b, hyp_c):
        await graph_ops.add_edge(str(hyp.id), str(question.id), "answers", {})

    seeded = SeededGraph(
        question_id=str(question.id),
        ids={
            "Q": str(question.id),
            "A": str(hyp_a.id),
            "B": str(hyp_b.id),
            "C": str(hyp_c.id),
            "E_qa": str(ev_qa.id),
            "E_email": str(ev_email.id),
            "K": str(conclusion_k.id),
        },
        labels={
            "Q": QUESTION_TEXT,
            "A": hyp_a.statement,
            "B": hyp_b.statement,
            "C": hyp_c.statement,
            "E_qa": ev_qa.claim,
            "E_email": ev_email.claim,
            "K": conclusion_k.statement,
        },
    )
    logger.info("Seeded investigation graph: %s", seeded.ids)
    return seeded


async def build_diamond_investigation() -> SeededGraph:
    """Build the Session-1 graph WITH a diamond dependency (Conclusion K2).

    This is the same investigation as :func:`build_investigation`, plus one extra
    Conclusion K2 that *critically depends on BOTH* E_qa and E_email. It exists to
    demonstrate that FALSIFY's propagation is a grounded least-fixpoint, not a naive
    cascade: a conclusion with two critical supports survives losing one of them, and
    only collapses when the LAST support dies.

    Two-phase story the demo drives on top of this graph:

        Phase 1 — refute E_qa (NEW_FACT):
            E_qa -> refuted ; K (single dep) -> invalidated -> forgotten
            K2 -> STILL ALIVE (E_email keeps it grounded) ; A superseded, B promoted
            E_qa -> retained (K2 alive still depends on it, and it anchors NEW_FACT)

        Phase 2 — refute E_email (NEW_FACT_2):
            E_email -> refuted ; K2 (last dep now dead) -> invalidated -> forgotten
            B -> superseded ; only C ("did not know") may remain

    The base nodes are re-created here (rather than shared with build_investigation)
    so the proven single-contradiction demo path stays untouched.
    """
    from cognee.tasks.storage import add_data_points

    # ------------------------------------------------------------------ nodes
    question = InvestigationQuestion(question=QUESTION_TEXT, source_id="investigation")

    hyp_a = Hypothesis(
        statement="Company X knew via the QA report dated March 2021.",
        question_id=str(question.id), prior=0.5, confidence=0.6, source_id="analyst",
    )
    hyp_b = Hypothesis(
        statement="Company X knew via a supplier email in January 2021.",
        question_id=str(question.id), prior=0.5, confidence=0.55, source_id="analyst",
    )
    hyp_c = Hypothesis(
        statement="Company X did not know about the defect before the recall.",
        question_id=str(question.id), prior=0.5, confidence=0.4, source_id="analyst",
    )

    ev_qa = Evidence(
        claim="A QA report dated March 2021 documented the defect internally.",
        source_id="qa_report_2021_03",
        quote="Internal QA report, dated 2021-03-15, flags the defect.",
        stance="supports", asserted_at="2021-03-15", confidence=0.8,
    )
    ev_email = Evidence(
        claim="A supplier email in January 2021 warned Company X about the defect.",
        source_id="supplier_email_2021_01",
        quote="Supplier email, 2021-01-20: 'we have observed the defect in test units.'",
        stance="supports", asserted_at="2021-01-20", confidence=0.7,
    )

    conclusion_k = Conclusion(
        statement="Company X knew about the defect by March 2021.",
        confidence=0.8, depends_on_ids=[str(ev_qa.id)], source_id="analyst",
    )
    # THE DIAMOND: K2 rests on two independent critical supports.
    conclusion_k2 = Conclusion(
        statement="Multiple independent sources confirm Company X had pre-recall "
                  "knowledge of the defect.",
        confidence=0.85,
        depends_on_ids=[str(ev_qa.id), str(ev_email.id)],
        source_id="analyst",
    )

    nodes: List = [question, hyp_a, hyp_b, hyp_c, ev_qa, ev_email, conclusion_k, conclusion_k2]

    logger.info("Persisting %d belief nodes (diamond) via add_data_points", len(nodes))
    await add_data_points(nodes)

    # ------------------------------------------------------------------ edges
    await graph_ops.add_edge(str(ev_qa.id), str(hyp_a.id), SUPPORTS, {"weight": 0.8})
    await graph_ops.add_edge(str(ev_email.id), str(hyp_b.id), SUPPORTS, {"weight": 0.7})

    # K depends on E_qa only (single-leg — dies in phase 1).
    await graph_ops.add_edge(
        str(conclusion_k.id), str(ev_qa.id), DEPENDS_ON, {"critical": True}
    )
    # K2's two critical legs — the diamond. Survives phase 1, collapses in phase 2.
    await graph_ops.add_edge(
        str(conclusion_k2.id), str(ev_qa.id), DEPENDS_ON, {"critical": True}
    )
    await graph_ops.add_edge(
        str(conclusion_k2.id), str(ev_email.id), DEPENDS_ON, {"critical": True}
    )

    for hyp in (hyp_a, hyp_b, hyp_c):
        await graph_ops.add_edge(str(hyp.id), str(question.id), "answers", {})

    seeded = SeededGraph(
        question_id=str(question.id),
        ids={
            "Q": str(question.id),
            "A": str(hyp_a.id), "B": str(hyp_b.id), "C": str(hyp_c.id),
            "E_qa": str(ev_qa.id), "E_email": str(ev_email.id),
            "K": str(conclusion_k.id), "K2": str(conclusion_k2.id),
        },
        labels={
            "Q": QUESTION_TEXT,
            "A": hyp_a.statement, "B": hyp_b.statement, "C": hyp_c.statement,
            "E_qa": ev_qa.claim, "E_email": ev_email.claim,
            "K": conclusion_k.statement, "K2": conclusion_k2.statement,
        },
    )
    logger.info("Seeded diamond investigation graph: %s", seeded.ids)
    return seeded