Custom query component

[sumitbindra]

How do I create a custom query component that returns both the query engine and context?

I tried it the following way but there are errors regarding the index not being initialized when I run the query engine itself:

This function converts the survey yaml file to a vector store that can be used as context with the SQL

class LoadQuestions(CustomQueryComponent):

def __init__(self):
    super().__init__()
    self.__dict__['vector_retriever'] = None
    self.__dict__['index'] = None


def _load_or_create_index(self, table_name):
    print(f"\n\nRunning load or create index {table_name}\n\n")
    dir_name = table_name.replace("_responses", "") if table_name else "default"
    PERSIST_DIR = f"./surveys/{dir_name}/RAGFiles/quesindex"

    if not os.path.exists(PERSIST_DIR):
        reader = SimpleDirectoryReader(f"./surveys/{dir_name}")
        documents = reader.load_data()
        self.index = VectorStoreIndex.from_documents(documents)
        self.index.storage_context.persist(persist_dir=PERSIST_DIR)
    else:
        storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
        self.index = cast(VectorStoreIndex, load_index_from_storage(storage_context))

    self.vector_retriever = self.index.as_retriever(similarity_top_k=3)

def _get_question_context(self, query_str):
    relevant_nodes = self.vector_retriever.retrieve(query_str) if self.vector_retriever else []
    question_context = ""
    if len(relevant_nodes) > 0:
        question_context = "\nHere are questions that are relevant to this query\n"
        for node in relevant_nodes:
            question_context += str(node.get_content()) + "\n"
    return question_context

def get_query_engine(self):
    if self.index is None:
        raise ValueError("Index has not been loaded. Run the component first.")
    return self.index.as_query_engine()
  
def _run_component(self, **kwargs):
    query_str = kwargs.get('query_str')
    table_name = kwargs.get('table_name')
    if query_str is None:
        raise ValueError("query_str cannot be None")

    return {"quesindex": question_context}

@property
def _input_keys(self):
    return {"query_str", "table_name"}

@property
def _optional_input_keys(self):
    return set()

@property
def _output_keys(self):
    return {"quesindex"}

[dosubot[bot]]

To create a custom query component that returns both the query engine and context, and to resolve errors regarding the index not being initialized, you can follow these steps:

1. Ensure Proper Initialization of the Index: Make sure the index is properly initialized before running the query engine.
2. Use the QueryEngineComponent Class: This class encapsulates the query engine and provides methods to run the component synchronously and asynchronously.

Here's an example of how you can modify your LoadQuestions class to achieve this:

import os
from typing import Optional, Dict, Any, cast
from llama_index.core.base.base_query_engine import BaseQueryEngine, QueryEngineComponent
from llama_index.core.callbacks.base import CallbackManager
from llama_index.core.readers.file.base import SimpleDirectoryReader
from llama_index.core.storage import StorageContext, load_index_from_storage
from llama_index.core.indices.vector_store import VectorStoreIndex

class LoadQuestions(QueryEngineComponent):
    
    def __init__(self):
        super().__init__()
        self.vector_retriever = None
        self.index = None

    def _load_or_create_index(self, table_name: Optional[str]):
        print(f"\n\nRunning load or create index {table_name}\n\n")
        dir_name = table_name.replace("_responses", "") if table_name else "default"
        PERSIST_DIR = f"./surveys/{dir_name}/RAGFiles/quesindex"

        if not os.path.exists(PERSIST_DIR):
            reader = SimpleDirectoryReader(f"./surveys/{dir_name}")
            documents = reader.load_data()
            self.index = VectorStoreIndex.from_documents(documents)
            self.index.storage_context.persist(persist_dir=PERSIST_DIR)
        else:
            storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
            self.index = cast(VectorStoreIndex, load_index_from_storage(storage_context))

        self.vector_retriever = self.index.as_retriever(similarity_top_k=3)

    def _get_question_context(self, query_str: str) -> str:
        relevant_nodes = self.vector_retriever.retrieve(query_str) if self.vector_retriever else []
        question_context = ""
        if len(relevant_nodes) > 0:
            question_context = "\nHere are questions that are relevant to this query\n"
            for node in relevant_nodes:
                question_context += str(node.get_content()) + "\n"
        return question_context
    
    def get_query_engine(self) -> BaseQueryEngine:
        if self.index is None:
            raise ValueError("Index has not been loaded. Run the component first.")
        return self.index.as_query_engine()
      
    def _run_component(self, **kwargs: Any) -> Dict[str, Any]:
        query_str = kwargs.get('query_str')
        table_name = kwargs.get('table_name')
        if query_str is None:
            raise ValueError("query_str cannot be None")

        self._load_or_create_index(table_name)
        question_context = self._get_question_context(query_str)
        return {"quesindex": question_context}

    @property
    def input_keys(self) -> set:
        return {"query_str", "table_name"}

    @property
    def optional_input_keys(self) -> set:
        return set()

    @property
    def output_keys(self) -> set:
        return {"quesindex"}

Key Points:

• Initialization: The _load_or_create_index method ensures the index is loaded or created before any queries are run.
• Query Engine: The get_query_engine method returns the query engine if the index is initialized.
• Component Execution: The _run_component method handles the execution of the component, ensuring the index is loaded and the query context is retrieved.

This approach ensures that the index is properly initialized and the query engine can be used without errors ^[1][2].

To continue talking to Dosu, mention @dosu.

[sumitbindra]

I have this custom query functions.

This function converts the survey yaml file to a vector store that can be used as context with the SQL

class LoadQuestions(CustomQueryComponent):

def __init__(self):
    super().__init__()
    self.__dict__['_index'] = None
    self.__dict__['_dir_name'] = None


def _run_component(self, **kwargs) -> Dict[str, Any]:
    query_str = kwargs.get('query_str')
    table_name = kwargs.get('table_name')
    if query_str is None:
        raise ValueError("query_str cannot be None")

    self._dir_name = table_name.replace("_responses", "") if table_name else "default"
    PERSIST_DIR = f"./surveys/{self._dir_name}/RAGFiles/quesindex"

    try:
        if not os.path.exists(PERSIST_DIR):
            reader = SimpleDirectoryReader(f"./surveys/{self._dir_name}")
            documents = reader.load_data()
            self._index = VectorStoreIndex.from_documents(documents)
            self._index.storage_context.persist(persist_dir=PERSIST_DIR)
        else:
            storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
            self._index = cast(VectorStoreIndex, load_index_from_storage(storage_context))

        # Retrieve relevant nodes for context
        vector_retriever = self._index.as_retriever(similarity_top_k=3)
        relevant_nodes = vector_retriever.retrieve(query_str)

        question_context = ""
        if len(relevant_nodes) > 0:
            question_context = "\nHere are questions that are relevant to this query\n"
            for node in relevant_nodes:
                question_context += str(node.get_content()) + "\n"

        return {"quesindex": question_context}

    except Exception as e:
        print(f"Error in LoadQuestions: {str(e)}")
        return {"quesindex": ""}

def get_query_engine(self):
    if self._index is None:
        raise ValueError("Index has not been initialized. Run _run_component first.")

    class QuestionQueryEngine(CustomQueryEngine):
        def __init__(self, _index):
            self._index = _index

        def custom_query(self, query_str: str):
            # Retrieve relevant nodes
            retriever = self._index.as_retriever(similarity_top_k=3)
            nodes = retriever.retrieve(query_str)

            # Prepare context from retrieved nodes
            context = "\n".join([node.get_content() for node in nodes])

            # Prepare prompt for the LLM
            prompt = f"Based on the following survey questions:\n\n{context}\n\nPlease answer the following question: {query_str}"

            # Query LLM
            #llm = OpenAI(temperature=0, model="gpt-3.5-turbo")
            llm_response = llm.complete(prompt)

            # Create and return Response object
            return str(llm_response)

    return QuestionQueryEngine(self._index)

@property
def _input_keys(self):
    return {"query_str", "table_name"}

@property
def _optional_input_keys(self):
    return set()

@property
def _output_keys(self):
    return {"quesindex"}

Why am I getting the error:

File "/home/runner/surveyanalysis/llama.py", line 141, in _run_component
self._dir_name = table_name.replace("_responses", "") if table_name else "default"
File "/home/runner/surveyanalysis/.pythonlibs/lib/python3.10/site-packages/pydantic/v1/main.py", line 357, in setattr
raise ValueError(f'"{self.class.name}" object has no field "{name}"')
ValueError: "LoadQuestions" object has no field "_dir_name"