Integrating Sources into RAG Answers

The Challenge of Preserving State in LCEL

The chain you built in the previous lesson ends with StrOutputParser(). This is effective for getting a simple string answer, but it means all prior information, including the retrieved documents, is lost at the final step.

The goal is to restructure the chain so that its final output is a dictionary containing both the generated answer and the documents used as sources. The key is to design the pipeline to carry the documents from the retrieval step all the way to the end, alongside the main data flow that generates the answer.

The Solution: Combining Retrieval and Generation with .assign()

We can achieve this by creating a main data-flow that prepares the necessary inputs, and then using the .assign() method to append the result of a sub-chain as a new key. This is a common and powerful pattern in LCEL for building chains that have structured outputs.

Here's the logical flow we will implement:

1. A RunnableParallel step will be the entry point. It will take the user's question and execute two tasks in parallel:
   • Retrieve the relevant documents.
   • Pass the original question through unchanged.
     The output will be a dictionary: {'documents': [...], 'question': '...'}.
2. We will then .assign() a new key, answer, to this dictionary.
3. The value for answer will be generated by a sub-chain that takes the dictionary from step 1, formats the documents into a context string, populates the prompt, and calls the LLM.

This pattern keeps the raw documents available in the main dictionary while the sub-chain works on a temporary, formatted version of the context to generate the answer.

A Stack Overflow discussion provides an excellent and concise example of this technique. While the accepted answer shows a very verbose, step-by-step assignment, a comment below it offers a more elegant and idiomatic approach that we will adapt.

Implementing the Chain with Sources

Let's modify our src/lcel_quickstart/chain.py to implement this. We will create a new function, get_rag_with_sources_chain, to distinguish it from the previous version.

First, ensure your imports are up to date. We will use itemgetter from the operator module, which is a convenient utility for pulling specific keys from a dictionary within an LCEL chain.

# src/lcel_quickstart/chain.py (updated imports)

import os
from operator import itemgetter

from langchain_community.vectorstores import FAISS
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda, RunnableParallel, RunnablePassthrough
from langchain_core.vectorstores import VectorStoreRetriever
from langchain_openai import ChatOpenAI, OpenAIEmbeddings

from .data_pipeline import load_vector_store, VECTOR_STORE_PATH




# ... (get_retriever and format_docs functions remain the same) ...

Now, let's define the new chain builder function. The structure is slightly different, so pay close attention to how we compose the final chain.

# src/lcel_quickstart/chain.py (new chain function)

def get_retriever() -> VectorStoreRetriever:
    """Gets the FAISS vector store retriever."""
    embedding_model = OpenAIEmbeddings(model="text-embedding-3-small")
    if not os.path.exists(VECTOR_STORE_PATH):
        raise FileNotFoundError(
            "FAISS index not found. Please run data_pipeline.py to create it."
        )
    vector_store = load_vector_store(embedding_model)
    return vector_store.as_retriever(search_kwargs={"k": 3})

def format_docs(docs: list) -> str:
    """Converts a list of documents into a single string."""
    return "\n\n".join(doc.page_content for doc in docs)




# Previous chain for reference (no changes needed)
def get_rag_chain():



    # ...
    return chain




# New chain implementation
def get_rag_with_sources_chain():
    """
    Builds and returns a RAG chain that includes source documents in the output.
    """
    retriever = get_retriever()
    
    template = """Answer the question based only on the following context:
{context}

Question: {question}
"""
    prompt = ChatPromptTemplate.from_template(template)

    model = ChatOpenAI(model="gpt-3.5-turbo")




    # This sub-chain is responsible for generating the answer
    answer_generator = (
        {
            "context": itemgetter("documents") | RunnableLambda(format_docs),
            "question": itemgetter("question"),
        }
        | prompt
        | model
        | StrOutputParser()
    )




    # This is the final chain that orchestrates the process
    chain = (
        RunnableParallel(
            {
                "documents": retriever,
                "question": RunnablePassthrough(),
            }
        ).assign(answer=answer_generator)
    )
    
    return chain

Let's break down this new structure:

1. answer_generator: This is a self-contained chain that assumes its input is a dictionary with documents and question keys.

   • itemgetter("documents") | RunnableLambda(format_docs): This part pulls the documents list from the input dictionary and passes it to our format_docs function to create the context string.
   • itemgetter("question"): This pulls the question from the input dictionary.
   • The rest is familiar: the resulting dictionary is passed to the prompt, model, and parser to generate the final answer string.

2. chain: This is the main, user-facing chain.

   • RunnableParallel: It takes the initial string query. retriever gets the documents, and RunnablePassthrough carries the query forward. The output is { "documents": [...], "question": "..." }.
   • .assign(answer=answer_generator): This is the key step. It takes the output from RunnableParallel and passes it as input to the answer_generator chain. The result of answer_generator (the answer string) is then added to the dictionary under the key answer.

The final output of the entire chain is a dictionary that looks like this: {'documents': [...], 'question': '...', 'answer': '...'}.

Testing the Modified Chain

Finally, let's update the if __name__ == '__main__': block to test our new chain and inspect its output.

# src/lcel_quickstart/chain.py (updated main block)

if __name__ == '__main__':
    rag_with_sources_chain = get_rag_with_sources_chain()
    query = "What is LCEL?"
    
    print("--- RAG Chain with Sources Demonstration ---")
    print(f"Query: '{query}'\n")
    
    result = rag_with_sources_chain.invoke(query)
    
    print("Answer:")
    print(result["answer"])
    print("\n--- Sources ---")
    for doc in result["documents"]:
        print(f"Source: {doc.metadata.get('source', 'N/A')}")



        # print(f"Content: {doc.page_content[:150]}...") # Optional: print a snippet
    print("---------------")

Now, run the script from your terminal:

uv run src/lcel_quickstart/chain.py

You will see the generated answer, followed by a "Sources" section that lists the metadata of the documents retrieved from your vector store. Since we loaded a single blog post, the source will be the same for all chunks, but this demonstrates the principle that will be vital when working with multiple documents.

Conclusion

In this lesson, you have significantly enhanced our RAG pipeline by modifying it to return source documents alongside the generated answer. This is not just a cosmetic change; it's a fundamental step toward building a transparent and production-worthy AI application.

The key takeaways are:

• Structured Outputs: LCEL is not limited to simple string inputs and outputs. You can design chains that operate on and produce complex dictionary objects.
• The .assign() Method: This is a powerful tool for enriching a data dictionary as it flows through a chain. It allows you to add new computed values while preserving the original data.
• itemgetter: A useful utility for cleanly extracting data from a dictionary at a specific point in a chain.

You now have a complete, source-citing RAG pipeline. The next major step in building our chatbot is to give it memory. In the next module, we will dive into managing conversational context, allowing our application to handle follow-up questions and maintain a coherent dialogue.