Add Upstash Vector support (#2493)
This commit is contained in:
ytkimirti
384
tests/vector_stores/test_upstash_vector.py
Normal file
384
tests/vector_stores/test_upstash_vector.py
Normal file
@@ -0,0 +1,384 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Optional
|
||||
from unittest.mock import MagicMock, call, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from mem0.vector_stores.upstash_vector import UpstashVector
|
||||
|
||||
|
||||
@dataclass
|
||||
class QueryResult:
|
||||
id: str
|
||||
score: Optional[float]
|
||||
vector: Optional[List[float]] = None
|
||||
metadata: Optional[Dict] = None
|
||||
data: Optional[str] = None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_index():
|
||||
with patch("upstash_vector.Index") as mock_index:
|
||||
yield mock_index
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def upstash_instance(mock_index):
|
||||
return UpstashVector(client=mock_index.return_value, collection_name="ns")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def upstash_instance_with_embeddings(mock_index):
|
||||
return UpstashVector(
|
||||
client=mock_index.return_value, collection_name="ns", enable_embeddings=True
|
||||
)
|
||||
|
||||
|
||||
def test_insert_vectors(upstash_instance, mock_index):
|
||||
vectors = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
||||
payloads = [{"name": "vector1"}, {"name": "vector2"}]
|
||||
ids = ["id1", "id2"]
|
||||
|
||||
upstash_instance.insert(vectors=vectors, payloads=payloads, ids=ids)
|
||||
|
||||
upstash_instance.client.upsert.assert_called_once_with(
|
||||
vectors=[
|
||||
{"id": "id1", "vector": [0.1, 0.2, 0.3], "metadata": {"name": "vector1"}},
|
||||
{"id": "id2", "vector": [0.4, 0.5, 0.6], "metadata": {"name": "vector2"}},
|
||||
],
|
||||
namespace="ns",
|
||||
)
|
||||
|
||||
|
||||
def test_search_vectors(upstash_instance, mock_index):
|
||||
mock_result = [
|
||||
QueryResult(
|
||||
id="id1", score=0.1, vector=None, metadata={"name": "vector1"}, data=None
|
||||
),
|
||||
QueryResult(
|
||||
id="id2", score=0.2, vector=None, metadata={"name": "vector2"}, data=None
|
||||
),
|
||||
]
|
||||
|
||||
upstash_instance.client.query_many.return_value = [mock_result]
|
||||
|
||||
vectors = [[0.1, 0.2, 0.3]]
|
||||
results = upstash_instance.search(
|
||||
query="hello world",
|
||||
vectors=vectors,
|
||||
limit=2,
|
||||
filters={"age": 30, "name": "John"},
|
||||
)
|
||||
|
||||
upstash_instance.client.query_many.assert_called_once_with(
|
||||
queries=[
|
||||
{
|
||||
"vector": vectors[0],
|
||||
"top_k": 2,
|
||||
"namespace": "ns",
|
||||
"include_metadata": True,
|
||||
"filter": 'age = 30 AND name = "John"',
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
assert len(results) == 2
|
||||
assert results[0].id == "id1"
|
||||
assert results[0].score == 0.1
|
||||
assert results[0].payload == {"name": "vector1"}
|
||||
|
||||
|
||||
def test_delete_vector(upstash_instance):
|
||||
vector_id = "id1"
|
||||
|
||||
upstash_instance.delete(vector_id=vector_id)
|
||||
|
||||
upstash_instance.client.delete.assert_called_once_with(
|
||||
ids=[vector_id], namespace="ns"
|
||||
)
|
||||
|
||||
|
||||
def test_update_vector(upstash_instance):
|
||||
vector_id = "id1"
|
||||
new_vector = [0.7, 0.8, 0.9]
|
||||
new_payload = {"name": "updated_vector"}
|
||||
|
||||
upstash_instance.update(vector_id=vector_id, vector=new_vector, payload=new_payload)
|
||||
|
||||
upstash_instance.client.update.assert_called_once_with(
|
||||
id="id1",
|
||||
vector=new_vector,
|
||||
data=None,
|
||||
metadata={"name": "updated_vector"},
|
||||
namespace="ns",
|
||||
)
|
||||
|
||||
|
||||
def test_get_vector(upstash_instance):
|
||||
mock_result = [
|
||||
QueryResult(
|
||||
id="id1", score=None, vector=None, metadata={"name": "vector1"}, data=None
|
||||
)
|
||||
]
|
||||
upstash_instance.client.fetch.return_value = mock_result
|
||||
|
||||
result = upstash_instance.get(vector_id="id1")
|
||||
|
||||
upstash_instance.client.fetch.assert_called_once_with(
|
||||
ids=["id1"], namespace="ns", include_metadata=True
|
||||
)
|
||||
|
||||
assert result.id == "id1"
|
||||
assert result.payload == {"name": "vector1"}
|
||||
|
||||
|
||||
def test_list_vectors(upstash_instance):
|
||||
mock_result = [
|
||||
QueryResult(
|
||||
id="id1", score=None, vector=None, metadata={"name": "vector1"}, data=None
|
||||
),
|
||||
QueryResult(
|
||||
id="id2", score=None, vector=None, metadata={"name": "vector2"}, data=None
|
||||
),
|
||||
QueryResult(
|
||||
id="id3", score=None, vector=None, metadata={"name": "vector3"}, data=None
|
||||
),
|
||||
]
|
||||
handler = MagicMock()
|
||||
|
||||
upstash_instance.client.info.return_value.dimension = 10
|
||||
upstash_instance.client.resumable_query.return_value = (mock_result[0:1], handler)
|
||||
handler.fetch_next.side_effect = [mock_result[1:2], mock_result[2:3], []]
|
||||
|
||||
filters = {"age": 30, "name": "John"}
|
||||
print("filters", filters)
|
||||
[results] = upstash_instance.list(filters=filters, limit=15)
|
||||
|
||||
upstash_instance.client.info.return_value = {
|
||||
"dimension": 10,
|
||||
}
|
||||
|
||||
upstash_instance.client.resumable_query.assert_called_once_with(
|
||||
vector=[1.0] * 10,
|
||||
filter='age = 30 AND name = "John"',
|
||||
include_metadata=True,
|
||||
namespace="ns",
|
||||
top_k=100,
|
||||
)
|
||||
|
||||
handler.fetch_next.assert_has_calls([call(100), call(100), call(100)])
|
||||
handler.__exit__.assert_called_once()
|
||||
|
||||
assert len(results) == len(mock_result)
|
||||
assert results[0].id == "id1"
|
||||
assert results[0].payload == {"name": "vector1"}
|
||||
|
||||
|
||||
def test_insert_vectors_with_embeddings(upstash_instance_with_embeddings, mock_index):
|
||||
vectors = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
||||
payloads = [
|
||||
{"name": "vector1", "data": "data1"},
|
||||
{"name": "vector2", "data": "data2"},
|
||||
]
|
||||
ids = ["id1", "id2"]
|
||||
|
||||
upstash_instance_with_embeddings.insert(vectors=vectors, payloads=payloads, ids=ids)
|
||||
|
||||
upstash_instance_with_embeddings.client.upsert.assert_called_once_with(
|
||||
vectors=[
|
||||
{
|
||||
"id": "id1",
|
||||
# Uses the data field instead of using vectors
|
||||
"data": "data1",
|
||||
"metadata": {"name": "vector1", "data": "data1"},
|
||||
},
|
||||
{
|
||||
"id": "id2",
|
||||
"data": "data2",
|
||||
"metadata": {"name": "vector2", "data": "data2"},
|
||||
},
|
||||
],
|
||||
namespace="ns",
|
||||
)
|
||||
|
||||
|
||||
def test_search_vectors_with_embeddings(upstash_instance_with_embeddings, mock_index):
|
||||
mock_result = [
|
||||
QueryResult(
|
||||
id="id1", score=0.1, vector=None, metadata={"name": "vector1"}, data="data1"
|
||||
),
|
||||
QueryResult(
|
||||
id="id2", score=0.2, vector=None, metadata={"name": "vector2"}, data="data2"
|
||||
),
|
||||
]
|
||||
|
||||
upstash_instance_with_embeddings.client.query.return_value = mock_result
|
||||
|
||||
results = upstash_instance_with_embeddings.search(
|
||||
query="hello world",
|
||||
vectors=[],
|
||||
limit=2,
|
||||
filters={"age": 30, "name": "John"},
|
||||
)
|
||||
|
||||
upstash_instance_with_embeddings.client.query.assert_called_once_with(
|
||||
# Uses the data field instead of using vectors
|
||||
data="hello world",
|
||||
top_k=2,
|
||||
filter='age = 30 AND name = "John"',
|
||||
include_metadata=True,
|
||||
namespace="ns",
|
||||
)
|
||||
|
||||
assert len(results) == 2
|
||||
assert results[0].id == "id1"
|
||||
assert results[0].score == 0.1
|
||||
assert results[0].payload == {"name": "vector1"}
|
||||
|
||||
|
||||
def test_update_vector_with_embeddings(upstash_instance_with_embeddings):
|
||||
vector_id = "id1"
|
||||
new_payload = {"name": "updated_vector", "data": "updated_data"}
|
||||
|
||||
upstash_instance_with_embeddings.update(vector_id=vector_id, payload=new_payload)
|
||||
|
||||
upstash_instance_with_embeddings.client.update.assert_called_once_with(
|
||||
id="id1",
|
||||
vector=None,
|
||||
data="updated_data",
|
||||
metadata={"name": "updated_vector", "data": "updated_data"},
|
||||
namespace="ns",
|
||||
)
|
||||
|
||||
|
||||
def test_insert_vectors_with_embeddings_missing_data(upstash_instance_with_embeddings):
|
||||
vectors = [[0.1, 0.2, 0.3]]
|
||||
payloads = [{"name": "vector1"}] # Missing data field
|
||||
ids = ["id1"]
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="When embeddings are enabled, all payloads must contain a 'data' field",
|
||||
):
|
||||
upstash_instance_with_embeddings.insert(
|
||||
vectors=vectors, payloads=payloads, ids=ids
|
||||
)
|
||||
|
||||
|
||||
def test_update_vector_with_embeddings_missing_data(upstash_instance_with_embeddings):
|
||||
# Should still work, data is not required for update
|
||||
vector_id = "id1"
|
||||
new_payload = {"name": "updated_vector"} # Missing data field
|
||||
|
||||
upstash_instance_with_embeddings.update(vector_id=vector_id, payload=new_payload)
|
||||
|
||||
upstash_instance_with_embeddings.client.update.assert_called_once_with(
|
||||
id="id1",
|
||||
vector=None,
|
||||
data=None,
|
||||
metadata={"name": "updated_vector"},
|
||||
namespace="ns",
|
||||
)
|
||||
|
||||
|
||||
def test_list_cols(upstash_instance):
|
||||
mock_namespaces = ["ns1", "ns2", "ns3"]
|
||||
upstash_instance.client.list_namespaces.return_value = mock_namespaces
|
||||
|
||||
result = upstash_instance.list_cols()
|
||||
|
||||
upstash_instance.client.list_namespaces.assert_called_once()
|
||||
assert result == mock_namespaces
|
||||
|
||||
|
||||
def test_delete_col(upstash_instance):
|
||||
upstash_instance.delete_col()
|
||||
upstash_instance.client.reset.assert_called_once_with(namespace="ns")
|
||||
|
||||
|
||||
def test_col_info(upstash_instance):
|
||||
mock_info = {
|
||||
"dimension": 10,
|
||||
"total_vectors": 100,
|
||||
"pending_vectors": 0,
|
||||
"disk_size": 1024,
|
||||
}
|
||||
upstash_instance.client.info.return_value = mock_info
|
||||
|
||||
result = upstash_instance.col_info()
|
||||
|
||||
upstash_instance.client.info.assert_called_once()
|
||||
assert result == mock_info
|
||||
|
||||
|
||||
def test_get_vector_not_found(upstash_instance):
|
||||
upstash_instance.client.fetch.return_value = []
|
||||
|
||||
result = upstash_instance.get(vector_id="nonexistent")
|
||||
|
||||
upstash_instance.client.fetch.assert_called_once_with(
|
||||
ids=["nonexistent"], namespace="ns", include_metadata=True
|
||||
)
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_search_vectors_empty_filters(upstash_instance):
|
||||
mock_result = [
|
||||
QueryResult(
|
||||
id="id1", score=0.1, vector=None, metadata={"name": "vector1"}, data=None
|
||||
)
|
||||
]
|
||||
upstash_instance.client.query_many.return_value = [mock_result]
|
||||
|
||||
vectors = [[0.1, 0.2, 0.3]]
|
||||
results = upstash_instance.search(
|
||||
query="hello world",
|
||||
vectors=vectors,
|
||||
limit=1,
|
||||
filters=None,
|
||||
)
|
||||
|
||||
upstash_instance.client.query_many.assert_called_once_with(
|
||||
queries=[
|
||||
{
|
||||
"vector": vectors[0],
|
||||
"top_k": 1,
|
||||
"namespace": "ns",
|
||||
"include_metadata": True,
|
||||
"filter": "",
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].id == "id1"
|
||||
|
||||
|
||||
def test_insert_vectors_no_payloads(upstash_instance):
|
||||
vectors = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
||||
ids = ["id1", "id2"]
|
||||
|
||||
upstash_instance.insert(vectors=vectors, ids=ids)
|
||||
|
||||
upstash_instance.client.upsert.assert_called_once_with(
|
||||
vectors=[
|
||||
{"id": "id1", "vector": [0.1, 0.2, 0.3], "metadata": None},
|
||||
{"id": "id2", "vector": [0.4, 0.5, 0.6], "metadata": None},
|
||||
],
|
||||
namespace="ns",
|
||||
)
|
||||
|
||||
|
||||
def test_insert_vectors_no_ids(upstash_instance):
|
||||
vectors = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
||||
payloads = [{"name": "vector1"}, {"name": "vector2"}]
|
||||
|
||||
upstash_instance.insert(vectors=vectors, payloads=payloads)
|
||||
|
||||
upstash_instance.client.upsert.assert_called_once_with(
|
||||
vectors=[
|
||||
{"id": None, "vector": [0.1, 0.2, 0.3], "metadata": {"name": "vector1"}},
|
||||
{"id": None, "vector": [0.4, 0.5, 0.6], "metadata": {"name": "vector2"}},
|
||||
],
|
||||
namespace="ns",
|
||||
)
|
||||
Reference in New Issue
Block a user