interactive-story/backend/tests/test_v11_b2_memory_ranking.py

"""v1.1 WP-B.2 (B2.1): what memory retrieval searches for, and how it scores.

B.1 found the planting-era memory created and retained but ranked out of
`memory_top_k`, because the query was three turns of narration with the player's
question at the end. The query is now the player's input plus a short scene
context, and the score adds one transparent lexical term over the input.

These tests pin the pieces. The end-to-end fixture tests (crowded bank,
context-dependent question, negative controls) are in
`test_v11_b1_memory_diagnostic.py`, beside the diagnostic they use.

    python -m pytest tests/test_v11_b2_memory_ranking.py -v
"""

import asyncio
import math

import pytest
from sqlalchemy import event

from app import memorybank, models, tree
from app.context import builder
from app.database import Base, SessionLocal, engine

# ------------------------------------------------------------------ lexical


def test_terms_are_folded_but_not_stemmed():
    terms = memorybank.lexical_terms("The tavern's teapots, the glass and the SUNDIAL")
    assert {"tavern", "teapot", "glass", "sundial"} <= terms
    assert "the" not in terms  # the knowledge path's stop list
    assert "glas" not in terms  # a word ending in "ss" is not a plural


def test_a_word_every_candidate_holds_weighs_nothing():
    scores = memorybank.lexical_scores(
        frozenset({"travellers"}), {1: frozenset({"travellers", "road"}), 2: frozenset({"travellers"})})
    assert scores == {1: 0.0, 2: 0.0}


def test_a_rarer_word_weighs_more_than_a_common_one():
    scores = memorybank.lexical_scores(
        frozenset({"sundial", "road"}),
        {1: frozenset({"sundial"}), 2: frozenset({"road"}), 3: frozenset({"road"}),
         4: frozenset({"gate"})})
    assert scores[1] > scores[2] == scores[3] > scores[4] == 0.0


def test_a_single_rare_word_of_a_longer_question_is_only_its_share():
    """The share is over the whole question, so one incidental word match
    cannot score like a memory that answers it."""
    question = frozenset({"where", "amber", "sundial", "fish"})
    scores = memorybank.lexical_scores(
        question, {1: frozenset({"fish"}), 2: frozenset({"amber", "sundial"}), 3: frozenset({"road"})})
    assert 0.0 < scores[1] < scores[2] <= 1.0
    assert scores[1] < 0.5


def test_scores_are_bounded_and_empty_inputs_score_zero():
    candidates = {1: frozenset({"a1", "b2"}), 2: frozenset({"a1"})}
    assert all(0.0 <= v <= 1.0 for v in memorybank.lexical_scores(frozenset({"a1", "b2"}), candidates).values())
    assert memorybank.lexical_scores(frozenset(), candidates) == {1: 0.0, 2: 0.0}
    assert memorybank.lexical_scores(frozenset({"a1"}), {}) == {}


# ------------------------------------------------------------------ scoring


def _unit(angle):
    return [math.cos(angle), math.sin(angle)]


def test_ties_are_broken_by_id_not_by_row_order():
    held = {7: [1.0, 0.0], 3: [1.0, 0.0], 5: [1.0, 0.0]}
    rows = memorybank.score_candidates([7, 3, 5], held, {}, [1.0, 0.0], None, [])
    assert [row[1] for row in rows] == [3, 5, 7]


def test_the_semantic_score_mixes_input_and_context_by_the_fixed_weight():
    held = {1: [1.0, 0.0]}
    [(final, _, semantic, lexical)] = memorybank.score_candidates(
        [1], held, {}, [1.0, 0.0], [0.0, 1.0], [])
    assert semantic == pytest.approx(memorybank.INPUT_WEIGHT)
    assert final == semantic and lexical == 0.0
    # Either part alone is used as it is.
    [(_, _, only_context, _)] = memorybank.score_candidates([1], held, {}, None, [0.0, 1.0], [])
    assert only_context == pytest.approx(0.0)


@pytest.mark.parametrize("margin, relevant_first", [(0.01, True), (-0.01, False)])
def test_a_lexical_match_moves_a_memory_at_most_the_lexical_weight(margin, relevant_first):
    """The bound that keeps rarity from overruling meaning: a memory more than
    `LEXICAL_WEIGHT` behind semantically cannot pass one ahead of it, however
    rare the word it shares."""
    decoy_cos = 1.0 - memorybank.LEXICAL_WEIGHT - margin
    held = {1: [1.0, 0.0], 2: _unit(math.acos(decoy_cos))}
    terms = {1: frozenset(), 2: frozenset({"zeppelin"})}
    rows = memorybank.score_candidates([1, 2], held, terms, [1.0, 0.0], None, ["zeppelin"])
    order = [row[1] for row in rows]
    assert (order[0] == 1) is relevant_first


# ---------------------------------------------------- the query, on real rows


@pytest.fixture()
def db():
    Base.metadata.create_all(bind=engine)
    memorybank._vector_cache.clear()
    memorybank._terms_cache.clear()
    session = SessionLocal()
    try:
        yield session
    finally:
        session.close()
        memorybank._vector_cache.clear()
        memorybank._terms_cache.clear()
        Base.metadata.drop_all(bind=engine)


@pytest.fixture(autouse=True)
def restore_embedding_provider():
    real = memorybank.embedding_provider
    try:
        yield
    finally:
        memorybank.embedding_provider = real


@pytest.fixture()
def settings(db):
    user = models.User(is_guest=False, email="b2-rank@example.com")
    db.add(user)
    db.flush()
    row = models.Settings(user_id=user.id, model="m", embedding_model="stub-embed",
                          memory_top_k=2, memory_bank_capacity=80)
    db.add(row)
    db.commit()
    return row


SCENE_STATE = {
    "entities": {"mara": {"type": "character", "name": "Mara"},
                 "tavern": {"type": "location", "name": "The Crooked Lantern"}},
    "scene": {"summary": "Closing time", "location": "tavern", "present": ["mara"]},
}


def make_adventure(db, settings, texts, state=None):
    """`texts` is `[(type, text)]`, oldest first, each placed on the tree."""
    adventure = models.Adventure(user_id=settings.user_id, title="Rank", script_state={},
                                 memory_bank_enabled=True, narrative_state=state or {})
    db.add(adventure)
    db.flush()
    placed = []
    for kind, text in texts:
        action = models.Action(adventure_id=adventure.id, type=kind, text=text)
        tree.place_action(db, adventure, action)
        db.add(action)
        db.flush()
        placed.append(action)
    db.commit()
    return adventure, placed


def test_the_query_is_the_players_input_and_the_scene(db, settings):
    adventure, _ = make_adventure(db, settings, [
        ("start", "Rain over the harbour."),
        ("ai", "Mara wipes down the counter and glances up at the shelf."),
        ("do", "> You ask Mara about the brass dial."),
    ], state=SCENE_STATE)
    query = memorybank.retrieval_query(adventure)
    assert query["input"] == "> You ask Mara about the brass dial."
    assert "The Crooked Lantern" in query["context"] and "Mara" in query["context"]
    assert "glances up at the shelf" in query["context"]
    assert "brass" in query["input_terms"] and "dial" in query["input_terms"]


def test_a_continue_turn_has_no_input_and_searches_by_the_scene(db, settings):
    adventure, _ = make_adventure(db, settings, [
        ("do", "> You sit down."),
        ("ai", "The fire burns low in the grate."),
    ])
    query = memorybank.retrieval_query(adventure)
    assert query["input"] == "" and query["input_terms"] == []
    assert "fire burns low" in query["context"]


def test_a_retry_searches_with_the_input_it_is_retrying(db, settings):
    adventure, placed = make_adventure(db, settings, [
        ("ai", "The market is quiet."),
        ("do", "> You ask about the sundial."),
        ("ai", "A discarded attempt about lanterns."),
    ])
    query = memorybank.retrieval_query(adventure, exclude_action_id=placed[-1].id)
    assert query["input"] == "> You ask about the sundial."
    assert "lanterns" not in query["context"]
    assert "market is quiet" in query["context"]


def test_the_query_is_bounded_however_long_the_story(db, settings):
    long = "The travellers walked the long grey road north past the salt market. " * 400
    adventure, _ = make_adventure(db, settings, [
        ("ai", long), ("story", long)], state=SCENE_STATE)
    query = memorybank.retrieval_query(adventure)
    assert builder.count_tokens(query["input"]) <= memorybank.QUERY_INPUT_TOKENS
    assert builder.count_tokens(query["context"]) <= (
        memorybank.QUERY_SCENE_TOKENS + memorybank.QUERY_NARRATION_TOKENS + 2)


# ------------------------------------------------ retrieval, end to end


class SameVector:
    """Every text embeds the same, so only the lexical term separates memories."""

    async def embed(self, texts):
        return [[1.0, 0.0, 0.0] for _ in texts]


def add_memory(db, adventure, text, **kwargs):
    memory = models.Memory(adventure_id=adventure.id, text=text, **kwargs)
    db.add(memory)
    db.flush()
    memorybank.set_vector(memory, [1.0, 0.0, 0.0])
    db.commit()
    return memory


def retrieve(adventure, settings, **kwargs):
    memorybank.embedding_provider = lambda s: SameVector()
    return asyncio.run(memorybank.retrieve_memories(adventure, settings, **kwargs))


@pytest.fixture()
def played(db, settings):
    adventure, _ = make_adventure(db, settings, [
        ("ai", "The tavern is warm."),
        ("do", "> You ask Mara where the amber sundial went."),
    ], state=SCENE_STATE)
    bank = {
        "road": add_memory(db, adventure, "Aldric walked the north road."),
        "sundial": add_memory(db, adventure, "Mara hid the amber sundial in the teapot."),
        "gate": add_memory(db, adventure, "The gate guard asked for a toll."),
    }
    return adventure, bank


def test_every_used_memory_reports_the_parts_of_its_score(db, settings, played):
    adventure, bank = played
    result = retrieve(adventure, settings)
    first = result["used"][0]
    assert first["id"] == bank["sundial"].id
    assert first["similarity"] == first["semantic_score"]
    assert first["lexical_score"] > 0
    assert first["final_score"] == pytest.approx(
        first["semantic_score"] + memorybank.LEXICAL_WEIGHT * first["lexical_score"], abs=2e-4)
    assert result["query"]["input"] == "> You ask Mara where the amber sundial went."
    assert result["query"]["lexical_weight"] == memorybank.LEXICAL_WEIGHT
    assert result["query"]["input_weight"] == memorybank.INPUT_WEIGHT


def test_a_pin_is_still_always_used_and_counts_toward_top_k(db, settings, played):
    adventure, bank = played
    settings.memory_top_k = 1
    bank["gate"].pinned = True
    db.commit()
    used = retrieve(adventure, settings)["used"]
    assert [m["id"] for m in used] == [bank["gate"].id]
    assert used[0]["pinned"] is True


def memory_text_reads(statements):
    return [s for s in statements
            if s.lstrip().upper().startswith("SELECT") and "FROM memories" in s
            and "memories.text" in s]


@pytest.fixture()
def sql_log():
    statements: list[str] = []

    def record(conn, cursor, statement, parameters, context, executemany):
        statements.append(statement)

    event.listen(engine, "before_cursor_execute", record)
    try:
        yield statements
    finally:
        event.remove(engine, "before_cursor_execute", record)


def test_memory_text_is_read_once_and_then_held(db, settings, played, sql_log):
    adventure, _ = played
    retrieve(adventure, settings)
    sql_log.clear()
    result = retrieve(adventure, settings)
    reads = memory_text_reads(sql_log)
    # Only the detail read of the memories chosen remains. (Every memory here
    # embeds identically, so redundancy suppression keeps just one of them.)
    assert len(reads) == 1 and reads[0].count("?") == len(result["used"])


def test_a_continue_turn_reads_no_memory_text_to_rank(db, settings, sql_log):
    adventure, _ = make_adventure(db, settings, [("do", "> You wait."), ("ai", "Night falls.")])
    for text in ("one", "two", "three"):
        add_memory(db, adventure, f"memory {text}")
    sql_log.clear()
    result = retrieve(adventure, settings)
    assert all(m["lexical_score"] == 0.0 for m in result["used"])
    assert len(memory_text_reads(sql_log)) == 1  # the top-k detail read only


def test_an_edited_memory_is_matched_on_its_new_text(db, settings, played):
    adventure, bank = played
    assert retrieve(adventure, settings)["used"][0]["id"] == bank["sundial"].id
    # An edit clears the vector (the route calls set_vector(None)); re-embedding
    # sets it again. Both go through set_vector, which drops the held terms.
    bank["road"].text = "The amber sundial was traded for the road toll."
    memorybank.set_vector(bank["road"], None)
    memorybank.set_vector(bank["road"], [1.0, 0.0, 0.0])
    bank["sundial"].text = "Mara hid a bottle in the cellar."
    memorybank.set_vector(bank["sundial"], None)
    memorybank.set_vector(bank["sundial"], [1.0, 0.0, 0.0])
    db.commit()
    assert retrieve(adventure, settings)["used"][0]["id"] == bank["road"].id