""" Tests for core/sufficient_context.py — Google ICLR 2025 abstention layer. """ from __future__ import annotations from unittest.mock import MagicMock from models import QueryMode, RetrievalContext, RetrievalResult def _make_result(sim: float, text: str = "test chunk") -> RetrievalResult: return RetrievalResult( chunk_text=text, source="test.txt", similarity_score=sim, chunk_index=0, ) def _make_context(results: list[RetrievalResult]) -> RetrievalContext: return RetrievalContext( query="test question", results=results, query_mode=QueryMode.HYBRID, ) class TestSufficientContextChecker: def test_empty_context_is_not_sufficient(self) -> None: from core.sufficient_context import SufficientContextChecker checker = SufficientContextChecker() ctx = _make_context([]) result = checker.score("what is X?", ctx) assert result.is_sufficient is False assert result.overall_score == 0.0 assert result.recommendation in ("abstain", "web_search") def test_high_similarity_context_is_sufficient(self) -> None: from core.sufficient_context import SufficientContextChecker checker = SufficientContextChecker(sufficiency_threshold=0.4) ctx = _make_context( [ _make_result(0.85, "very relevant chunk about X"), _make_result(0.80, "another highly relevant chunk"), _make_result(0.75, "also relevant to the question"), ] ) result = checker.score("what is X?", ctx) assert result.is_sufficient is True assert result.overall_score > 0.4 assert result.recommendation == "generate" def test_low_similarity_context_is_not_sufficient(self) -> None: from core.sufficient_context import SufficientContextChecker checker = SufficientContextChecker(sufficiency_threshold=0.6) ctx = _make_context([_make_result(0.1)]) result = checker.score("what is X?", ctx) assert result.is_sufficient is False def test_density_score_is_mean_similarity(self) -> None: from core.sufficient_context import SufficientContextChecker checker = SufficientContextChecker() ctx = _make_context([_make_result(0.6), _make_result(0.8)]) result = checker.score("test", ctx) assert abs(result.density_score - 0.7) < 0.01 def test_crag_score_included_in_ensemble(self) -> None: from core.sufficient_context import SufficientContextChecker checker = SufficientContextChecker(sufficiency_threshold=0.3) ctx = _make_context([_make_result(0.1)]) # low density # Without CRAG score without_crag = checker.score("test", ctx) # With high CRAG score — should raise overall with_crag = checker.score("test", ctx, crag_score=0.9) assert with_crag.overall_score > without_crag.overall_score def test_self_rating_called_when_enabled(self) -> None: from core.sufficient_context import SufficientContextChecker checker = SufficientContextChecker() ctx = _make_context([_make_result(0.7, "relevant content")]) mock_llm = MagicMock(return_value="0.8") result = checker.score("test", ctx, llm_fn=mock_llm, enable_self_rating=True) assert result.self_rating is not None assert abs(result.self_rating - 0.8) < 0.01 mock_llm.assert_called_once() def test_self_rating_not_called_when_disabled(self) -> None: from core.sufficient_context import SufficientContextChecker checker = SufficientContextChecker() ctx = _make_context([_make_result(0.7)]) mock_llm = MagicMock(return_value="0.8") result = checker.score("test", ctx, llm_fn=mock_llm, enable_self_rating=False) assert result.self_rating is None mock_llm.assert_not_called() def test_recommendation_abstain_when_very_low(self) -> None: from config import settings from core.sufficient_context import SufficientContextChecker # Save original original = settings.web_search_fallback settings.web_search_fallback = False try: checker = SufficientContextChecker(sufficiency_threshold=0.9) ctx = _make_context([_make_result(0.1)]) result = checker.score("test", ctx) assert result.recommendation == "abstain" finally: settings.web_search_fallback = original def test_component_scores_present(self) -> None: from core.sufficient_context import SufficientContextChecker checker = SufficientContextChecker() ctx = _make_context([_make_result(0.7)]) result = checker.score("test", ctx) assert "density" in result.component_scores assert "coverage" in result.component_scores class TestCheckSufficiencyConvenience: def test_module_function_works(self) -> None: from core.sufficient_context import check_sufficiency ctx = _make_context([_make_result(0.7)]) result = check_sufficiency("test question", ctx) assert isinstance(result.overall_score, float) assert 0.0 <= result.overall_score <= 1.0 class TestAbstentionResponse: def test_abstention_response_contains_score(self) -> None: from core.sufficient_context import SufficiencyResult, abstention_response result = SufficiencyResult( is_sufficient=False, overall_score=0.2, density_score=0.1, coverage_score=0.3, crag_score=None, self_rating=None, num_chunks=1, recommendation="abstain", explanation="Low context quality", ) response = abstention_response(result) assert "20%" in response assert "Low context quality" in response