INTPYTHON-701 - Use PyMongo Asynchronous API instead of Motor

llama-index-integrations/readers/llama-index-readers-mongodb/llama_index/readers/mongodb/base.py

@@ -2,6 +2,7 @@
 
 from collections.abc import Callable
 from typing import Dict, Iterable, List, Optional
+from importlib.metadata import version
 
 from llama_index.core.readers.base import BaseReader
 from llama_index.core.schema import Document
@@ -27,11 +28,11 @@ def __init__(
     ) -> None:
         """Initialize with parameters."""
         try:
-            from motor.motor_asyncio import AsyncIOMotorClient
-            from pymongo import MongoClient
+            from pymongo import MongoClient, AsyncMongoClient
+            from pymongo.driver_info import DriverInfo
         except ImportError as err:
             raise ImportError(
-                "`pymongo / motor` package not found, please run `pip install pymongo motor`"
+                "`pymongo` package not found, please run `pip install pymongo`"
             ) from err
 
         if uri:
@@ -42,7 +43,17 @@ def __init__(
             raise ValueError("Either `host` and `port` or `uri` must be provided.")
 
         self.client = MongoClient(*client_args)
-        self.async_client = AsyncIOMotorClient(*client_args)
+        self.async_client = AsyncMongoClient(*client_args)
+
+        # append_metadata was added in PyMongo 4.14.0, but is a valid database name on earlier versions
+        if callable(self.client.append_metadata):
+            self.client.append_metadata(
+                DriverInfo(name="llama-index", version=version("llama-index"))
+            )
+        if callable(self.async_client.append_metadata):
+            self.async_client.append_metadata(
+                DriverInfo(name="llama-index", version=version("llama-index"))
+            )
 
     def lazy_load_data(
         self,

llama-index-integrations/readers/llama-index-readers-mongodb/pyproject.toml

@@ -26,16 +26,15 @@ dev = [
 
 [project]
 name = "llama-index-readers-mongodb"
-version = "0.4.1"
+version = "0.5.0"
 description = "llama-index readers mongodb integration"
 authors = [{name = "Your Name", email = "[email protected]"}]
 requires-python = ">=3.9,<4.0"
 readme = "README.md"
 license = "MIT"
 maintainers = [{name = "jerryjliu"}]
 dependencies = [
-    "pymongo>=4.6.1,<5",
-    "motor",
+    "pymongo>=4.13.0,<5",
     "llama-index-core>=0.13.0,<0.15",
 ]
 

llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-mongo/README.md

@@ -22,14 +22,13 @@ chat_store = MongoChatStore(
 )
 ```
 
-You can also initialize the chat store with a `MongoClient` or `AsyncIOMotorClient` and a database name and collection name.
+You can also initialize the chat store with a `MongoClient` or `AsyncMongoClient` and a database name and collection name.
 
 ```python
-from pymongo import MongoClient
-from motor.motor_asyncio import AsyncIOMotorClient
+from pymongo import MongoClient, AsyncMongoClient
 
 client = MongoClient("mongodb://localhost:27017/")
-async_client = AsyncIOMotorClient("mongodb://localhost:27017/")
+async_client = AsyncMongoClient("mongodb://localhost:27017/")
 
 chat_store = MongoChatStore(
     client=client,
@@ -39,14 +38,14 @@ chat_store = MongoChatStore(
 )
 ```
 
-You can also initialize the chat store with a `Collection` or `AsyncIOMotorCollection`.
+You can also initialize the chat store with a `Collection` or `AsyncCollection`.
 
 ```python
 from pymongo import Collection
-from motor.motor_asyncio import AsyncIOMotorCollection
+from pymongo.asynchronous import AsyncCollection
 
 client = MongoClient("mongodb://localhost:27017/")
-async_client = AsyncIOMotorClient("mongodb://localhost:27017/")
+async_client = AsyncMongoClient("mongodb://localhost:27017/")
 
 collection = client["llama_index"]["chat_sessions"]
 async_collection = async_client["llama_index"]["chat_sessions"]

llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-mongo/llama_index/storage/chat_store/mongo/base.py

@@ -1,12 +1,14 @@
+from importlib.metadata import version
 from typing import Any, List, Optional
 from datetime import datetime
 
 from llama_index.core.bridge.pydantic import Field, PrivateAttr
 from llama_index.core.llms import ChatMessage
 from llama_index.core.storage.chat_store.base import BaseChatStore
-from pymongo import MongoClient
+from pymongo import MongoClient, AsyncMongoClient
+from pymongo.driver_info import DriverInfo
 from pymongo.collection import Collection
-from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorCollection
+from pymongo.asynchronous.collection import AsyncCollection
 
 
 def _message_to_dict(message: ChatMessage) -> dict:
@@ -33,18 +35,18 @@ class MongoChatStore(BaseChatStore):
         default=None, description="Time to live in seconds."
     )
     _mongo_client: Optional[MongoClient] = PrivateAttr()
-    _async_client: Optional[AsyncIOMotorClient] = PrivateAttr()
+    _async_client: Optional[AsyncMongoClient] = PrivateAttr()
 
     def __init__(
         self,
         mongo_uri: str = "mongodb://localhost:27017",
         db_name: str = "default",
         collection_name: str = "sessions",
         mongo_client: Optional[MongoClient] = None,
-        amongo_client: Optional[AsyncIOMotorClient] = None,
+        amongo_client: Optional[AsyncMongoClient] = None,
         ttl_seconds: Optional[int] = None,
         collection: Optional[Collection] = None,
-        async_collection: Optional[AsyncIOMotorCollection] = None,
+        async_collection: Optional[AsyncCollection] = None,
         **kwargs: Any,
     ) -> None:
         """
@@ -63,7 +65,17 @@ def __init__(
         super().__init__(ttl=ttl_seconds)
 
         self._mongo_client = mongo_client or MongoClient(mongo_uri, **kwargs)
-        self._async_client = amongo_client or AsyncIOMotorClient(mongo_uri, **kwargs)
+        self._async_client = amongo_client or AsyncMongoClient(mongo_uri, **kwargs)
+
+        # append_metadata was added in PyMongo 4.14.0, but is a valid database name on earlier versions
+        if callable(self._mongo_client.append_metadata):
+            self._mongo_client.append_metadata(
+                DriverInfo(name="llama-index", version=version("llama-index"))
+            )
+        if callable(self._async_client.append_metadata):
+            self._async_client.append_metadata(
+                DriverInfo(name="llama-index", version=version("llama-index"))
+            )
 
         if collection:
             self._collection = collection

llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-mongo/pyproject.toml

@@ -24,20 +24,20 @@ dev = [
     "black[jupyter]<=23.9.1,>=23.7.0",
     "codespell[toml]>=v2.2.6",
     "diff-cover>=9.2.0",
+    "llama-index>=0.13.0,<0.15",
 ]
 
 [project]
 name = "llama-index-storage-chat-store-mongo"
-version = "0.2.1"
+version = "0.3.0"
 description = "llama-index storage-chat-store mongo integration"
 authors = [{name = "Vrushab Ghodke", email = "[email protected]"}]
 requires-python = ">=3.9,<4.0"
 readme = "README.md"
 license = "MIT"
 dependencies = [
     "llama-index-core>=0.13.0,<0.15",
-    "pymongo>=4.11.1,<5",
-    "motor>=3.7.0,<4",
+    "pymongo>=4.13.0,<5",
 ]
 
 [tool.codespell]

llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-mongo/tests/test_chat_store_mongo_chat_store.py

@@ -10,7 +10,7 @@
 
 try:
     import pymongo  # noqa: F401
-    import motor.motor_asyncio  # noqa: F401
+    import llama_index  # noqa: F401
 
     no_packages = False
 except ImportError:
@@ -54,7 +54,7 @@ def mongo_container() -> Generator[Dict[str, Union[str, Container]], None, None]
 
 
 @pytest.fixture()
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 def mongo_chat_store(
     mongo_container: Dict[str, Union[str, Container]],
 ) -> Generator[MongoChatStore, None, None]:
@@ -73,7 +73,7 @@ def mongo_chat_store(
             chat_store._collection.drop()
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 def test_create_chat(mongo_chat_store: MongoChatStore):
     """Test creating a chat session."""
     # Create a chat with metadata
@@ -89,7 +89,7 @@ def test_create_chat(mongo_chat_store: MongoChatStore):
     assert key in keys
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 def test_get_messages(mongo_chat_store: MongoChatStore):
     """Test retrieving messages from a chat session."""
     # Create a chat with messages
@@ -114,7 +114,7 @@ def test_get_messages(mongo_chat_store: MongoChatStore):
     assert retrieved_messages[2].content == "I need information about databases."
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 def test_add_message(mongo_chat_store: MongoChatStore):
     """Test adding a message to an existing chat session."""
     # Create a chat with initial messages
@@ -137,7 +137,7 @@ def test_add_message(mongo_chat_store: MongoChatStore):
     assert messages[1].content == "Response to initial message"
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 def test_delete_messages(mongo_chat_store: MongoChatStore):
     """Test deleting all messages from a chat session."""
     # Create a chat with messages
@@ -164,7 +164,7 @@ def test_delete_messages(mongo_chat_store: MongoChatStore):
     assert key not in keys
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 def test_delete_message(mongo_chat_store: MongoChatStore):
     """Test deleting a specific message from a chat session."""
     # Create a chat with multiple messages
@@ -190,7 +190,7 @@ def test_delete_message(mongo_chat_store: MongoChatStore):
     assert remaining[1].content == "Last message"
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 def test_delete_last_message(mongo_chat_store: MongoChatStore):
     """Test deleting the last message from a chat session."""
     # Create a chat with messages
@@ -214,7 +214,7 @@ def test_delete_last_message(mongo_chat_store: MongoChatStore):
     assert remaining[0].content == "First message"
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 @pytest.mark.asyncio
 async def test_async_get_messages(mongo_chat_store: MongoChatStore):
     """Test retrieving messages asynchronously."""
@@ -235,7 +235,7 @@ async def test_async_get_messages(mongo_chat_store: MongoChatStore):
     assert retrieved[1].content == "Async response"
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 @pytest.mark.asyncio
 async def test_async_add_message(mongo_chat_store: MongoChatStore):
     """Test adding a message asynchronously."""
@@ -253,7 +253,7 @@ async def test_async_add_message(mongo_chat_store: MongoChatStore):
     assert messages[1].content == "Async response"
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 @pytest.mark.asyncio
 async def test_async_set_messages(mongo_chat_store: MongoChatStore):
     """Test setting messages asynchronously."""
@@ -273,7 +273,7 @@ async def test_async_set_messages(mongo_chat_store: MongoChatStore):
     assert retrieved[1].content == "Second async set message"
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 @pytest.mark.asyncio
 async def test_async_delete_messages(mongo_chat_store: MongoChatStore):
     """Test deleting all messages asynchronously."""
@@ -300,7 +300,7 @@ async def test_async_delete_messages(mongo_chat_store: MongoChatStore):
     assert key not in keys
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 @pytest.mark.asyncio
 async def test_async_delete_message(mongo_chat_store: MongoChatStore):
     """Test deleting a specific message asynchronously."""
@@ -326,7 +326,7 @@ async def test_async_delete_message(mongo_chat_store: MongoChatStore):
     assert remaining[1].content == "Async last message"
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 @pytest.mark.asyncio
 async def test_async_delete_last_message(mongo_chat_store: MongoChatStore):
     """Test deleting the last message asynchronously."""
@@ -350,7 +350,7 @@ async def test_async_delete_last_message(mongo_chat_store: MongoChatStore):
     assert remaining[0].content == "Async first message"
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 @pytest.mark.asyncio
 async def test_async_get_keys(mongo_chat_store: MongoChatStore):
     """Test getting all keys asynchronously."""
@@ -372,7 +372,7 @@ async def test_async_get_keys(mongo_chat_store: MongoChatStore):
     assert "async_keys_test2" in keys
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 def test_nonexistent_key(mongo_chat_store: MongoChatStore):
     """Test behavior with nonexistent keys."""
     # Try to get messages for nonexistent key
@@ -394,7 +394,7 @@ def test_nonexistent_key(mongo_chat_store: MongoChatStore):
     assert deleted is None
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 def test_ttl_configuration(mongo_container):
     """Test TTL configuration is applied correctly."""
     # Create chat store with TTL
@@ -416,7 +416,7 @@ def test_ttl_configuration(mongo_container):
     chat_store._collection.drop()
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 def test_invalid_message_index(mongo_chat_store: MongoChatStore):
     """Test behavior when trying to delete a message with invalid index."""
     key = "test_invalid_index"
@@ -435,7 +435,7 @@ def test_invalid_message_index(mongo_chat_store: MongoChatStore):
     assert remaining[0].content == "Only message"
 
 
-@pytest.mark.skipif(no_packages, reason="pymongo and motor not installed")
+@pytest.mark.skipif(no_packages, reason="pymongo not installed")
 def test_multiple_clients(mongo_container):
     """Test using multiple chat store instances with the same database."""
     # Create two chat store instances