fix(vertex): pass required vectors arg in list and similarity search (#5627)

This commit is contained in:
Yash Singh
2026-06-19 12:44:48 +05:30
committed by GitHub
parent c5971193a2
commit f0ccd99924
2 changed files with 28 additions and 2 deletions
@@ -487,7 +487,7 @@ class GoogleMatchingEngine(VectorStoreBase):
# Use a large top_k if none specified
search_limit = top_k if top_k is not None else 10000
results = self.search(query=zero_vector, top_k=search_limit, filters=filters)
results = self.search(query="", vectors=zero_vector, top_k=search_limit, filters=filters)
logger.debug("Found %d results", len(results))
return [results] # Wrap in extra array to match interface
@@ -620,7 +620,7 @@ class GoogleMatchingEngine(VectorStoreBase):
logger.debug("Filter: %s", filter)
embedding = self.embedder.embed_query(query)
results = self.search(query=embedding, top_k=k, filters=filter)
results = self.search(query=query, vectors=embedding, top_k=k, filters=filter)
docs_and_scores = [
(Document(page_content=result.payload.get("text", ""), metadata=result.payload), result.score)
@@ -116,6 +116,32 @@ def test_search_vectors(vector_store, mock_vertex_ai):
assert results[0].payload == {"user_id": "test_user"}
def test_list_does_not_raise_type_error(vector_store, mock_vertex_ai):
"""list() must call search() with the required positional `vectors` arg."""
mock_vertex_ai["endpoint"].find_neighbors.return_value = [[]]
results = vector_store.list(filters={"user_id": "test_user"}, top_k=5)
mock_vertex_ai["endpoint"].find_neighbors.assert_called_once()
queries = mock_vertex_ai["endpoint"].find_neighbors.call_args[1]["queries"]
assert queries == [[0.0] * 768]
assert results == [[]]
def test_similarity_search_with_score_passes_embedding(vector_store, mock_vertex_ai):
"""similarity_search_with_score() must pass the embedding as `vectors`."""
embedding = [0.1, 0.2, 0.3]
vector_store.embedder = Mock()
vector_store.embedder.embed_query.return_value = embedding
mock_vertex_ai["endpoint"].find_neighbors.return_value = [[]]
vector_store.similarity_search_with_score(query="hello", k=3)
mock_vertex_ai["endpoint"].find_neighbors.assert_called_once()
queries = mock_vertex_ai["endpoint"].find_neighbors.call_args[1]["queries"]
assert queries == [embedding]
def test_delete(vector_store, mock_vertex_ai):
"""Test deleting vectors"""
vector_id = "test-id"