Chat engine with fusion retriever

[d3buggerdan]

@dosu

I made a FusionRetriever with the following code but instead of using a query engine id like to use a chat engine. How can i do it?

Function to generate queries

def generate_queries(llm, query_str: str, num_queries: int = 4):
fmt_prompt = query_gen_prompt.format(
num_queries=num_queries - 1, query=query_str
)
print(f"Formatted prompt for query generation:\n{fmt_prompt}\n") # Debugging
response = llm.complete(fmt_prompt)
print(f"Response from LLM for query generation:\n{response.text}\n") # Debugging
queries = response.text.split("\n")
return queries

Function to run queries synchronously

def run_queries_sync(queries, retrievers):
results_dict = {}
for query in queries:
print(f"Running query: {query}") # Debugging
for i, retriever in enumerate(retrievers):
query_result = retriever.retrieve(query)
results_dict[(query, i)] = query_result
print(f"Results for retriever {i} and query '{query}': {query_result}\n") # Debugging
return results_dict

Function to fuse results

def fuse_results(results_dict, similarity_top_k: int = 2):
k = 60.0 # Parameter used to control the impact of outlier rankings.
fused_scores = {}
text_to_node = {}

print(f"Results before fusing: {results_dict}\n")  # Debugging

# Compute reciprocal rank scores
for nodes_with_scores in results_dict.values():
    for rank, node_with_score in enumerate(
        sorted(
            nodes_with_scores, key=lambda x: x.score or 0.0, reverse=True
        )
    ):
        text = node_with_score.node.get_content()
        text_to_node[text] = node_with_score
        if text not in fused_scores:
            fused_scores[text] = 0.0
        fused_scores[text] += 1.0 / (rank + k)

# Sort results
reranked_results = dict(
    sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)
)

print(f"Fused scores: {fused_scores}\n")  # Debugging

# Adjust node scores
reranked_nodes: List[NodeWithScore] = []
for text, score in reranked_results.items():
    reranked_nodes.append(text_to_node[text])
    reranked_nodes[-1].score = score
    print(f"Reranked node: {text} with score: {score}\n")  # Debugging

return reranked_nodes[:similarity_top_k]

Function to determine whether to use the reranked result or fallback to LLM

def should_fallback_to_llm(reranked_nodes, cutoff_score=0.05):
if not reranked_nodes:
return True
top_score = reranked_nodes[0].score
return top_score < cutoff_score

Define FusionRetriever class

class FusionRetriever(BaseRetriever):
def init(
self,
llm,
retrievers: List[BaseRetriever],
similarity_top_k: int = 2,
cutoff_score: float = 0.05, # Added cutoff_score parameter
) -> None:
self._retrievers = retrievers
self._similarity_top_k = similarity_top_k
self._llm = llm
self._cutoff_score = cutoff_score # Store the cutoff score
super().init()

def _retrieve(self, query_bundle: QueryBundle):
    print(f"Query Bundle: {query_bundle}\n")  # Debugging
    queries = generate_queries(self._llm, query_bundle.query_str, num_queries=4)
    print(f"Generated Queries: {queries}\n")  # Debugging
    results = run_queries_sync(queries, self._retrievers)
    final_results = fuse_results(results, similarity_top_k=self._similarity_top_k)
    
    # Decide whether to use the reranked results or fallback to LLM
    if should_fallback_to_llm(final_results, cutoff_score=self._cutoff_score):
        print("Fallback to LLM for a more general response.\n")  # Debugging
        # Directly return the LLM response as a string
        return self._llm.complete(query_bundle.query_str).text
    
    return final_results

Example usage

vector_retriever = index.as_retriever(similarity_top_k=2)

fusion_retriever = FusionRetriever(llm, [vector_retriever], similarity_top_k=2, cutoff_score=0.05)
query_engine = RetrieverQueryEngine(fusion_retriever)

[dosubot[bot]]

To use a chat engine instead of a query engine with your FusionRetriever, you can follow these steps:

1. Initialize the Chat Engine: Convert your index to a chat engine.
2. Modify the FusionRetriever: Ensure it works seamlessly with the chat engine.
3. Use the Chat Engine: Replace the query engine usage with the chat engine.

Here's how you can modify your code:

Step 1: Initialize the Chat Engine

# Assuming `index` is already defined and set up
chat_engine = index.as_chat_engine()

Step 2: Modify the FusionRetriever

Ensure your FusionRetriever class is compatible with the chat engine. The main change is to handle chat messages instead of query bundles.

Step 3: Use the Chat Engine

Replace the query engine usage with the chat engine in your example usage.

Here is the complete modified code:

# Function to generate queries
def generate_queries(llm, query_str: str, num_queries: int = 4):
    fmt_prompt = query_gen_prompt.format(
        num_queries=num_queries - 1, query=query_str
    )
    print(f"Formatted prompt for query generation:\n{fmt_prompt}\n")  # Debugging
    response = llm.complete(fmt_prompt)
    print(f"Response from LLM for query generation:\n{response.text}\n")  # Debugging
    queries = response.text.split("\n")
    return queries

# Function to run queries synchronously
def run_queries_sync(queries, retrievers):
    results_dict = {}
    for query in queries:
        print(f"Running query: {query}")  # Debugging
        for i, retriever in enumerate(retrievers):
            query_result = retriever.retrieve(query)
            results_dict[(query, i)] = query_result
            print(f"Results for retriever {i} and query '{query}': {query_result}\n")  # Debugging
    return results_dict

# Function to fuse results
def fuse_results(results_dict, similarity_top_k: int = 2):
    k = 60.0  # Parameter used to control the impact of outlier rankings.
    fused_scores = {}
    text_to_node = {}

    print(f"Results before fusing: {results_dict}\n")  # Debugging

    # Compute reciprocal rank scores
    for nodes_with_scores in results_dict.values():
        for rank, node_with_score in enumerate(
            sorted(
                nodes_with_scores, key=lambda x: x.score or 0.0, reverse=True
            )
        ):
            text = node_with_score.node.get_content()
            text_to_node[text] = node_with_score
            if text not in fused_scores:
                fused_scores[text] = 0.0
            fused_scores[text] += 1.0 / (rank + k)

    # Sort results
    reranked_results = dict(
        sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)
    )

    print(f"Fused scores: {fused_scores}\n")  # Debugging

    # Adjust node scores
    reranked_nodes: List[NodeWithScore] = []
    for text, score in reranked_results.items():
        reranked_nodes.append(text_to_node[text])
        reranked_nodes[-1].score = score
        print(f"Reranked node: {text} with score: {score}\n")  # Debugging

    return reranked_nodes[:similarity_top_k]

# Function to determine whether to use the reranked result or fallback to LLM
def should_fallback_to_llm(reranked_nodes, cutoff_score=0.05):
    if not reranked_nodes:
        return True
    top_score = reranked_nodes[0].score
    return top_score < cutoff_score

# Define FusionRetriever class
class FusionRetriever(BaseRetriever):
    def __init__(
        self,
        llm,
        retrievers: List[BaseRetriever],
        similarity_top_k: int = 2,
        cutoff_score: float = 0.05,  # Added cutoff_score parameter
    ) -> None:
        self._retrievers = retrievers
        self._similarity_top_k = similarity_top_k
        self._llm = llm
        self._cutoff_score = cutoff_score  # Store the cutoff score
        super().__init__()

    def _retrieve(self, query_bundle: QueryBundle):
        print(f"Query Bundle: {query_bundle}\n")  # Debugging
        queries = generate_queries(self._llm, query_bundle.query_str, num_queries=4)
        print(f"Generated Queries: {queries}\n")  # Debugging
        results = run_queries_sync(queries, self._retrievers)
        final_results = fuse_results(results, similarity_top_k=self._similarity_top_k)
        
        # Decide whether to use the reranked results or fallback to LLM
        if should_fallback_to_llm(final_results, cutoff_score=self._cutoff_score):
            print("Fallback to LLM for a more general response.\n")  # Debugging
            # Directly return the LLM response as a string
            return self._llm.complete(query_bundle.query_str).text
        
        return final_results

# Example usage
vector_retriever = index.as_retriever(similarity_top_k=2)

fusion_retriever = FusionRetriever(llm, [vector_retriever], similarity_top_k=2, cutoff_score=0.05)
chat_engine = index.as_chat_engine()

# Use the chat engine to have a conversation
response = chat_engine.chat("Tell me a joke.")
print(response)

# To stream the response
streaming_response = chat_engine.stream_chat("Tell me a joke.")
for token in streaming_response.response_gen:
    print(token, end="")

This code initializes the chat engine from the index and uses it to have a conversation, either by getting a complete response or by streaming the response token by token ^[1][2].

To continue talking to Dosu, mention @dosu.

[d3buggerdan]

@dosu i want the chat engine to use the fusion retriever not the index

[dosubot[bot]]

To replace the query engine with a chat engine that uses the FusionRetriever, you can configure the chat engine using the low-level composition API. Here is an example of how you can do this:

from llama_index.legacy.prompts.base import PromptTemplate
from llama_index.legacy.core.llms.types import ChatMessage, MessageRole
from llama_index.legacy.chat_engine.condense_question import CondenseQuestionChatEngine
from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.core.query_engine import RetrieverQueryEngine

# Define a custom prompt template
custom_prompt = PromptTemplate(
    """\
Given a conversation (between Human and Assistant) and a follow up message from Human, \
rewrite the message to be a standalone question that captures all relevant context \
from the conversation.

<Chat History>
{chat_history}

<Follow Up Message>
{question}

<Standalone question>
"""
)

# List of `ChatMessage` objects for initial chat history
custom_chat_history = [
    ChatMessage(
        role=MessageRole.USER,
        content="Hello assistant, we are having an insightful discussion about Paul Graham today.",
    ),
    ChatMessage(role=MessageRole.ASSISTANT, content="Okay, sounds good."),
]

# Initialize the FusionRetriever
documents_1 = SimpleDirectoryReader(input_files=["../../community/integrations/vector_stores.md"]).load_data()
documents_2 = SimpleDirectoryReader(input_files=["../../module_guides/storing/vector_stores.md"]).load_data()

index_1 = VectorStoreIndex.from_documents(documents_1)
index_2 = VectorStoreIndex.from_documents(documents_2)

fusion_retriever = QueryFusionRetriever(
    [index_1.as_retriever(), index_2.as_retriever()],
    similarity_top_k=2,
    num_queries=4,  # set this to 1 to disable query generation
    use_async=True,
    verbose=True,
)

# Create the query engine using the FusionRetriever
query_engine = RetrieverQueryEngine.from_args(fusion_retriever)

# Create the chat engine using the custom prompt and query engine
chat_engine = CondenseQuestionChatEngine.from_defaults(
    query_engine=query_engine,
    condense_question_prompt=custom_prompt,
    chat_history=custom_chat_history,
    verbose=True,
)

In this example, the FusionRetriever is initialized by combining retrieval results from multiple indexes and generating additional queries to enhance the retrieval process. This setup configures the chat engine to use the FusionRetriever for handling queries.

The CondenseQuestionChatEngine class is designed to generate a standalone question from the conversation context and the last message, then query the query engine for a response. The from_defaults method allows you to initialize the chat engine with default parameters, including the query engine, prompt template, chat history, and other optional settings ^[1][2][3][4].

To continue talking to Dosu, mention @dosu.

[d3buggerdan]

@dosu
{
"name": "ValidationError",
"message": "1 validation error for LLMPredictStartEvent
template
Can't instantiate abstract class BasePromptTemplate with abstract methods format, format_messages, get_template, partial_format (type=type_error)",
"stack": "---------------------------------------------------------------------------
ValidationError Traceback (most recent call last)
in
125 # response = query_engine.query(query_str)
126 # print(str(response))
--> 127 response = chat_engine.chat(
128 "what day of the week is it"
129 )

~\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.9_qbz5n2kfra8p0\LocalCache\local-packages\Python39\site-packages\llama_index\legacy\callbacks\utils.py in wrapper(self, *args, **kwargs)
39 callback_manager = cast(CallbackManager, callback_manager)
40 with callback_manager.as_trace(trace_id):
---> 41 return func(self, *args, **kwargs)
42
43 @functools.wraps(func) # preserve signature, name, etc. of func

~\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.9_qbz5n2kfra8p0\LocalCache\local-packages\Python39\site-packages\llama_index\legacy\chat_engine\condense_question.py in chat(self, message, chat_history)
169
170 # Generate standalone question from conversation context and last message
--> 171 condensed_question = self._condense_question(chat_history, message)
172
173 log_str = f"Querying with: {condensed_question}"

~\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.9_qbz5n2kfra8p0\LocalCache\local-packages\Python39\site-packages\llama_index\legacy\chat_engine\condense_question.py in _condense_question(self, chat_history, last_message)
123 logger.debug(chat_history_str)
124
--> 125 return self._llm.predict(
126 self._condense_question_prompt,
127 question=last_message,

~\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.9_qbz5n2kfra8p0\LocalCache\local-packages\Python39\site-packages\llama_index\core\instrumentation\dispatcher.py in wrapper(func, instance, args, kwargs)
233 )
234 try:
--> 235 result = func(*args, **kwargs)
236 except BaseException as e:
237 self.event(SpanDropEvent(span_id=id_, err_str=str(e)))

~\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.9_qbz5n2kfra8p0\LocalCache\local-packages\Python39\site-packages\llama_index\core\llms\llm.py in predict(self, prompt, **prompt_args)
551 """
552 dispatcher.event(
--> 553 LLMPredictStartEvent(template=prompt, template_args=prompt_args)
554 )
555 self._log_template_data(prompt, **prompt_args)

~\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.9_qbz5n2kfra8p0\LocalCache\local-packages\Python39\site-packages\pydantic\v1\main.py in init(pydantic_self, **data)
339 values, fields_set, validation_error = validate_model(pydantic_self.class, data)
340 if validation_error:
--> 341 raise validation_error
342 try:
343 object_setattr(pydantic_self, 'dict', values)

ValidationError: 1 validation error for LLMPredictStartEvent
template
Can't instantiate abstract class BasePromptTemplate with abstract methods format, format_messages, get_template, partial_format (type=type_error)"
}

[dosubot[bot]]

The ValidationError you're encountering indicates that the BasePromptTemplate class is being instantiated directly, which is not allowed because it contains abstract methods that need to be implemented.

To resolve this issue, you need to ensure that your custom prompt template class implements all the required abstract methods and properties from the BasePromptTemplate class. Here are the methods and properties you need to implement:

1. Methods:

   • partial_format(self, **kwargs: Any) -> "BasePromptTemplate"
   • format(self, llm: Optional[BaseLLM] = None, **kwargs: Any) -> str
   • format_messages(self, llm: Optional[BaseLLM] = None, **kwargs: Any) -> List[ChatMessage]
   • get_template(self, llm: Optional[BaseLLM] = None) -> str

2. Properties:

   • metadata: Dict[str, Any]
   • template_vars: List[str]
   • kwargs: Dict[str, str]
   • output_parser: Optional[BaseOutputParser]
   • template_var_mappings: Optional[Dict[str, Any]]
   • function_mappings: Optional[Dict[str, Callable]]

Here is an example of how you can implement a custom prompt template class:

from llama_index.legacy.prompts.base import BasePromptTemplate
from typing import Any, Dict, List, Optional
from llama_index.legacy.core.llms.types import ChatMessage, BaseLLM

class CustomPromptTemplate(BasePromptTemplate):
    metadata: Dict[str, Any] = {}
    template_vars: List[str] = []
    kwargs: Dict[str, str] = {}
    output_parser: Optional[BaseOutputParser] = None
    template_var_mappings: Optional[Dict[str, Any]] = {}
    function_mappings: Optional[Dict[str, Callable]] = {}

    def partial_format(self, **kwargs: Any) -> "BasePromptTemplate":
        # Implement your logic here
        pass

    def format(self, llm: Optional[BaseLLM] = None, **kwargs: Any) -> str:
        # Implement your logic here
        return "formatted string"

    def format_messages(self, llm: Optional[BaseLLM] = None, **kwargs: Any) -> List[ChatMessage]:
        # Implement your logic here
        return []

    def get_template(self, llm: Optional[BaseLLM] = None) -> str:
        # Implement your logic here
        return "template string"

# Example usage
custom_prompt = CustomPromptTemplate()

Make sure to replace the placeholder logic with your actual implementation. This should resolve the ValidationError and allow you to use the custom prompt template with the CondenseQuestionChatEngine ^[1][2].

To continue talking to Dosu, mention @dosu.