googleapis / googleapis/python-aiplatform

gRPC Metadata not being passed in VertexRagDataServiceClient when polling for LRO

Aperta
#5,286 0 commenti 1 reazione 0 assegnatari Vedi su GitHub
api: vertex-ai
Lingua principale
Python
Stelle
905
Fork
465
Merge medio
1g 13h
PR unite (30g)
44

Descrizione

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

Guida per i contributori

Apri la guida per i contributori

Valutazione

Questa issue non è ancora stata valutata.

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.