mem0ai--mem0/tests/vector_stores/test_upstash_vector.py

from dataclasses import dataclass
from typing import Dict, List, Optional
from unittest.mock import MagicMock, call, patch

import pytest

from mem0.configs.vector_stores.upstash_vector import UpstashVectorConfig
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,
        top_k=2,
        filters={"age": 30, "name": "John"},
    )

    upstash_instance.client.query_many.assert_called_once_with(
        queries=[
            {
                "vector": vectors[0],
                "top_k": 2,
                "include_metadata": True,
                "filter": 'age = 30 AND name = "John"',
            }
        ],
        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_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, top_k=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=[],
        top_k=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,
        top_k=1,
        filters=None,
    )

    upstash_instance.client.query_many.assert_called_once_with(
        queries=[
            {
                "vector": vectors[0],
                "top_k": 1,
                "include_metadata": True,
                "filter": "",
            }
        ],
        namespace="ns",
    )

    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",
    )


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"}


def test_env_var_only_config_builds_provider(monkeypatch):
    """Regression: an env-var-only config (no url/token/client) must build.

    VectorStoreFactory does ``UpstashVector(**config.model_dump())``. The config
    validator read ``UPSTASH_VECTOR_REST_URL``/``UPSTASH_VECTOR_REST_TOKEN`` only
    to pass its presence check, then returned the config unchanged, so
    ``model_dump()`` still carried ``url=token=None`` and construction raised
    "Either a client or URL and token must be provided." — even though the docs
    advertise env-var setup. The resolved credentials must reach the ctor.
    """
    monkeypatch.setenv("UPSTASH_VECTOR_REST_URL", "https://example.upstash.io")
    monkeypatch.setenv("UPSTASH_VECTOR_REST_TOKEN", "tok_123")

    dumped = UpstashVectorConfig(collection_name="mem0").model_dump()
    assert dumped["url"] == "https://example.upstash.io"
    assert dumped["token"] == "tok_123"

    # Mirrors VectorStoreFactory.create: instance(**config.model_dump()).
    with patch("mem0.vector_stores.upstash_vector.Index") as mock_index:
        UpstashVector(**dumped)

    mock_index.assert_called_once_with("https://example.upstash.io", "tok_123")