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,
|
"top_k": top_k,
|
||||||
"filter": filters_str or "",
|
"filter": filters_str or "",
|
||||||
"include_metadata": True,
|
"include_metadata": True,
|
||||||
"namespace": self.collection_name,
|
|
||||||
}
|
}
|
||||||
for v in vectors
|
for v in vectors
|
||||||
]
|
]
|
||||||
responses = self.client.query_many(queries=queries)
|
responses = self.client.query_many(queries=queries, namespace=self.collection_name)
|
||||||
# flatten
|
# flatten
|
||||||
response = [res for res_list in responses for res in res_list]
|
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],
|
"vector": vectors[0],
|
||||||
"top_k": 2,
|
"top_k": 2,
|
||||||
"namespace": "ns",
|
|
||||||
"include_metadata": True,
|
"include_metadata": True,
|
||||||
"filter": 'age = 30 AND name = "John"',
|
"filter": 'age = 30 AND name = "John"',
|
||||||
}
|
}
|
||||||
]
|
],
|
||||||
|
namespace="ns",
|
||||||
)
|
)
|
||||||
|
|
||||||
assert len(results) == 2
|
assert len(results) == 2
|
||||||
@@ -311,11 +311,11 @@ def test_search_vectors_empty_filters(upstash_instance):
|
|||||||
{
|
{
|
||||||
"vector": vectors[0],
|
"vector": vectors[0],
|
||||||
"top_k": 1,
|
"top_k": 1,
|
||||||
"namespace": "ns",
|
|
||||||
"include_metadata": True,
|
"include_metadata": True,
|
||||||
"filter": "",
|
"filter": "",
|
||||||
}
|
}
|
||||||
]
|
],
|
||||||
|
namespace="ns",
|
||||||
)
|
)
|
||||||
|
|
||||||
assert len(results) == 1
|
assert len(results) == 1
|
||||||
@@ -350,3 +350,46 @@ def test_insert_vectors_no_ids(upstash_instance):
|
|||||||
],
|
],
|
||||||
namespace="ns",
|
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