Fixing _InactiveRpcError in MatchingEngine vectorstore (#10056)

- Description: There was an issue with the MatchingEngine VectorStore,
preventing from using it with a public endpoint. In the Google Cloud
library there are two similar methods for private or public endpoints :
`match()` and `find_neighbors()`.
  - Issue: Fixes #8378 
- This uses the `google.cloud.aiplatform` library :
https://github.com/googleapis/python-aiplatform/blob/main/google/cloud/aiplatform/matching_engine/matching_engine_index_endpoint.py
This commit is contained in:
Maxime Bourliatoux 2023-09-20 01:16:04 +02:00 committed by GitHub
parent 4f19ba3065
commit 21b236e5e4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -174,6 +174,14 @@ class MatchingEngine(VectorStore):
logger.debug(f"Embedding query {query}.") logger.debug(f"Embedding query {query}.")
embedding_query = self.embedding.embed_documents([query]) embedding_query = self.embedding.embed_documents([query])
# If the endpoint is public we use the find_neighbors function.
if self.endpoint._public_match_client:
response = self.endpoint.find_neighbors(
deployed_index_id=self._get_index_id(),
queries=embedding_query,
num_neighbors=k,
)
else:
response = self.endpoint.match( response = self.endpoint.match(
deployed_index_id=self._get_index_id(), deployed_index_id=self._get_index_id(),
queries=embedding_query, queries=embedding_query,