googleapis / googleapis/python-aiplatform

gRPC Metadata not being passed in VertexRagDataServiceClient when polling for LRO

Đang mở
#5,286 0 bình luận 1 reaction 0 người được giao Xem trên GitHub
api: vertex-ai
Ngôn ngữ chính
Python
Star
905
Fork
465
Merge trung bình
1 ngày 13 giờ
Pull request đã merge (30 ngày)
44

Mô tả

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

Hướng dẫn đóng góp

Mở hướng dẫn đóng góp

Đánh giá

Issue này chưa được đánh giá.

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.