fix: Increase default timeout to 300 seconds for ask_contexts and async_retrieve_contexts in VertexRagServiceClient.

PiperOrigin-RevId: 892413253

Author: vertex-sdk-bot

tests/unit/vertex_rag/test_rag_retrieval.py

@@ -113,6 +113,19 @@ def test_ask_contexts_rag_resources_success(self):
         )
         retrieve_contexts_eq(response, tc.TEST_RETRIEVAL_RESPONSE)
 
+    @pytest.mark.usefixtures("ask_contexts_mock")
+    def test_ask_contexts_with_timeout(self, ask_contexts_mock):
+        response = rag.ask_contexts(
+            rag_resources=[tc.TEST_RAG_RESOURCE],
+            text=tc.TEST_QUERY_TEXT,
+            rag_retrieval_config=tc.TEST_RAG_RETRIEVAL_CONFIG,
+            timeout=300,
+        )
+        ask_contexts_mock.assert_called_once()
+        args, kwargs = ask_contexts_mock.call_args
+        assert kwargs["timeout"] == 300
+
+
     @pytest.mark.usefixtures("ask_contexts_mock")
     def test_ask_contexts_multiple_rag_resources_success(self):
         response = rag.ask_contexts(
@@ -132,6 +145,22 @@ async def test_async_retrieve_contexts_rag_resources_success(self):
         )
         retrieve_contexts_eq(response, tc.TEST_RETRIEVAL_RESPONSE)
 
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures("async_retrieve_contexts_mock")
+    async def test_async_retrieve_contexts_with_timeout(
+        self, async_retrieve_contexts_mock
+    ):
+        response = await rag.async_retrieve_contexts(
+            rag_resources=[tc.TEST_RAG_RESOURCE],
+            text=tc.TEST_QUERY_TEXT,
+            rag_retrieval_config=tc.TEST_RAG_RETRIEVAL_CONFIG,
+            timeout=300,
+        )
+        async_retrieve_contexts_mock.assert_called_once()
+        args, kwargs = async_retrieve_contexts_mock.call_args
+        assert kwargs["timeout"] == 300
+
+
     @pytest.mark.asyncio
     @pytest.mark.usefixtures("async_retrieve_contexts_mock")
     async def test_async_retrieve_contexts_multiple_rag_resources_success(self):

tests/unit/vertex_rag/test_rag_retrieval_preview.py

@@ -118,6 +118,19 @@ def test_ask_contexts_rag_resources_success(self):
         )
         retrieve_contexts_eq(response, tc.TEST_RETRIEVAL_RESPONSE)
 
+    @pytest.mark.usefixtures("ask_contexts_mock")
+    def test_ask_contexts_with_timeout(self, ask_contexts_mock):
+        response = rag.ask_contexts(
+            rag_resources=[tc.TEST_RAG_RESOURCE],
+            text=tc.TEST_QUERY_TEXT,
+            rag_retrieval_config=tc.TEST_RAG_RETRIEVAL_CONFIG_ALPHA,
+            timeout=300,
+        )
+        ask_contexts_mock.assert_called_once()
+        args, kwargs = ask_contexts_mock.call_args
+        assert kwargs["timeout"] == 300
+
+
     @pytest.mark.usefixtures("ask_contexts_mock")
     def test_ask_contexts_multiple_rag_resources_success(self):
         response = rag.ask_contexts(
@@ -147,6 +160,22 @@ async def test_async_retrieve_contexts_rag_resources_success(self):
         )
         retrieve_contexts_eq(response, tc.TEST_RETRIEVAL_RESPONSE)
 
+    @pytest.mark.asyncio
+    @pytest.mark.usefixtures("async_retrieve_contexts_mock")
+    async def test_async_retrieve_contexts_with_timeout(
+        self, async_retrieve_contexts_mock
+    ):
+        response = await rag.async_retrieve_contexts(
+            rag_resources=[tc.TEST_RAG_RESOURCE],
+            text=tc.TEST_QUERY_TEXT,
+            rag_retrieval_config=tc.TEST_RAG_RETRIEVAL_CONFIG_ALPHA,
+            timeout=300,
+        )
+        async_retrieve_contexts_mock.assert_called_once()
+        args, kwargs = async_retrieve_contexts_mock.call_args
+        assert kwargs["timeout"] == 300
+
+
     @pytest.mark.asyncio
     @pytest.mark.usefixtures("async_retrieve_contexts_mock")
     async def test_async_retrieve_contexts_multiple_rag_resources_success(self):

vertexai/preview/rag/rag_retrieval.py

@@ -290,6 +290,7 @@ async def async_retrieve_contexts(
     vector_distance_threshold: Optional[float] = None,
     vector_search_alpha: Optional[float] = None,
     rag_retrieval_config: Optional[resources.RagRetrievalConfig] = None,
+    timeout: int = 600,
 ) -> aiplatform_v1beta1.RetrieveContextsResponse:
     """Retrieve top k relevant docs/chunks asynchronously.
 
@@ -523,7 +524,9 @@ async def async_retrieve_contexts(
         tools=[tool],
     )
     try:
-        response_lro = await client.async_retrieve_contexts(request=request)
+        response_lro = await client.async_retrieve_contexts(
+            request=request, timeout=timeout
+        )
         response = await response_lro.result()
     except Exception as e:
         raise RuntimeError(
@@ -541,6 +544,7 @@ def ask_contexts(
     vector_distance_threshold: Optional[float] = None,
     vector_search_alpha: Optional[float] = None,
     rag_retrieval_config: Optional[resources.RagRetrievalConfig] = None,
+    timeout: int = 600,
 ) -> aiplatform_v1beta1.AskContextsResponse:
     """Ask questions on top k relevant docs/chunks.
 
@@ -774,7 +778,7 @@ def ask_contexts(
         tools=[tool],
     )
     try:
-        response = client.ask_contexts(request=request)
+        response = client.ask_contexts(request=request, timeout=timeout)
     except Exception as e:
         raise RuntimeError("Failed in asking contexts due to: ", e) from e