ThothII/harness/tests/memory/test_solved_lifecycle.py

from datetime import UTC, datetime
from types import SimpleNamespace

import pytest

from tht.decisions import DecisionRecord
from tht.memory import (
    SolvedIndexError,
    index_solved_question,
    index_solved_question_best_effort,
    search_solved_questions,
)
from tht.session.models import SessionManifest, SessionSnapshot


def _decision(seq: int, type_: str, subject: str, detail: str = "") -> DecisionRecord:
    return DecisionRecord(
        seq=seq,
        ts=datetime(2026, 8, 24, tzinfo=UTC),
        type=type_,
        subject=subject,
        detail=detail,
    )


def _snapshot(*, sql: str | None = "SELECT 1\n", approved: bool = True) -> SessionSnapshot:
    decisions = [_decision(1, "question_rewritten", "question", "rewritten question")]
    if approved:
        decisions.append(_decision(2, "sql_approved", "phase:7"))
    decisions.extend(
        _decision(index + 2, "phase_approved", f"phase:{index}")
        for index in range(1, 9)
    )
    return SessionSnapshot(
        manifest=SessionManifest(
            id="s1",
            created_at=datetime(2026, 8, 24, tzinfo=UTC),
            status="finalized",
            question="original question",
            database="analytics",
            schema="mart",
        ),
        artifacts={} if sql is None else {"sql_final": sql},
        decisions=decisions,
    )


class Store:
    def __init__(self):
        self.hashes = {}
        self.upserts = []

    def existing_hashes(self, collection, kinds):
        assert (collection, kinds) == ("memory", ["solved_question"])
        return self.hashes

    def upsert(self, collection, records):
        self.upserts.append((collection, records))
        return len(records)


class Embedder:
    def __init__(self):
        self.documents = []
        self.queries = []

    def embed_documents(self, documents):
        self.documents.append(documents)
        return [[0.1, 0.2]]

    def embed_query(self, question):
        self.queries.append(question)
        return [0.3, 0.4]


def test_memory_facade_indexes_finalized_question_with_compatible_record_and_dedup():
    store = Store()
    embedder = Embedder()
    snapshot = _snapshot()

    assert index_solved_question(
        snapshot,
        {"fact_z", "dim_a"},
        store=store,
        embedder=embedder,
    ) == 1

    collection, rows = store.upserts[0]
    assert collection == "memory"
    assert len(rows) == 1
    row = rows[0]
    assert row.record.model_dump() == {
        "id": "solved:s1",
        "kind": "solved_question",
        "ref": "s1",
        "title": "rewritten question",
        "content": "rewritten question",
        "metadata": {
            "question": "rewritten question",
            "sql": "SELECT 1",
            "tables": ["dim_a", "fact_z"],
            "session_id": "s1",
        },
    }
    assert embedder.documents == [["rewritten question"]]

    store.hashes = {row.record.id: row.content_hash}
    assert index_solved_question(
        snapshot,
        {"fact_z", "dim_a"},
        store=store,
        embedder=embedder,
    ) == 0
    assert len(store.upserts) == 1
    assert embedder.documents == [["rewritten question"]]

    changed = snapshot.model_copy(
        update={"artifacts": {"sql_final": "SELECT 2\n"}},
    )
    assert index_solved_question(
        changed,
        {"fact_z", "dim_a"},
        store=store,
        embedder=embedder,
    ) == 1
    assert store.upserts[-1][1][0].record.metadata["sql"] == "SELECT 2"
    assert embedder.documents == [["rewritten question"], ["rewritten question"]]


def test_memory_facade_falls_back_to_manifest_question():
    snapshot = _snapshot().model_copy(
        update={"decisions": [
            decision
            for decision in _snapshot().decisions
            if decision.type != "question_rewritten"
        ]},
    )
    store = Store()

    assert index_solved_question(
        snapshot,
        None,
        store=store,
        embedder=Embedder(),
    ) == 1
    assert store.upserts[0][1][0].record.content == "original question"
    assert store.upserts[0][1][0].record.metadata["tables"] == []


@pytest.mark.parametrize(
    ("snapshot", "message"),
    [
        (_snapshot(sql=None), "sql_final.sql assente"),
        (_snapshot(approved=False), "decisione sql_approved assente"),
    ],
)
def test_memory_facade_rejects_incomplete_solved_question(snapshot, message):
    with pytest.raises(SolvedIndexError, match=message):
        index_solved_question(snapshot, set(), store=Store(), embedder=Embedder())


def test_memory_facade_keeps_solved_indexing_best_effort_after_finalization():
    snapshot = _snapshot()

    outcome = index_solved_question_best_effort(
        snapshot,
        set(),
        store_factory=lambda: (_ for _ in ()).throw(RuntimeError("vector unavailable")),
        embedder_factory=Embedder,
    )

    assert outcome.upserted is None
    assert outcome.error == "vector unavailable"
    assert snapshot.manifest.status == "finalized"


def test_memory_facade_searches_solved_questions_in_rank_order_with_compatible_payload():
    hits = [
        SimpleNamespace(
            ref="s2",
            content="second fallback question",
            metadata={
                "session_id": "s2",
                "question": "second question",
                "sql": "SELECT 2",
                "tables": ["fact_two"],
            },
            similarity=0.93456,
        ),
        SimpleNamespace(
            ref="s1",
            content="first fallback question",
            metadata={},
            similarity=0.81234,
        ),
    ]

    class Searcher:
        def __init__(self):
            self.calls = []

        def search(self, embedding, *, top_n, kinds):
            self.calls.append((embedding, top_n, kinds))
            return hits

    searcher = Searcher()
    embedder = Embedder()

    results = search_solved_questions(
        "similar question",
        searcher=searcher,
        embedder=embedder,
        top=3,
    )

    assert results == [
        {
            "session_id": "s2",
            "question": "second question",
            "sql": "SELECT 2",
            "tables": ["fact_two"],
            "score": 0.9346,
        },
        {
            "session_id": "s1",
            "question": "first fallback question",
            "sql": "",
            "tables": [],
            "score": 0.8123,
        },
    ]
    assert embedder.queries == ["similar question"]
    assert searcher.calls == [([0.3, 0.4], 3, ["solved_question"])]