Optimizing Search Precision With Self-Querying Retrieval (SQR) and Langchain
This tutorial describes the merits of SQR and how to use it with the Langchain library to improve retrieval precision and relevance within a RAG framework.
Join the DZone community and get the full member experience.
Join For FreeWhat Is Self-Querying Retrieval?
Self-querying retrieval (SQR) is a method ensuring that LLMs are at the core of understanding the user's intent against a document collection. Key ideas include:
- Document representation: Word embedding provides a numerical vector for every document. This helps in fast comparison between the documents.
- User query: The user submits a natural language query expressing their need for information.
- LLM-driven retrieval: The query and the document representations are fed into the LLM, which then retrieves documents that maximize the user's intent.
- Refine and repeat: The user is now able to refine his query or ask follow-up questions to narrow the search based on the retrieved documents.
Why Self-Querying Retrieval?
Traditional retrieval systems usually require complex query languages or predefined search facets. However, self-querying retrieval would provide a much more natural and user-friendly approach. Here is why:
Natural Language Queries
SQR allows the user to frame questions in natural language, making the interface in a way self-explainable and more usable. In other words, this is comparable to visiting a librarian who asks you, "What are you looking for?" Users can formulate questions in their own parlance without studying complicated query languages or specifically search facets.
Advanced Retrieval Capabilities
SQR does not apply to basic keyword searching: it supports advanced search and retrieval facilities. One can conduct content and metadata searches by author, genre, and year. This provides two layers of filtering that can obtain pretty accurate results.
Imagine looking for a certain kind of legal document. SQR will find you documents that have the proper keywords and then further refine them by author, perhaps the judge, or even by year, like the case date.
Flexibility
SQR systems can adjust the search results about user intentions by modifying them according to further questions or inputs. This will allow the search results to be narrowed down to more appropriately fit the needs of the user.
Step-By-Step Implementation
The following are the steps to implement the naive RAG self-querying retrieval (SQR) using the Langchain library:
1. Import Libraries
We’ll import the required modules from the installed libraries to implement SQR:
import os
from langchain.schema import Document
from langchain_openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain_openai import OpenAI
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain.chains.query_constructor.base import AttributeInfoThese libraries and modules are essential for the subsequent steps in the process.
2. Set up the OpenAI API Key
The following code snippet retrieves your OpenAI API key from an environment variable. We’ll need a valid API key to interact with the OpenAI language model:
OPENAI_API_KEY = os.environ["OPENAI_API_KEY"] = "" # Add your OpenAI API key
if OPENAI_API_KEY == "":
raise ValueError("Please set the OPENAI_API_KEY environment variable")3. Example Data With Metadata
Here, we define a list of documents. Each document is a Document object containing page content and metadata like title, author, year, etc:
docs = [
Document(
page_content="A complex, layered narrative exploring themes of identity and belonging",
metadata={"title":"The Namesake", "author": "Jhumpa Lahiri", "year": 2003, "genre": "Fiction", "rating": 4.5, "language":"English", "country":"USA", "ISBN": "0618485228", "publisher": "Mariner Books"},
),
Document(
page_content="A luxurious, heartfelt novel with themes of love and loss set against a historical backdrop",
metadata={"title":"The Nightingale", "author": "Kristin Hannah", "year": 2015, "genre": "Historical Fiction", "rating": 4.8, "language":"English", "country":"France", "ISBN": "0312577222", "publisher": "St. Martin's Press"},
),
Document(
page_content="A full-bodied epic with rich characters and a sprawling plot",
metadata={"title":"War and Peace", "author": "Leo Tolstoy", "year": 1869, "genre": "Historical Fiction", "rating": 4.7, "language":"Russian", "country":"Russia", "ISBN": "0199232768", "publisher": "Oxford University Press"},
),
Document(
page_content="An elegant, balanced narrative with intricate character development and subtle themes",
metadata={"title":"Pride and Prejudice", "author": "Jane Austen", "year": 1813, "genre": "Romance", "rating": 4.6, "language":"English", "country":"UK", "ISBN": "0141439513", "publisher": "Penguin Classics"},
),
Document(
page_content="A highly regarded novel with deep themes and a nuanced exploration of human nature",
metadata={"title":"To Kill a Mockingbird", "author": "Harper Lee", "year": 1960, "genre": "Fiction", "rating": 4.9, "language":"English", "country":"USA", "ISBN": "0061120081", "publisher": "Harper Perennial Modern Classics"},
),
Document(
page_content="A crisp, engaging story with vibrant characters and a compelling plot",
metadata={"title":"The Alchemist", "author": "Paulo Coelho", "year": 1988, "genre": "Adventure", "rating": 4.4, "language":"Portuguese", "country":"Brazil", "ISBN": "0061122416", "publisher": "HarperOne"},
),
Document(
page_content="A rich, complex narrative set in a dystopian future with strong thematic elements",
metadata={"title":"1984", "author": "George Orwell", "year": 1949, "genre": "Dystopian", "rating": 4.7, "language":"English", "country":"UK", "ISBN": "0451524934", "publisher": "Signet Classics"},
),
Document(
page_content="An intense, gripping story with dark themes and intricate plot twists",
metadata={"title":"Gone Girl", "author": "Gillian Flynn", "year": 2012, "genre": "Thriller", "rating": 4.3, "language":"English", "country":"USA", "ISBN": "0307588378", "publisher": "Crown Publishing Group"},
),
Document(
page_content="An exotic, enchanting tale with rich descriptions and an intricate plot",
metadata={"title":"One Hundred Years of Solitude", "author": "Gabriel García Márquez", "year": 1967, "genre": "Magical Realism", "rating": 4.8, "language":"Spanish", "country":"Colombia", "ISBN": "0060883286", "publisher": "Harper Perennial Modern Classics"},
),
# ... (add more book documents as needed)
] 4. Define the Embedding Function
Create an instance of OpenAIEmbeddings which converts document text into numerical representations suitable for retrieval:
embeddings = OpenAIEmbeddings()5. Initialize Vector Store
Now that we have document embeddings, let’s create a vector store using the Chroma library. This store will hold the document embeddings for efficient retrieval:
vectorstore = Chroma.from_documents(docs, embeddings)6. Create LLM and Retriever
Here, we’ll establish communication with the OpenAI LLM and create the core component for SQR — the retriever object:
metadata_field_info = [
AttributeInfo(
name="title",
description="The title of the book",
type="string or list[string]",
),
AttributeInfo(
name="author",
description="The author of the book",
type="string or list[string]",
),
AttributeInfo(
name="year",
description="The year the book was published",
type="integer",
),
AttributeInfo(
name="genre",
description="The genre of the book",
type="string or list[string]",
),
AttributeInfo(
name="rating",
description="The rating of the book (1-5 scale)",
type="float",
),
AttributeInfo(
name="language",
description="The language the book is written in",
type="string",
),
AttributeInfo(
name="country",
description="The country the author is from",
type="string",
),
]
document_content_description = "Brief description of the book"
llm = OpenAI(temperature=0)
retriever = SelfQueryRetriever.from_llm(
llm,
vectorstore,
document_content_description,
metadata_field_info,
verbose=True
)7. Example Queries
Finally, we can explore various SQR queries you can use with the retriever object to interact with your document collection:
# Basic query
retriever.invoke("What are some highly rated historical fiction books")
# Query with filtering
retriever.invoke("I want a book with deep themes and a rating above 4.5")
# Query with composite filter
retriever.invoke("I want a book with complex characters and a gripping plot")
# Query with country filter
retriever.invoke("What books come from the USA?")
# Query with year range and theme filter
retriever.invoke("What's a book published after 2003 but before 2015 with deep themes and a high rating")
# Retrieval with limiting
retriever = SelfQueryRetriever.from_llm(
llm,
vectorstore,
document_content_description,
metadata_field_info,
enable_limit=True,
verbose=True
)
retriever.invoke("What are two books that have a rating above 4.8")
retriever.invoke("What are two books that come from USA or UK")This way we can effectively use SQR to perform detailed and specific book searches.
Conclusion
The self-query retriever significantly improves the precision and relevance of what is retrieved within a Retrieval-Augmented Generation framework. SQR supports dynamic and context-aware querying that personalizes retrieval in accordance with certain needs, hence improving the quality of the responses generated. In other words, with SQR, RAG models will be better placed to provide accurate and contextually relevant information, therefore improving overall performance and the user experience.
Opinions expressed by DZone contributors are their own.
Comments