| import unittest |
| from unittest.mock import MagicMock, patch |
|
|
| from agents.rag_graph import ( |
| build_initial_graph_state, |
| decide_next_action_node, |
| refine_query_node, |
| retrieve_node, |
| ) |
|
|
|
|
| class RefineQueryNodeTests(unittest.TestCase): |
| @patch("agents.rag_graph.generate_sub_queries") |
| def test_stores_refined_query_from_shared_query_transform(self, mock_generate_sub_queries): |
| state = build_initial_graph_state("What is machine learning?") |
| mock_generate_sub_queries.return_value = ["machine learning definition and types"] |
|
|
| updated_state = refine_query_node(state, llm=MagicMock()) |
|
|
| self.assertEqual(updated_state["refined_query"], "machine learning definition and types") |
|
|
| @patch("agents.rag_graph.generate_sub_queries") |
| def test_increments_retry_count(self, mock_generate_sub_queries): |
| state = build_initial_graph_state("What is machine learning?") |
| mock_generate_sub_queries.return_value = ["machine learning definition"] |
|
|
| updated_state = refine_query_node(state, llm=MagicMock()) |
|
|
| self.assertEqual(updated_state["retry_count"], 1) |
|
|
| @patch("agents.rag_graph.generate_sub_queries") |
| def test_increments_retry_count_on_second_retry(self, mock_generate_sub_queries): |
| state = build_initial_graph_state("What is machine learning?") |
| state["retry_count"] = 1 |
| mock_generate_sub_queries.return_value = ["supervised learning algorithms"] |
|
|
| updated_state = refine_query_node(state, llm=MagicMock()) |
|
|
| self.assertEqual(updated_state["retry_count"], 2) |
|
|
| @patch("agents.rag_graph.generate_sub_queries") |
| def test_falls_back_to_original_query_when_query_transform_returns_empty(self, mock_generate_sub_queries): |
| state = build_initial_graph_state("What is machine learning?") |
| mock_generate_sub_queries.return_value = [] |
|
|
| updated_state = refine_query_node(state, llm=MagicMock()) |
|
|
| self.assertEqual(updated_state["refined_query"], "What is machine learning?") |
|
|
| @patch("agents.rag_graph.generate_sub_queries") |
| def test_records_refined_query_in_debug_data_when_debug_mode(self, mock_generate_sub_queries): |
| state = build_initial_graph_state("What is ML?", debug_mode=True) |
| state["debug_data"] = {"stage_counts": {}} |
| mock_generate_sub_queries.return_value = ["machine learning overview"] |
|
|
| updated_state = refine_query_node(state, llm=MagicMock()) |
|
|
| self.assertEqual(updated_state["debug_data"]["refined_query"], "machine learning overview") |
| self.assertEqual(updated_state["debug_data"]["retry_count"], 1) |
|
|
|
|
| class RetryCapTests(unittest.TestCase): |
| def _state_with_partial_evidence(self, retry_count): |
| state = build_initial_graph_state("What is deep learning?") |
| state["retry_count"] = retry_count |
| state["retrieved_documents"] = [MagicMock()] |
| state["expanded_documents"] = [MagicMock()] |
| state["grounding"] = {"passed": False, "reason": "low_rerank_score"} |
| return state |
|
|
| def test_retries_when_partial_evidence_and_under_cap(self): |
| state = self._state_with_partial_evidence(retry_count=0) |
|
|
| updated_state = decide_next_action_node(state) |
|
|
| self.assertEqual(updated_state["next_action"], "retry_retrieval") |
|
|
| def test_retries_when_partial_evidence_and_one_retry_done(self): |
| state = self._state_with_partial_evidence(retry_count=1) |
|
|
| updated_state = decide_next_action_node(state) |
|
|
| self.assertEqual(updated_state["next_action"], "retry_retrieval") |
|
|
| def test_falls_back_when_retry_cap_reached(self): |
| state = self._state_with_partial_evidence(retry_count=2) |
|
|
| updated_state = decide_next_action_node(state) |
|
|
| self.assertEqual(updated_state["next_action"], "fallback") |
| self.assertEqual(updated_state["decision_reason"], "local_retry_cap_reached_research_disabled") |
|
|
|
|
| class RetrieveNodeUsesRefinedQueryTests(unittest.TestCase): |
| def test_uses_refined_query_when_available(self): |
| state = build_initial_graph_state("original query") |
| state["refined_query"] = "refined query" |
|
|
| with patch("agents.rag_graph.retrieve_documents_with_query_transform") as mock_retrieve: |
| mock_retrieve.return_value = [] |
| retrieve_node( |
| state, |
| vectorstore=MagicMock(), |
| reranker=None, |
| bm25_index=None, |
| llm=MagicMock(), |
| retrieval_k=4, |
| rerank_candidate_k=8, |
| bm25_candidate_k=8, |
| enable_query_transform=False, |
| ) |
| used_query = mock_retrieve.call_args[0][1] |
| self.assertEqual(used_query, "refined query") |
| self.assertFalse(mock_retrieve.call_args.kwargs["enable_query_transform"]) |
|
|
| def test_uses_original_query_when_no_refined_query(self): |
| state = build_initial_graph_state("original query") |
|
|
| with patch("agents.rag_graph.retrieve_documents_with_query_transform") as mock_retrieve: |
| mock_retrieve.return_value = [] |
| retrieve_node( |
| state, |
| vectorstore=MagicMock(), |
| reranker=None, |
| bm25_index=None, |
| llm=MagicMock(), |
| retrieval_k=4, |
| rerank_candidate_k=8, |
| bm25_candidate_k=8, |
| enable_query_transform=False, |
| ) |
| used_query = mock_retrieve.call_args[0][1] |
| self.assertEqual(used_query, "original query") |
| self.assertFalse(mock_retrieve.call_args.kwargs["enable_query_transform"]) |
|
|
| def test_keeps_query_transform_enabled_when_no_refined_query_and_flag_is_true(self): |
| state = build_initial_graph_state("original query") |
|
|
| with patch("agents.rag_graph.retrieve_documents_with_query_transform") as mock_retrieve: |
| mock_retrieve.return_value = [] |
| retrieve_node( |
| state, |
| vectorstore=MagicMock(), |
| reranker=None, |
| bm25_index=None, |
| llm=MagicMock(), |
| retrieval_k=4, |
| rerank_candidate_k=8, |
| bm25_candidate_k=8, |
| enable_query_transform=True, |
| ) |
| self.assertTrue(mock_retrieve.call_args.kwargs["enable_query_transform"]) |
|
|
