Open
Description
Environment details
- OS type and version: macOS Sonoma Version 14.3.1
- Python version: Python 3.12.9
- pip version: pip 25.1
google-cloud-aiplatformversion: 1.91.0
Steps to reproduce
- Do vertex.init with custom metadata
- Create RAG Corpus
Code example
def create_rag_corpus(
display_name: str,
description: str,
embedding_model: str,
) -> rag.RagCorpus:
embedding_model_config = rag.EmbeddingModelConfig(publisher_model=embedding_model)
corpus = rag.create_corpus(
display_name=display_name,
description=description,
embedding_model_config=embedding_model_config,
)
return corpus
vertexai_metadata = [
("client-id", CLIENT_ID),
("client-secret", CLIENT_SECRET),
("user-agent", "test"),
]
# Initialize vertexai with the correct configuration
vertexai.init(
project=GOOGLE_CLOUD_PROJECT,
location=GOOGLE_CLOUD_LOCATION,
request_metadata=vertexai_metadata,
api_endpoint="http://url.com",
api_transport="rest",
)Issue
While initializing Vertex AI with custom metadata, necessary for routing API calls through a proxy, I noticed that although the RAG corpus is created successfully, it appears to get stuck when polling for the status of the creation request.
Reason
This seems to be due to how the Operation future is currently being constructed in VertexRagDataServiceClient. At the moment, the gRPC metadata isn't being passed into the following code:
response = gac_operation.from_gapic(
response,
self._transport.operations_client,
vertex_rag_data.RagCorpus,
metadata_type=vertex_rag_data_service.CreateRagCorpusOperationMetadata,
)Resolution
A small tweak here resolves the issue nicely:
response = gac_operation.from_gapic(
response,
self._transport.operations_client,
vertex_rag_data.RagCorpus,
grpc_metadata=metadata, # Added
metadata_type=vertex_rag_data_service.CreateRagCorpusOperationMetadata,
)A similar fix would be needed for other VertexRagDataServiceClient operations like delete_rag_corpus, update_rag_corpus, etc.
I’d be happy to open a PR with the changes, if the changes seems fine.
Thanks!