How to add question rewriter function to pandas_dataframe_agent?

[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

agent = create_pandas_dataframe_agent(
    llm,
    df,
    agent_type=AgentType.OPENAI_FUNCTIONS,
    verbose=True,
    return_intermediate_steps=True,
    allow_dangerous_code=True,
    allow_dangerous_requests=True,
)
agent.invoke({"input": query_str})
---------------------------------------
def create_question_rewriter(model):
    """
    Creates a question rewriter chain that rewrites a given question to improve its clarity and relevance.

    Returns:
        A callable function that takes a question as input and returns the rewritten question as a string.
    """
    re_write_prompt = hub.pull("efriis/self-rag-question-rewriter")
    question_rewriter = re_write_prompt | self.model | StrOutputParser()

    return question_rewriter

Description

I have two requests:

1. I am trying to in corporate 'create_question_rewriter' with pandas_dataframe_agent. Is that possible? can you give me a sample code?
2. And also I like to build 'Define the Query Generator' for my pandas_dataframe agent as well. Like to have a sample code.
   In the given example, it is using vector store. However, in my pandas_dataframe_agent, I am using 'df' not a vector store.

Link for the query generator topic: https://github.com/langchain-ai/langchain/blob/master/cookbook/rag_fusion.ipynb?ref=blog.langchain.dev

System Info

Name: langchain
Version: 0.2.16

[dosubot[bot]]

To incorporate the create_question_rewriter function with the pandas_dataframe_agent and build a query generator, you can follow these steps:

1. Incorporate create_question_rewriter Function:

   Modify the agent's invocation to include the rewritten question as input. Here is a sample code:

   from langchain_community.chat_models import ChatOpenAI
   from langchain_experimental.agents import create_pandas_dataframe_agent
   import pandas as pd
   
   # Load your DataFrame
   df = pd.read_csv("your_data.csv")
   
   # Initialize your language model
   llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
   
   # Define the create_question_rewriter function
   def create_question_rewriter(model):
       """
       Creates a question rewriter chain that rewrites a given question to improve its clarity and relevance.
   
       Returns:
           A callable function that takes a question as input and returns the rewritten question as a string.
       """
       re_write_prompt = hub.pull("efriis/self-rag-question-rewriter")
       question_rewriter = re_write_prompt | model | StrOutputParser()
       return question_rewriter
   
   # Create the question rewriter
   question_rewriter = create_question_rewriter(llm)
   
   # Create the agent executor using the OPENAI_FUNCTIONS agent type
   agent_executor = create_pandas_dataframe_agent(
       llm,
       df,
       verbose=True,
       agent_type="openai-tools",
       allow_dangerous_code=True  # Ensure you understand the security implications
   )
   
   # Rewrite the query and invoke the agent
   query_str = "What is the average age of passengers?"
   rewritten_query = question_rewriter(query_str)
   response = agent_executor.invoke({"input": rewritten_query})
   print(response)

2. Build a Query Generator:

   You can create a function that generates queries based on the DataFrame's structure and content. Here is an example:

   from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
   from langchain_core.pydantic_v1 import BaseModel, Field
   from langchain_experimental.tools import PythonAstREPLTool
   from langchain.agents import AgentExecutor, OpenAIFunctionsAgent
   
   class PythonInputs(BaseModel):
       query: str = Field(description="code snippet to run")
   
   df = pd.read_csv("your_data.csv")
   template = """
   You are working with a pandas dataframe in Python. The name of the dataframe is `df`.
   It is important to understand the attributes of the dataframe before working with it. This is the result of running `df.head().to_markdown()`
   
   <df>
   {dhead}
   </df>
   
   You are not meant to use only these rows to answer questions - they are meant as a way of telling you about the shape and schema of the dataframe.
   You also do not have use only the information here to answer questions - you can run intermediate queries to do exploratory data analysis to give you more information as needed.
   
   For example:
   
   <question>How old is Jane?</question>
   <logic>Use `person_name_search` since you can use the query `Jane`</logic>
   <question>Who has id 320</question>
   <logic>Use `python_repl` since even though the question is about a person, you don't know their name so you can't include it.</logic>
   """  # noqa: E501
   
   template = template.format(dhead=df.head().to_markdown())
   
   prompt = ChatPromptTemplate.from_messages(
       [
           ("system", template),
           MessagesPlaceholder(variable_name="agent_scratchpad"),
           ("human", "{input}"),
       ]
   )
   
   repl = PythonAstREPLTool(
       locals={"df": df},
       name="python_repl",
       description="Runs code and returns the output of the final line",
       args_schema=PythonInputs,
   )
   tools = [repl]
   agent = OpenAIFunctionsAgent(
       llm=ChatOpenAI(temperature=0, model="gpt-4"), prompt=prompt, tools=tools
   )
   agent_executor = AgentExecutor(
       agent=agent, tools=tools, max_iterations=5, early_stopping_method="generate"
   ) | (lambda x: x["output"])
   
   # Typing for playground inputs
   class AgentInputs(BaseModel):
       input: str
   
   agent_executor = agent_executor.with_types(input_type=AgentInputs)
   
   # Example usage
   query = "What is the average age of passengers?"
   response = agent_executor.invoke({"input": query})
   print(response)

This code demonstrates how to set up the pandas_dataframe_agent with a DataFrame and use it to generate queries, incorporating both the question rewriter and a custom query generator ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[mraguth]

@dosu I am getting this error at execution:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[71], [line 38](vscode-notebook-cell:?execution_count=71&line=38)
     [36](vscode-notebook-cell:?execution_count=71&line=36) # Rewrite the query and invoke the agent
     [37](vscode-notebook-cell:?execution_count=71&line=37) query_str = output_query_str
---> [38](vscode-notebook-cell:?execution_count=71&line=38) rewritten_query = question_rewriter(query_str)
     [39](vscode-notebook-cell:?execution_count=71&line=39) response = agent_executor.invoke({"input": rewritten_query})
     [40](vscode-notebook-cell:?execution_count=71&line=40) print(response)

TypeError: 'RunnableSequence' object is not callable

My code:

from langchain_community.chat_models import ChatOpenAI
from langchain_experimental.agents import create_pandas_dataframe_agent
import pandas as pd
from langchain import hub
from langchain_core.output_parsers import StrOutputParser


# Define the create_question_rewriter function
def create_question_rewriter(model):
    """
    Creates a question rewriter chain that rewrites a given question to improve its clarity and relevance.

    Returns:
        A callable function that takes a question as input and returns the rewritten question as a string.
    """
    re_write_prompt = hub.pull("efriis/self-rag-question-rewriter")
    question_rewriter = re_write_prompt | model | StrOutputParser()
    return question_rewriter


# Create the question rewriter
question_rewriter = create_question_rewriter(llm)

# Create the agent executor using the OPENAI_FUNCTIONS agent type
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,
    agent_executor_kwargs={"handle_parsing_errors": True},
)

# Rewrite the query and invoke the agent
query_str = output_query_str
rewritten_query = question_rewriter(query_str)
response = agent_executor.invoke({"input": rewritten_query})
print(response)

[dosubot[bot]]

The error you're encountering, TypeError: 'RunnableSequence' object is not callable, occurs because RunnableSequence objects should be invoked using the invoke method instead of being called directly.

Here's how you can modify your code to use the invoke method:

from langchain_community.chat_models import ChatOpenAI
from langchain_experimental.agents import create_pandas_dataframe_agent
import pandas as pd
from langchain import hub
from langchain_core.output_parsers import StrOutputParser

# Load your DataFrame
df = pd.read_csv("your_data.csv")

# Initialize your language model
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)

# Define the create_question_rewriter function
def create_question_rewriter(model):
    """
    Creates a question rewriter chain that rewrites a given question to improve its clarity and relevance.

    Returns:
        A callable function that takes a question as input and returns the rewritten question as a string.
    """
    re_write_prompt = hub.pull("efriis/self-rag-question-rewriter")
    question_rewriter = re_write_prompt | model | StrOutputParser()
    return question_rewriter

# Create the question rewriter
question_rewriter = create_question_rewriter(llm)

# Create the agent executor using the OPENAI_FUNCTIONS agent type
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,
    agent_executor_kwargs={"handle_parsing_errors": True},
)

# Rewrite the query and invoke the agent
query_str = "What is the average age of passengers?"
rewritten_query = question_rewriter.invoke(query_str)  # Use invoke method here
response = agent_executor.invoke({"input": rewritten_query})
print(response)

By using the invoke method on the question_rewriter, you should be able to avoid the TypeError and successfully rewrite the query before passing it to the agent ^[1][2].

To continue talking to Dosu, mention @dosu.

[mraguth]

@dosu. Both scenarios are working!!! I have one last question. When it comes to the create_question_rewriter function, what model should I use as the base? Is 'self-rag-question-rewriter' the best option, or are there other models you’d recommend?

Also, is it possible to build custom templates or give the model clear insights about my data before it rewrites the questions? If so, could you provide some sample code on how to add custom templates and give the model insights about my data?

Thanks a lot!

[mraguth]

@dosu I mean how I customize create_question_rewriter by adding my own prompt?

[dosubot[bot]]

To customize the create_question_rewriter function by adding your own prompt and providing the model with clear insights about your data, you can follow these steps:

1. Define a Custom Prompt: Create a PromptTemplate with your custom prompt.
2. Create an LLMChain: Use the custom prompt to create an LLMChain.
3. Instantiate the RePhraseQueryRetriever: Use the LLMChain to create a RePhraseQueryRetriever.

Here is the code to achieve this:

from langchain.chains import LLMChain
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langchain.retrievers import RePhraseQueryRetriever

# Define your custom prompt
QUERY_PROMPT = PromptTemplate(
    input_variables=["question"],
    template="""You are an assistant tasked with taking a natural language query from a user
    and converting it into a query for a vectorstore. In the process, strip out all 
    information that is not relevant for the retrieval task and return a new, simplified
    question for vectorstore retrieval. The new user query should be in pirate speech.
    Here is the user query: {question} """,
)

# Create an LLMChain with the custom prompt
llm = ChatOpenAI(temperature=0)
llm_chain = LLMChain(llm=llm, prompt=QUERY_PROMPT)

# Instantiate the RePhraseQueryRetriever with the custom LLMChain
retriever_from_llm_chain = RePhraseQueryRetriever(
    retriever=vectorstore.as_retriever(), llm_chain=llm_chain
)

# Use the customized retriever
docs = retriever_from_llm_chain.invoke(
    "Hi I'm Lance. What is Maximum Inner Product Search?"
)

Additionally, you can refer to the Rewrite-Retrieve-Read method proposed in the paper Query Rewriting for Retrieval-Augmented Large Language Models for more advanced query rewriting techniques. This method involves rewriting the search query to improve retrieval performance, especially for complex or "distracted" queries. Here is an example of how to implement this method:

from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI

# Define the rewrite prompt
template = """Provide a better search query for \
web search engine to answer the given question, end \
the queries with ’**’. Question: \
{x} Answer:"""
rewrite_prompt = ChatPromptTemplate.from_template(template)

# Create the rewriter
rewriter = rewrite_prompt | ChatOpenAI(temperature=0) | StrOutputParser() | (lambda text: text.strip('"').strip("**"))

# Define the retriever
search = DuckDuckGoSearchAPIWrapper()
def retriever(query):
    return search.run(query)

# Create the Rewrite-Retrieve-Read chain
rewrite_retrieve_read_chain = (
    {
        "context": {"x": RunnablePassthrough()} | rewriter | retriever,
        "question": RunnablePassthrough(),
    }
    | ChatPromptTemplate.from_template("Answer the users question based only on the following context:\n\n<context>\n{context}\n</context>\n\nQuestion: {question}")
    | ChatOpenAI(temperature=0)
    | StrOutputParser()
)

# Use the chain with a distracted query
distracted_query = "man that sam bankman fried trial was crazy! what is langchain?"
rewrite_retrieve_read_chain.invoke(distracted_query)

For more advanced customization, you can use the FewShotPromptTemplate to provide examples and structure for your prompts, as shown in the libs/langchain/langchain/chains/constitutional_ai/prompts.py file. This allows you to include examples of critiques and revisions to guide the model's behavior more precisely.

from langchain_core.prompts.few_shot import FewShotPromptTemplate
from langchain_core.prompts.prompt import PromptTemplate

# Define critique and revision examples
critique_example = PromptTemplate(
    template="""Human: {input_prompt}

Model: {output_from_model}

Critique Request: {critique_request}

Critique: {critique}""",
    input_variables=[
        "input_prompt",
        "output_from_model",
        "critique_request",
        "critique",
    ],
)

revision_example = PromptTemplate(
    template="""Human: {input_prompt}

Model: {output_from_model}

Critique Request: {critique_request}

Critique: {critique}

Revision Request: {revision_request}

Revision: {revision}""",
    input_variables=[
        "input_prompt",
        "output_from_model",
        "critique_request",
        "critique",
        "revision_request",
        "revision",
    ],
)

# Create FewShotPromptTemplate with examples
CRITIQUE_PROMPT = FewShotPromptTemplate(
    example_prompt=critique_example,
    examples=[
        {k: v for k, v in e.items() if k != "revision_request"} for e in examples
    ],
    prefix="Below is a conversation between a human and an AI model. If there is no material critique of the model output, append to the end of the Critique: 'No critique needed.' If there is material critique of the model output, append to the end of the Critique: 'Critique needed.'",
    suffix="""Human: {input_prompt}
Model: {output_from_model}

Critique Request: {critique_request}

Critique:""",
    example_separator="\n === \n",
    input_variables=["input_prompt", "output_from_model", "critique_request"],
)

REVISION_PROMPT = FewShotPromptTemplate(
    example_prompt=revision_example,
    examples=examples,
    prefix="Below is a conversation between a human and an AI model.",
    suffix="""Human: {input_prompt}

Model: {output_from_model}

Critique Request: {critique_request}

Critique: {critique}

If the critique does not identify anything worth changing, ignore the Revision Request and do not make any revisions. Instead, return "No revisions needed".

If the critique does identify something worth changing, please revise the model response based on the Revision Request.

Revision Request: {revision_request}

Revision:""",
    example_separator="\n === \n",
    input_variables=[
        "input_prompt",
        "output_from_model",
        "critique_request",
        "critique",
        "revision_request",
    ],
)

This approach ensures that even complex or distracted queries are effectively rewritten and retrieved, providing more accurate and relevant results ^[1][2][3].

To continue talking to Dosu, mention @dosu.