fix(vector_stores): pass namespace as top-level kwarg to UpstashVector query_many (#5202)
This commit is contained in:
@@ -132,11 +132,10 @@ class UpstashVector(VectorStoreBase):
|
||||
"top_k": top_k,
|
||||
"filter": filters_str or "",
|
||||
"include_metadata": True,
|
||||
"namespace": self.collection_name,
|
||||
}
|
||||
for v in vectors
|
||||
]
|
||||
responses = self.client.query_many(queries=queries)
|
||||
responses = self.client.query_many(queries=queries, namespace=self.collection_name)
|
||||
# flatten
|
||||
response = [res for res_list in responses for res in res_list]
|
||||
|
||||
|
||||
@@ -69,11 +69,11 @@ def test_search_vectors(upstash_instance, mock_index):
|
||||
{
|
||||
"vector": vectors[0],
|
||||
"top_k": 2,
|
||||
"namespace": "ns",
|
||||
"include_metadata": True,
|
||||
"filter": 'age = 30 AND name = "John"',
|
||||
}
|
||||
]
|
||||
],
|
||||
namespace="ns",
|
||||
)
|
||||
|
||||
assert len(results) == 2
|
||||
@@ -311,11 +311,11 @@ def test_search_vectors_empty_filters(upstash_instance):
|
||||
{
|
||||
"vector": vectors[0],
|
||||
"top_k": 1,
|
||||
"namespace": "ns",
|
||||
"include_metadata": True,
|
||||
"filter": "",
|
||||
}
|
||||
]
|
||||
],
|
||||
namespace="ns",
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
@@ -350,3 +350,46 @@ def test_insert_vectors_no_ids(upstash_instance):
|
||||
],
|
||||
namespace="ns",
|
||||
)
|
||||
|
||||
|
||||
def test_search_vectors_multi_query_namespace_at_top_level(upstash_instance):
|
||||
"""Regression test for #4207.
|
||||
|
||||
The Upstash client's `query_many` takes `namespace` as a top-level kwarg
|
||||
and the per-query `QueryRequest` TypedDict has no `namespace` field. If we
|
||||
pass `namespace` inside each query dict it is silently dropped (multi-query
|
||||
branch) or causes `TypeError: got multiple values for keyword argument
|
||||
'namespace'` (single-query branch, where `query_many` internally calls
|
||||
`self.query(**query, namespace=namespace)`). This test locks in that
|
||||
`namespace` is forwarded only at the top level and that multi-query
|
||||
results are flattened across all per-query response lists.
|
||||
"""
|
||||
mock_results = [
|
||||
[QueryResult(id="id1", score=0.9, vector=None, metadata={"name": "vector1"}, data=None)],
|
||||
[
|
||||
QueryResult(id="id2", score=0.8, vector=None, metadata={"name": "vector2"}, data=None),
|
||||
QueryResult(id="id3", score=0.7, vector=None, metadata={"name": "vector3"}, data=None),
|
||||
],
|
||||
]
|
||||
upstash_instance.client.query_many.return_value = mock_results
|
||||
|
||||
vectors = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
||||
results = upstash_instance.search(query="hello world", vectors=vectors, top_k=5, filters=None)
|
||||
|
||||
upstash_instance.client.query_many.assert_called_once_with(
|
||||
queries=[
|
||||
{"vector": vectors[0], "top_k": 5, "include_metadata": True, "filter": ""},
|
||||
{"vector": vectors[1], "top_k": 5, "include_metadata": True, "filter": ""},
|
||||
],
|
||||
namespace="ns",
|
||||
)
|
||||
|
||||
# No per-query dict should carry a `namespace` key.
|
||||
sent_queries = upstash_instance.client.query_many.call_args.kwargs["queries"]
|
||||
for q in sent_queries:
|
||||
assert "namespace" not in q
|
||||
|
||||
# Results must be flattened across both per-query response lists.
|
||||
assert [r.id for r in results] == ["id1", "id2", "id3"]
|
||||
assert results[0].score == 0.9
|
||||
assert results[1].payload == {"name": "vector2"}
|
||||
|
||||
Reference in New Issue
Block a user