From 5aabd30b88de06f1c4d73cfd3e6d02bf4f73c943 Mon Sep 17 00:00:00 2001 From: A Vertex SDK engineer Date: Thu, 14 May 2026 16:52:02 -0700 Subject: [PATCH] fix: Remove SDK workaround for async_retrieve_contexts. PiperOrigin-RevId: 915686364 --- vertexai/preview/rag/rag_retrieval.py | 19 +------------------ 1 file changed, 1 insertion(+), 18 deletions(-) diff --git a/vertexai/preview/rag/rag_retrieval.py b/vertexai/preview/rag/rag_retrieval.py index d3d879509a..0b60a2e668 100644 --- a/vertexai/preview/rag/rag_retrieval.py +++ b/vertexai/preview/rag/rag_retrieval.py @@ -25,8 +25,6 @@ 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, @@ -530,22 +528,7 @@ async def async_retrieve_contexts( response_lro = await client.async_retrieve_contexts( request=request, timeout=timeout ) - try: - response = await response_lro.result(timeout=timeout) - 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 + response = await response_lro.result(timeout=timeout) except Exception as e: raise RuntimeError( "Failed in retrieving contexts asynchronously due to: ", e