Fix CI issues related to missing dependency (#3096)

tests/vector_stores/test_baidu.py

@@ -2,10 +2,9 @@ from unittest.mock import Mock, patch, PropertyMock
 
 import pytest
 
-from mem0.vector_stores.baidu import BaiduDB, OutputData
-from pymochow.model.enum import MetricType, TableState, ServerErrCode
-from pymochow.model.schema import Field, Schema, VectorIndex, FilteringIndex, HNSWParams, AutoBuildRowCountIncrement
-from pymochow.model.table import Partition, Row, VectorSearchConfig, VectorTopkSearchRequest, FloatVector, Table
+from mem0.vector_stores.baidu import BaiduDB
+from pymochow.model.enum import TableState, ServerErrCode
+from pymochow.model.table import VectorSearchConfig, VectorTopkSearchRequest, FloatVector, Table
 from pymochow.exception import ServerError
 
 

tests/vector_stores/test_mongodb.py

@@ -1,7 +1,7 @@
 import pytest
 from unittest.mock import MagicMock, patch
 from mem0.vector_stores.mongodb import MongoDB
-from pymongo.operations import SearchIndexModel
+
 
 @pytest.fixture
 @patch("mem0.vector_stores.mongodb.MongoClient")
@@ -19,10 +19,11 @@ def mongo_vector_fixture(mock_mongo_client):
         db_name="test_db",
         collection_name="test_collection",
         embedding_model_dims=1536,
-        mongo_uri="mongodb://username:password@localhost:27017"
+        mongo_uri="mongodb://username:password@localhost:27017",
     )
     return mongo_vector, mock_collection, mock_db
 
+
 def test_initalize_create_col(mongo_vector_fixture):
     mongo_vector, mock_collection, mock_db = mongo_vector_fixture
     assert mongo_vector.collection_name == "test_collection"
@@ -49,12 +50,13 @@ def test_initalize_create_col(mongo_vector_fixture):
                         "dimensions": 1536,
                         "similarity": "cosine",
                     }
-                }
+                },
             }
-        }
+        },
     }
     assert mongo_vector.collection == mock_collection
 
+
 def test_insert(mongo_vector_fixture):
     mongo_vector, mock_collection, _ = mongo_vector_fixture
     vectors = [[0.1] * 1536, [0.2] * 1536]
@@ -62,12 +64,13 @@ def test_insert(mongo_vector_fixture):
     ids = ["id1", "id2"]
 
     mongo_vector.insert(vectors, payloads, ids)
-    expected_records=[
+    expected_records = [
         ({"_id": ids[0], "embedding": vectors[0], "payload": payloads[0]}),
-        ({"_id": ids[1], "embedding": vectors[1], "payload": payloads[1]})
+        ({"_id": ids[1], "embedding": vectors[1], "payload": payloads[1]}),
     ]
     mock_collection.insert_many.assert_called_once_with(expected_records)
 
+
 def test_search(mongo_vector_fixture):
     mongo_vector, mock_collection, _ = mongo_vector_fixture
     query_vector = [0.1] * 1536
@@ -79,25 +82,28 @@ def test_search(mongo_vector_fixture):
 
     results = mongo_vector.search("query_str", query_vector, limit=2)
     mock_collection.list_search_indexes.assert_called_with(name="test_collection_vector_index")
-    mock_collection.aggregate.assert_called_once_with([
-        {
-            "$vectorSearch": {
-                "index": "test_collection_vector_index",
-                "limit": 2,
-                "numCandidates": 2,
-                "queryVector": query_vector,
-                "path": "embedding",
+    mock_collection.aggregate.assert_called_once_with(
+        [
+            {
+                "$vectorSearch": {
+                    "index": "test_collection_vector_index",
+                    "limit": 2,
+                    "numCandidates": 2,
+                    "queryVector": query_vector,
+                    "path": "embedding",
+                },
             },
-        },
-        {"$set": {"score": {"$meta": "vectorSearchScore"}}},
-        {"$project": {"embedding": 0}},
-    ])
+            {"$set": {"score": {"$meta": "vectorSearchScore"}}},
+            {"$project": {"embedding": 0}},
+        ]
+    )
     assert len(results) == 2
     assert results[0].id == "id1"
     assert results[0].score == 0.9
     assert results[1].id == "id2"
     assert results[1].score == 0.8
 
+
 def test_delete(mongo_vector_fixture):
     mongo_vector, mock_collection, _ = mongo_vector_fixture
     mock_delete_result = MagicMock()
@@ -107,6 +113,7 @@ def test_delete(mongo_vector_fixture):
     mongo_vector.delete("id1")
     mock_collection.delete_one.assert_called_with({"_id": "id1"})
 
+
 def test_update(mongo_vector_fixture):
     mongo_vector, mock_collection, _ = mongo_vector_fixture
     mock_update_result = MagicMock()
@@ -122,6 +129,7 @@ def test_update(mongo_vector_fixture):
         {"$set": {"embedding": vectorValue, "payload": payloadValue}},
     )
 
+
 def test_get(mongo_vector_fixture):
     mongo_vector, mock_collection, _ = mongo_vector_fixture
     mock_collection.find_one.return_value = {"_id": "id1", "payload": {"key": "value1"}}
@@ -131,6 +139,7 @@ def test_get(mongo_vector_fixture):
     assert result.id == "id1"
     assert result.payload == {"key": "value1"}
 
+
 def test_list_cols(mongo_vector_fixture):
     mongo_vector, _, mock_db = mongo_vector_fixture
     mock_db.list_collection_names.return_value = ["col1", "col2"]
@@ -138,12 +147,14 @@ def test_list_cols(mongo_vector_fixture):
     collections = mongo_vector.list_cols()
     assert collections == ["col1", "col2"]
 
+
 def test_delete_col(mongo_vector_fixture):
     mongo_vector, mock_collection, _ = mongo_vector_fixture
 
     mongo_vector.delete_col()
     mock_collection.drop.assert_called_once()
 
+
 def test_col_info(mongo_vector_fixture):
     mongo_vector, _, mock_db = mongo_vector_fixture
     mock_db.command.return_value = {"count": 10, "size": 1024}
@@ -154,6 +165,7 @@ def test_col_info(mongo_vector_fixture):
     assert info["count"] == 10
     assert info["size"] == 1024
 
+
 def test_list(mongo_vector_fixture):
     mongo_vector, mock_collection, _ = mongo_vector_fixture
     mock_cursor = MagicMock()

tests/vector_stores/test_pinecone.py

@@ -88,7 +88,7 @@ def test_get_vector_found(pinecone_db):
     # or a list of dictionaries, not a dictionary with an 'id' field
 
     # Create a mock Vector object
-    from pinecone.data.dataclasses.vector import Vector
+    from pinecone import Vector
 
     mock_vector = Vector(id="id1", values=[0.1] * 128, metadata={"name": "vector1"})