Skip to content

Commit 23fe6e9

Browse files
vertex-sdk-botcopybara-github
authored andcommitted
fix: Add SDK workaround for double-wrapped Any response in async_retrieve_contexts.
PiperOrigin-RevId: 896867061
1 parent 9ed3759 commit 23fe6e9

1 file changed

Lines changed: 1 addition & 16 deletions

File tree

vertexai/preview/rag/rag_retrieval.py

Lines changed: 1 addition & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -530,22 +530,7 @@ async def async_retrieve_contexts(
530530
response_lro = await client.async_retrieve_contexts(
531531
request=request, timeout=timeout
532532
)
533-
try:
534-
response = await response_lro.result(timeout=timeout)
535-
except Exception as e:
536-
if response_lro.done():
537-
raw_op = response_lro.operation
538-
if raw_op.WhichOneof("result") == "response":
539-
any_response = raw_op.response
540-
inner_any = any_pb2.Any()
541-
if any_response.Unpack(inner_any):
542-
inner_any.type_url = "type.googleapis.com/google.cloud.aiplatform.v1beta1.RagContexts"
543-
rag_contexts = aiplatform_v1beta1.RagContexts()
544-
if inner_any.Unpack(rag_contexts._pb):
545-
return aiplatform_v1beta1.AsyncRetrieveContextsResponse(
546-
contexts=rag_contexts
547-
)
548-
raise e
533+
response = await response_lro.result(timeout=timeout)
549534
except Exception as e:
550535
raise RuntimeError(
551536
"Failed in retrieving contexts asynchronously due to: ", e

0 commit comments

Comments
 (0)