ThothII/harness/tests/test_semantic_kind_isolation.py

from __future__ import annotations

import hashlib
from dataclasses import dataclass
from datetime import UTC, datetime

from tht.adapters.vector.qdrant import point_id
from tht.cli.vector_cmd import sync_canonical_records
from tht.corpus.chunk import ChunkPolicy
from tht.corpus.models import CanonicalChunk
from tht.corpus.pipeline import CorpusPipeline
from tht.corpus.store import CorpusStore
from tht.memory import MemoryRecord, save_one_memory
from tht.mschema.models import (
    Annotations,
    ColumnPhysical,
    PhysicalSchema,
    TablePhysical,
)
from tht.ports.vector import VectorCapabilities, VectorHealth
from tht.vectorstore.records import qdrant_payload, schema_records


def _sha(content: str) -> str:
    return f"sha256:{hashlib.sha256(content.encode('utf-8')).hexdigest()}"


class _Embedder:
    def embed_documents(self, documents):
        return [[float(index + 1)] * 4 for index, _ in enumerate(documents)]


@dataclass
class _Point:
    point_id: str
    payload: dict
    embedding: list[float]


class FakeVectorStore:
    def __init__(self, workspace_id="psd-clinical", workspace_revision=None):
        self.workspace_id = workspace_id
        self.workspace_revision = workspace_revision or "a" * 40
        self.points: dict[str, _Point] = {}
        self.search_calls: list[dict] = []

    @property
    def capabilities(self):
        return VectorCapabilities(
            search=True,
            existing_hashes=True,
            upsert=True,
            metadata_filter=True,
            delete_generation=True,
            list_evidence_generations=True,
        )

    def health(self):
        return VectorHealth(ok=True)

    def search(self, collections, embedding, *, limit, kinds=None, metadata_filter=None):
        self.search_calls.append(
            {
                "collections": collections,
                "embedding": embedding,
                "limit": limit,
                "kinds": kinds,
                "metadata_filter": metadata_filter,
            }
        )
        return []

    def existing_hashes(self, collection, kinds):
        allowed = set(kinds)
        return {
            point.payload["record_key"]: point.payload["content_hash"]
            for point in self.points.values()
            if point.payload["record_kind"] in allowed
        }

    def upsert(self, collection, records):
        for row in records:
            semantic_kind = qdrant_payload(
                row.record,
                content_hash=row.content_hash,
                workspace_id=self.workspace_id,
                workspace_revision=self.workspace_revision,
            )["kind"]
            payload = qdrant_payload(
                row.record,
                content_hash=row.content_hash,
                workspace_id=self.workspace_id,
                workspace_revision=self.workspace_revision,
            )
            self.points[point_id(self.workspace_id, semantic_kind, row.record.id)] = _Point(
                point_id=point_id(self.workspace_id, semantic_kind, row.record.id),
                payload=payload,
                embedding=row.embedding,
            )
        return len(records)

    def delete_generation(self, collection, generation, workspace_id):
        doomed = [
            key
            for key, point in self.points.items()
            if point.payload.get("record_kind") == "evidence"
            and point.payload.get("vector_generation") == generation
            and point.payload.get("workspace_id") == workspace_id
        ]
        for key in doomed:
            self.points.pop(key)
        return len(doomed)

    def list_evidence_generations(self, collection, workspace_id):
        return sorted(
            {
                point.payload["vector_generation"]
                for point in self.points.values()
                if point.payload.get("record_kind") == "evidence"
                and point.payload.get("workspace_id") == workspace_id
            }
        )


def _schema_records():
    return schema_records(
        PhysicalSchema(
            database="analytics",
            schema="mart",
            introspected_at=datetime.now(UTC),
            tables={
                "fact_patient": TablePhysical(
                    comment="Patients",
                    columns={"id": ColumnPhysical(type="bigint", comment="pk")},
                )
            },
        ),
        Annotations(),
    )


def _memory_records():
    return [
        MemoryRecord(
            id="mem-0001",
            ts=datetime(2026, 1, 1, tzinfo=UTC),
            session_id="s1",
            decision_seq=7,
            type="concept_clarified",
            subject="paziente attivo",
            detail="flag_attivo = true",
            rationale="r",
            question_context="dammi i pazienti attivi",
            tables=[],
            concepts=["paziente attivo"],
        )
    ]


def test_schema_and_memory_use_expected_semantic_kinds_and_shared_identity():
    store = FakeVectorStore()
    embedder = _Embedder()

    schema_stats = sync_canonical_records(
        "schema_records",
        _schema_records(),
        store=store,
        embedder=embedder,
    )
    memory_count = save_one_memory(_memory_records(), 7, store=store, embedder=embedder)

    assert schema_stats.added == 2
    assert memory_count == 1
    payloads = {point.payload["record_kind"]: point.payload for point in store.points.values()}
    assert payloads["schema_table"]["kind"] == "schema"
    assert payloads["schema_column"]["kind"] == "schema"
    assert payloads["memory"]["kind"] == "memory"
    assert {payload["workspace_id"] for payload in payloads.values()} == {"psd-clinical"}
    assert {payload["workspace_revision"] for payload in payloads.values()} == {"a" * 40}


def test_corpus_vector_records_keep_exact_generation_and_retry_is_idempotent(tmp_path):
    store = FakeVectorStore()
    pipeline = CorpusPipeline(
        store=CorpusStore(tmp_path / "corpus"),
        sources=[],
        embedder=None,
        vector_store=store,
        embedding_model="qwen3-embedding:0.6b",
        embedding_dimensions=1024,
        chunk_policy=ChunkPolicy(version="chunk-v1", max_chars=4000),
        pipeline_version="evidence-v1",
        workspace_id="psd-clinical",
    )
    chunk = CanonicalChunk(
        chunk_id="chunk:1",
        document_id="doc:patient-guide",
        ordinal=0,
        content="Patient evidence",
        content_hash=_sha("Patient evidence"),
        source_uri="file:///tmp/patient-guide.md",
        pipeline_version="evidence-v1",
    )
    row = pipeline._vector_record(
        chunk,
        [0.1, 0.2, 0.3, 0.4],
        "gen:" + "1" * 32,
        "psd-clinical",
    )

    assert row.record.kind == "evidence"
    assert row.record.metadata["vector_generation"] == "gen:" + "1" * 32

    store.upsert("evidence", [row])
    store.upsert("evidence", [row])

    assert len(store.points) == 1
    point = next(iter(store.points.values()))
    assert point.payload["kind"] == "evidence"
    assert point.payload["vector_generation"] == "gen:" + "1" * 32
    assert point.payload["workspace_id"] == "psd-clinical"
    assert point.payload["workspace_revision"] == "a" * 40