Spaces:
Runtime error
Runtime error
| from typing import Iterator, List, Optional, Sequence, Tuple | |
| from langchain_core.documents import Document | |
| from langchain_core.stores import BaseStore | |
| class MongoDBStore(BaseStore[str, Document]): | |
| """BaseStore implementation using MongoDB as the underlying store. | |
| Examples: | |
| Create a MongoDBStore instance and perform operations on it: | |
| .. code-block:: python | |
| # Instantiate the MongoDBStore with a MongoDB connection | |
| from langchain.storage import MongoDBStore | |
| mongo_conn_str = "mongodb://localhost:27017/" | |
| mongodb_store = MongoDBStore(mongo_conn_str, db_name="test-db", | |
| collection_name="test-collection") | |
| # Set values for keys | |
| doc1 = Document(...) | |
| doc2 = Document(...) | |
| mongodb_store.mset([("key1", doc1), ("key2", doc2)]) | |
| # Get values for keys | |
| values = mongodb_store.mget(["key1", "key2"]) | |
| # [doc1, doc2] | |
| # Iterate over keys | |
| for key in mongodb_store.yield_keys(): | |
| print(key) | |
| # Delete keys | |
| mongodb_store.mdelete(["key1", "key2"]) | |
| """ | |
| def __init__( | |
| self, | |
| connection_string: str, | |
| db_name: str, | |
| collection_name: str, | |
| *, | |
| client_kwargs: Optional[dict] = None, | |
| ) -> None: | |
| """Initialize the MongoDBStore with a MongoDB connection string. | |
| Args: | |
| connection_string (str): MongoDB connection string | |
| db_name (str): name to use | |
| collection_name (str): collection name to use | |
| client_kwargs (dict): Keyword arguments to pass to the Mongo client | |
| """ | |
| try: | |
| from pymongo import MongoClient | |
| except ImportError as e: | |
| raise ImportError( | |
| "The MongoDBStore requires the pymongo library to be " | |
| "installed. " | |
| "pip install pymongo" | |
| ) from e | |
| 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 = MongoClient(connection_string, **(client_kwargs or {})) | |
| self.collection = self.client[db_name][collection_name] | |
| def mget(self, keys: Sequence[str]) -> List[Optional[Document]]: | |
| """Get the list of documents associated with the given keys. | |
| Args: | |
| keys (list[str]): A list of keys representing Document IDs.. | |
| Returns: | |
| list[Document]: A list of Documents corresponding to the provided | |
| keys, where each Document is either retrieved successfully or | |
| represented as None if not found. | |
| """ | |
| result = self.collection.find({"_id": {"$in": keys}}) | |
| result_dict = {doc["_id"]: Document(**doc["value"]) for doc in result} | |
| return [result_dict.get(key) for key in keys] | |
| def mset(self, key_value_pairs: Sequence[Tuple[str, Document]]) -> None: | |
| """Set the given key-value pairs. | |
| Args: | |
| key_value_pairs (list[tuple[str, Document]]): A list of id-document | |
| pairs. | |
| Returns: | |
| None | |
| """ | |
| from pymongo import UpdateOne | |
| updates = [{"_id": k, "value": v.__dict__} for k, v in key_value_pairs] | |
| self.collection.bulk_write( | |
| [UpdateOne({"_id": u["_id"]}, {"$set": u}, upsert=True) for u in updates] | |
| ) | |
| def mdelete(self, keys: Sequence[str]) -> None: | |
| """Delete the given ids. | |
| Args: | |
| keys (list[str]): A list of keys representing Document IDs.. | |
| """ | |
| self.collection.delete_many({"_id": {"$in": keys}}) | |
| def yield_keys(self, prefix: Optional[str] = None) -> Iterator[str]: | |
| """Yield keys in the store. | |
| Args: | |
| prefix (str): prefix of keys to retrieve. | |
| """ | |
| if prefix is None: | |
| for doc in self.collection.find(projection=["_id"]): | |
| yield doc["_id"] | |
| else: | |
| for doc in self.collection.find( | |
| {"_id": {"$regex": f"^{prefix}"}}, projection=["_id"] | |
| ): | |
| yield doc["_id"] | |