help with custom chat engine

[SarahSohana27]

I have a query pipeline that I pass to a custom chat engine. I want to add memory in the chat engine. how to do it ?

this is my Custom Chat Engine:

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
        print(query_pipeline)

    def chat(self, message, chat_history):
        try:
            # Retrieve chat history
            chat_history = self._memory.get()
        
            # Process the message with the query pipeline
            response = self.query_pipeline.run(query=message, chat_history=chat_history)

            print(response)

            # Add the new message to the chat history
            self._memory.put(ChatMessage(role="user", content=message))
            # Add the response to the chat history
            self._memory.put(ChatMessage(role="assistant", content=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.reset()

    async def achat(
        self, message: str, chat_history: Optional[List[ChatMessage]] = None
    ):
        try:
            # Retrieve chat history
            chat_history = self._memory.get()
            
            chat_response = await self.query_pipeline.arun(query= message)
            print('response', chat_response)

            # Add the new message to the chat history
            chat_history.append(ChatMessage(role="user", content=message))
            chat_history.append(ChatMessage(role="assistant", content=chat_response))

            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)

    def stream_chat(
        self, message: str, chat_history: Optional[List[ChatMessage]] = None
    ) -> StreamingAgentChatResponse:
        # Implement your streaming chat logic here
        return StreamingAgentChatResponse(response="Streaming response")

    async def astream_chat(
        self, message: str, chat_history: Optional[List[ChatMessage]] = None
    ) -> StreamingAgentChatResponse:
        # Implement your async streaming chat logic here
        return StreamingAgentChatResponse(response="Async streaming response")

    @property
    def chat_history(self) -> List[ChatMessage]:
        # Return chat history if applicable
        return []

below is the query pipeline I am using:

# LlamaIndex query pipeline
def initialize_chat_engine():
    # memory
    chat_store = SimpleChatStore()
    memory = ChatMemoryBuffer.from_defaults(
        token_limit=15000, chat_store=chat_store, chat_store_key=uuid.uuid4().hex
    )
  
    # handles sql generation
    vn_component_netstat = create_vn_component('netstat')
    vn_component_crowdsource = create_vn_component('location_insight')

    # handles final response
    response_component = create_response_component()
    vn_chain_netstat = QP(chain=[vn_component_netstat, response_component], verbose=True)
    vn_chain_crowdsource = QP(chain=[vn_component_crowdsource, response_component], verbose=True)

    # Step 3: Define Router and Pipelines
    choices = [
    """This tool generates SQL queries for evaluating daily network performance metrics in mobile telecommunications. 
    (the tool does not analyze the demographic and location-based insights related to mobile network coverage and signal quality )
    """,
    """This tool generates SQL queries for analyzing demographic and location-based insights related to mobile network coverage and signal quality. """]

    selector = LLMSingleSelector.from_defaults()

    router_c = RouterComponent(
        selector=selector,
        choices=choices,
        components=[vn_chain_netstat, vn_chain_crowdsource],
        verbose=True,
    )
    # top-level pipeline

    qp = QP(chain=[InputComponent(), router_c, router_response()], verbose=True)
    chat_engine = CustomChatEngine(qp, memory)
    if chat_engine:
        st.session_state.chat_engine = chat_engine
        st.session_state.chat_engine_status = True
    else:
        print("Chat engine not created! Check for errors!")

def custom_query_pipeline():
    if "chat_engine_status" not in st.session_state:
        st.session_state.chat_engine_status = False
    if not st.session_state.chat_engine_status:
        initialize_chat_engine()

[dosubot[bot]]

To add memory to your custom chat engine using a query pipeline, you can utilize the ChatMemoryBuffer from the llama-index library. Your current implementation already includes the necessary components to manage chat history. Here's a brief overview of how you can ensure the memory is effectively integrated:

1. Initialize the Memory Buffer: You have already initialized the ChatMemoryBuffer with a token limit and a SimpleChatStore. This setup is correct for managing chat history.

2. Retrieve and Update Chat History: In your chat and achat methods, you are correctly retrieving the chat history using self._memory.get() and updating it with new messages and responses using self._memory.put(). This ensures that the chat history is maintained across interactions.

3. Handling Asynchronous Interactions: In the achat method, you are appending new messages and responses to the chat history and updating the memory. This is crucial for maintaining context in asynchronous interactions.

4. Resetting Memory: The reset method is correctly implemented to clear the chat history when needed.

5. Ensure Consistency: Make sure that the chat_history parameter in your methods is consistently used to manage the chat history. You might want to remove the chat_history parameter from the method signatures if it's not being used, as you are directly managing the history with the memory buffer.

Your implementation seems to be on the right track. If you encounter any specific issues or errors, ensure that the ChatMemoryBuffer is correctly initialized and that the chat history is being updated as expected after each interaction. If you need further customization, you can adjust the token limit or the way messages are stored in the SimpleChatStore to better fit your application's needs.

To continue talking to Dosu, mention @dosu.

[SarahSohana27]

@dosu can you modify my code above? achat function is not working properly after adding the chat history

[dosubot[bot]]

To modify your achat function so that it works properly with the chat history, you should ensure that the chat history is correctly set and updated within the function. Here's an updated version of your achat function:

async def achat(
    self, message: str, chat_history: Optional[List[ChatMessage]] = None
):
    try:
        # Set chat history if provided
        if chat_history is not None:
            self._memory.set(chat_history)
        
        # Add the new user message to the memory
        self._memory.put(ChatMessage(role="user", content=message))

        # Retrieve the updated chat history
        all_messages = self._memory.get()

        # Process the message with the query pipeline
        chat_response = await self.query_pipeline.arun(query=message, chat_history=all_messages)
        print('response', chat_response)

        # Add the response to the chat history
        self._memory.put(ChatMessage(role="assistant", content=chat_response))

        return chat_response
    except Exception as e:
        return "An error occurred while processing your query: {}".format(e)

Key Changes:

1. Set Chat History: If a chat_history is provided, it is set in the memory using self._memory.set(chat_history).
2. Update Memory: The new user message is added to the memory with self._memory.put(ChatMessage(role="user", content=message)).
3. Retrieve Updated History: The updated chat history is retrieved using self._memory.get().
4. Process with Pipeline: The message is processed with the query pipeline, passing the updated chat history.
5. Store AI Response: The AI's response is added to the memory with self._memory.put(ChatMessage(role="assistant", content=chat_response)).

This ensures that the chat history is maintained and updated correctly during asynchronous interactions.

To continue talking to Dosu, mention @dosu.

[SarahSohana27]

@dosu, how can I add a router component in this pipeline?

pipeline = QueryPipeline(
    modules={
        "input": input_component,
        "rewrite_template": rewrite_template,
        "llm": llm,
        "rewrite_retriever": retriever,
        "query_retriever": retriever,
        "join": argpack_component,
        "reranker": reranker,
        "response_component": response_component,
    },
    verbose=False,
)

# run both retrievers -- once with the hallucinated query, once with the real query
pipeline.add_link(
    "input", "rewrite_template", src_key="query_str", dest_key="query_str"
)
pipeline.add_link(
    "input",
    "rewrite_template",
    src_key="chat_history_str",
    dest_key="chat_history_str",
)
pipeline.add_link("rewrite_template", "llm")
pipeline.add_link("llm", "rewrite_retriever")
pipeline.add_link("input", "query_retriever", src_key="query_str")

# each input to the argpack component needs a dest key -- it can be anything
# then, the argpack component will pack all the inputs into a single list
pipeline.add_link("rewrite_retriever", "join", dest_key="rewrite_nodes")
pipeline.add_link("query_retriever", "join", dest_key="query_nodes")

# reranker needs the packed nodes and the query string
pipeline.add_link("join", "reranker", dest_key="nodes")
pipeline.add_link(
    "input", "reranker", src_key="query_str", dest_key="query_str"
)

# synthesizer needs the reranked nodes and query str
pipeline.add_link("reranker", "response_component", dest_key="nodes")
pipeline.add_link(
    "input", "response_component", src_key="query_str", dest_key="query_str"
)
pipeline.add_link(
    "input",
    "response_component",
    src_key="chat_history",
    dest_key="chat_history",
)

[dosubot[bot]]

To add a RouterComponent to your existing query pipeline, you need to define a selector, a list of choices, and the corresponding components. Here's how you can integrate a RouterComponent into your pipeline:

1. Define the Selector and Components: You need a selector to decide which component to route to, and the components themselves.

2. Create the RouterComponent: Instantiate the RouterComponent with the selector, choices, and components.

3. Add the RouterComponent to the Pipeline: Insert the router into the appropriate place in your pipeline.

Here's an example of how you can do this:

from llama_index.legacy.core.base_selector import BaseSelector
from llama_index.legacy.core.query_pipeline.components.router import RouterComponent

# Define your selector and components
selector = MySelector()  # Replace with your actual selector
choices = ["component1", "component2"]  # Names corresponding to your components
components = [Component1(), Component2()]  # Replace with your actual components

# Create the RouterComponent
router_component = RouterComponent(
    selector=selector,
    choices=choices,
    components=components,
    verbose=True  # Set to True if you want verbose output
)

# Add the router component to your query pipeline
pipeline.add_component(router_component)

# Add links to and from the router component as needed
pipeline.add_link("input", "router_component", src_key="query_str", dest_key="query_str")
pipeline.add_link("router_component", "response_component", dest_key="response")

Make sure that MySelector, Component1, and Component2 are properly defined and imported in your code. The choices list should correspond to the components you are routing to. The RouterComponent will route queries to different components based on the selection made by the selector ^[1][2].

To continue talking to Dosu, mention @dosu.