TableTransactionError in AzureChatStore when using with ReActAgent

[MarkHmnv]

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

Version of llama-index: 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:

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

[dosubot[bot]]

I found a similar open issue: [Bug]: TableTransactionError in AzureChatStore. The suggested solution is to ensure that the transaction operations are correctly formed and that the entities being deleted or created do not conflict.

Here are the steps to consider:

1. Check for Entity Existence: Ensure that the entities you are trying to delete exist and are not already deleted in a previous transaction.
2. Unique Row Keys: Ensure that the RowKey values for the new messages are unique and do not conflict with existing ones.
3. 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].

Additionally, here is a more detailed approach to resolving the TableTransactionError:

1. Check Entity Existence: Ensure that the entities you are trying to delete exist in the table. If they do not exist, the delete operation will fail.
2. Correct Transaction Operations: Ensure that the transaction operations are correctly formed. Each operation should be a tuple of the form (TransactionOperation, entity).
3. Proper Serialization: Ensure that the entities are properly serialized before being included in the transaction.
4. Partition Key and Row Key: Ensure that the PartitionKey and RowKey values are correctly set and match the entities in the table.

Here is an example of how you might modify the set_messages method to include additional checks:

def set_messages(self, key: str, messages: List[ChatMessage]) -> None:
    """Set messages for a key."""
    chat_client = self._table_service_client.create_table_if_not_exists(
        self.chat_table_name
    )
    entities = chat_client.query_entities(f"PartitionKey eq '{key}'")
    
    # Ensure entities exist before attempting to delete
    existing_entities = list(entities)
    if not existing_entities:
        raise ResourceNotFoundError(f"No entities found for PartitionKey: {key}")

    delete_operations = (
        (TransactionOperation.DELETE, entity) for entity in existing_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,
    )

For a practical example of how to set up and use AzureChatStore with SimpleChatEngine and ChatMemoryBuffer in llama-index, you can refer to the AzureChatStoreDemo ^[2][3].

To continue talking to Dosu, mention @dosu.