fix: Add SDK workaround for double-wrapped Any response in async_retrieve_contexts.

PiperOrigin-RevId: 896867061

Author: vertex-sdk-bot

vertexai/preview/rag/rag_retrieval.py

@@ -25,7 +25,7 @@
 from vertexai.preview.rag.utils import _gapic_utils
 from vertexai.preview.rag.utils import resources
 
-from google.protobuf import any_pb2
+
 
 
 def retrieval_query(
@@ -537,14 +537,15 @@ async def async_retrieve_contexts(
                 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
-                            )
+                    # HACK: Override type_url to match SDK's expected type
+                    any_response.type_url = "type.googleapis.com/google.cloud.aiplatform.v1beta1.RagContexts"
+                    rag_contexts = aiplatform_v1beta1.RagContexts()
+                    if any_response.Unpack(rag_contexts._pb):
+                        if not rag_contexts.contexts:
+                            raise ValueError("No rag contexts were returned.")
+                        return aiplatform_v1beta1.AsyncRetrieveContextsResponse(
+                            contexts=rag_contexts
+                        )
             raise e
     except Exception as e:
         raise RuntimeError(