RetrievalQA chain return source_documents when using it as a Tool for an Agent

[damiam]

Hello everyone,
I am building a simple chatBot using Agent which uses two custom tools to answer the users questions.

Tools and Agent

qa = RetrievalQA.from_chain_type(llm=llm, chain_type='stuff', retriever=retriever, return_source_documents=True, verbose=True)
llm_math = LLMMathChain.from_llm(OpenAI(temperature=0), verbose=True)
tools = [
    Tool(
        name="Search summaries",
        func=qa.run,
        description="useful for when you need to answer questions. Input should be a fully formed question."
    ),
    Tool(
        name="Calculator",
        func=llm_math.run,
        description="useful for doing calculations"
    )
]

And the agent is defined as

agent = initialize_agent(
    tools,
    llm,
    agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
    verbose=True,
    memory=memory,
    return_intermediate_steps=True
)

Problem

When running the agent with

answer = agent({"input": query}, return_only_outputs=False)

if the first tool is selected ("Search summaries") I get the following error

ValueError: `run` not supported when there is not exactly one output key. Got ['result', 'source_documents'].

And when trying other solutions like calling directly the chain in the tool

tools = [
    Tool(
        name="Search summaries",
        func=qa,
        description="useful for when you need to answer questions. Input should be a fully formed question."
    ),

I get the a NotImplementedError.

Expected behaviour

When running the chain singularly (not as a Tool in an Agent) the answer works as expected. Note that when I set return_source_documents=False, the Agent works perfectly, but without returning the metadata of the source_documents.

Is it possible to have the source_documents also while using an agent?

[damiam]

EDIT:

Tool(
    name="Search summaries",
    func=lambda query: chain({"question": query}),
    description="useful for when you need to answer questions. Input should be a fully formed question."
),

This new func=lambda query: chain({"question": query}) worked for me, but it looks kinda clunky.

the chain is now

chain = RetrievalQA.from_chain_type(llm=llm, chain_type='stuff', retriever=retriever, return_source_documents=True, verbose=True, input_key="question")

[cnndabbler]

#4187 worked for me

[Sarthak1306]

Could you tell me how to print the Source Document for a particular query?

[Mansko]

Does anyone have a solution for async/coroutine implementation?

[swarnim-sawane]

I would like to know how can we cite sources while using an agent.