How to fit a multistep query decomposition with a custom chat engine built with a query pipeline

[SarahSohana27]

How can I fit multistep query decomposition like this one into my custom chat engine:

from llama_index.core.query_engine import MultiStepQueryEngine
query_engine = index.as_query_engine(llm=gpt4)
query_engine = MultiStepQueryEngine(
    query_engine=query_engine,
    query_transform=step_decompose_transform,
    index_summary=index_summary,
)
response_gpt4 = query_engine.query(
    "Who was in the first batch of the accelerator program the author"
    " started?",
)

this is my code:

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 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)

response_llm = await st.session_state.chat_engine.achat(user_input)

[dosubot[bot]]

To fit multistep query decomposition into your custom chat engine built with a query pipeline, you can integrate the MultiStepQueryEngine class. Here's how you can modify your CustomChatEngine to include multistep query decomposition:

1. Instantiate the MultiStepQueryEngine: Create an instance of MultiStepQueryEngine by passing the required parameters such as a base query engine, a query transform, and optionally a response synthesizer, number of steps, early stopping flag, index summary, and a stop function.

2. Modify the achat method: Use the MultiStepQueryEngine to process the message.

Here is 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
from llama_index.core.query_engine import MultiStepQueryEngine
from llama_index.legacy.indices.query.query_transform.base import StepDecomposeQueryTransform
from llama_index.legacy.response_synthesizers import BaseSynthesizer

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

        # Initialize MultiStepQueryEngine
        self.multi_step_query_engine = MultiStepQueryEngine(
            query_engine=self.query_pipeline,
            query_transform=StepDecomposeQueryTransform(),
            response_synthesizer=BaseSynthesizer(),
            num_steps=3,  # Number of steps for the multi-step query
            early_stopping=True,  # Whether to stop early if the stop function returns True
            index_summary="Used to answer questions about the author",
            stop_fn=None  # Optional custom stop function
        )

    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 MultiStepQueryEngine
            query_bundle = QueryBundle(query_str=message)
            chat_response = await self.multi_step_query_engine._aquery(query_bundle)
            # 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 and multistep query engine
chat_engine = CustomChatEngine(qp, memory)
print("chat_engine", chat_engine)

This code integrates the MultiStepQueryEngine into your custom chat engine, allowing it to handle multistep query decomposition seamlessly ^[1][2].

To continue talking to Dosu, mention @dosu.

[SarahSohana27]

@dosu I get the following error:

Unexpected error during initialization: Can't instantiate abstract class BaseSynthesizer with abstract methods _get_prompts, _update_prompts, aget_response, get_response

[dosubot[bot]]

To resolve the error "Can't instantiate abstract class BaseSynthesizer with abstract methods _get_prompts, _update_prompts, aget_response, get_response" and properly instantiate the MultiStepQueryEngine in your custom chat engine, you need to use a concrete implementation of the BaseSynthesizer class. The Generation class is a concrete implementation of BaseSynthesizer.

Here's how you can modify your CustomChatEngine to use the Generation class:

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
from llama_index.legacy.core.query_engine.multistep_query_engine import MultiStepQueryEngine
from llama_index.legacy.indices.query.query_transform.base import StepDecomposeQueryTransform
from llama_index.legacy.response_synthesizers.generation import Generation
from llama_index.legacy.service_context import ServiceContext
from llama_index.legacy.llm_predictor import LLMPredictor

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

        # Initialize the LLMPredictor and ServiceContext
        llm_predictor = LLMPredictor()  # Replace with your actual LLMPredictor instance
        service_context = ServiceContext(llm_predictor=llm_predictor)

        # Initialize the Generation synthesizer
        generation_synthesizer = Generation(service_context=service_context)

        # Initialize MultiStepQueryEngine
        self.multi_step_query_engine = MultiStepQueryEngine(
            query_engine=self.query_pipeline,
            query_transform=StepDecomposeQueryTransform(),
            response_synthesizer=generation_synthesizer,
            num_steps=3,  # Number of steps for the multi-step query
            early_stopping=True,  # Whether to stop early if the stop function returns True
            index_summary="Used to answer questions about the author",
            stop_fn=None  # Optional custom stop function
        )

    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 MultiStepQueryEngine
            query_bundle = QueryBundle(query_str=message)
            chat_response = await self.multi_step_query_engine._aquery(query_bundle)
            # 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 and multistep query engine
chat_engine = CustomChatEngine(qp, memory)
print("chat_engine", chat_engine)

This code uses the Generation class to instantiate the MultiStepQueryEngine properly ^[1][2][3][4].

To continue talking to Dosu, mention @dosu.

[SarahSohana27]

@dosu getting this error:

Unexpected error during initialization: ServiceContext.init() missing 5 required positional arguments: 'prompt_helper', 'embed_model', 'transformations', 'llama_logger', and 'callback_manager'

[dosubot[bot]]

To resolve the error and properly initialize the ServiceContext with the required arguments, you can use the from_defaults class method provided in the ServiceContext class. Here is an example of how to do it:

from llama_index.legacy.llm_predictor import LLMPredictor
from llama_index.legacy.indices.prompt_helper import PromptHelper
from llama_index.legacy.core.embeddings.base import BaseEmbedding
from llama_index.legacy.node_parser.interface import NodeParser
from llama_index.legacy.logger import LlamaLogger
from llama_index.legacy.callbacks.base import CallbackManager
from llama_index.legacy.service_context import ServiceContext

# Initialize the required components
llm_predictor = LLMPredictor()  # or provide a specific LLM predictor
prompt_helper = PromptHelper()  # or provide specific prompt helper settings
embed_model = BaseEmbedding()  # or provide a specific embedding model
node_parser = NodeParser()  # or provide a specific node parser
llama_logger = LlamaLogger()  # optional, as it is deprecated
callback_manager = CallbackManager()  # or provide a specific callback manager

# Create the ServiceContext using from_defaults
service_context = ServiceContext.from_defaults(
    llm_predictor=llm_predictor,
    prompt_helper=prompt_helper,
    embed_model=embed_model,
    node_parser=node_parser,
    llama_logger=llama_logger,
    callback_manager=callback_manager
)

# Now you can use the service_context as needed

This code initializes the ServiceContext with the required arguments by creating instances of the necessary components and passing them to the from_defaults method ^[1].

To continue talking to Dosu, mention @dosu.