diff --git a/tests/unit/test_retrieval/test_document_rag_service.py b/tests/unit/test_retrieval/test_document_rag_service.py index 2bdf3959..7d1b5637 100644 --- a/tests/unit/test_retrieval/test_document_rag_service.py +++ b/tests/unit/test_retrieval/test_document_rag_service.py @@ -130,4 +130,73 @@ class TestDocumentRagService: assert sent_response.response == "A document about cats." assert sent_response.end_of_stream is True, "Non-streaming response must have end_of_stream=True" assert sent_response.end_of_session is True - assert sent_response.error is None \ No newline at end of file + assert sent_response.error is None + + def test_fetch_chunk_timeout_defaults_to_120(self): + """Processor without fetch_chunk_timeout override uses default of 120.""" + processor = Processor( + taskgroup=MagicMock(), + id="test-processor", + doc_limit=10 + ) + assert processor.fetch_chunk_timeout == 120 + + def test_fetch_chunk_timeout_uses_overridden_value(self): + """Processor with fetch_chunk_timeout override stores that value.""" + processor = Processor( + taskgroup=MagicMock(), + id="test-processor", + doc_limit=10, + fetch_chunk_timeout=45 + ) + assert processor.fetch_chunk_timeout == 45 + + @patch('trustgraph.retrieval.document_rag.rag.DocumentRag') + @pytest.mark.asyncio + async def test_fetch_chunk_uses_configured_timeout(self, mock_document_rag_class): + """ + Test that the fetch_chunk closure built in on_request calls + flow.librarian.fetch_document_text with the configured + fetch_chunk_timeout, not the old hardcoded 120. + """ + processor = Processor( + taskgroup=MagicMock(), + id="test-processor", + doc_limit=10, + fetch_chunk_timeout=45 + ) + + mock_rag_instance = AsyncMock() + mock_document_rag_class.return_value = mock_rag_instance + mock_rag_instance.query.return_value = ( + "test response", {"in_token": None, "out_token": None, "model": None}) + + msg = MagicMock() + msg.value.return_value = DocumentRagQuery( + query="test query", + collection="default", + doc_limit=5 + ) + msg.properties.return_value = {"id": "test-id"} + + consumer = MagicMock() + flow = MagicMock() + flow.librarian.fetch_document_text = AsyncMock(return_value="chunk text") + + mock_producer = AsyncMock() + + def flow_router(service_name): + if service_name == "response": + return mock_producer + return AsyncMock() + flow.side_effect = flow_router + + await processor.on_request(msg, consumer, flow) + + # Retrieve the fetch_chunk callable that on_request built and passed into DocumentRag(...) + fetch_chunk = mock_document_rag_class.call_args.kwargs["fetch_chunk"] + await fetch_chunk("some-chunk-id") + + flow.librarian.fetch_document_text.assert_called_once_with( + document_id="some-chunk-id", timeout=45, + ) \ No newline at end of file diff --git a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py b/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py index f2c128db..d8ee661b 100755 --- a/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py +++ b/trustgraph-flow/trustgraph/retrieval/document_rag/rag.py @@ -36,6 +36,7 @@ class Processor(FlowProcessor): fetch_limit = params.get("fetch_limit", 0) rerank_diversity_mode = params.get("rerank_diversity_mode", "none") rerank_diversity_lambda = params.get("rerank_diversity_lambda", 0.7) + fetch_chunk_timeout = params.get("fetch_chunk_timeout", 120) retrieval_mode = params.get("retrieval_mode", "vector") vector_weight = params.get("vector_weight", 1.0) keyword_weight = params.get("keyword_weight", 1.0) @@ -47,6 +48,7 @@ class Processor(FlowProcessor): "fetch_limit": fetch_limit, "rerank_diversity_mode": rerank_diversity_mode, "rerank_diversity_lambda": rerank_diversity_lambda, + "fetch_chunk_timeout": fetch_chunk_timeout, "retrieval_mode": retrieval_mode, "vector_weight": vector_weight, "keyword_weight": keyword_weight, @@ -57,6 +59,7 @@ class Processor(FlowProcessor): self.fetch_limit = fetch_limit self.rerank_diversity_mode = rerank_diversity_mode self.rerank_diversity_lambda = rerank_diversity_lambda + self.fetch_chunk_timeout = fetch_chunk_timeout self.retrieval_mode = retrieval_mode self.vector_weight = vector_weight self.keyword_weight = keyword_weight @@ -139,7 +142,7 @@ class Processor(FlowProcessor): logger.info(f"Handling input {id}...") - async def fetch_chunk(chunk_id, timeout=120): + async def fetch_chunk(chunk_id, timeout=self.fetch_chunk_timeout): return await flow.librarian.fetch_document_text( document_id=chunk_id, timeout=timeout, ) @@ -329,6 +332,14 @@ class Processor(FlowProcessor): help='MMR relevance/diversity tradeoff, higher values prefer relevance' ) + parser.add_argument( + '--fetch-chunk-timeout', + type=int, + default=120, + help='Timeout in seconds for fetching a document chunk from the ' + 'librarian (default: 120)' + ) + parser.add_argument( '--retrieval-mode', choices=['vector', 'keyword', 'hybrid'],