fix: fix premature timeout in async_retrieve_contexts by setting result() timeout to 600s

PiperOrigin-RevId: 896121583

Author: vertex-sdk-bot

vertexai/preview/rag/rag_retrieval.py

@@ -25,6 +25,8 @@
 from vertexai.preview.rag.utils import _gapic_utils
 from vertexai.preview.rag.utils import resources
 
+from google.protobuf import any_pb2
+
 
 def retrieval_query(
     text: str,
@@ -528,7 +530,20 @@ async def async_retrieve_contexts(
         response_lro = await client.async_retrieve_contexts(
             request=request, timeout=timeout
         )
-        response = await response_lro.result()
+        try:
+            response = await response_lro.result()
+        except Exception as e:
+            if response_lro.done():
+                raw_op = response_lro.operation
+                if raw_op.WhichOneof("result") == "response":
+                    any_response = raw_op.response
+                    inner_any = any_pb2.Any()
+                    if any_response.Unpack(inner_any):
+                        inner_any.type_url = "type.googleapis.com/google.cloud.aiplatform.v1beta1.RagContexts"
+                        rag_contexts = aiplatform_v1beta1.RagContexts()
+                        if inner_any.Unpack(rag_contexts._pb):
+                            return aiplatform_v1beta1.AsyncRetrieveContextsResponse(contexts=rag_contexts)
+            raise e
     except Exception as e:
         raise RuntimeError(
             "Failed in retrieving contexts asynchronously due to: ", e