Multiple retrieval for seperate indexes

[regybean]

I want to create a system which can query two seperate indexes seperately and include their output in a contextualised query. The plan is that a question would be provided. The first index would get the relevant user docs nodes. Then the second index would be used to get context nodes. Then the query would be formatted as in text_qa_template_str to be finally sent to the LLM for a response.

text_qa_template_str = (
"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n"
"User documentation information is"
" below.\n---------------------\n{user_docs_str}\n---------------------\n"
"Context information is"
" below.\n---------------------\n{context_str}\n---------------------\n"
" You are an AI model specialized in providing advice on railway operational technology (OT) cybersecurity.\n"
" Your responses must be based solely on the context information provided to you.\n"
" If the information given is insufficient for you to provide a useful answer, explicitly state that you do not know the answer given the provided context.\n"
" Your goal is to offer precise, relevant, and actionable advice or insights.\n"
" Answer the question clearly and avoid sentences explaining that you have used the provided context'\n"
" <|eot_id|><|start_header_id|>user<|end_header_id|>\n\n"
" {query_str}\n"
"<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
)

This is my existing code which is partially implemented I don't fully understand how the links would work or if my modules are wrong, please can someone help out to get this working.

docs_path = "../docs2"

domain_documents = SimpleDirectoryReader(docs_path).load_data()

domain_index = VectorStoreIndex.from_documents(domain_documents)#, service_context=svc)

domain_index.storage_context.persist()

domain_engine = domain_index.as_query_engine(llm=llm, text_qa_template=text_qa_template)
domain_retriever = domain_index.as_retriever(llm=llm)

docs_path = "../pdfs"

user_documents = SimpleDirectoryReader(docs_path).load_data()

user_index = VectorStoreIndex.from_documents(user_documents)#, service_context=svc)

user_index.storage_context.persist()

user_engine = user_index.as_query_engine(llm=llm, text_qa_template=text_qa_template)
user_docs_retriever = user_index.as_retriever(llm=llm)

prompt_template = PromptTemplate(text_qa_template_str)

input_component = PromptTemplate("Does the SUC description include a high-level description and depiction of the SUC?")

p = QueryPipeline(verbose=True)

p.add_modules(
{
"input":input_component,
"user_docs_retriever": user_docs_retriever,
"domain_retriever": domain_retriever,
"prompt_template": prompt_template,
"llm":llm,
}

)
p.add_link(...

[regybean]

@dosu

[dosubot[bot]]

To create a system that queries two separate indexes and includes their output in a contextualized query, you can use the QueryPipeline class to chain the different modules together. Here is a step-by-step guide to achieve this:

1. Create Tool Configurations for Each Index:

from llama_index.core.base.base_query_engine import BaseQueryEngine
from llama_index.core.langchain_helpers.agents.tools import IndexToolConfig, LlamaIndexTool

# Assuming you have two query engines for two separate indexes
domain_engine = BaseQueryEngine()  # Replace with actual initialization
user_engine = BaseQueryEngine()  # Replace with actual initialization

# Create tool configurations for each index
tool_config_domain = IndexToolConfig(
    query_engine=domain_engine,
    name="DomainIndex",
    description="Query engine for the domain index"
)

tool_config_user = IndexToolConfig(
    query_engine=user_engine,
    name="UserIndex",
    description="Query engine for the user index"
)

2. Create Tools from Configurations:

# Create tools from the configurations
tool_domain = LlamaIndexTool.from_tool_config(tool_config_domain)
tool_user = LlamaIndexTool.from_tool_config(tool_config_user)

3. Create and Configure the Query Pipeline:

from llama_index.core.query_pipeline.query import QueryPipeline, Link

# Initialize the query pipeline
pipeline = QueryPipeline(verbose=True)

# Add tools as modules to the pipeline
pipeline.add_modules({
    "user_docs_retriever": tool_user,
    "domain_retriever": tool_domain,
    "prompt_template": prompt_template,
    "llm": llm,
})

# Link the tools in the pipeline
pipeline.add_link("user_docs_retriever", "domain_retriever")
pipeline.add_link("domain_retriever", "prompt_template")
pipeline.add_link("prompt_template", "llm")

4. Query the Pipeline and Combine Results:

def query_both_indexes(input_query: str) -> str:
    # Run the pipeline
    result = pipeline.run(input_query=input_query)
    
    # Combine the results
    combined_response = f"Response from User Index: {result['user_docs_retriever']}\nResponse from Domain Index: {result['domain_retriever']}"
    return combined_response

# Example usage
input_query = "Does the SUC description include a high-level description and depiction of the SUC?"
combined_response = query_both_indexes(input_query)
print(combined_response)

This setup ensures that the system queries two separate indexes and includes their output in a contextualized query. The QueryPipeline class allows for chaining different modules, making it flexible for various query workflows ^[1][2].

To continue talking to Dosu, mention @dosu.