| import unittest |
| from unittest.mock import patch |
|
|
| from agents.rag_graph import ( |
| build_initial_graph_state, |
| check_grounding_node, |
| decide_next_action_node, |
| expand_context_node, |
| fallback_answer_node, |
| generate_answer_node, |
| prepare_input_node, |
| route_after_grounding, |
| run_rag_graph_answer, |
| retrieve_node, |
| ) |
|
|
|
|
| class Phase2RagGraphTests(unittest.TestCase): |
| def test_build_initial_graph_state_sets_expected_defaults(self): |
| state = build_initial_graph_state( |
| "What is this roadmap about?", |
| selected_file="roadmap.pdf", |
| debug_mode=True, |
| ) |
|
|
| self.assertEqual(state["query"], "What is this roadmap about?") |
| self.assertEqual(state["selected_file"], "roadmap.pdf") |
| self.assertTrue(state["debug_mode"]) |
| self.assertEqual(state["retrieved_documents"], []) |
| self.assertIsNone(state["metadata_filter"]) |
| self.assertIsNone(state["next_action"]) |
|
|
| def test_prepare_input_node_builds_metadata_filter(self): |
| state = build_initial_graph_state( |
| "What is this roadmap about?", |
| selected_file="roadmap.pdf", |
| selected_file_type="pdf", |
| page_start="4", |
| page_end="2", |
| ) |
|
|
| updated_state = prepare_input_node(state) |
|
|
| self.assertEqual( |
| updated_state["metadata_filter"], |
| { |
| "file_name": "roadmap.pdf", |
| "file_type": "pdf", |
| "page_range": {"start": 2, "end": 4}, |
| }, |
| ) |
|
|
| @patch("agents.rag_graph.retrieve_documents_with_query_transform") |
| def test_retrieve_node_stores_retrieved_documents_and_debug_data(self, mock_retrieve): |
| state = build_initial_graph_state( |
| "What is this roadmap about?", |
| selected_file="roadmap.pdf", |
| debug_mode=True, |
| ) |
| state["metadata_filter"] = {"file_name": "roadmap.pdf"} |
| mock_retrieve.return_value = (["doc-a", "doc-b"], {"stage_counts": {"search_queries": 1}}) |
|
|
| updated_state = retrieve_node( |
| state, |
| vectorstore=object(), |
| reranker=object(), |
| bm25_index=object(), |
| llm=object(), |
| retrieval_k=4, |
| rerank_candidate_k=8, |
| bm25_candidate_k=8, |
| enable_query_transform=True, |
| ) |
|
|
| self.assertEqual(updated_state["retrieved_documents"], ["doc-a", "doc-b"]) |
| self.assertEqual(updated_state["debug_data"]["metadata_filter"], {"file_name": "roadmap.pdf"}) |
| self.assertEqual(updated_state["debug_data"]["grounding"]["reason"], "not_checked") |
|
|
| @patch("agents.rag_graph.expand_with_context_window") |
| def test_expand_context_node_stores_expanded_documents(self, mock_expand): |
| state = build_initial_graph_state("What is this roadmap about?", debug_mode=True) |
| state["retrieved_documents"] = ["doc-a"] |
| state["debug_data"] = {"stage_counts": {}, "grounding": {"reason": "not_checked"}} |
| mock_expand.return_value = ["doc-a", "doc-a-next"] |
|
|
| updated_state = expand_context_node( |
| state, |
| chunk_registry={"by_chunk_id": {}}, |
| context_window=1, |
| max_expanded_chunks=12, |
| ) |
|
|
| self.assertEqual(updated_state["expanded_documents"], ["doc-a", "doc-a-next"]) |
| self.assertEqual(updated_state["debug_data"]["expanded_hits"], ["doc-a", "doc-a-next"]) |
| self.assertEqual(updated_state["debug_data"]["stage_counts"]["expanded_context"], 2) |
|
|
| @patch("agents.rag_graph.check_grounding_evidence") |
| def test_check_grounding_node_stores_grounding_result(self, mock_check_grounding): |
| state = build_initial_graph_state("What is this roadmap about?", debug_mode=True) |
| state["retrieved_documents"] = ["doc-a"] |
| state["expanded_documents"] = ["doc-a", "doc-a-next"] |
| state["debug_data"] = { |
| "stage_counts": {}, |
| "grounding": {"stage": "retrieval", "passed": None, "reason": "not_checked"}, |
| } |
| mock_check_grounding.return_value = { |
| "passed": True, |
| "reason": "enough_evidence", |
| "top_rerank_score": 1.2, |
| "retrieved_count": 1, |
| "expanded_count": 2, |
| } |
|
|
| updated_state = check_grounding_node( |
| state, |
| min_grounded_rerank_score=1.0, |
| min_grounded_chunks=1, |
| ) |
|
|
| self.assertTrue(updated_state["grounding"]["passed"]) |
| self.assertEqual(updated_state["debug_data"]["grounding"]["reason"], "enough_evidence") |
| self.assertEqual(updated_state["debug_data"]["grounding"]["stage"], "retrieval") |
|
|
| def test_route_after_grounding_chooses_refine_query_when_retry_retrieval(self): |
| state = build_initial_graph_state("What is this roadmap about?") |
| state["next_action"] = "retry_retrieval" |
|
|
| next_node = route_after_grounding(state) |
|
|
| self.assertEqual(next_node, "refine_query") |
|
|
| def test_route_after_grounding_chooses_fallback_when_no_evidence(self): |
| state = build_initial_graph_state("What is this roadmap about?") |
| state["next_action"] = "fallback" |
|
|
| next_node = route_after_grounding(state) |
|
|
| self.assertEqual(next_node, "fallback_answer") |
|
|
| def test_decide_next_action_node_marks_retry_when_local_evidence_is_partial(self): |
| state = build_initial_graph_state("What is this roadmap about?", debug_mode=True) |
| state["retrieved_documents"] = ["doc-a"] |
| state["expanded_documents"] = ["doc-a", "doc-a-next"] |
| state["grounding"] = {"passed": False} |
| state["debug_data"] = {"grounding": {}, "stage_counts": {}} |
|
|
| updated_state = decide_next_action_node(state) |
|
|
| self.assertEqual(updated_state["next_action"], "retry_retrieval") |
| self.assertEqual(updated_state["decision_reason"], "partial_local_evidence_retry_1") |
| self.assertEqual(updated_state["debug_data"]["next_action"], "retry_retrieval") |
|
|
| def test_decide_next_action_node_marks_fallback_when_no_local_evidence_exists(self): |
| state = build_initial_graph_state("What is this roadmap about?", debug_mode=True) |
| state["grounding"] = {"passed": False} |
| state["debug_data"] = {"grounding": {}, "stage_counts": {}} |
|
|
| updated_state = decide_next_action_node(state) |
|
|
| self.assertEqual(updated_state["next_action"], "fallback") |
| self.assertEqual(updated_state["decision_reason"], "no_evidence_found_research_disabled") |
|
|
| def test_decide_next_action_node_marks_answer_when_grounding_passes(self): |
| state = build_initial_graph_state("What is this roadmap about?", debug_mode=True) |
| state["grounding"] = {"passed": True} |
| state["debug_data"] = {"grounding": {}, "stage_counts": {}} |
|
|
| updated_state = decide_next_action_node(state) |
|
|
| self.assertEqual(updated_state["next_action"], "answer") |
| self.assertEqual(updated_state["decision_reason"], "grounding_passed") |
|
|
| @patch("agents.rag_graph.build_sources") |
| def test_fallback_answer_node_stores_fallback_answer_and_sources(self, mock_build_sources): |
| state = build_initial_graph_state("What is this roadmap about?", debug_mode=True) |
| state["retrieved_documents"] = ["doc-a"] |
| state["debug_data"] = {"grounding": {}, "stage_counts": {}} |
| mock_build_sources.return_value = [{"chunk_id": "c1"}] |
|
|
| updated_state = fallback_answer_node( |
| state, |
| grounded_fallback_message="fallback", |
| ) |
|
|
| self.assertEqual(updated_state["answer"], "fallback") |
| self.assertEqual(updated_state["sources"], [{"chunk_id": "c1"}]) |
| self.assertEqual(updated_state["citations"], []) |
| self.assertEqual(updated_state["debug_data"]["graph_path"], "fallback_answer") |
|
|
| @patch("agents.rag_graph.build_citations") |
| @patch("agents.rag_graph.answer_is_grounded") |
| @patch("agents.rag_graph.generate_answer") |
| @patch("agents.rag_graph.format_context") |
| @patch("agents.rag_graph.build_sources") |
| def test_generate_answer_node_stores_answer_sources_and_citations( |
| self, |
| mock_build_sources, |
| mock_format_context, |
| mock_generate_answer, |
| mock_answer_is_grounded, |
| mock_build_citations, |
| ): |
| state = build_initial_graph_state("What is this roadmap about?", debug_mode=True) |
| state["retrieved_documents"] = ["doc-a"] |
| state["expanded_documents"] = ["doc-a", "doc-a-next"] |
| state["grounding"] = { |
| "passed": True, |
| "top_rerank_score": 1.2, |
| "retrieved_count": 1, |
| "expanded_count": 2, |
| } |
| state["debug_data"] = {"grounding": {}, "stage_counts": {}} |
| mock_build_sources.return_value = [{"chunk_id": "c1"}] |
| mock_format_context.return_value = "[1] expanded context" |
| mock_generate_answer.return_value = "Grounded answer [1]" |
| mock_answer_is_grounded.return_value = True |
| mock_build_citations.return_value = [{"number": 1}] |
|
|
| updated_state = generate_answer_node( |
| state, |
| llm=object(), |
| grounded_fallback_message="fallback", |
| ) |
|
|
| self.assertEqual(updated_state["answer"], "Grounded answer [1]") |
| self.assertEqual(updated_state["sources"], [{"chunk_id": "c1"}]) |
| self.assertEqual(updated_state["citations"], [{"number": 1}]) |
| self.assertEqual(updated_state["debug_data"]["graph_path"], "generate_answer") |
| self.assertEqual(updated_state["debug_data"]["grounding"]["reason"], "answer_is_grounded") |
|
|
| @patch("agents.rag_graph.build_rag_graph") |
| def test_run_rag_graph_answer_returns_answer_result_shape(self, mock_build_graph): |
| fake_graph = mock_build_graph.return_value |
| fake_graph.invoke.return_value = { |
| "answer": "Graph answer [1]", |
| "sources": [{"chunk_id": "c1"}], |
| "citations": [{"number": 1}], |
| "debug_data": {"grounding": {"reason": "answer_is_grounded"}}, |
| } |
|
|
| result = run_rag_graph_answer( |
| "What is this roadmap about?", |
| vectorstore=object(), |
| chunk_registry={"by_chunk_id": {}}, |
| reranker=object(), |
| bm25_index=object(), |
| llm=object(), |
| retrieval_k=4, |
| rerank_candidate_k=8, |
| bm25_candidate_k=8, |
| context_window=1, |
| max_expanded_chunks=12, |
| min_grounded_rerank_score=1.0, |
| min_grounded_chunks=1, |
| grounded_fallback_message="fallback", |
| enable_query_transform=True, |
| selected_file="roadmap.pdf", |
| debug_mode=True, |
| ) |
|
|
| self.assertEqual(result.answer, "Graph answer [1]") |
| self.assertEqual(result.sources, [{"chunk_id": "c1"}]) |
| self.assertEqual(result.citations, [{"number": 1}]) |
| self.assertEqual(result.debug_data["pipeline_mode"], "langgraph_rag") |
| self.assertEqual(result.debug_data["grounding"]["reason"], "answer_is_grounded") |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
