googleapis / googleapis/python-aiplatform

gRPC Metadata not being passed in VertexRagDataServiceClient when polling for LRO

Abierto
#5,286 0 comentarios 1 reacción 0 asignados Ver en GitHub
api: vertex-ai
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

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.