fix(vector_stores): return None from get() for missing IDs (milvus/weaviate/supabase) (#5562)

Co-authored-by: ly-wang19 <ly-wang19@users.noreply.github.com>

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,

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)

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(

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 = [

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

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