How to include document retrieval results in LangChain agent response

Other Questions langchain
1

I’m building a LangChain agent that uses document retrieval but I need help getting more than just the final answer. Right now my agent only returns the answer string, but I want to also see which documents were retrieved during the search process.

# Setting up my agent tools
my_tools = [
    StructuredTool.from_function(custom_function),
    Tool(
        name="DocumentSearch",
        func=(lambda query: perform_search(vector_db, query, org_name, lang, metadata=metadata)),
        description="use this tool when you need to find information from documents",
        return_direct=True,
        verbose=True
    )
]

# Creating the agent
my_agent = initialize_agent(
    my_tools,
    llm_model,
    agent=AgentType.OPENAI_FUNCTIONS,
    verbose=True,
    system_message=prompt_template,
    max_iterations=3,
    agent_kwargs={
        "system_message": prompt_template
    },
    memory=conversation_memory
)

My search function looks like this:

def perform_search(vector_db: Pinecone, query: str, org_name: str, lang: str = "english", metadata: str = ""):
    search_result = vector_db(inputs={"query": query, "text": query, "organization": org_name, "lang": lang})
    retrieved_info = search_result.get("documents", "")
    return search_result["response"]

Is there a way to modify this so the agent returns both the answer and the retrieved document information? I want to see what sources were used to generate the response.

2

Hit this exact problem last year building a knowledge base system. Your tool’s only returning the response string - you’re missing the full context.

Here’s what fixed it for me - change your perform_search function to return a structured object:

def perform_search(vector_db: Pinecone, query: str, org_name: str, lang: str = "english", metadata: str = ""):
    search_result = vector_db(inputs={"query": query, "text": query, "organization": org_name, "lang": lang})
    retrieved_docs = search_result.get("documents", [])
    
    # Return both answer and sources
    return {
        "answer": search_result["response"],
        "sources": retrieved_docs,
        "query": query
    }

Then update your tool definition to handle the structured response:

Tool(
    name="DocumentSearch",
    func=(lambda query: json.dumps(perform_search(vector_db, query, org_name, lang, metadata=metadata))),
    description="use this tool when you need to find information from documents. Returns JSON with answer and source documents",
    return_direct=False,  # Let agent process the full response
    verbose=True
)

The key is setting return_direct=False so the agent sees the full JSON response and can decide how to present it. I’ve found the agent will naturally mention sources when they’re available in the tool output.

3

To get both documents and answers from your LangChain agent, change your search function’s return structure. Rather than only returning search_result["response"], modify it to return a dictionary that includes both the response and the retrieved documents. Additionally, ensure that your tool can manage this new return structure. Consider using LangChain’s create_retriever_tool, which is specifically designed for these scenarios. Alternatively, you could employ class variables to retain document information after execution, allowing for easier access later.

4

Your perform_search function is throwing away the retrieved documents and only returning the response. You need to return both. I ran into the same thing last month - fixed it by restructuring the function to return a formatted string with both the answer and source metadata.

def perform_search(vector_db: Pinecone, query: str, org_name: str, lang: str = "english", metadata: str = ""):
    search_result = vector_db(inputs={"query": query, "text": query, "organization": org_name, "lang": lang})
    retrieved_docs = search_result.get("documents", "")
    response = search_result["response"]
    
    # Format the output to include both answer and sources
    formatted_output = f"Answer: {response}\n\nSources Retrieved:\n{retrieved_docs}"
    return formatted_output

Now your agent will return everything - the answer plus which documents were used. The Tool’s func parameter can return any string format you want, so you’ve got full control over structuring the output to include source attribution with the actual answer.

5

Had this exact problem building a research assistant. The issue isn’t just returning structured data - your tool setup needs to work with how the agent makes decisions. Keep return_direct=True but change your function to return one consolidated string with source references built right into the answer. Just append document IDs or titles directly to the response text. The agent gets everything at once without parsing JSON. If you’re using OpenAI Functions, make sure your system message tells it to cite sources when available. Once I embedded document references in the tool output, the agent started mentioning them naturally - but only after I explicitly told it to in the prompt. Your vector_db response probably already has metadata you can use for attribution without rebuilding your whole pipeline.

6

try changing your perform_search function to return a dict instead of just the response string. something like return {"answer": search_result["response"], "sources": retrieved_info} - then your agent will pass both parts through. you might need to adjust the tool description tho.

7

Been dealing with similar retrieval workflows for years - manual JSON formatting gets messy fast when you scale.

I automate the whole pipeline with Latenode instead of wrestling with LangChain’s tool structures. Build a workflow that handles document search, formats responses with sources, and logs everything for debugging.

Set up Pinecone search as one node, pipe results through a formatter that structures your answer and sources however you want. Add conditional logic for different response formats or route to different LLMs based on query type.

Switched our document Q&A system over last month - way cleaner than managing tool definitions and return structures. Built-in monitoring shows exactly what documents get retrieved for each query.

Scales better when you want to add source ranking or multi-step retrieval later.