Return_source_documents when RetrievalQA is an Agent tool

[pythonhacknewby]

I can return source documents fine when I run this code without being a tool in an agent. But when I run this code as part of an agent, below I get an error if I "return_source_documents" as a tool in an agent. I know it is from the qa.run since it returns more than one, but cant figure out how to fix it. I have commented out several tries of mine. Any help would be very much appreciated!

Doesn't work:
def qa_chat3(vector_store, query):
vectorstore = u.load_vectorstore(vector_store)
template="You are an expert test taker."
system_message_prompt = SystemMessagePromptTemplate.from_template(template)
human_template="Answer this question: {queries}"
human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])

qa = RetrievalQA.from_chain_type(

qa = RetrievalQAWithSourcesChain.from_chain_type(

    llm=OpenAI(temperature = 0),
    chain_type="stuff",
    retriever=vectorstore.as_retriever(),
    **_#return_source_documents = True,_**        
    )
print('Loading tools')
tools = [
	Tool.from_function(
    	name="Chat QA System",
    	#func=qa({"query":query}),
    	func=qa.run,
        #func=qa,
    	description="Useful for when you need to answer questions about the aspects asked. Input may be a partial or fully formed question.",
    	)
	]

prefix = """Have a conversation with a human, answering the following questions as best you can based on the context and memory available.
    You have access to a single tool:"""
suffix = """ Begin!
	{chat_history}
	Question: {input}
	{agent_scratchpad}"""

prompt = ZeroShotAgent.create_prompt(
    tools,
    prefix=prefix,
    suffix=suffix,
    input_variables=["input", "chat_history", "agent_scratchpad"],
)

print('loading chain')
llm_chain = LLMChain(
    	llm=OpenAI(
    	temperature=0,
    	model_name="gpt-3.5-turbo",
    ),
    prompt=prompt,
)

Declared outside of function

memory = ConversationBufferMemory(

memory_key="chat_history")

agent = ZeroShotAgent(llm_chain=llm_chain, tools=tools, verbose=True)
agent_chain = AgentExecutor.from_agent_and_tools(
    agent=agent,
    tools=tools,
    verbose=True,
    memory=memory,
    max_iterations=1
)

print('Running agent')
with get_openai_callback() as cb:
    result = agent_chain.run({"input":query})
    #result = agent_chain({"input":query})

return result, cb

This works fine:

Question and answer llm

def qa_llm(vector_store, queries):
vectorstore = u.load_vectorstore(vector_store)

if type(queries) == str:
    qa = RetrievalQA.from_chain_type(
        llm=OpenAI(temperature = 0),
        chain_type="stuff",
        retriever=vectorstore.as_retriever(),
        return_source_documents = True,
    )

    with get_openai_callback() as cb:
        result = qa({"query":queries})
    return result, cb

elif type(queries) == list:
    results = []
    with get_openai_callback() as cb:
        for query in queries:
            qa = RetrievalQA.from_chain_type(
                llm=OpenAI(temperature = 0),
                chain_type="stuff",
                retriever=vectorstore.as_retriever(),
                return_source_documents = True,
            )
            response = qa({"query":query})
            results.append(response)
    return results, cb

[pythonhacknewby]

Found the answer here. Sorry, originally missed it...
#5097