diff --git a/mem0/vector_stores/milvus.py b/mem0/vector_stores/milvus.py index 129a9675b..80a4d0ffa 100644 --- a/mem0/vector_stores/milvus.py +++ b/mem0/vector_stores/milvus.py @@ -290,7 +290,7 @@ class MilvusDB(VectorStoreBase): schema = {"id": vector_id, "vectors": vector, "metadata": payload, "text": text} self.client.upsert(collection_name=self.collection_name, data=schema) - def get(self, vector_id): + def get(self, vector_id) -> Optional[OutputData]: """ Retrieve a vector by ID. @@ -298,9 +298,11 @@ class MilvusDB(VectorStoreBase): vector_id (str): ID of the vector to retrieve. Returns: - OutputData: Retrieved vector. + Optional[OutputData]: Retrieved vector, or None if the ID is not found. """ result = self.client.get(collection_name=self.collection_name, ids=vector_id) + if not result: + return None output = OutputData( id=result[0].get("id", None), score=None, diff --git a/mem0/vector_stores/supabase.py b/mem0/vector_stores/supabase.py index 62c9054ec..6faa09b09 100644 --- a/mem0/vector_stores/supabase.py +++ b/mem0/vector_stores/supabase.py @@ -176,7 +176,7 @@ class Supabase(VectorStoreBase): """ result = self.collection.fetch([(vector_id,)]) if not result: - return [] + return None record = result[0] return OutputData(id=str(record.id), score=None, payload=record.metadata) diff --git a/mem0/vector_stores/weaviate.py b/mem0/vector_stores/weaviate.py index cfa580ca3..e3c71742e 100644 --- a/mem0/vector_stores/weaviate.py +++ b/mem0/vector_stores/weaviate.py @@ -305,7 +305,7 @@ class Weaviate(VectorStoreBase): existing_payload: Mapping[str, str] = existing_data collection.data.update(uuid=vector_id, properties=existing_payload, vector=vector) - def get(self, vector_id): + def get(self, vector_id) -> Optional[OutputData]: """ Retrieve a vector by ID. @@ -313,7 +313,7 @@ class Weaviate(VectorStoreBase): vector_id: ID of the vector to retrieve. Returns: - dict: Retrieved vector and metadata. + Optional[OutputData]: Retrieved vector, or None if the ID is not found. """ vector_id = get_valid_uuid(vector_id) collection = self.client.collections.get(str(self.collection_name)) @@ -322,9 +322,8 @@ class Weaviate(VectorStoreBase): uuid=vector_id, return_properties=["hash", "created_at", "updated_at", "user_id", "agent_id", "run_id", "data", "category"], ) - # results = {} - # print("reponse",response) - # for obj in response.objects: + if response is None: + return None payload = response.properties.copy() payload["id"] = str(response.uuid).split("'")[0] results = OutputData( diff --git a/tests/vector_stores/test_milvus.py b/tests/vector_stores/test_milvus.py index 637411a2e..b0109ca56 100644 --- a/tests/vector_stores/test_milvus.py +++ b/tests/vector_stores/test_milvus.py @@ -189,11 +189,17 @@ class TestMilvusDB: ] result = milvus_db.get(vector_id) - + assert result.id == vector_id assert result.payload == {"user_id": "alice"} assert result.score is None + def test_get_missing_returns_none(self, milvus_db, mock_milvus_client): + """get() must return None (not raise IndexError) for an unknown id.""" + mock_milvus_client.get.return_value = [] + + assert milvus_db.get("missing") is None + def test_list_with_filters(self, milvus_db, mock_milvus_client): """Test listing memories with filters.""" mock_milvus_client.query.return_value = [ diff --git a/tests/vector_stores/test_supabase.py b/tests/vector_stores/test_supabase.py index e00d2b279..d4d0d5bd8 100644 --- a/tests/vector_stores/test_supabase.py +++ b/tests/vector_stores/test_supabase.py @@ -111,6 +111,13 @@ def test_get_vector(supabase_instance, mock_collection): assert result.payload == {"name": "vector1"} +def test_get_missing_returns_none(supabase_instance, mock_collection): + # An unknown id yields an empty fetch; get() must return None (not []). + mock_collection.fetch.return_value = [] + + assert supabase_instance.get(vector_id="missing") is None + + def test_list_vectors(supabase_instance, mock_collection): mock_query_results = [("id1", 0.9, {}), ("id2", 0.8, {})] mock_fetch_results = [("id1", [0.1, 0.2, 0.3], {"name": "vector1"}), ("id2", [0.4, 0.5, 0.6], {"name": "vector2"})] diff --git a/tests/vector_stores/test_weaviate.py b/tests/vector_stores/test_weaviate.py index 76bcff043..d4ace1a1b 100644 --- a/tests/vector_stores/test_weaviate.py +++ b/tests/vector_stores/test_weaviate.py @@ -4,9 +4,7 @@ import uuid from unittest.mock import MagicMock, patch import dotenv -import httpx import weaviate -from weaviate.exceptions import UnexpectedStatusCodeException from mem0.vector_stores.weaviate import Weaviate @@ -112,11 +110,13 @@ class TestWeaviateDB(unittest.TestCase): assert result.payload == expected_payload def test_get_not_found(self): - mock_response = httpx.Response(status_code=404, json={"error": "Not found"}) + # fetch_object_by_id returns None for an unknown id; get() must return + # None rather than raising AttributeError on response.properties. + self.client_mock.collections.get.return_value.query.fetch_object_by_id.return_value = None - self.client_mock.collections.get.return_value.data.get_by_id.side_effect = UnexpectedStatusCodeException( - "Not found", mock_response - ) + result = self.weaviate_db.get(vector_id=str(uuid.uuid4())) + + assert result is None def test_search(self): mock_objects = [{"uuid": "id1", "properties": {"key1": "value1"}, "metadata": {"distance": 0.2}}]