async mongo document loader

docs/modules/indexes/document_loaders/examples/mongodb.ipynb

@@ -0,0 +1,148 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "vm8vn9t8DvC_"
+   },
+   "source": [
+    "# MongoDB"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "5WjXERXzFEhg"
+   },
+   "source": [
+    "## Overview"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "juAmbgoWD17u"
+   },
+   "source": [
+    "The MongoDB Document Loader returns a list of Langchain Documents from a MongoDB database.\n",
+    "\n",
+    "The Loader requires the following parameters:\n",
+    "\n",
+    "*   MongoDB connection string\n",
+    "*   MongoDB database name\n",
+    "*   MongoDB collection name\n",
+    "\n",
+    "The output takes the following format:\n",
+    "\n",
+    "- pageContent= Mongo Document\n",
+    "- metadata={'database': '[database_name]', 'collection': '[collection_name]'}"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load the Document Loader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# for running async code in jupyter notebook\n",
+    "import nest_asyncio\n",
+    "nest_asyncio.apply()\n",
+    "\n",
+    "\n",
+    "from langchain.document_loaders.mongodb import MongodbLoader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "loader = MongodbLoader(connection_string=\"mongodb://localhost:27017/\",\n",
+    "                       db_name=\"sample_restaurants\", \n",
+    "                       collection_name=\"restaurants\"\n",
+    "                       ) "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "25359"
+      ]
+     },
+     "execution_count": 26,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "docs = loader.load()\n",
+    "\n",
+    "len(docs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Document(page_content=\"{'_id': ObjectId('5eb3d668b31de5d588f4292a'), 'address': {'building': '2780', 'coord': [-73.98241999999999, 40.579505], 'street': 'Stillwell Avenue', 'zipcode': '11224'}, 'borough': 'Brooklyn', 'cuisine': 'American', 'grades': [{'date': datetime.datetime(2014, 6, 10, 0, 0), 'grade': 'A', 'score': 5}, {'date': datetime.datetime(2013, 6, 5, 0, 0), 'grade': 'A', 'score': 7}, {'date': datetime.datetime(2012, 4, 13, 0, 0), 'grade': 'A', 'score': 12}, {'date': datetime.datetime(2011, 10, 12, 0, 0), 'grade': 'A', 'score': 12}], 'name': 'Riviera Caterer', 'restaurant_id': '40356018'}\", metadata={'database': 'sample_restaurants', 'collection': 'restaurants'})"
+      ]
+     },
+     "execution_count": 27,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "docs[0]"
+   ]
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "collapsed_sections": [
+    "5WjXERXzFEhg"
+   ],
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

langchain/document_loaders/__init__.py

@@ -49,6 +49,7 @@
 from langchain.document_loaders.markdown import UnstructuredMarkdownLoader
 from langchain.document_loaders.mediawikidump import MWDumpLoader
 from langchain.document_loaders.modern_treasury import ModernTreasuryLoader
+from langchain.document_loaders.mongodb import MongodbLoader
 from langchain.document_loaders.notebook import NotebookLoader
 from langchain.document_loaders.notion import NotionDirectoryLoader
 from langchain.document_loaders.notiondb import NotionDBLoader
@@ -150,6 +151,7 @@
     "MWDumpLoader",
     "MathpixPDFLoader",
     "ModernTreasuryLoader",
+    "MongodbLoader",
     "NotebookLoader",
     "NotionDBLoader",
     "NotionDirectoryLoader",

langchain/document_loaders/mongodb.py

@@ -0,0 +1,61 @@
+import asyncio
+from typing import List
+
+from motor.motor_asyncio import AsyncIOMotorClient
+
+from langchain.docstore.document import Document
+from langchain.document_loaders.base import BaseLoader
+
+
+class MongodbLoader(BaseLoader):
+    """Load MongoDB documents."""
+
+    def __init__(
+        self,
+        connection_string: str,
+        db_name: str,
+        collection_name: str,
+    ):
+        if not connection_string:
+            raise ValueError("connection_string must be provided.")
+
+        if not db_name:
+            raise ValueError("db_name must be provided.")
+
+        if not collection_name:
+            raise ValueError("collection_name must be provided.")
+
+        self.client = AsyncIOMotorClient(connection_string)
+        self.db_name = db_name
+        self.collection_name = collection_name
+        self.db = self.client.get_database(db_name)
+        self.collection = self.db.get_collection(collection_name)
+
+    def load(self) -> List[Document]:
+        result = asyncio.run(self._async_load())
+        return result
+
+    async def _async_load(self) -> List[Document]:
+        result = []
+
+        try:
+            cursor = self.collection.find()
+
+            total_docs = self.collection.count_documents({})
+
+            async for doc in cursor:
+                content = str(doc)
+                metadata = {
+                    "database": self.db_name,
+                    "collection": self.collection_name,
+                }
+                result.append(Document(page_content=content, metadata=metadata))
+
+            if len(result) != total_docs:
+                raise Exception(f"Only partial collection of documents returned.")
+
+
+        except Exception as e:
+            print(f"Error: {e}")
+
+        return result

pyproject.toml

@@ -77,6 +77,7 @@ pexpect = {version = "^4.8.0", optional = true}
 pyvespa = {version = "^0.33.0", optional = true}
 O365 = {version = "^2.0.26", optional = true}
 jq = {version = "^1.4.1", optional = true}
+motor = "^3.1.2"
 
 [tool.poetry.group.docs.dependencies]
 autodoc_pydantic = "^1.8.0"

tests/integration_tests/document_loaders/test_mongodb.py

@@ -0,0 +1,109 @@
+import os
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from motor.motor_asyncio import AsyncIOMotorClient
+
+from langchain.docstore.document import Document
+from langchain.document_loaders import MongodbLoader
+
+if "MONGODB_API_KEY" in os.environ:
+    mongo_api_key_set = True
+    mongo_api_key = os.environ["MONGODB_API_KEY"]
+else:
+    mongo_api_key_set = False
+
+
+@pytest.fixture
+def raw_docs():
+    return [
+        {"_id": "1", "address": {"building": "1", "room": "1"}},
+        {"_id": "2", "address": {"building": "2", "room": "2"}},
+    ]
+
+
+@pytest.fixture
+def expected_documents():
+    return [
+        Document(
+            page_content="{'_id': '1', 'address': {'building': '1', 'room': '1'}}",
+            metadata={"database": "sample_restaurants", "collection": "restaurants"},
+        ),
+        Document(
+            page_content="{'_id': '2', 'address': {'building': '2', 'room': '2'}}",
+            metadata={"database": "sample_restaurants", "collection": "restaurants"},
+        ),
+    ]
+
+
+@pytest.fixture
+def mock_mongodb(mocker, raw_docs):
+    async def mock_find():
+        for doc in raw_docs:
+            yield doc
+
+    mock_collection = MagicMock()
+    mock_db = MagicMock()
+    mock_client = MagicMock()
+
+    mocker.patch.object(mock_collection, "find", return_value=mock_find())
+    mocker.patch.object(AsyncIOMotorClient, "get_database", return_value=mock_db)
+    mocker.patch.object(mock_db, "get_collection", return_value=mock_collection)
+    mocker.patch("motor.motor_asyncio.AsyncIOMotorClient", return_value=mock_client)
+
+    return mocker
+
+
+@pytest.mark.asyncio
+async def test_async_load_mocked(mock_mongodb, expected_documents):
+    loader = MongodbLoader(
+        "mongodb://localhost:27017", "sample_restaurants", "restaurants"
+    )
+    documents = await loader._async_load()
+    assert documents == expected_documents
+
+
+@pytest.mark.asyncio
+async def test_async_partial_load_mocked(expected_documents):
+    loader = MongodbLoader(
+        "mongodb://localhost:27017", "sample_restaurants", "restaurants"
+    )
+    expected_documents.remove(expected_documents[1])
+
+    with pytest.raises(Exception) as error_PartialLoad:
+        await loader._async_load()
+
+        assert (
+            str(error_PartialLoad.value)
+            == "Error: Only partial collection of documents returned."
+        )
+
+
+def test_load_mocked(expected_documents):
+    mock_async_load = AsyncMock()
+    mock_async_load.return_value = expected_documents
+
+    with patch("langchain.document_loaders.MongodbLoader", MagicMock()):
+        loader = MongodbLoader(
+            "mongodb://localhost:27017", "test_db", "test_collection"
+        )
+
+    loader._async_load = mock_async_load
+
+    documents = loader.load()
+
+    assert documents == expected_documents
+
+
+@pytest.mark.skipif(not mongo_api_key_set, reason="MONGODB_API_KEY not provided.")
+def test_load_actual_db() -> None:
+    if "MONGODB_API_KEY" in os.environ:
+        mongo_api_key = os.environ["MONGODB_API_KEY"]
+        db_name = os.environ["MONGODB_DB_NAME"]
+        collection_name = os.environ["MONGODB_COLLECTION_NAME"]
+
+        loader = MongodbLoader(mongo_api_key, db_name, collection_name)
+
+        result = loader.load()
+
+        assert len(result) > 0