feat: make fetch timeout for chunks configurable (#1031)

* tests: processor accepts default and overrides of fetch_chunk_timeout

* feat: configure fetch timeout for chunks from librarian

tests/unit/test_retrieval/test_document_rag_service.py

@@ -131,3 +131,72 @@ class TestDocumentRagService:
         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
+
+    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,
+        )

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'],