Applying MultiQueryRetriever Concepts to Pandas DataFrames

[mraguth]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

Code I am using for pandas dataframe agent:

agent_executor = create_pandas_dataframe_agent(
    llm,
    df,
    agent_type="openai-tools",
    verbose=True,
    return_intermediate_steps=True,
    allow_dangerous_code=True,
    allow_dangerous_requests=True,
)

Description

The concept I am interested in: Web link:
https://python.langchain.com/v0.1/docs/modules/data_connection/retrievers/MultiQueryRetriever/#supplying-your-own-prompt

I like the concept of MultiQueryRetriever; however, most of my work involves structured data and utilizes the pandas DataFrame agent concept. My question is: is it possible to apply a similar concept to a DataFrame instead of a vector database? Thank you for your help.

System Info

Name: langchain
Version: 0.2.16

[dosubot[bot]]

Yes, it is possible to apply the concept of MultiQueryRetriever to a pandas DataFrame instead of a vector database. The Retriever class in LangChain is designed to return documents given a text query and does not need to store documents, making it more general than a vector store. This flexibility allows it to be adapted to different data sources, including pandas DataFrames ^[1].

To implement this, you can create a custom retriever that processes DataFrame queries similarly to how MultiQueryRetriever handles vector database queries. Here is an example of how you can adapt the MultiQueryRetriever concept for a pandas DataFrame:

1. Create a Custom Retriever:

   • Implement a retrieval mechanism that can interpret and execute the generated queries against the DataFrame.

2. Generate Multiple Queries:

   • Use a language model (LLM) to generate multiple queries from different perspectives for a given user input query.

3. Retrieve Data from DataFrame:

   • Execute the generated queries against the DataFrame and combine the results.

Here is a sample implementation:

from typing import List
from langchain_core.output_parsers import BaseOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langchain.retrievers.multi_query import MultiQueryRetriever
import pandas as pd

# Output parser to split the LLM result into a list of queries
class LineListOutputParser(BaseOutputParser[List[str]]):
    def parse(self, text: str) -> List[str]:
        lines = text.strip().split("\n")
        return list(filter(None, lines))  # Remove empty lines

output_parser = LineListOutputParser()

QUERY_PROMPT = PromptTemplate(
    input_variables=["question"],
    template="""You are an AI language model assistant. Your task is to generate five 
    different versions of the given user question to retrieve relevant data from a pandas DataFrame. 
    Provide these alternative questions separated by newlines.
    Original question: {question}""",
)

llm = ChatOpenAI(temperature=0)

# Chain
llm_chain = QUERY_PROMPT | llm | output_parser

# Sample DataFrame
data = {
    'Name': ['Alice', 'Bob', 'Charlie', 'David'],
    'Age': [24, 30, 22, 35],
    'City': ['New York', 'Los Angeles', 'Chicago', 'Houston']
}
df = pd.DataFrame(data)

# Custom retriever for DataFrame
class DataFrameRetriever:
    def __init__(self, df: pd.DataFrame):
        self.df = df

    def retrieve(self, query: str) -> pd.DataFrame:
        # Implement your query execution logic here
        # For simplicity, let's assume the query is a column name
        if query in self.df.columns:
            return self.df[[query]]
        else:
            return pd.DataFrame()

# MultiQueryRetriever adapted for DataFrame
class MultiQueryDataFrameRetriever(MultiQueryRetriever):
    def __init__(self, retriever, llm_chain, include_original=False):
        super().__init__(retriever, llm_chain, include_original)

    def retrieve_documents(self, queries: List[str]) -> pd.DataFrame:
        results = [self.retriever.retrieve(query) for query in queries]
        return pd.concat(results).drop_duplicates()

# Instantiate the custom retriever
df_retriever = DataFrameRetriever(df)

# Instantiate the MultiQueryDataFrameRetriever
multi_query_retriever = MultiQueryDataFrameRetriever(
    retriever=df_retriever, llm_chain=llm_chain, include_original=True
)

# Example usage
question = "What are the ages of the people?"
queries = multi_query_retriever.generate_queries(question)
result_df = multi_query_retriever.retrieve_documents(queries)
print(result_df)

This code demonstrates how to set up a custom retriever for a pandas DataFrame and use it with the MultiQueryRetriever concept to generate and execute multiple queries, providing a more comprehensive set of results ^[1].

To continue talking to Dosu, mention @dosu.

[mraguth]

@dosu can you hlep on this error? I am getting error on this statement...

multi_query_retriever = MultiQueryDataFrameRetriever(
    retriever=df_retriever, llm_chain=llm_chain, include_original=True
)

**TypeError: BaseModel.__init__() takes 1 positional argument but 4 were given**

[dosubot[bot]]

The error "TypeError: BaseModel.init() takes 1 positional argument but 4 were given" occurs because the MultiQueryDataFrameRetriever class is likely inheriting from a class that does not accept the arguments you are passing to it.

To resolve this issue, you should ensure that the class you are trying to instantiate (MultiQueryDataFrameRetriever) is defined correctly and accepts the arguments retriever, llm_chain, and include_original.

Here is an example of how you might define such a class:

from langchain_core.retrievers import BaseRetriever
from langchain_core.pydantic_v1 import Field
from langchain_core.runnables import Runnable

class MultiQueryDataFrameRetriever(BaseRetriever):
    retriever: BaseRetriever
    llm_chain: Runnable
    include_original: bool = False

    class Config:
        allow_population_by_field_name = True
        arbitrary_types_allowed = True

    def __init__(self, retriever: BaseRetriever, llm_chain: Runnable, include_original: bool = False, **kwargs):
        super().__init__(**kwargs)
        self.retriever = retriever
        self.llm_chain = llm_chain
        self.include_original = include_original

Ensure that the MultiQueryDataFrameRetriever class is correctly defined and accepts the arguments you are passing to it ^[1].

To continue talking to Dosu, mention @dosu.

[mraguth]

@dosu....can you give me full code please?

[mraguth]

@dosu the main problem is ... with pandas dataframe. I am getting this error with df_retriever from the code:

retriever
  value is not a valid dict (type=type_error.dict)

[dosubot[bot]]

To address the issues you're encountering, I'll provide a complete and corrected implementation of the MultiQueryDataFrameRetriever class and its usage with a pandas DataFrame.

Here is the full code:

from typing import List
from langchain_core.output_parsers import BaseOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain_core.retrievers import BaseRetriever
from langchain_core.pydantic_v1 import Field
from langchain_core.runnables import Runnable
import pandas as pd

# Output parser to split the LLM result into a list of queries
class LineListOutputParser(BaseOutputParser[List[str]]):
    def parse(self, text: str) -> List[str]:
        lines = text.strip().split("\n")
        return list(filter(None, lines))  # Remove empty lines

output_parser = LineListOutputParser()

QUERY_PROMPT = PromptTemplate(
    input_variables=["question"],
    template="""You are an AI language model assistant. Your task is to generate five 
    different versions of the given user question to retrieve relevant data from a pandas DataFrame. 
    Provide these alternative questions separated by newlines.
    Original question: {question}""",
)

llm = ChatOpenAI(temperature=0)

# Chain
llm_chain = QUERY_PROMPT | llm | output_parser

# Sample DataFrame
data = {
    'Name': ['Alice', 'Bob', 'Charlie', 'David'],
    'Age': [24, 30, 22, 35],
    'City': ['New York', 'Los Angeles', 'Chicago', 'Houston']
}
df = pd.DataFrame(data)

# Custom retriever for DataFrame
class DataFrameRetriever(BaseRetriever):
    def __init__(self, df: pd.DataFrame):
        self.df = df

    def _get_relevant_documents(self, query: str) -> pd.DataFrame:
        # Implement your query execution logic here
        # For simplicity, let's assume the query is a column name
        if query in self.df.columns:
            return self.df[[query]]
        else:
            return pd.DataFrame()

# MultiQueryRetriever adapted for DataFrame
class MultiQueryDataFrameRetriever(BaseRetriever):
    retriever: DataFrameRetriever
    llm_chain: Runnable
    include_original: bool = False

    class Config:
        allow_population_by_field_name = True
        arbitrary_types_allowed = True

    def __init__(self, retriever: DataFrameRetriever, llm_chain: Runnable, include_original: bool = False, **kwargs):
        super().__init__(**kwargs)
        self.retriever = retriever
        self.llm_chain = llm_chain
        self.include_original = include_original

    def generate_queries(self, question: str) -> List[str]:
        return self.llm_chain.invoke({"question": question})

    def retrieve_documents(self, queries: List[str]) -> pd.DataFrame:
        results = [self.retriever._get_relevant_documents(query) for query in queries]
        return pd.concat(results).drop_duplicates()

# Instantiate the custom retriever
df_retriever = DataFrameRetriever(df)

# Instantiate the MultiQueryDataFrameRetriever
multi_query_retriever = MultiQueryDataFrameRetriever(
    retriever=df_retriever, llm_chain=llm_chain, include_original=True
)

# Example usage
question = "What are the ages of the people?"
queries = multi_query_retriever.generate_queries(question)
result_df = multi_query_retriever.retrieve_documents(queries)
print(result_df)

This code defines a DataFrameRetriever class that retrieves data from a pandas DataFrame and a MultiQueryDataFrameRetriever class that generates multiple queries and retrieves data based on those queries. The MultiQueryDataFrameRetriever class is correctly defined to accept the necessary arguments and handle the retrieval process.

This should resolve the errors you were encountering and allow you to apply the MultiQueryRetriever concept to a pandas DataFrame.

To continue talking to Dosu, mention @dosu.

[mraguth]

@how I got queries like below... How I refine below queries output and pass those in to my dataframe agent agent.invoke({"input":xxxxxx)

INFO:langchain.retrievers.multi_query:Generated queries: ['1. How can I gather the average turnaround hours (TAT_hours) by year and quarter for the case record type data intake across all methods of intake?', '2. What is the process to analyze and compare the average TAT_hour values for phone versus other methods of intake on a quarterly basis?', '3. Can you provide insights on how to extract and compare average TAT_hours for different methods of intake, specifically focusing on phone and others, by quarter?', '4. What steps should I follow to derive data insights for each quarter based on the average TAT_hours for various methods of intake, particularly comparing phone intake?', '5. How do I calculate and interpret the average TAT_hours by quarter for case record types, and what insights can be drawn from comparing phone intake to other methods?']
[Document(page_content='data intake'), Document(page_content='fax'), Document(page_content='api'), Document(page_content='benefits verification'), Document(page_content='phone')]