fix(vertex): pass required vectors arg in list and similarity search (#5627)
This commit is contained in:
@@ -487,7 +487,7 @@ class GoogleMatchingEngine(VectorStoreBase):
|
|||||||
# Use a large top_k if none specified
|
# Use a large top_k if none specified
|
||||||
search_limit = top_k if top_k is not None else 10000
|
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))
|
logger.debug("Found %d results", len(results))
|
||||||
return [results] # Wrap in extra array to match interface
|
return [results] # Wrap in extra array to match interface
|
||||||
@@ -620,7 +620,7 @@ class GoogleMatchingEngine(VectorStoreBase):
|
|||||||
logger.debug("Filter: %s", filter)
|
logger.debug("Filter: %s", filter)
|
||||||
|
|
||||||
embedding = self.embedder.embed_query(query)
|
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 = [
|
docs_and_scores = [
|
||||||
(Document(page_content=result.payload.get("text", ""), metadata=result.payload), result.score)
|
(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"}
|
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):
|
def test_delete(vector_store, mock_vertex_ai):
|
||||||
"""Test deleting vectors"""
|
"""Test deleting vectors"""
|
||||||
vector_id = "test-id"
|
vector_id = "test-id"
|
||||||
|
|||||||
Reference in New Issue
Block a user