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"])]