Azure CosmosDB Mongo vCore Storage Integrations

llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosmongovcore/BUILD

@@ -0,0 +1,3 @@
+poetry_requirements(
+    name="poetry",
+)

llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosmongovcore/Makefile

@@ -0,0 +1,17 @@
+GIT_ROOT ?= $(shell git rev-parse --show-toplevel)
+
+help:	## Show all Makefile targets.
+	@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}'
+
+format:	## Run code autoformatters (black).
+	pre-commit install
+	git ls-files | xargs pre-commit run black --files
+
+lint:	## Run linters: pre-commit (black, ruff, codespell) and mypy
+	pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files
+
+test:	## Run tests via pytest.
+	pytest tests
+
+watch-docs:	## Build and watch documentation.
+	sphinx-autobuild docs/ docs/_build/html --open-browser --watch $(GIT_ROOT)/llama_index/

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

@@ -0,0 +1 @@
+# LlamaIndex Chat_Store Integration: Azure CosmosDB Mongo vCore Chat Store

llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosmongovcore/llama_index/storage/chat_store/azurecosmosmongovcore/BUILD

@@ -0,0 +1 @@
+python_sources()

llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosmongovcore/llama_index/storage/chat_store/azurecosmosmongovcore/__init__.py

@@ -0,0 +1,5 @@
+from llama_index.storage.chat_store.azurecosmosmongovcore.base import (
+    AzureCosmosMongoVCoreChatStore,
+)
+
+__all__ = ["AzureCosmosMongoVCoreChatStore"]

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

@@ -0,0 +1,157 @@
+import logging
+from abc import ABC
+from typing import List, Optional
+
+from llama_index.core.llms import ChatMessage
+from llama_index.core.storage.chat_store import BaseChatStore
+from pymongo import MongoClient
+from pymongo.database import Database
+from pymongo.collection import Collection
+
+logger = logging.getLogger(__name__)
+
+APP_NAME = "Llama-Index-CDBMongoVCore-ChatStore-Python"
+
+
+# Convert a ChatMessage to a JSON object
+def _message_to_dict(message: ChatMessage) -> dict:
+    return message.dict()
+
+
+# Convert a list of ChatMessages to a list of JSON objects
+def _messages_to_dict(messages: List[ChatMessage]) -> List[dict]:
+    return [_message_to_dict(message) for message in messages]
+
+
+# Convert a JSON object to a ChatMessage
+def _dict_to_message(d: dict) -> ChatMessage:
+    return ChatMessage.model_validate(d)
+
+
+class AzureCosmosMongoVCoreChatStore(BaseChatStore, ABC):
+    """Creates an Azure Cosmos DB NoSql Chat Store."""
+
+    _mongo_client = MongoClient
+    _database = Database
+    _collection = Collection
+
+    def __init__(
+        self,
+        mongo_client: MongoClient,
+        uri: Optional[str] = None,
+        host: Optional[str] = None,
+        port: Optional[int] = None,
+        db_name: Optional[str] = None,
+        collection_name: Optional[str] = None,
+    ):
+        super().__init__(
+            mongo_client=mongo_client,
+            uri=uri,
+            host=host,
+            port=port,
+            db_name=db_name,
+        )
+
+        self._mongo_client = mongo_client
+        self._uri = uri
+        self._host = host
+        self._port = port
+        self._database = self._mongo_client[db_name]
+        self._collection = self._mongo_client[db_name][collection_name]
+
+    @classmethod
+    def from_connection_string(
+        cls,
+        connection_string: str,
+        db_name: Optional[str] = None,
+        collection_name: Optional[str] = None,
+    ):
+        """Creates an instance of AzureCosmosMongoVCoreChatStore using a connection string."""
+        mongo_client = MongoClient(connection_string, appname=APP_NAME)
+
+        return cls(
+            mongo_client=mongo_client,
+            db_name=db_name,
+            collection_name=collection_name,
+        )
+
+    @classmethod
+    def from_host_and_port(
+        cls,
+        host: str,
+        port: int,
+        db_name: Optional[str] = None,
+        collection_name: Optional[str] = None,
+    ) -> "AzureCosmosMongoVCoreChatStore":
+        """Initializes AzureCosmosMongoVCoreChatStore from an endpoint url and key."""
+        mongo_client = MongoClient(host=host, port=port, appname=APP_NAME)
+
+        return cls(
+            mongo_client=mongo_client,
+            host=host,
+            port=port,
+            db_name=db_name,
+            collection_name=collection_name,
+        )
+
+    def set_messages(self, key: str, messages: List[ChatMessage]) -> None:
+        """Set messages for a key."""
+        self._collection.updateOne(
+            {"_id": key},
+            {"$set": {"messages": _messages_to_dict(messages)}},
+            upsert=True,
+        )
+
+    def get_messages(self, key: str) -> List[ChatMessage]:
+        """Get messages for a key."""
+        response = self._collection.find_one({"_id": key})
+        if response is not None:
+            message_history = response["messages"]
+        else:
+            message_history = []
+        return [_dict_to_message(message) for message in message_history]
+
+    def add_message(self, key: str, message: ChatMessage) -> None:
+        """Add a message for a key."""
+        current_messages = _messages_to_dict(self.get_messages(key))
+        current_messages.append(_message_to_dict(message))
+
+        self._collection.insert_one(
+            {
+                "id": key,
+                "messages": current_messages,
+            }
+        )
+
+    def delete_messages(self, key: str) -> Optional[List[ChatMessage]]:
+        """Delete messages for a key."""
+        messages_to_delete = self.get_messages(key)
+        self._collection.delete_one({"_id": key})
+        return messages_to_delete
+
+    def delete_message(self, key: str, idx: int) -> Optional[ChatMessage]:
+        """Delete specific message for a key."""
+        current_messages = self.get_messages(key)
+        try:
+            message_to_delete = current_messages[idx]
+            del current_messages[idx]
+            self.set_messages(key, current_messages)
+            return message_to_delete
+        except IndexError:
+            logger.error(
+                IndexError(f"No message exists at index, {idx}, for key {key}")
+            )
+            return None
+
+    def delete_last_message(self, key: str) -> Optional[ChatMessage]:
+        """Delete last message for a key."""
+        return self.delete_message(key, -1)
+
+    def get_keys(self) -> List[str]:
+        """Get all keys."""
+        return self._collection.find({}, {"_id": 1})
+
+    @classmethod
+    def class_name(cls) -> str:
+        """Get class name."""
+        return "AzureCosmosMongoVCoreChatStore"

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

@@ -0,0 +1,63 @@
+[build-system]
+build-backend = "poetry.core.masonry.api"
+requires = ["poetry-core"]
+
+[tool.codespell]
+check-filenames = true
+check-hidden = true
+skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb"
+
+[tool.llamahub]
+contains_example = false
+import_path = "llama_index.storage.chat_store.azurecosmosmongovcore"
+
+[tool.llamahub.class_authors]
+AzureCosmosNoSqlChatStore = "Aayush"
+
+[tool.mypy]
+disallow_untyped_defs = true
+exclude = ["_static", "build", "examples", "notebooks", "venv"]
+ignore_missing_imports = true
+python_version = "3.8"
+
+[tool.poetry]
+authors = ["Aayush Kataria <[email protected]>"]
+description = "llama-index storage-chat-store azure cosmosdb mongo vcore integration"
+exclude = ["**/BUILD"]
+license = "MIT"
+name = "llama-index-storage-chat-store-azurecosmosmongovcore"
+readme = "README.md"
+version = "1.0.0"
+
+[tool.poetry.dependencies]
+python = ">=3.8.1,<4.0"
+pymongo = "^4.9.1"
+llama-index-core = "^0.11.0"
+
+[tool.poetry.group.dev.dependencies]
+ipython = "8.10.0"
+jupyter = "^1.0.0"
+mypy = "0.991"
+pre-commit = "3.2.0"
+pylint = "2.15.10"
+pytest = "7.2.1"
+pytest-mock = "3.11.1"
+ruff = "0.0.292"
+tree-sitter-languages = "^1.8.0"
+types-Deprecated = ">=0.1.0"
+types-PyYAML = "^6.0.12.12"
+types-protobuf = "^4.24.0.4"
+types-redis = "4.5.5.0"
+types-requests = "2.28.11.8"
+types-setuptools = "67.1.0.0"
+
+[tool.poetry.group.dev.dependencies.black]
+extras = ["jupyter"]
+version = "<=23.9.1,>=23.7.0"
+
+[tool.poetry.group.dev.dependencies.codespell]
+extras = ["toml"]
+version = ">=v2.2.6"
+
+[[tool.poetry.packages]]
+include = "llama_index/"

llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosmongovcore/tests/BUILD

@@ -0,0 +1 @@
+python_tests()

llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosmongovcore/tests/test_storage_azurecosmosmongovcore_chat_store.py

@@ -0,0 +1,9 @@
+from llama_index.core.storage.chat_store.base import BaseChatStore
+from llama_index.storage.chat_store.azurecosmosmongovcore import (
+    AzureCosmosMongoVCoreChatStore,
+)
+
+
+def test_class():
+    names_of_base_classes = [b.__name__ for b in AzureCosmosMongoVCoreChatStore.__mro__]
+    assert BaseChatStore.__name__ in names_of_base_classes

llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosmongovcore/BUILD

@@ -0,0 +1,3 @@
+poetry_requirements(
+    name="poetry",
+)

llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosmongovcore/Makefile

@@ -0,0 +1,17 @@
+GIT_ROOT ?= $(shell git rev-parse --show-toplevel)
+
+help:	## Show all Makefile targets.
+	@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}'
+
+format:	## Run code autoformatters (black).
+	pre-commit install
+	git ls-files | xargs pre-commit run black --files
+
+lint:	## Run linters: pre-commit (black, ruff, codespell) and mypy
+	pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files
+
+test:	## Run tests via pytest.
+	pytest tests
+
+watch-docs:	## Build and watch documentation.
+	sphinx-autobuild docs/ docs/_build/html --open-browser --watch $(GIT_ROOT)/llama_index/

llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosmongovcore/README.md

@@ -0,0 +1 @@
+# LlamaIndex Document Store Integration: Azure CosmosDB Mongo vCore Document Store

llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosmongovcore/llama_index/storage/docstore/azurecosmosmongovcore/BUILD

@@ -0,0 +1 @@
+python_sources()

llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosmongovcore/llama_index/storage/docstore/azurecosmosmongovcore/__init__.py

@@ -0,0 +1,5 @@
+from llama_index.storage.docstore.azurecosmosmongovcore.base import (
+    AzureCosmosMongoVCoreDocumentStore,
+)
+
+__all__ = ["AzureCosmosMongoVCoreDocumentStore"]

llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosmongovcore/llama_index/storage/docstore/azurecosmosmongovcore/base.py

@@ -0,0 +1,62 @@
+from typing import Optional
+
+from llama_index.core.storage.docstore.keyval_docstore import KVDocumentStore
+from llama_index.storage.kvstore.azurecosmosmongovcore import (
+    AzureCosmosMongoVCoreKVStore,
+)
+
+
+APP_NAME = "Llama-Index-CDBMongoVCore-DocumentStore-Python"
+
+
+class AzureCosmosMongoVCoreDocumentStore(KVDocumentStore):
+    """Creates an AzureCosmosMongoVCoreDocumentStore."""
+
+    def __init__(
+        self,
+        azure_cosmos_mongo_vcore_kvstore: AzureCosmosMongoVCoreKVStore,
+        namespace: Optional[str] = None,
+        collection_suffix: Optional[str] = None,
+    ) -> None:
+        """Initializes the Azure Cosmos Mongo vCore Index Store."""
+        super().__init__(
+            azure_cosmos_mongo_vcore_kvstore,
+            namespace=namespace,
+            collection_suffix=collection_suffix,
+        )
+
+    @classmethod
+    def from_connection_string(
+        cls,
+        connection_string: str,
+        db_name: Optional[str] = None,
+        collection_name: Optional[str] = None,
+    ) -> "AzureCosmosMongoVCoreDocumentStore":
+        """Creates an instance of AzureCosmosMongoVCoreDocumentStore using a connection string."""
+        azure_cosmos_mongo_vcore_kvstore = (
+            AzureCosmosMongoVCoreKVStore.from_connection_string(
+                connection_string, db_name=db_name, collection_name=collection_name
+            )
+        )
+        namespace = db_name + "." + collection_name
+        return cls(azure_cosmos_mongo_vcore_kvstore, namespace)
+
+    @classmethod
+    def from_host_and_port(
+        cls,
+        host: str,
+        port: int,
+        db_name: Optional[str] = None,
+        collection_name: Optional[str] = None,
+    ) -> "AzureCosmosMongoVCoreDocumentStore":
+        """Initializes AzureCosmosMongoVCoreDocumentStore from an endpoint url and key."""
+        azure_cosmos_mongo_vcore_kvstore = (
+            AzureCosmosMongoVCoreKVStore.from_host_and_port(
+                host, port, db_name=db_name, collection_name=collection_name
+            )
+        )
+        namespace = db_name + "." + collection_name
+        return cls(
+            azure_cosmos_mongo_vcore_kvstore,
+            namespace,
+        )