Vector store agent with sources

[courtney-wright]

Hi!
I'm wondering if there's a way to use a vector store chain that provides sources with an agent?

For example, I'm trying to use a RetrievalQAWithSourcesChain with an executor agent.

I've tried these two main concepts from your documentation with a few different variants including switching the chain type and providing parameters to the chain.

1. creating the chain as a tool
2. creating a vector store agent

For 1. I get an error that the tool cannot return more than one parameter. For 2. I can't manage to get the sources that are accessed in the vectorstore.

Any help would be much appreciated!

[oddrationale]

So the tricky part is that the RetrievalQAwithSourcesChain chain does not receive and return a single input and output. But on the other hand, Tools require a single string input and single string output.

We can work around this by wrapping the RetrievalQAwithSourcesChain in a function that takes a single string input and single string output. For example, we can return the answer and source as a string.

from langchain.agents import (
    AgentType,
    initialize_agent,
    Tool,
)
from langchain.chains import RetrievalQAWithSourcesChain
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import TextLoader
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma


text = TextLoader(
    "../state_of_the_union_2023.txt"
).load()

docs = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=0,
).split_documents(text)

vectorstore = Chroma.from_documents(
    documents=docs,
    embedding=OpenAIEmbeddings(),
)

qa_chain = RetrievalQAWithSourcesChain.from_chain_type(
    llm=ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo"),
    chain_type="stuff",
    retriever=vectorstore.as_retriever(),
)

def run_qa_chain(question):
    results = qa_chain({"question":question},return_only_outputs=True)
    return str(results)

Then we setup our agent to use this function as a Tool:

agent_exe = initialize_agent(
    agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    llm=ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo"),
    tools=[
        Tool(
            name="State of the Union 2023 QA System",
            func=run_qa_chain,
            description="Useful for when you need to answer questions about "
                        "the most recent state of the union address. Input "
                        "should be a fully formed question.",
            return_direct=True,
        ),
    ],
    verbose=True,
)

We can set return_direct=True in order to return the exact result of the qa_chain including the sources.

Hope that helps put you in the right direction!

[courtney-wright]

This worked great. I then parsed the string back into a dictionary and am now able to explain to the user where the information comes from much better!

Thanks so much 😄

[meet1919]

You can also define a custom tool and modify the _run() function to give a single output containing answer and sources.

def _run(self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None) -> str:
        """Use the tool."""
        retrieval_chain = self.create_retrieval_chain()
        answer = retrieval_chain(query)['answer']
        sources = '\n'.join(retrieval_chain(query)['sources'].split(', '))

        return f'{answer}\nSources:\n{sources}'

Also use return_direct=True

[cryptocoinserver]

Thanks this helped. Maybe useful for others:

class CustomRetrievalTool(BaseTool):
    retriever: Any
    return_direct = True

    def _run(
        self, query: str, run_manager: Optional[CallbackManagerForToolRun] = None
    ) -> str:
        """Use the tool."""
        result = self.retriever({"query": query})
        answer = result["result"]
        # Convert the list of source_documents to a string.
        sources = str(result["source_documents"])
        # Concatenate the answer and the sources into a single string.
        return f"Answer: {answer}\nSources:\n{sources}"

    async def _arun(
        self, query: str, run_manager: Optional[AsyncCallbackManagerForToolRun] = None
    ) -> str:
        """Use the tool asynchronously."""
        raise NotImplementedError("CustomRetrieval does not support async")
        
custom_literature_tool = CustomRetrievalTool(
    name="Information Extractor",
    description="This tool is capable of extracting factual data related to a provided query.",
    retriever=literature_chain,
)

[fabmeyer]

@oddrationale I am getting the following error when using your solution but with two different tools:

code:

def load_qa_chain(docsearch):
    llm = OpenAI(
        client=OpenAI,
        model='text-davinci-003',
        streaming=True,
        callbacks=[StreamingStdOutCallbackHandler()],
        temperature=0,
        max_tokens=1024,
        verbose=True
    )

    FindallQAWithSourcesChain = RetrievalQAWithSourcesChain.from_chain_type(
        llm=llm,
        chain_type='refine',
        retriever=docsearch.as_retriever(),
        verbose=True,
    )
    return FindallQAWithSourcesChain
    
def run_qa_chain():
    results = st.session_state['qa_chain']({"question": st.session_state['user_input'][-1]}, return_only_outputs=True)
    return str(results)

def load_agent(companies):
    with st.spinner('Loading model...'):

        # Pandas agent
        PandasAgent = create_pandas_dataframe_agent(
            llm=ChatOpenAI(temperature=0, model='gpt-3.5-turbo-0613'),
            df=companies,
            verbose=True,
            agent_type=AgentType.OPENAI_FUNCTIONS,
            callback_manager=CallbackManager(handlers=[StreamingStdOutCallbackHandler()])
        )

        tools = [
            Tool(
                name="Findall QA with Sources System",
                func=run_qa_chain(),
                description=f"useful for when you need to answer questions about {st.session_state['thematic']['term']}. Will return text in natural language.",
                return_direct=True
            ),
            Tool(
                name="Pandas Agent",
                func=PandasAgent.run,
                description="useful for when you need to query Pandas DataFrames about companies. Will return a Pandas DataFrame.",
                return_direct=True
            ),
        ]

        agent = initialize_agent(
            agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
            llm=ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo"),
            tools=tools,
            verbose=True
        )

        return agent

2023-07-13 15:23:01.662 Uncaught app exception Traceback (most recent call last): File "/home/fabmeyer/.local/share/virtualenvs/personal-investment-assistant-Y3o7tQcK/lib/python3.9/site-packages/streamlit/runtime/scriptrunner/script_runner.py", line 552, in _run_script exec(code, module.__dict__) File "/home/fabmeyer/Dev/Python/personal-investment-assistant/personal_assistant.py", line 274, in <module> st.session_state['agent'] = load_agent(st.session_state['companies_df']) File "/home/fabmeyer/Dev/Python/personal-investment-assistant/personal_assistant.py", line 200, in load_agent Tool( File "/home/fabmeyer/.local/share/virtualenvs/personal-investment-assistant-Y3o7tQcK/lib/python3.9/site-packages/langchain/tools/base.py", line 480, in __init__ super(Tool, self).__init__( File "pydantic/main.py", line 341, in pydantic.main.BaseModel.__init__ pydantic.error_wrappers.ValidationError: 1 validation error for Tool func

Any idea why is that?