how to add chat memory in a a custom chat engine?

[SarahSohana27]

This my custom chat engine that wraps a query pipeline:

from llama_index.legacy.chat_engine.types import (
    AgentChatResponse,
    BaseChatEngine,
    StreamingAgentChatResponse,
    ChatMessage,
)
from typing import AsyncGenerator, Generator, List, Optional, Union
from llama_index.core.memory import BaseMemory, ChatMemoryBuffer


class CustomChatEngine(BaseChatEngine):
    def __init__(self, query_pipeline, memory: BaseMemory, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.query_pipeline = query_pipeline
        self._memory = memory

    def chat(self, message):
        try:
            response = self.query_pipeline.run(query=message)
            print(response)
            return response
        except Exception as e:
            return "An error occurred while processing your query: {}".format(e)

    def reset(self) -> None:
        # Reset conversation state if applicable
        pass

    async def achat(
        self, message: str, chat_history: Optional[List[ChatMessage]] = None
    ):
        chat_response = await self.query_pipeline.arun(query=message)
        return chat_response

this is my query pipeline:

    # handles sql generation
    vn_component = create_vn_component()
    # handles final response
    response_component = create_response_component()

    qp = QP(
        modules={
            "input": InputComponent(),
            "vn_component": vn_component,
            "response_component": response_component,
        },
        verbose=True,
    )

    qp.add_link("input", "vn_component", dest_key="input_str")
    qp.add_link("vn_component", "response_component", dest_key="response")
    memory = []
    chat_engine = CustomChatEngine(qp, memory)
    print("chat_engine", chat_engine)

how to add memory in this code?

[dosubot[bot]]

To add chat memory functionality to your custom chat engine, you can use the ChatMemoryBuffer class. Here's how you can modify your CustomChatEngine to include chat memory:

1. Initialize the ChatMemoryBuffer: Create an instance of ChatMemoryBuffer and pass it to your CustomChatEngine.

2. Update the chat method: Modify the chat method to store the chat messages in the memory buffer and retrieve the chat history when processing a new message.

Here's the updated code:

from llama_index.legacy.chat_engine.types import (
    AgentChatResponse,
    BaseChatEngine,
    StreamingAgentChatResponse,
    ChatMessage,
)
from typing import AsyncGenerator, Generator, List, Optional, Union
from llama_index.core.memory import BaseMemory, ChatMemoryBuffer
from llama_index.core.storage.chat_store import SimpleChatStore

class CustomChatEngine(BaseChatEngine):
    def __init__(self, query_pipeline, memory: BaseMemory, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.query_pipeline = query_pipeline
        self._memory = memory

    def chat(self, message):
        try:
            # Retrieve chat history
            chat_history = self._memory.get()
            # Add the new message to the chat history
            chat_history.append(ChatMessage(role="user", content=message))
            # Process the message with the query pipeline
            response = self.query_pipeline.run(query=message)
            # Add the response to the chat history
            chat_history.append(ChatMessage(role="assistant", content=response))
            # Update the memory with the new chat history
            self._memory.chat_store.set_messages(self._memory.chat_store_key, chat_history)
            print(response)
            return response
        except Exception as e:
            return "An error occurred while processing your query: {}".format(e)

    def reset(self) -> None:
        # Reset conversation state if applicable
        self._memory.chat_store.set_messages(self._memory.chat_store_key, [])

    async def achat(
        self, message: str, chat_history: Optional[List[ChatMessage]] = None
    ):
        chat_response = await self.query_pipeline.arun(query=message)
        return chat_response

# Initialize the query pipeline
vn_component = create_vn_component()
response_component = create_response_component()

qp = QP(
    modules={
        "input": InputComponent(),
        "vn_component": vn_component,
        "response_component": response_component,
    },
    verbose=True,
)

qp.add_link("input", "vn_component", dest_key="input_str")
qp.add_link("vn_component", "response_component", dest_key="response")

# Initialize the chat memory buffer
chat_store = SimpleChatStore()
memory = ChatMemoryBuffer.from_defaults(
    token_limit=3000,
    chat_store=chat_store,
    chat_store_key="user1",
)

# Create the custom chat engine with memory
chat_engine = CustomChatEngine(qp, memory)
print("chat_engine", chat_engine)

In this updated code:

• The ChatMemoryBuffer is initialized with a SimpleChatStore and a token limit.
• The chat method retrieves the chat history, processes the new message, updates the chat history with the new message and response, and stores the updated chat history back in the memory buffer.
• The reset method clears the chat history.

This will enable your custom chat engine to maintain a memory of the conversation history.

To continue talking to Dosu, mention @dosu.

[SarahSohana27]

@dosu I am using achat. Is it same for achat method?

[dosubot[bot]]

To add chat memory functionality to the achat method, you can follow a similar approach as with the chat method. Here's how you can modify the achat method to include chat memory:

1. Retrieve chat history: Get the chat history from the memory buffer.
2. Add the new message to the chat history: Append the new user message to the chat history.
3. Process the message with the query pipeline: Use the query pipeline to process the message.
4. Add the response to the chat history: Append the assistant's response to the chat history.
5. Update the memory: Store the updated chat history back in the memory buffer.

Here's the updated code:

from llama_index.legacy.chat_engine.types import (
    AgentChatResponse,
    BaseChatEngine,
    StreamingAgentChatResponse,
    ChatMessage,
)
from typing import AsyncGenerator, Generator, List, Optional, Union
from llama_index.core.memory import BaseMemory, ChatMemoryBuffer
from llama_index.core.storage.chat_store import SimpleChatStore

class CustomChatEngine(BaseChatEngine):
    def __init__(self, query_pipeline, memory: BaseMemory, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.query_pipeline = query_pipeline
        self._memory = memory

    def chat(self, message):
        try:
            # Retrieve chat history
            chat_history = self._memory.get()
            # Add the new message to the chat history
            chat_history.append(ChatMessage(role="user", content=message))
            # Process the message with the query pipeline
            response = self.query_pipeline.run(query=message)
            # Add the response to the chat history
            chat_history.append(ChatMessage(role="assistant", content=response))
            # Update the memory with the new chat history
            self._memory.chat_store.set_messages(self._memory.chat_store_key, chat_history)
            print(response)
            return response
        except Exception as e:
            return "An error occurred while processing your query: {}".format(e)

    def reset(self) -> None:
        # Reset conversation state if applicable
        self._memory.chat_store.set_messages(self._memory.chat_store_key, [])

    async def achat(
        self, message: str, chat_history: Optional[List[ChatMessage]] = None
    ):
        try:
            # Retrieve chat history
            chat_history = self._memory.get()
            # Add the new message to the chat history
            chat_history.append(ChatMessage(role="user", content=message))
            # Process the message with the query pipeline
            chat_response = await self.query_pipeline.arun(query=message)
            # Add the response to the chat history
            chat_history.append(ChatMessage(role="assistant", content=chat_response))
            # Update the memory with the new chat history
            self._memory.chat_store.set_messages(self._memory.chat_store_key, chat_history)
            return chat_response
        except Exception as e:
            return "An error occurred while processing your query: {}".format(e)

# Initialize the query pipeline
vn_component = create_vn_component()
response_component = create_response_component()

qp = QP(
    modules={
        "input": InputComponent(),
        "vn_component": vn_component,
        "response_component": response_component,
    },
    verbose=True,
)

qp.add_link("input", "vn_component", dest_key="input_str")
qp.add_link("vn_component", "response_component", dest_key="response")

# Initialize the chat memory buffer
chat_store = SimpleChatStore()
memory = ChatMemoryBuffer.from_defaults(
    token_limit=3000,
    chat_store=chat_store,
    chat_store_key="user1",
)

# Create the custom chat engine with memory
chat_engine = CustomChatEngine(qp, memory)
print("chat_engine", chat_engine)

In this updated code:

• The achat method retrieves the chat history, processes the new message, updates the chat history with the new message and response, and stores the updated chat history back in the memory buffer.
• This ensures that the chat memory is maintained for asynchronous interactions as well.

To continue talking to Dosu, mention @dosu.