Retrieval-Augmented Generation with LangChain
In this demo, we’ll implement a hybrid search using sparse vector embedding algorithms from LangChain. Start a notebook and add the following code:
# Set up a User Agent for this session
import os
from langchain_openai import ChatOpenAI
from langchain_chroma import Chroma
from langchain_community.document_loaders import WikipediaLoader
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
os.environ['USER_AGENT'] = 'sports-buddy-advanced'
llm = ChatOpenAI(model="gpt-4o-mini")
loader = WikipediaLoader("2024_Summer_Olympics",)
docs = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000,
chunk_overlap=0)
splits = text_splitter.split_documents(docs)
database = Chroma.from_documents(documents=splits,
embedding=OpenAIEmbeddings())
retriever = database.as_retriever()
There’s nothing new here. It’s the same initial setup for a basic RAG.
Create a cell. In this cell, create a retriever based on the Best Match 25 algorithm. This will allow you to do a sparse-vector search:
from langchain.retrievers import BM25Retriever, EnsembleRetriever
keyword_retriever = BM25Retriever.from_documents(splits)
keyword_retriever.k = 3
ensemble_retriever = EnsembleRetriever(retrievers=[retriever,
keyword_retriever], weights=[0.3, 0.7])
In the next cell, create simple chains for both the dense (semantic) and sparse (keyword) retrievers:
from langchain.chains import RetrievalQA
dense_chain = RetrievalQA.from_chain_type(
llm=llm, chain_type="stuff", retriever=retriever
)
sparse_chain = RetrievalQA.from_chain_type(
llm=llm, chain_type="stuff", retriever=ensemble_retriever
)
The dense_chain is the same familiar chain you’ve been using with the Chroma database retriever. The sparse_chain uses the sparse retriever for keyword-based search.
Run the dense_chain in the next cell:
normal_response = normal_chain.invoke("What happened at the opening
ceremony of the 2024 Summer Olympics")
print(normal_response['result'])
Observe the output:
The opening ceremony of the 2024 Summer Olympics was held outside of a stadium for the first time in modern Olympic history.
Athletes were paraded by boat along the Seine River in Paris.
Finally, run the sparse_chain with:
sparse_response = sparse_chain.invoke("What happened at the
opening ceremony of the 2024 Summer Olympics")
print(hybrid_response['result'])
And note its output:
The opening ceremony of the 2024 Summer Olympics took place outside of a stadium for the first time in modern Olympic history, with athletes being paraded by boat along the Seine River in Paris. This unique setting was part of the ceremony, making it a significant and memorable event in
Olympic history.
Notice how the keywords in the query contribute to a more elaborate response in the hybrid search.
Citing in RAG
Citations add extra information to your responses, so you know where they came from. Open a new notebook to learn how to add citations to SportsBuddy. In the notebook, start with the following code:
from langchain_community.retrievers import WikipediaRetriever
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
import os
llm = ChatOpenAI(model="gpt-4o-mini")
system_prompt = (
"You're a helpful AI assistant. Given a user question "
"and some Wikipedia article snippets, answer the user "
"question. If none of the articles answer the question, "
"just say you don't know."
"\n\nHere are the Wikipedia articles: "
"{context}"
)
retriever = WikipediaRetriever(top_k_results=6, doc_content_chars_max=2000)
prompt = ChatPromptTemplate.from_messages(
[
("system", system_prompt),
("human", "{input}"),
]
)
This prompt instructs the WikipediaRetriever to fetch relevant articles based on the given context. Heads up: It can get things wildly wrong. Because it’s limited to Wikipedia articles, it will fetch articles it thinks best answer the semantic understanding it has for your query. In the next cell, create a chain:
from typing import List
from langchain_core.documents import Document
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
def format_docs(docs: List[Document]):
return "\n\n".join(doc.page_content for doc in docs)
rag_chain = (
RunnablePassthrough.assign(context=(lambda x: format_docs(x["context"])))
| prompt
| llm
| StrOutputParser()
)
retrieve_docs = (lambda x: x["input"]) | retriever
chain = RunnablePassthrough.assign(context=retrieve_docs).assign(
answer=rag_chain
)
Execute a simple query and examine the structure of the response object:
result = chain.invoke({"input": "How did the USA fair at the 2024
Summer Olympics"})
print(result.keys())
dict_keys(['input', 'context', 'answer'])
The response contains input (the query), context (the reference document), and answer. This information is accessible due to OpenAI’s tool-calling support. We can format this neatly into a citation model.
Let’s use the CitedAnswer model:
from typing import List
from langchain_core.pydantic_v1 import BaseModel, Field
class CitedAnswer(BaseModel):
"""Answer the user question based only on the given sources, and cite
the sources used."""
answer: str = Field(
...,
description="The answer to the user question, which is based only on
the given sources.",
)
citations: List[int] = Field(
...,
description="The integer IDs of the SPECIFIC sources which justify
the answer.",
)
To use the citation model, search with the following:
structured_llm = llm.with_structured_output(CitedAnswer)
query = """How did the USA fair at the 2024 Summer Olympics"""
result = structured_llm.invoke(query)
result
The model assigns the values to answer and citations by interpreting the description. The response is wrapped in a CitedAnswer class.
You could modify the citation to reference source URLs instead of integer IDs:
citations: List[str] = Field(
...,
description="The string URLs of the SPECIFIC sources which justify
the answer.",
)
However, take note that because the documents aren’t retrieved verbatim, these URLs might be generated and lead to 404 errors. If you instead want to cite a portion of the retrieved document, consider using a model like this:
class Citation(BaseModel):
source_id: int = Field(
...,
description="The integer ID of a SPECIFIC source which
justifies the answer.",
)
quote: str = Field(
...,
description="The VERBATIM quote from the specified source that
justifies the answer.",
)
class QuotedAnswer(BaseModel):
"""Answer the user question based only on the given sources, and
cite the sources used."""
answer: str = Field(
...,
description="The answer to the user question, which is based
only on the given sources.",
)
citations: List[Citation] = Field(
..., description="Citations from the given sources that
justify the answer."
)
You can use it similarly:
rag_chain = (
RunnablePassthrough.assign(context=(lambda x:
format_docs_with_id(x["context"])))
| prompt
| llm.with_structured_output(QuotedAnswer)
)
retrieve_docs = (lambda x: x["input"]) | retriever
chain = RunnablePassthrough.assign(context=retrieve_docs).assign(
answer=rag_chain
)
chain.invoke({"input": "How did the USA fair at the 2024 Summer
Olympics"})
And that’s it! You’ve now added citations to your RAG’s responses.