pdf-quiz-generator/backend/tests/test_hybrid_search.py

"""Hybrid retrieval: lexical and semantic rankers fused, never intersected.

Run: DATABASE_URL=sqlite:///:memory: PYTHONPATH=backend python -m unittest discover -s backend/tests
"""
import os
os.environ["DATABASE_URL"] = "sqlite:///:memory:"

import unittest
from unittest.mock import patch

from sqlalchemy import create_engine
from sqlalchemy.orm import Session
from sqlalchemy.pool import StaticPool

from app.database import Base
from app.models.media import MediaAsset  # noqa — the health report walks every embeddable table.
from app.models.question import Question
from app.models.user import User
from app.services import search_service


class HybridSearchTests(unittest.TestCase):
    def setUp(self):
        self.engine = create_engine("sqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool)
        Base.metadata.create_all(self.engine)
        self.db = Session(self.engine)
        self.db.add(User(id=1, name="Mod", email="mod@example.test", hashed_password="unused", role="moderator"))
        for qid, text in [
            (1, "A child with fever and a seizure"),
            (2, "An infant with jaundice on day three"),
            (3, "A toddler with a febrile convulsion"),
        ]:
            self.db.add(Question(id=qid, user_id=1, question_text=text,
                                 question_type="mcq", options=["yes", "no"], correct_answer="yes"))
        self.db.commit()

    def tearDown(self):
        self.db.close()
        self.engine.dispose()

    def test_result_is_the_union_of_both_rankers(self):
        # "convulsion" never appears in question 1, and "fever" never in 3, so an
        # intersection would return one of them; the union must return both.
        with patch.object(search_service, "_semantic_ranked", return_value=[3]):
            ids, semantic = search_service.hybrid_question_ids(self.db, "fever")
        self.assertEqual(set(ids), {1, 3})
        self.assertEqual(semantic, {3})

    def test_agreement_between_rankers_outranks_a_single_hit(self):
        with patch.object(search_service, "_semantic_ranked", return_value=[3, 1]):
            ids, _ = search_service.hybrid_question_ids(self.db, "fever")
        # Question 1 is first lexically and second semantically; 3 only appears once
        # in the lexical list, so fusion should put 1 ahead of the semantic-only hit.
        self.assertEqual(ids[0], 1)

    def test_a_failing_ranker_degrades_instead_of_erroring(self):
        with patch.object(search_service, "_semantic_ranked", side_effect=RuntimeError("no pgvector")):
            ids, semantic = search_service.hybrid_question_ids(self.db, "fever")
        self.assertEqual(ids, [1])
        self.assertEqual(semantic, set())

        with patch.object(search_service, "_lexical_ranked", side_effect=RuntimeError("no index")), \
             patch.object(search_service, "_semantic_ranked", return_value=[2]):
            ids, _ = search_service.hybrid_question_ids(self.db, "fever")
        self.assertEqual(ids, [2])

    def test_blank_query_returns_nothing(self):
        self.assertEqual(search_service.hybrid_question_ids(self.db, "   "), ([], set()))


class EmbeddingProvenanceTests(unittest.TestCase):
    """A model change must be visible, not a silent quality regression."""

    def setUp(self):
        self.engine = create_engine("sqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool)
        Base.metadata.create_all(self.engine)
        self.db = Session(self.engine)
        self.db.add(User(id=1, name="Mod", email="mod@example.test", hashed_password="unused", role="moderator"))
        for qid in (1, 2, 3):
            self.db.add(Question(id=qid, user_id=1, question_text=f"Question {qid}",
                                 question_type="mcq", options=["yes", "no"], correct_answer="yes"))
        self.db.commit()

    def tearDown(self):
        self.db.close()
        self.engine.dispose()

    def test_embedding_records_the_model_that_produced_it(self):
        from app.services import embedding_service

        question = self.db.get(Question, 1)
        with patch.object(embedding_service, "_get_embedding_model", return_value="model-a"), \
             patch.object(embedding_service, "generate_embedding", return_value=[0.1] * 1024):
            self.assertTrue(embedding_service.embed_question(question))
        self.assertEqual(question.embedding_model, "model-a")
        self.assertIsNotNone(question.embedded_at)

    def test_a_failed_embedding_leaves_no_stale_provenance(self):
        from app.services import embedding_service

        question = self.db.get(Question, 1)
        with patch.object(embedding_service, "_get_embedding_model", return_value="model-a"), \
             patch.object(embedding_service, "generate_embedding", return_value=None):
            self.assertFalse(embedding_service.embed_question(question))
        self.assertIsNone(question.embedding_model)

    def test_health_separates_missing_from_stale(self):
        from app.services import embedding_service

        # One current, one from a retired model, one never embedded.
        self.db.query(Question).filter(Question.id == 1).update(
            {"embedding": [0.1] * 1024, "embedding_model": "model-a"})
        self.db.query(Question).filter(Question.id == 2).update(
            {"embedding": [0.2] * 1024, "embedding_model": "model-old"})
        self.db.commit()

        with patch.object(embedding_service, "_get_embedding_model", return_value="model-a"):
            health = embedding_service.stale_embedding_counts(self.db)
        self.assertEqual(health["active_model"], "model-a")
        self.assertEqual((health["current"], health["stale"], health["missing"]), (1, 1, 1))
        self.assertTrue(health["needs_regeneration"])

        # Switching the model makes every stored vector stale, and says so.
        with patch.object(embedding_service, "_get_embedding_model", return_value="model-b"):
            after = embedding_service.stale_embedding_counts(self.db)
        self.assertEqual((after["current"], after["stale"]), (0, 2))


class ProvenanceMatchesGeneratorTests(unittest.TestCase):
    """The stamp must name the model that actually produced the vector."""

    def test_generator_and_stamp_resolve_the_same_model(self):
        from app.config import settings
        from app.services import embedding_service

        calls = {}

        def fake_post(url, **kwargs):
            calls["model"] = kwargs["content"] if "content" in kwargs else kwargs.get("json")
            raise RuntimeError("stop after capturing the request")

        with patch.object(settings, "LITELLM_EMBEDDING_MODEL", "from-env"), \
             patch.object(settings, "LITELLM_API_KEY", "k"), \
             patch.object(settings, "LITELLM_API_BASE", "https://proxy.test"), \
             patch.object(embedding_service, "_get_embedding_model", return_value="from-resolver") as resolver, \
             patch("httpx.post", side_effect=fake_post):
            embedding_service.generate_embedding("some question text")

        # The request must carry the resolver's model, not the raw env value:
        # otherwise embed_question would stamp a different name than it embedded with.
        self.assertIn("from-resolver", str(calls.get("model")))
        self.assertNotIn("from-env", str(calls.get("model")))
        self.assertTrue(resolver.called)


if __name__ == "__main__":
    unittest.main()