from __future__ import annotations import json from config import DEFAULT_QWEN_MODEL, OPENAI_REASONING_EFFORT from services.providers import ProviderClient from services.providers import _PROVIDER_MODEL_MAP from services.rag_pipeline import ( AnswerResult, AnalysisResult, Citation, ScanMatch, _anthropic_thinking_kwargs, _analysis_answer_schema, _full_document_budget, _analysis_schema, _extract_chat_completion_text, _merge_scan_matches, _openai_reasoning_kwargs, _retrieve_context_from_index, answer_query, answer_query_from_full_document, build_analysis_snippets, generate_analysis, generate_analysis_progress, split_into_chunks, split_into_scan_chunks, ) class DummyBlock: def __init__(self, block_type: str, text: str) -> None: self.type = block_type self.text = text class DummyOutputItem: def __init__(self, block: DummyBlock) -> None: self.content = [block] class DummyResponse: def __init__(self, block: DummyBlock) -> None: self.output = [DummyOutputItem(block)] self.output_text = block.text class DummyResponsesClient: def __init__(self, payload: str) -> None: self._payload = payload def create(self, *args, **kwargs): # noqa: ANN002, ANN003 return DummyResponse(DummyBlock("output_json_schema", self._payload)) class DummyOpenAIClient: def __init__(self, payload: str) -> None: self.responses = DummyResponsesClient(payload) class DummyChatCompletionsClient: def __init__(self, responses: list[object]) -> None: self._responses = responses self.calls: list[str] = [] def create(self, *args, **kwargs): # noqa: ANN002, ANN003 self.calls.append(kwargs["model"]) response = self._responses.pop(0) if isinstance(response, Exception): raise response return response class DummyChatClient: def __init__(self, responses: list[object]) -> None: self.completions = DummyChatCompletionsClient(responses) class DummyOpenAIChatClient: def __init__(self, responses: list[object]) -> None: self.chat = DummyChatClient(responses) class DummyTokenizer: def encode(self, text: str, *args, **kwargs): # noqa: ANN002, ANN003 return list(range(len(text.split()))) class FakeVectorStore: def __init__(self, documents: list[str]) -> None: self._documents = documents def similarity_search(self, question: str, k: int): # noqa: ARG002 limit = min(k, len(self._documents)) return [type("Doc", (), {"page_content": doc})() for doc in self._documents[:limit]] def test_analysis_schema_contains_required_fields(): schema = _analysis_schema() assert set(schema["required"]) == { "executive_summary", "bill_summary", "implementation", "critique", "swot", } implementation_item = schema["properties"]["implementation"]["items"] critique_item = schema["properties"]["critique"]["items"] assert set(implementation_item["required"]) == { "stakeholder", "obligation", "implementation_burden", "risk_or_note", } assert set(critique_item["required"]) == {"issue", "why_it_matters", "recommendation"} def test_generate_analysis_with_dummy_client(): payload = json.dumps( { "executive_summary": "Point", "bill_summary": ["Impact"], "implementation": [ { "stakeholder": "Agency", "obligation": "Report", "implementation_burden": "Medium", "risk_or_note": "Needs guidance", } ], "critique": [ { "issue": "Broad powers", "why_it_matters": "Could be overbroad", "recommendation": "Narrow scope", } ], "swot": { "strengths": ["Strong"], "weaknesses": ["Weak"], "opportunities": ["Opp"], "threats": ["Threat"], }, } ) provider_client = ProviderClient( name="openai", client=DummyOpenAIClient(payload), default_model="dummy", api_key="dummy-api-key", ) result = generate_analysis(provider_client, "Sample document text") assert isinstance(result, AnalysisResult) assert result.executive_summary == "Point" assert result.implementation[0].stakeholder == "Agency" def test_split_into_chunks_respects_size(monkeypatch): monkeypatch.setattr("services.rag_pipeline._embedding_tokenizer", lambda: DummyTokenizer()) text = "Lorem ipsum " * 100 chunks = split_into_chunks(text, chunk_size=100, chunk_overlap=10) tokenizer = DummyTokenizer() assert all(len(tokenizer.encode(chunk)) <= 100 for chunk in chunks) assert len(chunks) > 1 def test_split_into_scan_chunks_respects_size(monkeypatch): monkeypatch.setattr("services.rag_pipeline._embedding_tokenizer", lambda: DummyTokenizer()) text = "Alpha beta gamma delta " * 400 chunks = split_into_scan_chunks(text, chunk_size=120, chunk_overlap=20) tokenizer = DummyTokenizer() assert chunks assert all(len(tokenizer.encode(chunk.text)) <= 120 for chunk in chunks) def test_merge_scan_matches_prefers_highest_score(): matches = [ ScanMatch(chunk_id=2, relevance_score=1, evidence_snippet="short"), ScanMatch(chunk_id=2, relevance_score=3, evidence_snippet="better"), ScanMatch(chunk_id=1, relevance_score=2, evidence_snippet="first"), ] merged = _merge_scan_matches(matches) assert [item.chunk_id for item in merged] == [2, 1] assert merged[0].evidence_snippet == "better" def test_full_document_budget_is_provider_aware(): openai_client = ProviderClient(name="openai", client=object(), default_model="dummy", api_key="dummy-api-key") qwen_client = ProviderClient(name="qwen", client=object(), default_model="dummy", api_key="dummy-api-key") assert _full_document_budget(openai_client) > _full_document_budget(qwen_client) assert _full_document_budget(qwen_client) == 24_000 def test_openai_reasoning_kwargs_only_for_reasoning_models(): gpt41_client = ProviderClient(name="openai", client=object(), default_model="gpt-4.1", api_key="dummy-api-key") gpt5_client = ProviderClient(name="openai", client=object(), default_model="gpt-5.1", api_key="dummy-api-key") assert _openai_reasoning_kwargs(gpt41_client) == {} assert _openai_reasoning_kwargs(gpt5_client) == {"reasoning": {"effort": OPENAI_REASONING_EFFORT}} def test_anthropic_thinking_kwargs_only_for_supported_models(): sonnet35_client = ProviderClient( name="anthropic", client=object(), default_model="claude-3-5-sonnet-20240620", api_key="dummy-api-key", ) sonnet37_client = ProviderClient( name="anthropic", client=object(), default_model="claude-3-7-sonnet-20250219", api_key="dummy-api-key", ) assert _anthropic_thinking_kwargs(sonnet35_client) == {} assert "extra_body" in _anthropic_thinking_kwargs(sonnet37_client) def test_gemini_default_model_is_thinking_capable(): assert _PROVIDER_MODEL_MAP["gemini"] == "gemini-2.5-flash" def test_provider_defaults_use_thinking_capable_models(): assert _PROVIDER_MODEL_MAP["openai"] == "gpt-5.5" assert _PROVIDER_MODEL_MAP["anthropic"] == "claude-sonnet-4-6" assert _PROVIDER_MODEL_MAP["cohere"] == "command-a-reasoning-08-2025" assert _PROVIDER_MODEL_MAP["qwen"] == DEFAULT_QWEN_MODEL def test_extract_chat_completion_text_strips_qwen_thinking_block(): message = type("Message", (), {"content": "internal reasoning\nFinal answer"})() choice = type("Choice", (), {"message": message})() response = type("Response", (), {"choices": [choice]})() assert _extract_chat_completion_text(response) == "Final answer" def test_answer_query_uses_scan_context(monkeypatch): provider_client = ProviderClient(name="openai", client=object(), default_model="dummy", api_key="dummy-api-key") analysis = AnalysisResult(executive_summary="Creates a new authority.") monkeypatch.setattr( "services.rag_pipeline.search_analysis_snippets", lambda snippets, question: ( "[ref1] Executive Summary\nCreates a new authority.", [Citation(ref_id=1, snippet="Creates a new authority.")], ), ) monkeypatch.setattr( "services.rag_pipeline._answer_question_from_analysis_context", lambda provider_client, context_text, question: {"answer": f"Scanned answer using {context_text}", "is_sufficient": True}, ) result = answer_query(provider_client, analysis, None, "What does the bill create?", doc_text="Long document") assert isinstance(result, AnswerResult) assert "Scanned answer" in result.answer assert result.citations[0].snippet == "Creates a new authority." assert result.provenance == "analysis_based" def test_answer_query_from_full_document_prefers_full_document(monkeypatch): provider_client = ProviderClient(name="openai", client=object(), default_model="dummy", api_key="dummy-api-key") monkeypatch.setattr( "services.rag_pipeline._answer_question_from_full_document", lambda provider_client, doc_text, question: "Full-document answer", ) monkeypatch.setattr( "services.rag_pipeline._scan_document_for_context", lambda provider_client, doc_text, question: (_ for _ in ()).throw(AssertionError("scan path should not run")), ) result = answer_query_from_full_document(provider_client, None, "What does the bill do?", doc_text="Small document") assert result.answer == "Full-document answer" assert result.citations == [] assert result.provenance == "full_document" def test_answer_query_from_full_document_falls_back_when_scan_errors(monkeypatch): provider_client = ProviderClient(name="openai", client=object(), default_model="dummy", api_key="dummy-api-key") vector_store = FakeVectorStore(["The bill establishes a new data protection authority."]) monkeypatch.setattr( "services.rag_pipeline._answer_question_from_full_document", lambda provider_client, doc_text, question: None, ) monkeypatch.setattr("services.rag_pipeline._scan_document_for_context", lambda provider_client, doc_text, question: None) monkeypatch.setattr( "services.rag_pipeline._generate_answer_from_context", lambda provider_client, context_text, question: f"Fallback answer from {context_text}", ) result = answer_query_from_full_document(provider_client, vector_store, "What does the bill establish?", doc_text="Long document") assert "Fallback answer" in result.answer assert result.citations assert result.citations[0].snippet == "The bill establishes a new data protection authority." def test_answer_query_from_full_document_falls_back_when_scan_has_no_hits(monkeypatch): provider_client = ProviderClient(name="openai", client=object(), default_model="dummy", api_key="dummy-api-key") vector_store = FakeVectorStore(["The bill outlines penalties for non-compliance."]) monkeypatch.setattr( "services.rag_pipeline._answer_question_from_full_document", lambda provider_client, doc_text, question: None, ) monkeypatch.setattr("services.rag_pipeline._scan_document_for_context", lambda provider_client, doc_text, question: None) monkeypatch.setattr( "services.rag_pipeline._generate_answer_from_context", lambda provider_client, context_text, question: context_text, ) result = answer_query_from_full_document(provider_client, vector_store, "What are the penalties?", doc_text="Long document") assert "[1]" in result.answer assert result.citations[0].ref_id == 1 def test_answer_query_from_full_document_supports_vector_store_only(monkeypatch): provider_client = ProviderClient(name="openai", client=object(), default_model="dummy", api_key="dummy-api-key") vector_store = FakeVectorStore(["Only fallback retrieval is available."]) monkeypatch.setattr( "services.rag_pipeline._generate_answer_from_context", lambda provider_client, context_text, question: context_text, ) result = answer_query_from_full_document(provider_client, vector_store, "What is available?") assert result.citations[0].snippet == "Only fallback retrieval is available." def test_answer_query_from_full_document_uses_scan_when_full_document_too_large(monkeypatch): provider_client = ProviderClient(name="openai", client=object(), default_model="dummy", api_key="dummy-api-key") monkeypatch.setattr( "services.rag_pipeline._answer_question_from_full_document", lambda provider_client, doc_text, question: None, ) monkeypatch.setattr( "services.rag_pipeline._scan_document_for_context", lambda provider_client, doc_text, question: ( "[ref1] Chunk 7\nEvidence: licensing rule\nFull context:\nThe bill sets licensing rules.", [Citation(ref_id=1, snippet="licensing rule")], ), ) monkeypatch.setattr( "services.rag_pipeline._generate_answer_from_context", lambda provider_client, context_text, question: "Scanned fallback answer", ) result = answer_query_from_full_document(provider_client, None, "How does licensing work?", doc_text="Very long document") assert result.answer == "Scanned fallback answer" assert result.citations[0].snippet == "licensing rule" def test_answer_query_requests_deeper_answer_when_analysis_is_insufficient(monkeypatch): provider_client = ProviderClient(name="openai", client=object(), default_model="dummy", api_key="dummy-api-key") analysis = AnalysisResult(executive_summary="A short summary.") monkeypatch.setattr( "services.rag_pipeline.search_analysis_snippets", lambda snippets, question: ( "[ref1] Executive Summary\nA short summary.", [Citation(ref_id=1, snippet="A short summary.")], ), ) monkeypatch.setattr( "services.rag_pipeline._answer_question_from_analysis_context", lambda provider_client, context_text, question: { "answer": "The summary only partly answers this.", "is_sufficient": False, }, ) result = answer_query(provider_client, analysis, None, "What is missing?", doc_text="Long document") assert result.provenance == "analysis_based" assert result.needs_deeper_consent is True assert result.deeper_answer_available is True def test_generate_analysis_progress_merges_sections(monkeypatch): payloads = iter( [ json.dumps({"executive_summary": "Summary", "bill_summary": ["Point"]}), json.dumps({"implementation": [{"stakeholder": "Users", "obligation": "Register", "implementation_burden": "Medium", "risk_or_note": "Needs outreach"}]}), json.dumps({"critique": [{"issue": "Ambiguity", "why_it_matters": "Creates uncertainty", "recommendation": "Clarify wording"}]}), json.dumps({"swot": {"strengths": ["Strong"], "weaknesses": ["Weak"], "opportunities": ["Opp"], "threats": ["Threat"]}}), ] ) monkeypatch.setattr("services.rag_pipeline._generate_json_payload", lambda *args, **kwargs: next(payloads)) progress = list(generate_analysis_progress(ProviderClient(name="openai", client=object(), default_model="dummy", api_key="dummy-api-key"), "Doc")) assert len(progress) == 4 assert progress[-1][1].swot.strengths == ["Strong"] def test_build_analysis_snippets_includes_sections(): analysis = AnalysisResult( executive_summary="Summary", bill_summary=["Point"], implementation=[{"stakeholder": "SMEs", "obligation": "File reports", "implementation_burden": "Low", "risk_or_note": "Admin cost"}], ) snippets = build_analysis_snippets(analysis) assert snippets[0].section == "Executive Summary" assert any(snippet.section == "Implementation" for snippet in snippets)