diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/BUILD b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/BUILD new file mode 100644 index 0000000000000..0896ca890d8bf --- /dev/null +++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/BUILD @@ -0,0 +1,3 @@ +poetry_requirements( + name="poetry", +) diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/Makefile b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/Makefile new file mode 100644 index 0000000000000..b9eab05aa3706 --- /dev/null +++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/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/ diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/README.md b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/README.md new file mode 100644 index 0000000000000..1fe6513fb84f8 --- /dev/null +++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/README.md @@ -0,0 +1 @@ +# LlamaIndex Chat_Store Integration: Azure CosmosDB NoSQL Chat Store diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/llama_index/storage/chat_store/azurecosmosnosql/BUILD b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/llama_index/storage/chat_store/azurecosmosnosql/BUILD new file mode 100644 index 0000000000000..db46e8d6c978c --- /dev/null +++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/llama_index/storage/chat_store/azurecosmosnosql/BUILD @@ -0,0 +1 @@ +python_sources() diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/llama_index/storage/chat_store/azurecosmosnosql/__init__.py b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/llama_index/storage/chat_store/azurecosmosnosql/__init__.py new file mode 100644 index 0000000000000..ae7f15fec2f77 --- /dev/null +++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/llama_index/storage/chat_store/azurecosmosnosql/__init__.py @@ -0,0 +1,5 @@ +from llama_index.storage.chat_store.azurecosmosnosql.base import ( + AzureCosmosNoSqlChatStore, +) + +__all__ = ["AzureCosmosNoSqlChatStore"] diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/llama_index/storage/chat_store/azurecosmosnosql/base.py b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/llama_index/storage/chat_store/azurecosmosnosql/base.py new file mode 100644 index 0000000000000..94f2e5daae2d1 --- /dev/null +++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/llama_index/storage/chat_store/azurecosmosnosql/base.py @@ -0,0 +1,233 @@ +import logging +from typing import Any, Dict, List, Optional + +from azure.cosmos import CosmosClient, DatabaseProxy, ContainerProxy +from llama_index.core.llms import ChatMessage +from llama_index.core.storage.chat_store import BaseChatStore + +DEFAULT_CHAT_DATABASE = "ChatMessagesDB" +DEFAULT_CHAT_CONTAINER = "ChatMessagesContainer" + + +logger = logging.getLogger(__name__) + + +# 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 AzureCosmosNoSqlChatStore(BaseChatStore): + """Creates an Azure Cosmos DB NoSql Chat Store.""" + + _cosmos_client = CosmosClient + _database = DatabaseProxy + _container = ContainerProxy + + def __init__( + self, + cosmos_client: CosmosClient, + chat_db_name: str = DEFAULT_CHAT_DATABASE, + chat_container_name: str = DEFAULT_CHAT_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + **kwargs, + ): + super().__init__( + cosmos_client=cosmos_client, + chat_db_name=chat_db_name, + chat_container_name=chat_container_name, + cosmos_container_properties=cosmos_container_properties, + cosmos_database_properties=cosmos_database_properties, + ) + + self._cosmos_client = cosmos_client + + # Create the database if it already doesn't exist + self._database = self._cosmos_client.create_database_if_not_exists( + id=chat_db_name, + offer_throughput=cosmos_database_properties.get("offer_throughput"), + session_token=cosmos_database_properties.get("session_token"), + initial_headers=cosmos_database_properties.get("initial_headers"), + etag=cosmos_database_properties.get("etag"), + match_condition=cosmos_database_properties.get("match_condition"), + ) + + # Create the collection if it already doesn't exist + self._container = self._database.create_container_if_not_exists( + id=chat_container_name, + partition_key=cosmos_container_properties["partition_key"], + indexing_policy=cosmos_container_properties.get("indexing_policy"), + default_ttl=cosmos_container_properties.get("default_ttl"), + offer_throughput=cosmos_container_properties.get("offer_throughput"), + unique_key_policy=cosmos_container_properties.get("unique_key_policy"), + conflict_resolution_policy=cosmos_container_properties.get( + "conflict_resolution_policy" + ), + analytical_storage_ttl=cosmos_container_properties.get( + "analytical_storage_ttl" + ), + computed_properties=cosmos_container_properties.get("computed_properties"), + etag=cosmos_container_properties.get("etag"), + match_condition=cosmos_container_properties.get("match_condition"), + session_token=cosmos_container_properties.get("session_token"), + initial_headers=cosmos_container_properties.get("initial_headers"), + ) + + @classmethod + def from_connection_string( + cls, + connection_string: str, + chat_db_name: str = DEFAULT_CHAT_DATABASE, + chat_container_name: str = DEFAULT_CHAT_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + ): + """Creates an instance of Azure Cosmos DB NoSql Chat Store using a connection string.""" + cosmos_client = CosmosClient.from_connection_string(connection_string) + + return cls( + cosmos_client, + chat_db_name, + chat_container_name, + cosmos_container_properties, + cosmos_database_properties, + ) + + @classmethod + def from_account_and_key( + cls, + endpoint: str, + key: str, + chat_db_name: str = DEFAULT_CHAT_DATABASE, + chat_container_name: str = DEFAULT_CHAT_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + ) -> "AzureCosmosNoSqlChatStore": + """Initializes AzureCosmosNoSqlChatStore from an endpoint url and key.""" + cosmos_client = CosmosClient(endpoint, key) + return cls( + cosmos_client, + chat_db_name, + chat_container_name, + cosmos_container_properties, + cosmos_database_properties, + ) + + @classmethod + def from_aad_token( + cls, + endpoint: str, + chat_db_name: str = DEFAULT_CHAT_DATABASE, + chat_container_name: str = DEFAULT_CHAT_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + ) -> "AzureCosmosNoSqlChatStore": + """Creates an AzureChatStore using an Azure Active Directory token.""" + from azure.identity import DefaultAzureCredential + + credential = DefaultAzureCredential() + return cls._from_clients( + endpoint, + credential, + chat_db_name, + chat_container_name, + cosmos_container_properties, + cosmos_database_properties, + ) + + def set_messages(self, key: str, messages: List[ChatMessage]) -> None: + """Set messages for a key.""" + if not self._container: + raise ValueError("Container not initialized") + self._container.upsert_item( + body={ + "id": key, + "messages": _messages_to_dict(messages), + } + ) + + def get_messages(self, key: str) -> List[ChatMessage]: + """Get messages for a key.""" + response = self._container.read_item(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._container.create_item( + body={ + "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._container.delete_item(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.""" + items = self._container.read_all_items() + return [item["id"] for item in items] + + @classmethod + def class_name(cls) -> str: + """Get class name.""" + return "AzureCosmosNoSqlChatStore" + + @classmethod + def _from_clients( + cls, + endpoint: str, + credential: Any, + chat_db_name: str = DEFAULT_CHAT_DATABASE, + chat_container_name: str = DEFAULT_CHAT_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + ) -> "AzureCosmosNoSqlChatStore": + """Create cosmos db service clients.""" + cosmos_client = CosmosClient(url=endpoint, credential=credential) + return cls( + cosmos_client, + chat_db_name, + chat_container_name, + cosmos_container_properties, + cosmos_database_properties, + ) diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/pyproject.toml b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/pyproject.toml new file mode 100644 index 0000000000000..9953c7db1a35b --- /dev/null +++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/pyproject.toml @@ -0,0 +1,64 @@ +[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.azurecosmosnosql" + +[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 "] +description = "llama-index storage-chat-store azure cosmosdb nosql integration" +exclude = ["**/BUILD"] +license = "MIT" +name = "llama-index-storage-chat-store-azurecosmosnosql" +readme = "README.md" +version = "1.0.0" + +[tool.poetry.dependencies] +python = ">=3.8.1,<4.0" +azure-cosmos = "^4.7.0" +azure-identity = "^1.7.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/" diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/tests/BUILD b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/tests/BUILD new file mode 100644 index 0000000000000..dabf212d7e716 --- /dev/null +++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/tests/BUILD @@ -0,0 +1 @@ +python_tests() diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/tests/__init__.py b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/tests/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/tests/test_storage_azurecosmosnosql_chat_store.py b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/tests/test_storage_azurecosmosnosql_chat_store.py new file mode 100644 index 0000000000000..5f1495926ec26 --- /dev/null +++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-azurecosmosnosql/tests/test_storage_azurecosmosnosql_chat_store.py @@ -0,0 +1,7 @@ +from llama_index.core.storage.chat_store.base import BaseChatStore +from llama_index.storage.chat_store.azurecosmosnosql import AzureCosmosNoSqlChatStore + + +def test_class(): + names_of_base_classes = [b.__name__ for b in AzureCosmosNoSqlChatStore.__mro__] + assert BaseChatStore.__name__ in names_of_base_classes diff --git a/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/BUILD b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/BUILD new file mode 100644 index 0000000000000..0896ca890d8bf --- /dev/null +++ b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/BUILD @@ -0,0 +1,3 @@ +poetry_requirements( + name="poetry", +) diff --git a/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/Makefile b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/Makefile new file mode 100644 index 0000000000000..b9eab05aa3706 --- /dev/null +++ b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/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/ diff --git a/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/README.md b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/README.md new file mode 100644 index 0000000000000..08bd00a6de0ec --- /dev/null +++ b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/README.md @@ -0,0 +1 @@ +# LlamaIndex Index_Store Integration: Azure CosmosDB NoSQL Index Store diff --git a/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/llama_index/storage/docstore/azurecosmosnosql/BUILD b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/llama_index/storage/docstore/azurecosmosnosql/BUILD new file mode 100644 index 0000000000000..db46e8d6c978c --- /dev/null +++ b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/llama_index/storage/docstore/azurecosmosnosql/BUILD @@ -0,0 +1 @@ +python_sources() diff --git a/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/llama_index/storage/docstore/azurecosmosnosql/__init__.py b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/llama_index/storage/docstore/azurecosmosnosql/__init__.py new file mode 100644 index 0000000000000..834628c663af4 --- /dev/null +++ b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/llama_index/storage/docstore/azurecosmosnosql/__init__.py @@ -0,0 +1,5 @@ +from llama_index.storage.docstore.azurecosmosnosql.base import ( + AzureCosmosNoSqlDocumentStore, +) + +__all__ = ["AzureCosmosNoSqlDocumentStore"] diff --git a/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/llama_index/storage/docstore/azurecosmosnosql/base.py b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/llama_index/storage/docstore/azurecosmosnosql/base.py new file mode 100644 index 0000000000000..85a93d748250e --- /dev/null +++ b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/llama_index/storage/docstore/azurecosmosnosql/base.py @@ -0,0 +1,82 @@ +from typing import Any, Dict, Optional + +from llama_index.core.storage.docstore.keyval_docstore import KVDocumentStore +from llama_index.storage.kvstore.azurecosmosnosql import AzureCosmosNoSqlKVStore + +DEFAULT_DOCUMENT_DATABASE = "DocumentStoreDB" +DEFAULT_DOCUMENT_CONTAINER = "DocumentStoreContainer" + + +class AzureCosmosNoSqlDocumentStore(KVDocumentStore): + """Creates an AzureCosmosNoSqlDocumentStore.""" + + def __init__( + self, + azure_cosmos_nosql_kvstore: AzureCosmosNoSqlKVStore, + namespace: Optional[str] = None, + collection_suffix: Optional[str] = None, + ) -> None: + """Initializes the AzureCosmosNoSqlDocumentStore.""" + super().__init__(azure_cosmos_nosql_kvstore, namespace, collection_suffix) + + @classmethod + def from_connection_string( + cls, + connection_string: str, + document_db_name: str = DEFAULT_DOCUMENT_DATABASE, + document_container_name: str = DEFAULT_DOCUMENT_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + ) -> "AzureCosmosNoSqlDocumentStore": + """Creates an instance of AzureCosmosNoSqlDocumentStore using a connection string.""" + azure_cosmos_nosql_kvstore = AzureCosmosNoSqlKVStore.from_connection_string( + connection_string, + document_db_name, + document_container_name, + cosmos_container_properties, + cosmos_database_properties, + ) + namespace = document_db_name + "." + document_container_name + return cls(azure_cosmos_nosql_kvstore, namespace) + + @classmethod + def from_account_and_key( + cls, + endpoint: str, + key: str, + document_db_name: str = DEFAULT_DOCUMENT_DATABASE, + document_container_name: str = DEFAULT_DOCUMENT_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + ) -> "AzureCosmosNoSqlDocumentStore": + """Creates an instance of AzureCosmosNoSqlDocumentStore using an account endpoint and key.""" + azure_cosmos_nosql_kvstore = AzureCosmosNoSqlKVStore.from_account_and_key( + endpoint, + key, + document_db_name, + document_container_name, + cosmos_container_properties, + cosmos_database_properties, + ) + namespace = document_db_name + "." + document_container_name + return cls(azure_cosmos_nosql_kvstore, namespace) + + @classmethod + def from_aad_token( + cls, + endpoint: str, + document_db_name: str = DEFAULT_DOCUMENT_DATABASE, + document_container_name: str = DEFAULT_DOCUMENT_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + ) -> "AzureCosmosNoSqlDocumentStore": + """Creates an instance of AzureCosmosNoSqlDocumentStore using an aad token.""" + azure_cosmos_nosql_kvstore = AzureCosmosNoSqlKVStore.from_aad_token( + endpoint, + document_db_name, + document_container_name, + cosmos_container_properties, + cosmos_database_properties, + ) + namespace = document_db_name + "." + document_container_name + return cls(azure_cosmos_nosql_kvstore, namespace) diff --git a/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/pyproject.toml b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/pyproject.toml new file mode 100644 index 0000000000000..e374b397d6637 --- /dev/null +++ b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/pyproject.toml @@ -0,0 +1,64 @@ +[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.index_store.azurecosmosnosql" + +[tool.llamahub.class_authors] +AzureCosmosNoSqlIndexStore = "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 "] +description = "llama-index storage-index-store azure cosmosdb nosql integration" +exclude = ["**/BUILD"] +license = "MIT" +name = "llama-index-storage-index-store-azurecosmosnosql" +readme = "README.md" +version = "1.0.0" + +[tool.poetry.dependencies] +python = ">=3.8.1,<4.0" +azure-cosmos = "^4.7.0" +azure-identity = "^1.7.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/" diff --git a/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/tests/BUILD b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/tests/BUILD new file mode 100644 index 0000000000000..dabf212d7e716 --- /dev/null +++ b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/tests/BUILD @@ -0,0 +1 @@ +python_tests() diff --git a/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/tests/__init__.py b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/tests/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/tests/test_storage_document_store_azurecosmosnosql.py b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/tests/test_storage_document_store_azurecosmosnosql.py new file mode 100644 index 0000000000000..c27e94a0f6676 --- /dev/null +++ b/llama-index-integrations/storage/docstore/llama-index-storage-docstore-azurecosmosnosql/tests/test_storage_document_store_azurecosmosnosql.py @@ -0,0 +1,7 @@ +from llama_index.core.storage.docstore.keyval_docstore import KVDocumentStore +from llama_index.storage.docstore.azurecosmosnosql import AzureCosmosNoSqlDocumentStore + + +def test_class(): + names_of_base_classes = [b.__name__ for b in AzureCosmosNoSqlDocumentStore.__mro__] + assert KVDocumentStore.__name__ in names_of_base_classes diff --git a/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/BUILD b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/BUILD new file mode 100644 index 0000000000000..0896ca890d8bf --- /dev/null +++ b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/BUILD @@ -0,0 +1,3 @@ +poetry_requirements( + name="poetry", +) diff --git a/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/Makefile b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/Makefile new file mode 100644 index 0000000000000..b9eab05aa3706 --- /dev/null +++ b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/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/ diff --git a/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/README.md b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/README.md new file mode 100644 index 0000000000000..08bd00a6de0ec --- /dev/null +++ b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/README.md @@ -0,0 +1 @@ +# LlamaIndex Index_Store Integration: Azure CosmosDB NoSQL Index Store diff --git a/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/llama_index/storage/index_store/azurecosmosnosql/BUILD b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/llama_index/storage/index_store/azurecosmosnosql/BUILD new file mode 100644 index 0000000000000..db46e8d6c978c --- /dev/null +++ b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/llama_index/storage/index_store/azurecosmosnosql/BUILD @@ -0,0 +1 @@ +python_sources() diff --git a/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/llama_index/storage/index_store/azurecosmosnosql/__init__.py b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/llama_index/storage/index_store/azurecosmosnosql/__init__.py new file mode 100644 index 0000000000000..b404c9045555b --- /dev/null +++ b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/llama_index/storage/index_store/azurecosmosnosql/__init__.py @@ -0,0 +1,5 @@ +from llama_index.storage.index_store.azurecosmosnosql.base import ( + AzureCosmosNoSqlIndexStore, +) + +__all__ = ["AzureCosmosNoSqlIndexStore"] diff --git a/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/llama_index/storage/index_store/azurecosmosnosql/base.py b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/llama_index/storage/index_store/azurecosmosnosql/base.py new file mode 100644 index 0000000000000..1e989303466dd --- /dev/null +++ b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/llama_index/storage/index_store/azurecosmosnosql/base.py @@ -0,0 +1,82 @@ +from typing import Any, Dict, Optional + +from llama_index.core.storage.index_store.keyval_index_store import KVIndexStore +from llama_index.storage.kvstore.azurecosmosnosql import AzureCosmosNoSqlKVStore + +DEFAULT_INDEX_DATABASE = "IndexStoreDB" +DEFAULT_INDEX_CONTAINER = "IndexStoreContainer" + + +class AzureCosmosNoSqlIndexStore(KVIndexStore): + """Creates an Azure Cosmos DB NoSql Index Store.""" + + def __init__( + self, + azure_cosmos_nosql_kvstore: AzureCosmosNoSqlKVStore, + namespace: Optional[str] = None, + collection_suffix: Optional[str] = None, + ) -> None: + """Initializes the Azure Cosmos NoSql Index Store.""" + super().__init__(azure_cosmos_nosql_kvstore, namespace, collection_suffix) + + @classmethod + def from_connection_string( + cls, + connection_string: str, + index_db_name: str = DEFAULT_INDEX_DATABASE, + index_container_name: str = DEFAULT_INDEX_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + ) -> "AzureCosmosNoSqlIndexStore": + """Creates an instance of Azure Cosmos DB NoSql KV Store using a connection string.""" + azure_cosmos_nosql_kvstore = AzureCosmosNoSqlKVStore.from_connection_string( + connection_string, + index_db_name, + index_container_name, + cosmos_container_properties, + cosmos_database_properties, + ) + namespace = index_db_name + "." + index_container_name + return cls(azure_cosmos_nosql_kvstore, namespace) + + @classmethod + def from_account_and_key( + cls, + endpoint: str, + key: str, + index_db_name: str = DEFAULT_INDEX_DATABASE, + index_container_name: str = DEFAULT_INDEX_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + ) -> "AzureCosmosNoSqlIndexStore": + """Creates an instance of Azure Cosmos DB NoSql KV Store using an account endpoint and key.""" + azure_cosmos_nosql_kvstore = AzureCosmosNoSqlKVStore.from_account_and_key( + endpoint, + key, + index_db_name, + index_container_name, + cosmos_container_properties, + cosmos_database_properties, + ) + namespace = index_db_name + "." + index_container_name + return cls(azure_cosmos_nosql_kvstore, namespace) + + @classmethod + def from_aad_token( + cls, + endpoint: str, + index_db_name: str = DEFAULT_INDEX_DATABASE, + index_container_name: str = DEFAULT_INDEX_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + ) -> "AzureCosmosNoSqlIndexStore": + """Creates an instance of Azure Cosmos DB NoSql KV Store using an aad token.""" + azure_cosmos_nosql_kvstore = AzureCosmosNoSqlKVStore.from_aad_token( + endpoint, + index_db_name, + index_container_name, + cosmos_container_properties, + cosmos_database_properties, + ) + namespace = index_db_name + "." + index_container_name + return cls(azure_cosmos_nosql_kvstore, namespace) diff --git a/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/pyproject.toml b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/pyproject.toml new file mode 100644 index 0000000000000..e374b397d6637 --- /dev/null +++ b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/pyproject.toml @@ -0,0 +1,64 @@ +[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.index_store.azurecosmosnosql" + +[tool.llamahub.class_authors] +AzureCosmosNoSqlIndexStore = "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 "] +description = "llama-index storage-index-store azure cosmosdb nosql integration" +exclude = ["**/BUILD"] +license = "MIT" +name = "llama-index-storage-index-store-azurecosmosnosql" +readme = "README.md" +version = "1.0.0" + +[tool.poetry.dependencies] +python = ">=3.8.1,<4.0" +azure-cosmos = "^4.7.0" +azure-identity = "^1.7.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/" diff --git a/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/tests/BUILD b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/tests/BUILD new file mode 100644 index 0000000000000..dabf212d7e716 --- /dev/null +++ b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/tests/BUILD @@ -0,0 +1 @@ +python_tests() diff --git a/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/tests/__init__.py b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/tests/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/tests/test_storage_index_store_azurecosmosnosql.py b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/tests/test_storage_index_store_azurecosmosnosql.py new file mode 100644 index 0000000000000..fc30cca56c07e --- /dev/null +++ b/llama-index-integrations/storage/index_store/llama-index-storage-index-store-azurecosmosnosql/tests/test_storage_index_store_azurecosmosnosql.py @@ -0,0 +1,7 @@ +from llama_index.core.storage.index_store.keyval_index_store import KVIndexStore +from llama_index.storage.index_store.azurecosmosnosql import AzureCosmosNoSqlIndexStore + + +def test_class(): + names_of_base_classes = [b.__name__ for b in AzureCosmosNoSqlIndexStore.__mro__] + assert KVIndexStore.__name__ in names_of_base_classes diff --git a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/BUILD b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/BUILD new file mode 100644 index 0000000000000..0896ca890d8bf --- /dev/null +++ b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/BUILD @@ -0,0 +1,3 @@ +poetry_requirements( + name="poetry", +) diff --git a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/Makefile b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/Makefile new file mode 100644 index 0000000000000..b9eab05aa3706 --- /dev/null +++ b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/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/ diff --git a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/README.md b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/README.md new file mode 100644 index 0000000000000..6fa827342b93c --- /dev/null +++ b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/README.md @@ -0,0 +1 @@ +# LlamaIndex KV_Store Integration: Azure CosmosDB NoSQL KV Store diff --git a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/llama_index/storage/kvstore/azurecosmosnosql/BUILD b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/llama_index/storage/kvstore/azurecosmosnosql/BUILD new file mode 100644 index 0000000000000..db46e8d6c978c --- /dev/null +++ b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/llama_index/storage/kvstore/azurecosmosnosql/BUILD @@ -0,0 +1 @@ +python_sources() diff --git a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/llama_index/storage/kvstore/azurecosmosnosql/__init__.py b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/llama_index/storage/kvstore/azurecosmosnosql/__init__.py new file mode 100644 index 0000000000000..06ee6b9173d30 --- /dev/null +++ b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/llama_index/storage/kvstore/azurecosmosnosql/__init__.py @@ -0,0 +1,3 @@ +from llama_index.storage.kvstore.azurecosmosnosql.base import AzureCosmosNoSqlKVStore + +__all__ = ["AzureCosmosNoSqlKVStore"] diff --git a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/llama_index/storage/kvstore/azurecosmosnosql/base.py b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/llama_index/storage/kvstore/azurecosmosnosql/base.py new file mode 100644 index 0000000000000..55869951a9685 --- /dev/null +++ b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/llama_index/storage/kvstore/azurecosmosnosql/base.py @@ -0,0 +1,245 @@ +import logging +from abc import ABC +from typing import Any, Dict, Optional + +from azure.cosmos import CosmosClient, DatabaseProxy, ContainerProxy +from llama_index.core.bridge.pydantic import PrivateAttr +from llama_index.core.storage.kvstore.types import ( + BaseKVStore, + DEFAULT_COLLECTION, +) + +DEFAULT_CHAT_DATABASE = "KVStoreDB" +DEFAULT_CHAT_CONTAINER = "KVStoreContainer" + + +logger = logging.getLogger(__name__) + + +class AzureCosmosNoSqlKVStore(BaseKVStore, ABC): + """Creates an Azure Cosmos DB NoSql Chat Store.""" + + _cosmos_client: CosmosClient = PrivateAttr() + _database: DatabaseProxy = PrivateAttr() + _container: ContainerProxy = PrivateAttr() + + def __init__( + self, + cosmos_client: CosmosClient, + chat_db_name: str = DEFAULT_CHAT_DATABASE, + chat_container_name: str = DEFAULT_CHAT_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + **kwargs, + ): + self._cosmos_client = cosmos_client + + # Create the database if it already doesn't exist + self._database = self._cosmos_client.create_database_if_not_exists( + id=chat_db_name, + offer_throughput=cosmos_database_properties.get("offer_throughput"), + session_token=cosmos_database_properties.get("session_token"), + initial_headers=cosmos_database_properties.get("initial_headers"), + etag=cosmos_database_properties.get("etag"), + match_condition=cosmos_database_properties.get("match_condition"), + ) + + # Create the collection if it already doesn't exist + self._container = self._database.create_container_if_not_exists( + id=chat_container_name, + partition_key=cosmos_container_properties["partition_key"], + indexing_policy=cosmos_container_properties.get("indexing_policy"), + default_ttl=cosmos_container_properties.get("default_ttl"), + offer_throughput=cosmos_container_properties.get("offer_throughput"), + unique_key_policy=cosmos_container_properties.get("unique_key_policy"), + conflict_resolution_policy=cosmos_container_properties.get( + "conflict_resolution_policy" + ), + analytical_storage_ttl=cosmos_container_properties.get( + "analytical_storage_ttl" + ), + computed_properties=cosmos_container_properties.get("computed_properties"), + etag=cosmos_container_properties.get("etag"), + match_condition=cosmos_container_properties.get("match_condition"), + session_token=cosmos_container_properties.get("session_token"), + initial_headers=cosmos_container_properties.get("initial_headers"), + ) + + @classmethod + def from_connection_string( + cls, + connection_string: str, + chat_db_name: str = DEFAULT_CHAT_DATABASE, + chat_container_name: str = DEFAULT_CHAT_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + ) -> "AzureCosmosNoSqlKVStore": + """Creates an instance of Azure Cosmos DB NoSql KV Store using a connection string.""" + cosmos_client = CosmosClient.from_connection_string(connection_string) + + return cls( + cosmos_client, + chat_db_name, + chat_container_name, + cosmos_container_properties, + cosmos_database_properties, + ) + + @classmethod + def from_account_and_key( + cls, + endpoint: str, + key: str, + chat_db_name: str = DEFAULT_CHAT_DATABASE, + chat_container_name: str = DEFAULT_CHAT_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + ) -> "AzureCosmosNoSqlKVStore": + """Initializes AzureCosmosNoSqlKVStore from an endpoint url and key.""" + cosmos_client = CosmosClient(endpoint, key) + return cls( + cosmos_client, + chat_db_name, + chat_container_name, + cosmos_container_properties, + cosmos_database_properties, + ) + + @classmethod + def from_aad_token( + cls, + endpoint: str, + chat_db_name: str = DEFAULT_CHAT_DATABASE, + chat_container_name: str = DEFAULT_CHAT_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + ) -> "AzureCosmosNoSqlKVStore": + """Creates an AzureCosmosNoSqlKVStore using an Azure Active Directory token.""" + from azure.identity import DefaultAzureCredential + + credential = DefaultAzureCredential() + return cls._from_clients( + endpoint, + credential, + chat_db_name, + chat_container_name, + cosmos_container_properties, + cosmos_database_properties, + ) + + def put(self, key: str, val: dict, collection: str = DEFAULT_COLLECTION) -> None: + """Put a key-value pair into the store. + + Args: + key (str): key + val (dict): value + collection (str): collection name + """ + self._container.create_item( + body={ + "id": key, + "messages": val, + } + ) + + async def aput( + self, key: str, val: dict, collection: str = DEFAULT_COLLECTION + ) -> None: + """Put a key-value pair into the store. + + Args: + key (str): key + val (dict): value + collection (str): collection name + """ + raise NotImplementedError + + def get(self, key: str, collection: str = DEFAULT_COLLECTION) -> Optional[dict]: + """Get a value from the store. + + Args: + key (str): key + collection (str): collection name + + """ + response = self._container.read_item(key) + if response is not None: + messages = response.get("messages") + else: + messages = {} + return messages + + async def aget( + self, key: str, collection: str = DEFAULT_COLLECTION + ) -> Optional[dict]: + """Get a value from the store. + + Args: + key (str): key + collection (str): collection name + """ + raise NotImplementedError + + def get_all(self, collection: str = DEFAULT_COLLECTION) -> Dict[str, dict]: + """Get all values from the store. + + Args: + collection (str): collection name + + """ + items = self._container.read_all_items() + output = {} + for item in items: + key = item.get("id") + output[key] = item + return output + + async def aget_all(self, collection: str = DEFAULT_COLLECTION) -> Dict[str, dict]: + """Get all values from the store. + + Args: + collection (str): collection name + """ + raise NotImplementedError + + def delete(self, key: str, collection: str = DEFAULT_COLLECTION) -> bool: + try: + self._container.delete_item(key) + return True + except Exception as e: + logger.error(f"Error deleting item {e} with key {key}") + return False + + async def adelete(self, key: str, collection: str = DEFAULT_COLLECTION) -> bool: + """Delete a value from the store. + + Args: + key (str): key + collection (str): collection name + """ + raise NotImplementedError + + @classmethod + def class_name(cls) -> str: + """Get class name.""" + return "AzureCosmosNoSqlKVStore" + + @classmethod + def _from_clients( + cls, + endpoint: str, + credential: Any, + chat_db_name: str = DEFAULT_CHAT_DATABASE, + chat_container_name: str = DEFAULT_CHAT_CONTAINER, + cosmos_container_properties: Dict[str, Any] = None, + cosmos_database_properties: Dict[str, Any] = None, + ) -> "AzureCosmosNoSqlKVStore": + """Create cosmos db service clients.""" + cosmos_client = CosmosClient(url=endpoint, credential=credential) + return cls( + cosmos_client, + chat_db_name, + chat_container_name, + cosmos_container_properties, + cosmos_database_properties, + ) diff --git a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/pyproject.toml b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/pyproject.toml new file mode 100644 index 0000000000000..50a747e321a17 --- /dev/null +++ b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/pyproject.toml @@ -0,0 +1,64 @@ +[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.kvstore.azurecosmosnosql" + +[tool.llamahub.class_authors] +AzureCosmosNoSqlKVStore = "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 "] +description = "llama-index storage-kv-store azure cosmosdb nosql integration" +exclude = ["**/BUILD"] +license = "MIT" +name = "llama-index-storage-kv-store-azurecosmosnosql" +readme = "README.md" +version = "1.0.0" + +[tool.poetry.dependencies] +python = ">=3.8.1,<4.0" +azure-cosmos = "^4.7.0" +azure-identity = "^1.7.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/" diff --git a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/tests/BUILD b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/tests/BUILD new file mode 100644 index 0000000000000..dabf212d7e716 --- /dev/null +++ b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/tests/BUILD @@ -0,0 +1 @@ +python_tests() diff --git a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/tests/__init__.py b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/tests/__init__.py new file mode 100644 index 0000000000000..e69de29bb2d1d diff --git a/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/tests/test_storage_azurecosmosnosql_kv_store.py b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/tests/test_storage_azurecosmosnosql_kv_store.py new file mode 100644 index 0000000000000..cef55ea33e7d0 --- /dev/null +++ b/llama-index-integrations/storage/kvstore/llama-index-storage-kvstore-azurecosmosnosql/tests/test_storage_azurecosmosnosql_kv_store.py @@ -0,0 +1,7 @@ +from llama_index.core.storage.kvstore.types import BaseKVStore +from llama_index.storage.kvstore.azurecosmosnosql import AzureCosmosNoSqlKVStore + + +def test_class(): + names_of_base_classes = [b.__name__ for b in AzureCosmosNoSqlKVStore.__mro__] + assert BaseKVStore.__name__ in names_of_base_classes