googleapis / googleapis/python-aiplatform
gRPC Metadata not being passed in VertexRagDataServiceClient when polling for LRO
- Lenguaje dominante
- Python
- Estrellas
- 905
- Forks
- 465
- Merge medio
- 1 d 13 h
- PR fusionados (30 d)
- 44
Descripción
#### 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-aiplatform` version: 1.91.0
#### Steps to reproduce
1. Do vertex.init with custom metadata
2. Create RAG Corpus
#### Code example
```python
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:
```python
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:
```python
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!
Guía de contribución
Evaluación
Este issue todavía no se ha evaluado.