RAG with LangChain, pgvector and OpenAI
Taking RAG past a mock lookup: building and storing real embeddings in Postgres, retrieving on similarity, and handing the context to GPT.
- Published
- Reading
- 6 min
- Tags
- rag
- Originally
- Codemancers
In a previous post I looked at how retrieval-augmented generation extends what a GPT model can answer. That example used a mock lookup. This one takes it a step further: creating and storing real embeddings from a document set using LangChain and pgvector, then feeding them to OpenAI's GPT for contextually relevant responses.
The role of embeddings
Embeddings represent text in a dense numerical format, which lets a model capture the semantic meaning of that text. Storing them in a database like pgvector means you can efficiently retrieve the most relevant pieces of information and hand them to a model to answer against.
Installation and setup
Install the libraries
pip install langchain pgvector psycopg2-binary- langchain: a framework for working with LLMs and building AI applications.
- pgvector: a Postgres extension supporting vector embedding storage and similarity search.
- psycopg2-binary: a PostgreSQL adapter for Python.
Set up pgvector in PostgreSQL
Your PostgreSQL instance needs the extension installed:
psql -d your_database -c "CREATE EXTENSION IF NOT EXISTS vector;"Set up OpenAI API access
Sign up or log in at OpenAI, then put the key in an environment variable:
OPENAI_API_KEY="your_openai_api_key"With that in place you're ready to load, embed and query documents.
Loading the document data
from langchain.document_loaders import TextLoader
# Load the document from a local file on the device
with open("path_to_your_file.txt", "r", encoding="utf-8") as temp_file:
loader = TextLoader(temp_file.name, encoding="utf-8")
documents = loader.load()
# Optional: Split large documents for better embedding granularity
from langchain.text_splitter import RecursiveCharacterTextSplitter
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=50)
texts = text_splitter.split_documents(documents)Splitting documents into smaller chunks means each chunk is embedded individually, which improves retrieval accuracy.
Creating and storing embeddings
from langchain.embeddings.openai import OpenAIEmbeddings
from pgvector.vectorstore import PGVector
# Set up the embeddings model
embeddings_model = OpenAIEmbeddings(openai_api_key="your_openai_key")
# Set up the Pgvector database
collection_name = "my_collection"
connection_string = "postgres://user:password@localhost:5432/mydb"
# Save embeddings to the Pgvector database
pgvector = PGVector(
collection_name=collection_name,
connection_string=connection_string,
embedding_function=embeddings_model
)
pgvector.from_documents(texts)Querying the database
With the embeddings stored, retrieve the most relevant documents for a query:
# Retrieve documents based on similarity to the query
query = "Tell me about the Eiffel Tower."
retriever = pgvector.as_retriever(search_type="similarity", search_kwargs={"k": 5})
# Get the top 5 most relevant documents
relevant_docs = retriever.get_relevant_documents(query)
def format_docs(docs):
return "\n\n".join([d.page_content for d in docs])
context = format_docs(relevant_docs)Generating a response
With the context retrieved, hand it to the model:
from langchain.prompts import ChatPromptTemplate
from langchain.chat_models import ChatOpenAI
# Prepare the template for the prompt
template = """You are an AI assistant. Here's the context for the question:\n{context}\nNow answer the question below:\n\nQuestion: {question}\nAI: """
prompt = ChatPromptTemplate.from_template(template)
# Set up the GPT-4 model
model = ChatOpenAI(api_key="your_openai_key")
# Prepare the full chain
chain = prompt | model
# Generate the response
response = chain.invoke({
"context": context,
"question": query
})
print(response)Conclusion
That's LangChain and pgvector wired together to create and store document embeddings, query them for relevant context, and use that context to generate a response with GPT. It's a small amount of code for a system that answers from your own documents rather than from whatever the model happened to memorise.