Retrieval-Augmented Generation with LangChain
Query analysis is a query-optimization technique that refines queries to improve retrieval search. Install the following modules to begin:
pip install -qU langchain langchain-community langchain-openai
langchain-chroma wikipedia
Open the starter project. The first cell sets a session key for SportsBuddy. In the second cell, you load a list of sports articles from Wikipedia. Then, there are two helper functions - the first to round a number to the nearest thousand, and the other to print the title and number of approximate words in a document. Under # TODO: Add 'words' metadata, add the following to print these articles’ titles and their words metadata:
for doc in docs:
doc.metadata["words"] = round_to_nearest_thousand(len(
doc.page_content.split(" ")))
print_summary(doc)
print()
The words metadata approximates the number of words in the article to the nearest 1000. Execute the first two cells, and you’ll see the following output for the second cell:
Title: 2022 Ballon d'Or - Wikipedia
Approximate Word Count: 4000
Title: 2023 Ballon d'Or - Wikipedia
Approximate Word Count: 2000
Title: 2022–23 NBA season - Wikipedia
Approximate Word Count: 10000
Title: 2021–22 NBA season - Wikipedia
Approximate Word Count: 14000
Title: 2022–23 Premier League - Wikipedia
Approximate Word Count: 8000
Title: 2021–22 Premier League - Wikipedia
Approximate Word Count: 7000
Title: 2021–22 UEFA Champions League - Wikipedia
Approximate Word Count: 6000
Title: 2022–23 UEFA Champions League - Wikipedia
Approximate Word Count: 4000
Title: 2023 Cricket World Cup - Wikipedia
Approximate Word Count: 6000
In the third cell, you partition the loaded documents into sizable chunks and embed them in a Chroma database. Create a new cell to try out a query on the vector store:
search_results = database.similarity_search("Who won the 2022 ballon d'or?")
print_summary(search_results[0])
From previous lessons, you learned that similarity search by default returns a list of documents, with the most relevant at the top. So in this case, you printed only the first document from the results.
Execute the cell and observe the results:
Title: 2022 Ballon d'Or - Wikipedia
Approximate Word Count: 4000
From this result, you can tell it got the right document as the first item. Now, send a query that references information from the metadata, like the word count.
Add a new cell and execute the following:
search_results = database.similarity_search("Suggest a sports article with
approximately 14000 words")
print_summary(search_results[0])
You get:
Title: 2023 Cricket World Cup - Wikipedia
Approximate Word Count: 6000
You can tell that it ignored the “14000 words” part of your query. Otherwise, it should have returned the football article titled “2021–22 NBA season - Wikipedia” because it has the number of words requested.
You can use query analysis to fix this by generating a query that includes the words metadata as a filter.
With LangChain, you can achieve this by creating a structured output based on your initial prompt. Extending the BaseModel, you can add new fields or filters to your search query.
Add the following code in a new cell to define the structure for the generated query:
from typing import Optional
from pydantic import BaseModel, Field
class SportsSearch(BaseModel):
"""Search over a database of sports articles."""
query: str = Field(
...,
description="Similarity search query applied to sports articles.",
)
words: Optional[int] = Field(None, description="Number of words in article")
SportsSearch will contain your original search in its query property and the article’s word count in an optional words property.
Next, you’ll use the OpenAI LLM to generate the new prompt. Add the following to a new cell to create the prompt chain with LangChain:
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI
system = """You are an expert at converting user questions into
database queries. \
You have access to a database of sports articles. \
Given a question, return a list of database queries optimized to
retrieve the most relevant results.
If there are acronyms or words you are not familiar with, do
not try to rephrase them."""
prompt = ChatPromptTemplate.from_messages(
[
("system", system),
("human", "{question}"),
]
)
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
structured_llm = llm.with_structured_output(SportsSearch)
query_analyzer = {"question": RunnablePassthrough()} | prompt | structured_llm
Two important things to note here are the system prompt and the LLM’s temperature. The system prompt is carefully conditioned to return an improved query suitable for a database search. You order it not to try anything fancy if it isn’t sure about the task.
By setting the temperature to 0, you’ve told your LLM not to attempt to be creative with its responses. It must stick to the given query’s context strictly.
Next, create the filter function for the new search query:
from typing import List
from langchain_core.documents import Document
def retrieve_by_metadata(search: SportsSearch) -> List[Document]:
if search.words not None:
_filter = {"words": {"$eq": search.words}}
else:
_filter = None
return database.similarity_search(search.query, filter=_filter)
This function adds a filter to the database during the search. If the generated query includes the words property, it’s included in the filter. Otherwise, it’s not. The way this filter is constructed depends on the fact that you’re using Chroma. For a different database, you’d have to comply with its API.
Try out the query regeneration using your first query:
query_analyzer.invoke("Who won the 2022 ballon d'or?")
The response is:
SportsSearch(query="2022 Ballon d'Or winner", words=None)
The query analyzer didn’t find anything in your query that fit the “words” filter, hence words is None. Is this new query as good as the first? Try it directly on the database to find out:
search_results = database.similarity_search("2022 Ballon d'Or winner")
print_summary(search_results[0])
The results are the same as before:
Title: 2022 Ballon d'Or - Wikipedia
Approximate Word Count: 4000
Great. Now, try the query regeneration on your second query:
query_analyzer.invoke("Suggest a sports article with approximately
14000 words")
Execute it.
SportsSearch(query='sports article', words=14000)
This time, the regenerated query captured some extra context from your query. Hence, you see words=4000. Good. Now, run the entire pipeline, which includes the filter on the regenerated query:
retrieval_chain = query_analyzer | retrieve_by_metadata
search_results = retrieval_chain.invoke("Suggest a sports article with
approximately 14000 words")
print_summary(search_results[0])
Check the results:
Title: 2021–22 NBA season - Wikipedia
Approximate Word Count: 14000
Excellent. To show that it takes the full query into account and not just the word count filter, update your query to search for a football article with about 6,000 words. From the second cell’s output, you can tell two articles have the same approximate length. But one is for football and the other is for cricket:
search_results = retrieval_chain.invoke("Suggest a football article
with approximately 6000 words")
print_summary(search_results[0])
It shows the 2021-22 UEFA Champions League document, which is a football article and has approximately 6,000 words.
As you can see, query analysis can be a great way to boost the performance of your RAG. That’s all for this demo, continue on to learn more about RAG optimizations.