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,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,
@@ -537,14 +535,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(
@@ -564,7 +563,7 @@ def ask_contexts(
     rag_retrieval_config: Optional[resources.RagRetrievalConfig] = None,
     timeout: int = 600,
 ) -> aiplatform_v1beta1.AskContextsResponse:
-    """Ask questions on top k relevant docs/chunks.
+  """Ask questions on top k relevant docs/chunks.
 
     Example usage:
     ```
@@ -610,99 +609,99 @@ def ask_contexts(
     Returns:
         AskContextsResponse.
     """
-    parent = initializer.global_config.common_location_path()
+  parent = initializer.global_config.common_location_path()
 
-    client = _gapic_utils.create_rag_service_client()
+  client = _gapic_utils.create_rag_service_client()
 
-    if not rag_resources and not rag_corpora:
-        raise ValueError("rag_resources or rag_corpora must be specified.")
+  if not rag_resources and not rag_corpora:
+    raise ValueError("rag_resources or rag_corpora must be specified.")
 
-    data_client = _gapic_utils.create_rag_data_service_client()
+  data_client = _gapic_utils.create_rag_data_service_client()
 
-    gapic_rag_resources = []
-    gapic_rag_corpora = []
-    if rag_resources:
-        for rag_resource in rag_resources:
-            name = rag_resource.rag_corpus
-            if data_client.parse_rag_corpus_path(name):
-                rag_corpus_name = name
-            elif re.match(
+  gapic_rag_resources = []
+  gapic_rag_corpora = []
+  if rag_resources:
+    for rag_resource in rag_resources:
+      name = rag_resource.rag_corpus
+      if data_client.parse_rag_corpus_path(name):
+        rag_corpus_name = name
+      elif re.match(
                 "^{}$".format(
                     _gapic_utils._VALID_RESOURCE_NAME_REGEX  # pylint: disable=protected-access
                 ),
                 name,
             ):
-                rag_corpus_name = parent + "/ragCorpora/" + name
-            else:
-                raise ValueError(
+        rag_corpus_name = parent + "/ragCorpora/" + name
+      else:
+        raise ValueError(
                     f"Invalid RagCorpus name: {name}. Proper format should be:"
                     " projects/{project}/locations/{location}/ragCorpora/{rag_corpus_id}"
                 )
-            gapic_rag_resources.append(
+      gapic_rag_resources.append(
                 aiplatform_v1beta1.VertexRagStore.RagResource(
                     rag_corpus=rag_corpus_name,
                     rag_file_ids=rag_resource.rag_file_ids,
                 )
             )
-        vertex_rag_store = aiplatform_v1beta1.VertexRagStore(
+    vertex_rag_store = aiplatform_v1beta1.VertexRagStore(
             rag_resources=gapic_rag_resources,
         )
-    elif rag_corpora:
-        warnings.warn(
+  elif rag_corpora:
+    warnings.warn(
             "rag_corpora is deprecated. Please use rag_resources instead."
             f" After {resources.DEPRECATION_DATE} using"
             " rag_corpora will raise error",
             DeprecationWarning,
         )
-        for name in rag_corpora:
-            if data_client.parse_rag_corpus_path(name):
-                rag_corpus_name = name
-            elif re.match(
+    for name in rag_corpora:
+      if data_client.parse_rag_corpus_path(name):
+        rag_corpus_name = name
+      elif re.match(
                 "^{}$".format(
                     _gapic_utils._VALID_RESOURCE_NAME_REGEX  # pylint: disable=protected-access
                 ),
                 name,
             ):
-                rag_corpus_name = parent + "/ragCorpora/" + name
-            else:
-                raise ValueError(
+        rag_corpus_name = parent + "/ragCorpora/" + name
+      else:
+        raise ValueError(
                     f"Invalid RagCorpus name: {name}. Proper format should be:"
                     " projects/{project}/locations/{location}/ragCorpora/{rag_corpus_id}"
                 )
-            gapic_rag_corpora.append(rag_corpus_name)
-        vertex_rag_store = aiplatform_v1beta1.VertexRagStore(
+      gapic_rag_corpora.append(rag_corpus_name)
+    vertex_rag_store = aiplatform_v1beta1.VertexRagStore(
             rag_corpora=gapic_rag_corpora,
         )
 
-    # Check for deprecated parameters and raise warnings.
-    if similarity_top_k:
-        warnings.warn(
+  # Check for deprecated parameters and raise warnings.
+  if similarity_top_k:
+    warnings.warn(
             "similarity_top_k is deprecated. Please use"
             " rag_retrieval_config.top_k instead."
             f" After {resources.DEPRECATION_DATE} using"
             " similarity_top_k will raise error",
             DeprecationWarning,
         )
-    if vector_search_alpha:
-        warnings.warn(
+  if vector_search_alpha:
+    warnings.warn(
             "vector_search_alpha is deprecated. Please use"
             " rag_retrieval_config.alpha instead."
             f" After {resources.DEPRECATION_DATE} using"
             " vector_search_alpha will raise error",
             DeprecationWarning,
         )
-    if vector_distance_threshold:
-        warnings.warn(
+  if vector_distance_threshold:
+    warnings.warn(
             "vector_distance_threshold is deprecated. Please use"
             " rag_retrieval_config.filter.vector_distance_threshold instead."
             f" After {resources.DEPRECATION_DATE} using"
             " vector_distance_threshold will raise error",
             DeprecationWarning,
         )
 
-    # If rag_retrieval_config is not specified, set it to default values.
-    if not rag_retrieval_config:
-        api_retrival_config = aiplatform_v1beta1.RagRetrievalConfig(
+  # If rag_retrieval_config is not specified, set it to default values.
+  if not rag_retrieval_config:
+    api_retrival_config = aiplatform_v1beta1.RagRetrievalConfig(
             top_k=similarity_top_k,
             hybrid_search=aiplatform_v1beta1.RagRetrievalConfig.HybridSearch(
                 alpha=vector_search_alpha,
@@ -711,94 +710,101 @@ def ask_contexts(
                 vector_distance_threshold=vector_distance_threshold
             ),
         )
+  else:
+    api_retrival_config = aiplatform_v1beta1.RagRetrievalConfig()
+    if rag_retrieval_config.top_k:
+      api_retrival_config.top_k = rag_retrieval_config.top_k
     else:
-        api_retrival_config = aiplatform_v1beta1.RagRetrievalConfig()
-        if rag_retrieval_config.top_k:
-            api_retrival_config.top_k = rag_retrieval_config.top_k
-        else:
-            api_retrival_config.top_k = similarity_top_k
+      api_retrival_config.top_k = similarity_top_k
 
-        if (
+    if (
             rag_retrieval_config.hybrid_search
             and rag_retrieval_config.hybrid_search.alpha
         ):
-            api_retrival_config.hybrid_search.alpha = (
+      api_retrival_config.hybrid_search.alpha = (
                 rag_retrieval_config.hybrid_search.alpha
             )
-        else:
-            api_retrival_config.hybrid_search.alpha = vector_search_alpha
+    else:
+      api_retrival_config.hybrid_search.alpha = vector_search_alpha
 
-        if (
+    if (
             rag_retrieval_config.filter
             and rag_retrieval_config.filter.vector_distance_threshold
             and rag_retrieval_config.filter.vector_similarity_threshold
         ):
-            raise ValueError(
+      raise ValueError(
                 "Only one of vector_distance_threshold or"
                 " vector_similarity_threshold can be specified at a time"
                 " in rag_retrieval_config."
             )
 
-        if (
+    if (
             rag_retrieval_config.filter
             and rag_retrieval_config.filter.vector_distance_threshold
         ):
-            api_retrival_config.filter.vector_distance_threshold = (
+      api_retrival_config.filter.vector_distance_threshold = (
                 rag_retrieval_config.filter.vector_distance_threshold
             )
-        else:
-            api_retrival_config.filter.vector_distance_threshold = (
+    else:
+      api_retrival_config.filter.vector_distance_threshold = (
                 vector_distance_threshold
             )
 
-        if (
+    if (
             rag_retrieval_config.filter
             and rag_retrieval_config.filter.vector_similarity_threshold
         ):
-            api_retrival_config.filter.vector_similarity_threshold = (
+      api_retrival_config.filter.vector_similarity_threshold = (
                 rag_retrieval_config.filter.vector_similarity_threshold
             )
 
-        if (
+    if (
             rag_retrieval_config.ranking
             and rag_retrieval_config.ranking.rank_service
             and rag_retrieval_config.ranking.llm_ranker
         ):
-            raise ValueError("Only one of rank_service and llm_ranker can be set.")
-        if rag_retrieval_config.ranking and rag_retrieval_config.ranking.rank_service:
-            api_retrival_config.ranking.rank_service.model_name = (
+      raise ValueError("Only one of rank_service and llm_ranker can be set.")
+    if rag_retrieval_config.ranking and rag_retrieval_config.ranking.rank_service:
+      api_retrival_config.ranking.rank_service.model_name = (
                 rag_retrieval_config.ranking.rank_service.model_name
             )
-        elif rag_retrieval_config.ranking and rag_retrieval_config.ranking.llm_ranker:
-            api_retrival_config.ranking.llm_ranker.model_name = (
+    elif rag_retrieval_config.ranking and rag_retrieval_config.ranking.llm_ranker:
+      api_retrival_config.ranking.llm_ranker.model_name = (
                 rag_retrieval_config.ranking.llm_ranker.model_name
             )
-        if rag_retrieval_config.filter and rag_retrieval_config.filter.metadata_filter:
-            api_retrival_config.filter.metadata_filter = (
+    if rag_retrieval_config.filter and rag_retrieval_config.filter.metadata_filter:
+      api_retrival_config.filter.metadata_filter = (
                 rag_retrieval_config.filter.metadata_filter
             )
 
-    query = aiplatform_v1beta1.RagQuery(
+  query = aiplatform_v1beta1.RagQuery(
         text=text,
         rag_retrieval_config=api_retrival_config,
     )
 
-    vertex_rag_store.rag_retrieval_config = api_retrival_config
+  vertex_rag_store.rag_retrieval_config = api_retrival_config
 
-    tool = aiplatform_v1beta1.Tool(
+  tool = aiplatform_v1beta1.Tool(
         retrieval=aiplatform_v1beta1.Retrieval(
             vertex_rag_store=vertex_rag_store,
         )
     )
 
-    request = aiplatform_v1beta1.AskContextsRequest(
+  request = aiplatform_v1beta1.AskContextsRequest(
         parent=parent,
         query=query,
         tools=[tool],
     )
-    try:
-        response = client.ask_contexts(request=request, timeout=timeout)
-    except Exception as e:
-        raise RuntimeError("Failed in asking contexts due to: ", e) from e
-
-    return response
+  try:
+    response = client.ask_contexts(request=request, timeout=timeout)
+  except Exception as e:
+    print(f"DEBUG: ask_contexts failed with error: {e}")
+    if hasattr(e, "trailing_metadata"):
+      metadata = e.trailing_metadata()
+      print(f"DEBUG: trailing metadata: {metadata}")
+      for key, value in metadata:
+        if key == "x-google-trace-id" or "trace" in key:
+          print(f"DEBUG: Found trace header: {key}={value}")
+    raise RuntimeError("Failed in asking contexts due to: ", e) from e
+
+  return response