[Bug]: TableTransactionError in AzureChatStore when using with ReActAgent

Bug Description

When trying to use AzureChatStore together with ReActAgent, a TableTransactionError occurs

Version

0.11.1

Steps to Reproduce

Create agent:

from azure.core.credentials import AzureKeyCredential
from azure.search.documents.indexes.aio import (
    SearchIndexClient as AsyncSearchIndexClient,
)
from llama_index.core import (
    StorageContext,
    VectorStoreIndex,
    Settings,
)
from llama_index.core.agent import ReActAgent
from llama_index.core.base.llms.types import ChatMessage, MessageRole
from llama_index.core.chat_engine.types import AgentChatResponse
from llama_index.core.memory import ChatMemoryBuffer
from llama_index.core.objects import ObjectIndex
from llama_index.core.query_engine import CitationQueryEngine
from llama_index.core.tools import QueryEngineTool
from llama_index.storage.chat_store.azure import AzureChatStore
from llama_index.storage.docstore.azure import AzureDocumentStore
from llama_index.vector_stores.azureaisearch import (
    MetadataIndexFieldType,
    AzureAISearchVectorStore,
    IndexManagement,
)

search_client = AsyncSearchIndexClient(
            endpoint=ENDPOINT, credential=CREDENTIAL
)
vector_store = AzureAISearchVectorStore(
            index_name=INDEX_NAME,
            search_or_index_client=search_client,
            filterable_metadata_field_keys=METADATA_FIELDS,
            index_management=IndexManagement.CREATE_IF_NOT_EXISTS,
            id_field_key='id',
            chunk_field_key='chunk',
            embedding_field_key='embedding',
            embedding_dimensionality=1536,
            metadata_string_field_key='metadata',
            doc_id_field_key='doc_id'
)
chat_store = AzureChatStore.from_account_and_key(
    account_name=AZURE_STORAGE_ACCOUNT_NAME,
    account_key=AZURE_STORAGE_ACCOUNT_KEY
)
storage_context = StorageContext.from_defaults(
    vector_store=vector_store
)
index = VectorStoreIndex.from_documents([], storage_context=storage_context)
query_engine = index.as_query_engine()
tool = QueryEngineTool.from_defaults(
    name='vector_citation_tool',
    query_engine=query_engine,
    description='Useful for retrieving specific context from the file.',
)
object_index = ObjectIndex.from_objects(
    [tool],
    index_cls=VectorStoreIndex,
)
chat_store.metadata_partition_key = 'chat_id'
memory = ChatMemoryBuffer.from_defaults(
    token_limit=3000,
    chat_store=chat_store,
    chat_store_key='chat_id',
)
agent = ReActAgent.from_tools(
    llm=llm,
    tool_retriever=obj_index.as_retriever(similarity_top_k=1),
    verbose=True,
    memory=memory,
)

At the first call everything is ok, the data is entered into the table:

agent.achat('some query')

But if you try to use the agent for the second time, the TableTransactionError occurs

Relevant Logs/Tracbacks

File "\.venv\Lib\site-packages\llama_index\core\instrumentation\dispatcher.py", line 291, in async_wrapper
    result = await func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\.venv\Lib\site-packages\llama_index\core\callbacks\utils.py", line 56, in async_wrapper
    return await func(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\.venv\Lib\site-packages\llama_index\core\agent\runner\base.py", line 671, in achat
    chat_response = await self._achat(
                    ^^^^^^^^^^^^^^^^^^
  File "\.venv\Lib\site-packages\llama_index\core\instrumentation\dispatcher.py", line 291, in async_wrapper
    result = await func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\.venv\Lib\site-packages\llama_index\core\agent\runner\base.py", line 624, in _achat
    result = self.finalize_response(
             ^^^^^^^^^^^^^^^^^^^^^^^
  File "\.venv\Lib\site-packages\llama_index\core\instrumentation\dispatcher.py", line 261, in wrapper
    result = func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^
  File "\.venv\Lib\site-packages\llama_index\core\agent\runner\base.py", line 549, in finalize_response
    self.agent_worker.finalize_task(self.state.get_task(task_id))
  File "\.venv\Lib\site-packages\llama_index\core\instrumentation\dispatcher.py", line 261, in wrapper
    result = func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^
  File "\.venv\Lib\site-packages\llama_index\core\agent\react\step.py", line 806, in finalize_task
    task.memory.set(
  File "\.venv\Lib\site-packages\llama_index\core\memory\types.py", line 99, in set
    self.chat_store.set_messages(self.chat_store_key, messages)
  File "\.venv\Lib\site-packages\llama_index\storage\chat_store\azure\base.py", line 185, in set_messages
    chat_client.submit_transaction(chain(delete_operations, create_operations))
  File "\.venv\Lib\site-packages\azure\core\tracing\decorator.py", line 94, in wrapper_use_tracer
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "\.venv\Lib\site-packages\azure\data\tables\_table_client.py", line 718, in submit_transaction
    return self._batch_send(self.table_name, *batched_requests.requests, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\.venv\Lib\site-packages\azure\data\tables\_base_client.py", line 296, in _batch_send
    raise decoded
azure.data.tables._error.TableTransactionError: 2:The batch request contains multiple changes with same row key. An entity can appear only once in a batch request.
RequestId:56ae5e74-4002-0061-1021-f7b3b8000000
Time:2024-08-25T19:01:25.7243438Z
ErrorCode:InvalidDuplicateRow
Content: {"odata.error":{"code":"InvalidDuplicateRow","message":{"lang":"en-US","value":"2:The batch request contains multiple changes with same row key. An entity can appear only once in a batch request.\nRequestId:56ae5e74-4002-0061-1021-f7b3b8000000\nTime:2024-08-25T19:01:25.7243438Z"}}}

P.S. Everythinkg works fine with RedisChatStore, so this error is in the llama-index-storage-chat-store-azure library

Hello, @MarkHmnv! I'm here to assist you with any bugs, questions, or contributions you have.

To resolve the TableTransactionError that occurs when using AzureChatStore with ReActAgent on the second call, you need to ensure that the transaction operations are correctly formed and that the entities being deleted or created do not conflict.

The set_messages method in AzureChatStore performs a transaction that deletes existing messages and inserts new ones. If the second call results in a TableTransactionError, it could be due to conflicts or issues with the entities being processed.

Here is the relevant code snippet from the set_messages method:

def set_messages(self, key: str, messages: List[ChatMessage]) -> None:
    """Set messages for a key."""
    # Delete existing messages and insert new messages in one transaction
    chat_client = self._table_service_client.create_table_if_not_exists(
        self.chat_table_name
    )
    entities = chat_client.query_entities(f"PartitionKey eq '{key}'")
    delete_operations = (
        (TransactionOperation.DELETE, entity) for entity in entities
    )
    create_operations = (
        (
            TransactionOperation.CREATE,
            serialize(
                self.service_mode,
                {
                    "PartitionKey": key,
                    "RowKey": self._to_row_key(idx),
                    **message.dict(),
                },
            ),
        )
        for idx, message in enumerate(messages)
    )
    chat_client.submit_transaction(chain(delete_operations, create_operations))

    # Update metadata
    metadata_client = self._table_service_client.create_table_if_not_exists(
        self.metadata_table_name
    )
    messages_len = len(messages)
    metadata_client.upsert_entity(
        {
            "PartitionKey": self.metadata_partition_key,
            "RowKey": key,
            "LastMessageRowKey": self._to_row_key(messages_len - 1),
            "MessageCount": messages_len,
        },
        UpdateMode.REPLACE,
    )

To troubleshoot and resolve the TableTransactionError, consider the following steps:

Check for Entity Existence: Ensure that the entities you are trying to delete exist and are not already deleted in a previous transaction.
Unique Row Keys: Ensure that the RowKey values for the new messages are unique and do not conflict with existing ones.
Transaction Size: Azure Table Storage has a limit on the number of operations in a single transaction (100 operations). Ensure that your transaction does not exceed this limit.

If the issue persists, you may need to add logging or debugging statements to inspect the entities being processed in the transaction to identify any conflicts or issues ^[1].

To continue talking to Dosu, mention @dosu.

run-llama / llama_index