googleapis / googleapis/python-aiplatform

gRPC Metadata not being passed in VertexRagDataServiceClient when polling for LRO

Ouverte
#5,286 0 commentaires 1 réaction 0 personnes assignées Voir sur GitHub
api: vertex-ai
Langage dominant
Python
Étoiles
905
Forks
465
Merge moyen
1 j 13 h
PR mergées (30 j)
44

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-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!

Guide de contribution

Ouvrir le guide de contribution

Évaluation

Cette issue n'a pas encore été évaluée.

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.