chat/tests/test_memory_search.py

"""Task 23: FTS5 memory retrieval with witness filter and ranking boosts.

Verifies that ``search_memories`` applies recency + significance boosts on top
of the FTS5 BM25 rank so that newer / more significant memories surface above
older / less significant ones for the same match. Existing T8 behaviour
(witness filter, k limit, FTS match, role validation) is exercised again here
to lock the contract.
"""

from __future__ import annotations
import pytest

from chat.db.connection import open_db
from chat.db.migrate import apply_migrations
from chat.eventlog.log import append_event
from chat.eventlog.projector import project
from chat.state.memory import search_memories
import chat.state.memory  # noqa: F401  (registers memory_written handler)
import chat.state.embeddings  # noqa: F401  (registers embedding_indexed handler)


def _seed(db, *, memory_specs):
    """Apply migrations + project a list of memory_written events.

    memory_specs: list of dicts. Required key: ``pov_summary``. Optional keys
    override the defaults below.
    """
    apply_migrations(db)
    with open_db(db) as conn:
        for spec in memory_specs:
            payload = {
                "owner_id": spec.get("owner_id", "bot_a"),
                "chat_id": spec.get("chat_id", "chat_bot_a"),
                "pov_summary": spec["pov_summary"],
                "witness_you": spec.get("witness_you", 1),
                "witness_host": spec.get("witness_host", 1),
                "witness_guest": spec.get("witness_guest", 0),
                "source": "direct",
                "reliability": 1.0,
                "significance": spec.get("significance", 1),
                "pinned": 0,
                "auto_pinned": 0,
            }
            append_event(conn, kind="memory_written", payload=payload)
        project(conn)


def test_search_filters_by_witness_bit(tmp_path):
    db = tmp_path / "t.db"
    _seed(
        db,
        memory_specs=[
            {
                "pov_summary": "BotA mentioned her sister",
                "witness_you": 1,
                "witness_host": 1,
                "witness_guest": 0,
            },
        ],
    )
    with open_db(db) as conn:
        # Witnessed by host -> returned.
        out = search_memories(conn, "bot_a", "host", "sister", k=4)
        assert len(out) == 1
        # NOT witnessed by guest -> filtered out.
        out = search_memories(conn, "bot_a", "guest", "sister", k=4)
        assert out == []


def test_search_higher_significance_ranks_above_lower(tmp_path):
    db = tmp_path / "t.db"
    _seed(
        db,
        memory_specs=[
            # Both match "promise"; the third row carries significance 3 and
            # should outrank the first two, which carry the default of 1.
            {"pov_summary": "small promise"},
            {"pov_summary": "huge promise"},
            {"pov_summary": "tiny promise", "significance": 3},
        ],
    )
    with open_db(db) as conn:
        out = search_memories(conn, "bot_a", "host", "promise", k=3)
        assert len(out) == 3
        assert out[0]["pov_summary"] == "tiny promise"
        assert out[0]["significance"] == 3


def test_search_newer_memory_ranks_above_older_when_same_match(tmp_path):
    db = tmp_path / "t.db"
    _seed(
        db,
        memory_specs=[
            {"pov_summary": "BotA said hello"},
            {"pov_summary": "BotA said hello again"},
        ],
    )
    with open_db(db) as conn:
        out = search_memories(conn, "bot_a", "host", "hello", k=2)
        assert len(out) == 2
        # Newer (higher id, "again") wins on the recency boost when the BM25
        # rank and significance are otherwise comparable.
        assert out[0]["pov_summary"] == "BotA said hello again"


def test_search_respects_k_limit(tmp_path):
    db = tmp_path / "t.db"
    _seed(
        db,
        memory_specs=[
            {"pov_summary": "the cat sat"},
            {"pov_summary": "the cat ran"},
            {"pov_summary": "the cat slept"},
            {"pov_summary": "the cat ate"},
            {"pov_summary": "the cat purred"},
        ],
    )
    with open_db(db) as conn:
        out = search_memories(conn, "bot_a", "host", "cat", k=2)
        assert len(out) == 2


def test_search_invalid_witness_role_raises(tmp_path):
    db = tmp_path / "t.db"
    apply_migrations(db)
    with open_db(db) as conn:
        with pytest.raises(ValueError):
            search_memories(conn, "bot_a", "invalid_role", "anything", k=4)


def test_higher_significance_outranks_equal_rank(tmp_path):
    """T57: significance multiplier biases the SQL ORDER BY.

    Two memories with IDENTICAL FTS-matching text yield (effectively) equal
    BM25 ranks. The significance bias applied in the SQL ORDER BY must
    surface the higher-significance row first.
    """
    db = tmp_path / "t.db"
    _seed(
        db,
        memory_specs=[
            # Identical pov_summary text -> FTS BM25 rank is the same for both.
            {"pov_summary": "she swore an oath", "significance": 0},
            {"pov_summary": "she swore an oath", "significance": 3},
        ],
    )
    with open_db(db) as conn:
        out = search_memories(conn, "bot_a", "host", "oath", k=5)
        assert len(out) == 2
        # Higher significance wins despite tied FTS rank.
        assert out[0]["significance"] == 3
        assert out[1]["significance"] == 0


def test_significance_bias_is_constant_module_level():
    """T57: pin ``SIGNIFICANCE_RANK_BIAS`` as a tunable module-level numeric."""
    from chat.state.memory import SIGNIFICANCE_RANK_BIAS

    assert isinstance(SIGNIFICANCE_RANK_BIAS, (int, float))
    # Must be non-negative -- a negative bias would invert the desired
    # "higher significance ranks higher" semantics.
    assert SIGNIFICANCE_RANK_BIAS >= 0


# ---------------------------------------------------------------------------
# T96 (Phase 4): combined FTS + vector retrieval ranking via reciprocal-rank
# fusion. The fused path activates only when ``query_vector`` is provided —
# the no-vector path (above) is unchanged.
# ---------------------------------------------------------------------------


def _one_hot(dim: int, idx: int) -> list[float]:
    v = [0.0] * dim
    v[idx] = 1.0
    return v


def _seed_memories_with_optional_embeddings(db, *, memory_specs):
    """Like ``_seed`` but also projects ``embedding_indexed`` events for any
    spec carrying a ``vector`` key.

    Memory rows are assigned ids in the order their ``memory_written`` events
    were appended (the ``memories.id`` column is an autoincrementing primary
    key), so we predict ``memory_id = i + 1`` per spec and append both kinds
    of events back-to-back BEFORE projecting. Projecting only once keeps the
    INSERT-based ``memory_written`` handler from duplicating rows.
    """
    apply_migrations(db)
    with open_db(db) as conn:
        # First pass: append every memory_written event in order. The DB
        # assigns autoincrementing ids 1..N matching the order of these
        # events, so we can pair vectors to memories by index.
        for spec in memory_specs:
            payload = {
                "owner_id": spec.get("owner_id", "bot_a"),
                "chat_id": spec.get("chat_id", "chat_bot_a"),
                "pov_summary": spec["pov_summary"],
                "witness_you": spec.get("witness_you", 1),
                "witness_host": spec.get("witness_host", 1),
                "witness_guest": spec.get("witness_guest", 0),
                "source": "direct",
                "reliability": 1.0,
                "significance": spec.get("significance", 1),
                "pinned": 0,
                "auto_pinned": 0,
            }
            append_event(conn, kind="memory_written", payload=payload)
        # Second pass: append embedding_indexed events for any spec that
        # supplied a vector, using the predicted memory id.
        for i, spec in enumerate(memory_specs, start=1):
            if "vector" not in spec:
                continue
            vec = spec["vector"]
            append_event(
                conn,
                kind="embedding_indexed",
                payload={
                    "memory_id": i,
                    "vector": list(vec),
                    "model": "test-model",
                    "dim": len(vec),
                },
            )
        # Single projection — avoids the memory_written handler INSERTing
        # the same row twice on a re-projection.
        project(conn)


def test_search_memories_without_query_vector_uses_fts_only(tmp_path):
    """Regression: omitting ``query_vector`` keeps the existing FTS-only path.

    Identical seed to ``test_search_higher_significance_ranks_above_lower``
    but pinned to the no-vector code path explicitly (no kwarg passed).
    """
    db = tmp_path / "t.db"
    _seed(
        db,
        memory_specs=[
            {"pov_summary": "small promise"},
            {"pov_summary": "huge promise"},
            {"pov_summary": "tiny promise", "significance": 3},
        ],
    )
    with open_db(db) as conn:
        out = search_memories(conn, "bot_a", "host", "promise", k=3)
        assert len(out) == 3
        # The composite re-rank surfaces the high-significance row first.
        assert out[0]["pov_summary"] == "tiny promise"
        # Sanity: the row shape still carries ``fts_rank`` + ``composite_score``
        # like the FTS-only path always has.
        assert "fts_rank" in out[0]
        assert "composite_score" in out[0]


def test_search_memories_with_query_vector_includes_vector_hits(tmp_path):
    """RRF fuses FTS hits with vector hits — both kinds surface in the result.

    Memory 1 only matches FTS (keyword "rabbit", embedding far from query).
    Memory 2 only matches the vector (embedding identical to query, no
    keyword overlap). Memories 3-5 are unrelated. The fused top-K must
    contain BOTH memory 1 and memory 2.
    """
    db = tmp_path / "t.db"
    dim = 8
    # Query vector = one-hot at index 0. Memory 2 mirrors it exactly. The
    # FTS-only memory (memory 1) has NO embedding so it cannot leak into
    # the vector ranking; the filler memories (3-5) likewise have no
    # embeddings, so the vector ranking returns memory 2 alone.
    query_vec = _one_hot(dim, 0)
    _seed_memories_with_optional_embeddings(
        db,
        memory_specs=[
            # Memory 1: FTS-only match. No embedding indexed.
            {"pov_summary": "rabbit hopped over the fence"},
            # Memory 2: vector-only match. No keyword overlap with "rabbit".
            {
                "pov_summary": "completely unrelated narrative line",
                "vector": _one_hot(dim, 0),
            },
            # Memories 3-5: filler, irrelevant to both channels.
            {"pov_summary": "lighthouse keeper polished the lens"},
            {"pov_summary": "they discussed cartography for hours"},
            {"pov_summary": "she taught him semaphore signals"},
        ],
    )
    with open_db(db) as conn:
        out = search_memories(
            conn,
            "bot_a",
            "host",
            "rabbit",
            k=4,
            query_vector=query_vec,
        )
        summaries = [r["pov_summary"] for r in out]
        # FTS-only candidate (memory 1) made it through.
        assert "rabbit hopped over the fence" in summaries
        # Vector-only candidate (memory 2) also made it through despite
        # having no keyword overlap with the query string.
        assert "completely unrelated narrative line" in summaries


def test_search_memories_fusion_significance_bias_still_applies(tmp_path):
    """With two RRF-tied candidates, the higher-significance one ranks first.

    Two memories share the keyword "promise" AND share an identical
    embedding to the query — so their FTS rank and vector rank are both
    ties. RRF gives them the same fusion score. The Python-side
    significance + recency boost must break the tie in favour of the
    higher-significance memory.
    """
    db = tmp_path / "t.db"
    dim = 4
    shared_vec = _one_hot(dim, 0)
    _seed_memories_with_optional_embeddings(
        db,
        memory_specs=[
            {
                "pov_summary": "she made a promise",
                "significance": 0,
                "vector": list(shared_vec),
            },
            {
                "pov_summary": "she made a promise",
                "significance": 3,
                "vector": list(shared_vec),
            },
        ],
    )
    with open_db(db) as conn:
        out = search_memories(
            conn,
            "bot_a",
            "host",
            "promise",
            k=2,
            query_vector=list(shared_vec),
        )
        assert len(out) == 2
        # Higher significance breaks the RRF tie.
        assert out[0]["significance"] == 3
        assert out[1]["significance"] == 0


def test_search_memories_fusion_handles_empty_vector_results(tmp_path):
    """Vector path returning [] (no embeddings indexed) must not break FTS.

    No ``embedding_indexed`` events are projected, so ``vector_search``
    returns an empty list. The function should still return the FTS hits
    as if ``query_vector`` had not been supplied.
    """
    db = tmp_path / "t.db"
    _seed(
        db,
        memory_specs=[
            {"pov_summary": "the vault held an old promise"},
            {"pov_summary": "another promise was kept that night"},
        ],
    )
    with open_db(db) as conn:
        out = search_memories(
            conn,
            "bot_a",
            "host",
            "promise",
            k=4,
            query_vector=[0.0] * 384,  # No embeddings exist for this owner.
        )
        # Both FTS hits still come back — no error from the empty vector path.
        assert len(out) == 2
        summaries = {r["pov_summary"] for r in out}
        assert summaries == {
            "the vault held an old promise",
            "another promise was kept that night",
        }