Repository: run-llama/llama_index
Pull request: Feature Request #15810 :: Add DynamoDBChatStore
# Description

Adds support for DynamoDB as a chat store.  Follows AWS credential pattern of [Bedrock](https://github.com/run-llama/llama_index/blob/162f5a0523f5a4de33f8cc056ec2b24713d2ee9e/llama-index-integrations/llms/llama-index-llms-bedrock/llama_index/llms/bedrock/base.py#L66) and [Bedrock Converse](https://github.com/run-llama/llama_index/blob/162f5a0523f5a4de33f8cc056ec2b24713d2ee9e/llama-index-integrations/llms/llama-index-llms-bedrock-converse/llama_index/llms/bedrock_converse/base.py#L86) LLMs.

Has the following dependencies:
* `boto3`
* `boto3-stubs`

DEV:
*  `moto`

Fixes #15810 

## New Package?

Did I fill in the `tool.llamahub` section in the `pyproject.toml` and provide a detailed README.md for my new integration or package?

- [x] Yes
- [ ] No

## Version Bump?

Did I bump the version in the `pyproject.toml` file of the package I am updating? (Except for the `llama-index-core` package)

- [x] Yes - `0.1.0`
- [ ] No

## Type of Change

Please delete options that are not relevant.

- [x] New feature (non-breaking change which adds functionality)
- [x] This change requires a documentation update

## How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration

- [x] Added new unit/integration tests
- [x] Added new notebook (that tests end-to-end) - Not a notebook, but the updated README examples run end-to-end.
- [x] I stared at the code and made sure it makes sense

## Suggested Checklist:

- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my feature works
- [x] New and existing unit tests pass locally with my changes
- [x] I ran `make format; make lint` to appease the lint gods


Problem statement:
[Feature Request]: DynamoDBChatStore
### Feature Description

DynamoDB support for storing conversation histories.

### Reason

_No response_

### Value of Feature

Expands architecture options.  Currently, only two remote ChatStores are supported with LlamaIndex, which is fairly limiting (particularly with respect to LangChain.)

Hints from issue discussion:

Commit message: Merge branch 'main' into dynamodbchatstore

Patch to review:
diff --git a/docs/docs/module_guides/storing/chat_stores.md b/docs/docs/module_guides/storing/chat_stores.md
index 9e3c493bbf..22c3b94716 100644
--- a/docs/docs/module_guides/storing/chat_stores.md
+++ b/docs/docs/module_guides/storing/chat_stores.md
@@ -97,3 +97,67 @@ chat_engine = SimpleChatEngine(
 
 response = chat_engine.chat("Hello.")
 ```
+
+## DynamoDBChatStore
+
+Using `DynamoDBChatStore`, you can store your chat history in AWS DynamoDB.
+
+### Installation
+```bash
+pip install llama-index-storage-chat-store-dynamodb
+```
+
+### Usage
+Ensure you have a DynamoDB table created with the appropriate schema. By default, here is an example:
+```python
+import boto3
+
+# Get the service resource.
+dynamodb = boto3.resource("dynamodb")
+
+# Create the DynamoDB table.
+table = dynamodb.create_table(
+    TableName="EXAMPLE_TABLE",
+    KeySchema=[{"AttributeName": "SessionId", "KeyType": "HASH"}],
+    AttributeDefinitions=[
+        {"AttributeName": "SessionId", "AttributeType": "S"}
+    ],
+    BillingMode="PAY_PER_REQUEST",
+)
+```
+
+You can then use the `DynamoDBChatStore` class to persist and retrieve chat histories:
+```python
+import os
+from llama_index.core.llms import ChatMessage, MessageRole
+from llama_index.storage.chat_store.dynamodb.base import DynamoDBChatStore
+
+# Initialize DynamoDB chat store
+chat_store = DynamoDBChatStore(
+    table_name="EXAMPLE_TABLE", profile_name=os.getenv("AWS_PROFILE")
+)
+
+# A chat history, which doesn't exist yet, returns an empty array.
+print(chat_store.get_messages("123"))
+# >>> []
+
+# Initializing a chat history with a key of "SessionID = 123"
+messages = [
+    ChatMessage(role=MessageRole.USER, content="Who are you?"),
+    ChatMessage(
+        role=MessageRole.ASSISTANT, content="I am your helpful AI assistant."
+    ),
+]
+chat_store.set_messages(key="123", messages=messages)
+print(chat_store.get_messages("123"))
+# >>> [ChatMessage(role=<MessageRole.USER: 'user'>, content='Who are you?', additional_kwargs={}),
+#      ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>, content='I am your helpful AI assistant.', additional_kwargs={})]]
+
+# Appending a message to an existing chat history
+message = ChatMessage(role=MessageRole.USER, content="What can you do?")
+chat_store.add_message(key="123", message=message)
+print(chat_store.get_messages("123"))
+# >>> [ChatMessage(role=<MessageRole.USER: 'user'>, content='Who are you?', additional_kwargs={}),
+#      ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>, content='I am your helpful AI assistant.', additional_kwargs={})],
+#      ChatMessage(role=<MessageRole.USER: 'user'>, content='What can you do?', additional_kwargs={})]
+```
diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/.gitignore b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/.gitignore
new file mode 100644
index 0000000000..990c18de22
--- /dev/null
+++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/.gitignore
@@ -0,0 +1,153 @@
+llama_index/_static
+.DS_Store
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+bin/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+etc/
+include/
+lib/
+lib64/
+parts/
+sdist/
+share/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+.ruff_cache
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+notebooks/
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+pyvenv.cfg
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# Jetbrains
+.idea
+modules/
+*.swp
+
+# VsCode
+.vscode
+
+# pipenv
+Pipfile
+Pipfile.lock
+
+# pyright
+pyrightconfig.json
diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/BUILD b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/BUILD
new file mode 100644
index 0000000000..0896ca890d
--- /dev/null
+++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/BUILD
@@ -0,0 +1,3 @@
+poetry_requirements(
+    name="poetry",
+)
diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/Makefile b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/Makefile
new file mode 100644
index 0000000000..b9eab05aa3
--- /dev/null
+++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/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-dynamodb/README.md b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/README.md
new file mode 100644
index 0000000000..b6c9baf30e
--- /dev/null
+++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/README.md
@@ -0,0 +1,57 @@
+# LlamaIndex Chat_Store Integration: DynamoDB Chat Store
+
+This enables AWS DynamoDB to be used as a chat store.
+
+## Installation
+
+```bash
+pip install llama-index-storage-chat-store-dynamodb
+```
+
+## Usage
+
+### Assumptions
+
+- `SessionID`, a string, is used as the partition key.
+- The table used for the chat store already exists. Here is an example for that:
+
+  ```python
+  import boto3
+
+  dynamodb = boto3.resource("dynamodb")
+
+  # Create the DynamoDB table.
+  table = dynamodb.create_table(
+      TableName="EXAMPLE_TABLE",
+      KeySchema=[{"AttributeName": "SessionId", "KeyType": "HASH"}],
+      AttributeDefinitions=[
+          {"AttributeName": "SessionId", "AttributeType": "S"}
+      ],
+      BillingMode="PAY_PER_REQUEST",
+  )
+  ```
+
+### Using an AWS IAM Role
+
+You can use any of the following AWS arguments to setup the required `boto3` resource connection:
+
+- `profile_name`
+- `aws_access_key_id`
+- `aws_secret_access_key`
+- `aws_session_token`
+- `botocore_session` - A pre-existing existing Botocore session.
+- `botocore_config`
+
+As an example, if you have already assumed an AWS profile in your local environment or within an AWS compute
+environment, you can simply do the following:
+
+```python
+import os
+from llama_index.storage.chat_store.dynamodb.base import DynamoDBChatStore
+
+store = DynamoDBChatStore(
+    profile_name=os.getenv("AWS_PROFILE"),
+    table_name="EXAMPLE_TABLE",
+    session_id="123",
+)
+```
diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/llama_index/storage/chat_store/dynamodb/BUILD b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/llama_index/storage/chat_store/dynamodb/BUILD
new file mode 100644
index 0000000000..db46e8d6c9
--- /dev/null
+++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/llama_index/storage/chat_store/dynamodb/BUILD
@@ -0,0 +1 @@
+python_sources()
diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/llama_index/storage/chat_store/dynamodb/__init__.py b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/llama_index/storage/chat_store/dynamodb/__init__.py
new file mode 100644
index 0000000000..48dfc0c164
--- /dev/null
+++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/llama_index/storage/chat_store/dynamodb/__init__.py
@@ -0,0 +1,3 @@
+from llama_index.storage.chat_store.dynamodb.base import DynamoDBChatStore
+
+__all__ = ["DynamoDBChatStore"]
diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/llama_index/storage/chat_store/dynamodb/base.py b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/llama_index/storage/chat_store/dynamodb/base.py
new file mode 100644
index 0000000000..4723f57c12
--- /dev/null
+++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/llama_index/storage/chat_store/dynamodb/base.py
@@ -0,0 +1,282 @@
+import logging
+from typing import Any, Dict, List, Optional
+
+from llama_index.core.base.llms.types import ChatMessage
+from llama_index.core.bridge.pydantic import Field, PrivateAttr
+from llama_index.core.storage.chat_store.base import BaseChatStore
+from mypy_boto3_dynamodb import ServiceResource
+
+
+logger = logging.getLogger(__name__)
+
+
+# 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 ChatMessage to a JSON object
+def _message_to_dict(message: ChatMessage) -> dict:
+    return message.dict()
+
+
+# Convert a JSON object to a ChatMessage
+def _dict_to_message(d: dict) -> ChatMessage:
+    return ChatMessage.model_validate(d)
+
+
+class DynamoDBChatStore(BaseChatStore):
+    """DynamoDB Chat Store.
+
+    Args:
+        table_name (str): The name of the preexisting DynamoDB table.
+        primary_key (str, optional): The primary/partition key to use for the table.
+            Defaults to "SessionId".
+        profile_name (str, optional): The AWS profile to use. If not specified, then
+            the default AWS profile is used.
+        aws_access_key_id (str, optional): The AWS Access Key ID to use.
+        aws_secret_access_key (str, optional): The AWS Secret Access Key to use.
+        aws_session_token (str, optional): The AWS Session Token to use.
+        botocore_session (Any, optional): Use this Botocore session instead of creating a new default one.
+        botocore_config (Any, optional): Custom configuration object to use instead of the default generated one.
+        region_name (str, optional): The AWS region name to use. Uses the region configured in AWS CLI if not passed.
+        max_retries (int, optional): The maximum number of API retries. Defaults to 10.
+        timeout (float, optional): The timeout for API requests in seconds. Defaults to 60.0.
+        session_kwargs (Dict[str, Any], optional): Additional kwargs for the `boto3.Session` object.
+        resource_kwargs (Dict[str, Any], optional): Additional kwargs for the `boto3.Resource` object.
+
+    Returns:
+        DynamoDBChatStore: A DynamoDB chat store object.
+    """
+
+    table_name: str = Field(description="DynamoDB table")
+    primary_key: str = Field(
+        default="SessionId", description="Primary/partition key to use for the table."
+    )
+    profile_name: Optional[str] = Field(
+        description="AWS profile to use. If not specified, then the default AWS profile is used."
+    )
+    aws_access_key_id: Optional[str] = Field(
+        description="AWS Access Key ID to use.", exclude=True
+    )
+    aws_secret_access_key: Optional[str] = Field(
+        description="AWS Secret Access Key to use.", exclude=True
+    )
+    aws_session_token: Optional[str] = Field(
+        description="AWS Session Token to use.", exclude=True
+    )
+    botocore_session: Optional[Any] = Field(
+        description="Use this Botocore session instead of creating a new default one.",
+        exclude=True,
+    )
+    botocore_config: Optional[Any] = Field(
+        description="Custom configuration object to use instead of the default generated one.",
+        exclude=True,
+    )
+    region_name: Optional[str] = Field(
+        description="AWS region name to use. Uses the region configured in AWS CLI if not passed",
+        exclude=True,
+    )
+    max_retries: int = Field(
+        default=10, description="The maximum number of API retries.", gt=0
+    )
+    timeout: float = Field(
+        default=60.0,
+        description="The timeout for API requests in seconds.",
+    )
+    session_kwargs: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Additional kwargs for the `boto3.Session` object.",
+    )
+    resource_kwargs: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Additional kwargs for the `boto3.Resource` object.",
+    )
+
+    _client: ServiceResource = PrivateAttr()
+    _table: Any = PrivateAttr()
+
+    def __init__(
+        self,
+        table_name: str,
+        primary_key: str = "SessionId",
+        profile_name: Optional[str] = None,
+        region_name: Optional[str] = None,
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None,
+        aws_session_token: Optional[str] = None,
+        botocore_session: Optional[Any] = None,
+        botocore_config: Optional[Any] = None,
+        max_retries: int = 10,
+        timeout: float = 60.0,
+        session_kwargs: Optional[Dict[str, Any]] = None,
+        resource_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        session_kwargs = session_kwargs or {}
+        resource_kwargs = resource_kwargs or {}
+
+        super().__init__(
+            table_name=table_name,
+            primary_key=primary_key,
+            profile_name=profile_name,
+            region_name=region_name,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+            aws_session_token=aws_session_token,
+            botocore_session=botocore_session,
+            botocore_config=botocore_config,
+            max_retries=max_retries,
+            timeout=timeout,
+            session_kwargs=session_kwargs,
+            resource_kwargs=resource_kwargs,
+        )
+
+        session_kwargs = {
+            "profile_name": profile_name,
+            "region_name": region_name,
+            "aws_access_key_id": aws_access_key_id,
+            "aws_secret_access_key": aws_secret_access_key,
+            "aws_session_token": aws_session_token,
+            "botocore_session": botocore_session,
+            **session_kwargs,
+        }
+
+        try:
+            import boto3
+            from botocore.config import Config
+
+            config = (
+                Config(
+                    retries={"max_attempts": max_retries, "mode": "standard"},
+                    connect_timeout=timeout,
+                    read_timeout=timeout,
+                )
+                if botocore_config is None
+                else botocore_config
+            )
+            session = boto3.Session(**session_kwargs)
+        except ImportError:
+            raise ImportError(
+                "boto3 package not found, install with 'pip install boto3"
+            )
+
+        self._client = session.resource("dynamodb", config=config, **resource_kwargs)
+        self._table = self._client.Table(table_name)
+
+    @classmethod
+    def class_name(self) -> str:
+        return "DynamoDBChatStore"
+
+    def set_messages(self, key: str, messages: List[ChatMessage]) -> None:
+        """Assign all provided messages to the row with the given key.
+        Any pre-existing messages for that key will be overwritten.
+
+        Args:
+            key (str): The key specifying a row.
+            messages (List[ChatMessage]): The messages to assign to the key.
+
+        Returns:
+            None
+        """
+        self._table.put_item(
+            Item={self.primary_key: key, "History": _messages_to_dict(messages)}
+        )
+
+    def get_messages(self, key: str) -> List[ChatMessage]:
+        """Retrieve all messages for the given key.
+
+        Args:
+            key (str): The key specifying a row.
+
+        Returns:
+            List[ChatMessage]: The messages associated with the key.
+        """
+        response = self._table.get_item(Key={self.primary_key: key})
+
+        if response and "Item" in response:
+            message_history = response["Item"]["History"]
+        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 to the end of the chat history for the given key.
+        Creates a new row if the key does not exist.
+
+        Args:
+            key (str): The key specifying a row.
+            message (ChatMessage): The message to add to the chat history.
+
+        Returns:
+            None
+        """
+        current_messages = _messages_to_dict(self.get_messages(key))
+        current_messages.append(_message_to_dict(message))
+
+        self._table.put_item(Item={self.primary_key: key, "History": current_messages})
+
+    def delete_messages(self, key: str) -> Optional[List[ChatMessage]]:
+        """Deletes the entire chat history for the given key (i.e. the row).
+
+        Args:
+            key (str): The key specifying a row.
+
+        Returns:
+            Optional[List[ChatMessage]]: The messages that were deleted. None if the
+                deletion failed.
+        """
+        messages_to_delete = self.get_messages(key)
+        self._table.delete_item(Key={self.primary_key: key})
+        return messages_to_delete
+
+    def delete_message(self, key: str, idx: int) -> Optional[ChatMessage]:
+        """Deletes the message at the given index for the given key.
+
+        Args:
+            key (str): The key specifying a row.
+            idx (int): The index of the message to delete.
+
+        Returns:
+            Optional[ChatMessage]: The message that was deleted. None if the index
+                did not exist.
+        """
+        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 session {key}")
+            )
+            return None
+
+    def delete_last_message(self, key: str) -> Optional[ChatMessage]:
+        """Deletes the last message in the chat history for the given key.
+
+        Args:
+            key (str): The key specifying a row.
+
+        Returns:
+            Optional[ChatMessage]: The message that was deleted. None if the chat history
+                was empty.
+        """
+        return self.delete_message(key, -1)
+
+    def get_keys(self) -> List[str]:
+        """Retrieve all keys in the table.
+
+        Returns:
+            List[str]: The keys in the table.
+        """
+        response = self._table.scan(ProjectionExpression="SessionId")
+        keys = [item["SessionId"] for item in response["Items"]]
+        while "LastEvaluatedKey" in response:
+            response = self._table.scan(
+                ProjectionExpression="SessionId",
+                ExclusiveStartKey=response["LastEvaluatedKey"],
+            )
+            keys.extend([item["SessionId"] for item in response["Items"]])
+        return keys
diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/pyproject.toml b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/pyproject.toml
new file mode 100644
index 0000000000..cc77432a9e
--- /dev/null
+++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/pyproject.toml
@@ -0,0 +1,65 @@
+[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.dynamodb"
+
+[tool.llamahub.class_authors]
+DynamoDBChatStore = "Bryce Freshcorn"
+
+[tool.mypy]
+disallow_untyped_defs = true
+exclude = ["_static", "build", "examples", "notebooks", "venv"]
+ignore_missing_imports = true
+python_version = "3.8"
+
+[tool.poetry]
+authors = ["Bryce Freshcorn <26725654+brycecf@users.noreply.github.com>"]
+description = "llama-index storage-chat-store aws dynamodb integration"
+exclude = ["**/BUILD"]
+license = "MIT"
+name = "llama-index-storage-chat-store-dynamodb"
+readme = "README.md"
+version = "0.1.0"
+
+[tool.poetry.dependencies]
+python = ">=3.8.1,<4.0"
+llama-index-core = "^0.11.0"
+boto3 = "^1.35.13"
+boto3-stubs = {extras = ["dynamodb"], version = "^1.35.14"}
+
+[tool.poetry.group.dev.dependencies]
+ipython = "8.10.0"
+jupyter = "^1.0.0"
+moto = {extras = ["dynamodb"], version = "^5.0.14"}
+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-dynamodb/tests/BUILD b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/tests/BUILD
new file mode 100644
index 0000000000..dabf212d7e
--- /dev/null
+++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/tests/BUILD
@@ -0,0 +1 @@
+python_tests()
diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/tests/__init__.py b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/tests/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/tests/test_chat_store_dynamodb_chat_store.py b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/tests/test_chat_store_dynamodb_chat_store.py
new file mode 100644
index 0000000000..579fab5a08
--- /dev/null
+++ b/llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/tests/test_chat_store_dynamodb_chat_store.py
@@ -0,0 +1,89 @@
+import pytest
+from moto import mock_aws
+import boto3
+from llama_index.storage.chat_store.dynamodb.base import DynamoDBChatStore
+from llama_index.core.base.llms.types import ChatMessage
+
+
+@pytest.fixture()
+def aws_credentials():
+    """Mocked AWS Credentials for moto."""
+    import os
+
+    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
+    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
+    os.environ["AWS_SECURITY_TOKEN"] = "testing"
+    os.environ["AWS_SESSION_TOKEN"] = "testing"
+
+
+@pytest.fixture()
+def dynamo_db(aws_credentials):
+    with mock_aws():
+        yield boto3.resource("dynamodb", region_name="us-east-1")
+
+
+@pytest.fixture()
+def chat_store(dynamo_db):
+    dynamo_db.create_table(
+        TableName="TestTable",
+        KeySchema=[{"AttributeName": "SessionId", "KeyType": "HASH"}],
+        AttributeDefinitions=[{"AttributeName": "SessionId", "AttributeType": "S"}],
+        ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1},
+    )
+    return DynamoDBChatStore(table_name="TestTable", region_name="us-east-1")
+
+
+def test_set_get_messages(chat_store):
+    messages = [ChatMessage(content="Hello"), ChatMessage(content="World")]
+    chat_store.set_messages("TestSession", messages)
+    retrieved_messages = chat_store.get_messages("TestSession")
+    assert len(retrieved_messages) == 2
+    assert retrieved_messages[0].content == "Hello"
+    assert retrieved_messages[1].content == "World"
+
+
+def test_add_message(chat_store):
+    initial_message = ChatMessage(content="Initial")
+    chat_store.add_message("TestSession", initial_message)
+    added_message = ChatMessage(content="Added")
+    chat_store.add_message("TestSession", added_message)
+    messages = chat_store.get_messages("TestSession")
+    assert len(messages) == 2
+    assert messages[1].content == "Added"
+
+
+def test_delete_messages(chat_store):
+    messages = [ChatMessage(content="Hello"), ChatMessage(content="World")]
+    chat_store.set_messages("TestSession", messages)
+    deleted_messages = chat_store.delete_messages("TestSession")
+    assert len(deleted_messages) == 2
+    assert deleted_messages[0].content == "Hello"
+    assert chat_store.get_messages("TestSession") == []
+
+
+def test_delete_message(chat_store):
+    messages = [ChatMessage(content="First"), ChatMessage(content="Second")]
+    chat_store.set_messages("TestSession", messages)
+    chat_store.delete_message("TestSession", 0)
+    remaining_messages = chat_store.get_messages("TestSession")
+    assert len(remaining_messages) == 1
+    assert remaining_messages[0].content == "Second"
+
+
+def test_delete_last_message(chat_store):
+    messages = [ChatMessage(content="First"), ChatMessage(content="Second")]
+    chat_store.set_messages("TestSession", messages)
+    last_message = chat_store.delete_last_message("TestSession")
+    assert last_message.content == "Second"
+    remaining_messages = chat_store.get_messages("TestSession")
+    assert len(remaining_messages) == 1
+    assert remaining_messages[0].content == "First"
+
+
+def test_get_keys(chat_store):
+    chat_store._table.put_item(Item={"SessionId": "1", "History": []})
+    chat_store._table.put_item(Item={"SessionId": "2", "History": []})
+    keys = chat_store.get_keys()
+    assert len(keys) == 2
+    assert "1" in keys
+    assert "2" in keys


Code review comments:
diff_hunk: @@ -0,0 +1,65 @@
+[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.dynamodb"
+
+[tool.llamahub.class_authors]
+DynamoDBChatStore = "Bryce Freshcorn"
path: llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/pyproject.toml
text: @user1:
Please change this to your github username instead: `brycecf`

@author:
Done

diff_hunk: @@ -0,0 +1,282 @@
+import logging
+from typing import Any, Dict, List, Optional
+
+from llama_index.core.base.llms.types import ChatMessage
+from llama_index.core.bridge.pydantic import Field, PrivateAttr
+from llama_index.core.storage.chat_store.base import BaseChatStore
+from mypy_boto3_dynamodb import ServiceResource
+
+
+logger = logging.getLogger(__name__)
+
+
+# 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 ChatMessage to a JSON object
+def _message_to_dict(message: ChatMessage) -> dict:
+    return message.dict()
+
+
+# Convert a JSON object to a ChatMessage
+def _dict_to_message(d: dict) -> ChatMessage:
+    return ChatMessage.model_validate(d)
+
+
+class DynamoDBChatStore(BaseChatStore):
+    """DynamoDB Chat Store.
+
+    Args:
+        table_name (str): The name of the preexisting DynamoDB table.
+        primary_key (str, optional): The primary/partition key to use for the table.
+            Defaults to "SessionId".
+        profile_name (str, optional): The AWS profile to use. If not specified, then
+            the default AWS profile is used.
+        aws_access_key_id (str, optional): The AWS Access Key ID to use.
+        aws_secret_access_key (str, optional): The AWS Secret Access Key to use.
+        aws_session_token (str, optional): The AWS Session Token to use.
+        botocore_session (Any, optional): Use this Botocore session instead of creating a new default one.
+        botocore_config (Any, optional): Custom configuration object to use instead of the default generated one.
+        region_name (str, optional): The AWS region name to use. Uses the region configured in AWS CLI if not passed.
+        max_retries (int, optional): The maximum number of API retries. Defaults to 10.
+        timeout (float, optional): The timeout for API requests in seconds. Defaults to 60.0.
+        session_kwargs (Dict[str, Any], optional): Additional kwargs for the `boto3.Session` object.
+        resource_kwargs (Dict[str, Any], optional): Additional kwargs for the `boto3.Resource` object.
+
+    Returns:
+        DynamoDBChatStore: A DynamoDB chat store object.
+    """
+
+    table_name: str = Field(description="DynamoDB table")
+    primary_key: str = Field(
+        default="SessionId", description="Primary/partition key to use for the table."
+    )
+    profile_name: Optional[str] = Field(
+        description="AWS profile to use. If not specified, then the default AWS profile is used."
+    )
+    aws_access_key_id: Optional[str] = Field(
+        description="AWS Access Key ID to use.", exclude=True
+    )
+    aws_secret_access_key: Optional[str] = Field(
+        description="AWS Secret Access Key to use.", exclude=True
+    )
+    aws_session_token: Optional[str] = Field(
+        description="AWS Session Token to use.", exclude=True
+    )
+    botocore_session: Optional[Any] = Field(
+        description="Use this Botocore session instead of creating a new default one.",
+        exclude=True,
+    )
+    botocore_config: Optional[Any] = Field(
+        description="Custom configuration object to use instead of the default generated one.",
+        exclude=True,
+    )
+    region_name: Optional[str] = Field(
+        description="AWS region name to use. Uses the region configured in AWS CLI if not passed",
+        exclude=True,
+    )
+    max_retries: int = Field(
+        default=10, description="The maximum number of API retries.", gt=0
+    )
+    timeout: float = Field(
+        default=60.0,
+        description="The timeout for API requests in seconds.",
+    )
+    session_kwargs: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Additional kwargs for the `boto3.Session` object.",
+    )
+    resource_kwargs: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Additional kwargs for the `boto3.Resource` object.",
+    )
+
+    _client: ServiceResource = PrivateAttr()
+    _table: Any = PrivateAttr()
+
+    def __init__(
+        self,
+        table_name: str,
+        primary_key: str = "SessionId",
+        profile_name: Optional[str] = None,
+        region_name: Optional[str] = None,
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None,
+        aws_session_token: Optional[str] = None,
+        botocore_session: Optional[Any] = None,
+        botocore_config: Optional[Any] = None,
+        max_retries: int = 10,
+        timeout: float = 60.0,
+        session_kwargs: Optional[Dict[str, Any]] = None,
+        resource_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        session_kwargs = session_kwargs or {}
+        resource_kwargs = resource_kwargs or {}
+
+        super().__init__(
+            table_name=table_name,
+            primary_key=primary_key,
+            profile_name=profile_name,
+            region_name=region_name,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+            aws_session_token=aws_session_token,
+            botocore_session=botocore_session,
+            botocore_config=botocore_config,
+            max_retries=max_retries,
+            timeout=timeout,
+            session_kwargs=session_kwargs,
+            resource_kwargs=resource_kwargs,
+        )
+
+        session_kwargs = {
+            "profile_name": profile_name,
+            "region_name": region_name,
+            "aws_access_key_id": aws_access_key_id,
+            "aws_secret_access_key": aws_secret_access_key,
+            "aws_session_token": aws_session_token,
+            "botocore_session": botocore_session,
+            **session_kwargs,
+        }
+
+        try:
+            import boto3
+            from botocore.config import Config
+
+            config = (
+                Config(
+                    retries={"max_attempts": max_retries, "mode": "standard"},
+                    connect_timeout=timeout,
+                    read_timeout=timeout,
+                )
+                if botocore_config is None
+                else botocore_config
+            )
+            session = boto3.Session(**session_kwargs)
+        except ImportError:
+            raise ImportError(
+                "boto3 package not found, install with 'pip install boto3"
+            )
+
+        self._client = session.resource("dynamodb", config=config, **resource_kwargs)
+        self._table = self._client.Table(table_name)
+
+    @classmethod
+    def class_name(self) -> str:
+        return "DynamoDBChatStore"
+
+    def set_messages(self, key: str, messages: List[ChatMessage]) -> None:
+        """Assign all provided messages to the row with the given key.
+        Any pre-existing messages for that key will be overwritten.
+
+        Args:
+            key (str): The key specifying a row.
+            messages (List[ChatMessage]): The messages to assign to the key.
+
+        Returns:
+            None
+        """
+        self._table.put_item(
+            Item={self.primary_key: key, "History": _messages_to_dict(messages)}
+        )
+
+    def get_messages(self, key: str) -> List[ChatMessage]:
+        """Retrieve all messages for the given key.
+
+        Args:
+            key (str): The key specifying a row.
+
+        Returns:
+            List[ChatMessage]: The messages associated with the key.
+        """
+        response = self._table.get_item(Key={self.primary_key: key})
+
+        if response and "Item" in response:
+            message_history = response["Item"]["History"]
+        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 to the end of the chat history for the given key.
+        Creates a new row if the key does not exist.
+
+        Args:
+            key (str): The key specifying a row.
+            message (ChatMessage): The message to add to the chat history.
+
+        Returns:
+            None
+        """
+        current_messages = _messages_to_dict(self.get_messages(key))
+        current_messages.append(_message_to_dict(message))
+
+        self._table.put_item(Item={self.primary_key: key, "History": current_messages})
+
+    def delete_messages(self, key: str) -> Optional[List[ChatMessage]]:
+        """Deletes the entire chat history for the given key (i.e. the row).
+
+        Args:
+            key (str): The key specifying a row.
+
+        Returns:
+            Optional[List[ChatMessage]]: The messages that were deleted. None if the
+                deletion failed.
+        """
+        messages_to_delete = self.get_messages(key)
+        self._table.delete_item(Key={self.primary_key: key})
+        return messages_to_delete
+
+    def delete_message(self, key: str, idx: int) -> Optional[ChatMessage]:
+        """Deletes the message at the given index for the given key.
+
+        Args:
+            key (str): The key specifying a row.
+            idx (int): The index of the message to delete.
+
+        Returns:
+            Optional[ChatMessage]: The message that was deleted. None if the index
+                did not exist.
+        """
+        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 session {key}")
+            )
+            return None
+
+    def delete_last_message(self, key: str) -> Optional[ChatMessage]:
+        """Deletes the last message in the chat history for the given key.
+
+        Args:
+            key (str): The key specifying a row.
+
+        Returns:
+            Optional[ChatMessage]: The message that was deleted. None if the chat history
+                was empty.
+        """
+        return self.delete_message(key, -1)
+
+    def get_keys(self) -> List[str]:
+        """Retrieve all keys in the table.
+
+        Returns:
+            List[str]: The keys in the table.
+        """
+        response = self._table.scan(ProjectionExpression="SessionId")
path: llama-index-integrations/storage/chat_store/llama-index-storage-chat-store-dynamodb/llama_index/storage/chat_store/dynamodb/base.py
text: @user1:
should we be using `self.primary_key` instead of `SessionID` here?

@author:
Good catch. Corrected.