Vector Retriever

Vector search can be used in Retrieval Augmented Generation (RAG) applications to find relevant documents based on their content.

In this lesson, you will update the LangChain agent to use a vector retriever that will allow you to search for movies based on plot.

Open the genai-integration-langchain/vector_retriever.py file.

python

vector_retriever.py

import os
from dotenv import load_dotenv
load_dotenv()

from langchain_core.documents import Document
from langchain.chat_models import init_chat_model
from langgraph.graph import START, StateGraph
from langchain_core.prompts import PromptTemplate
from typing_extensions import List, TypedDict

# Initialize the LLM
model = init_chat_model("gpt-4o", model_provider="openai")

# Create a prompt
template = """Use the following pieces of context to answer the question at the end.
If you don't know the answer, just say that you don't know, don't try to make up an answer.

{context}

Question: {question}

Answer:"""

prompt = PromptTemplate.from_template(template)

# Define state for application
class State(TypedDict):
    question: str
    context: List[Document]
    answer: str

# Connect to Neo4j
# graph = 

# Create the embedding model
# embedding_model = 

# Create Vector
# plot_vector =

# Define functions for each step in the application

# Retrieve context 
def retrieve(state: State):
    # Use the vector to find relevant documents
    context = [Document("Nothing to see here")]
    return {"context": context}

# Generate the answer based on the question and context
def generate(state: State):
    messages = prompt.invoke({"question": state["question"], "context": state["context"]})
    response = model.invoke(messages)
    return {"answer": response.content}

# Define application steps
workflow = StateGraph(State).add_sequence([retrieve, generate])
workflow.add_edge(START, "retrieve")
app = workflow.compile()

# Run the application
question = "What is the movie with the pig who wants to be a sheep dog?"
response = app.invoke({"question": question})
print("Answer:", response["answer"])

The code is similar to the agent you explored in the previous module.

To add the vector RAG retriever, you will to:

Connect to Neo4j using Neo4jGraph
Create an embedding model
Create a vector store using Neo4jVector
Update the retrieve function to use the vector store

Neo4j connection

Connect to Neo4j using Neo4jGraph:

python

from langchain_neo4j import Neo4jGraph

# Connect to Neo4j
graph = Neo4jGraph(
    url=os.getenv("NEO4J_URI"),
    username=os.getenv("NEO4J_USERNAME"), 
    password=os.getenv("NEO4J_PASSWORD"),
)

Embedding model

Create the embedding model using OpenAIEmbeddings:

python

from langchain_openai import OpenAIEmbeddings

# Create the embedding model
embedding_model = OpenAIEmbeddings(model="text-embedding-ada-002")

Vector store

Create the vector store using Neo4jVector:

python

from langchain_neo4j import Neo4jVector

# Create Vector
plot_vector = Neo4jVector.from_existing_index(
    embedding_model,
    graph=graph,
    index_name="moviePlots",
    embedding_node_property="plotEmbedding",
    text_node_property="plot",
)

Retriever

Finally, you need to update the retrieve function to use the vector store.

The retrieve function will need to:

Use the similarity_search method of the Neo4jVector class to find similar nodes based on the question.
Store the results as the context for the agent.

python

# Retrieve context 
def retrieve(state: State):
    # Use the vector to find relevant documents
    context = plot_vector.similarity_search(
        state["question"], 
        k=6
    )
    return {"context": context}

Click to see the complete code

python

import os
from dotenv import load_dotenv
load_dotenv()

from langchain_core.documents import Document
from langchain.chat_models import init_chat_model
from langgraph.graph import START, StateGraph
from langchain_core.prompts import PromptTemplate
from typing_extensions import List, TypedDict
from langchain_neo4j import Neo4jGraph
from langchain_neo4j import Neo4jVector
from langchain_openai import OpenAIEmbeddings

# Initialize the LLM
model = init_chat_model("gpt-4o", model_provider="openai")

# Create a prompt
template = """Use the following pieces of context to answer the question at the end.
If you don't know the answer, just say that you don't know, don't try to make up an answer.

{context}

Question: {question}

Answer:"""

prompt = PromptTemplate.from_template(template)

# Define state for application
class State(TypedDict):
    question: str
    context: List[Document]
    answer: str

# Connect to Neo4j
graph = Neo4jGraph(
    url=os.getenv("NEO4J_URI"),
    username=os.getenv("NEO4J_USERNAME"), 
    password=os.getenv("NEO4J_PASSWORD"),
)

# Create the embedding model
embedding_model = OpenAIEmbeddings(model="text-embedding-ada-002")

# Create Vector
plot_vector = Neo4jVector.from_existing_index(
    embedding_model,
    graph=graph,
    index_name="moviePlots",
    embedding_node_property="plotEmbedding",
    text_node_property="plot",
)

# Define functions for each step in the application

# Retrieve context 
def retrieve(state: State):
    # Use the vector to find relevant documents
    context = plot_vector.similarity_search(
        state["question"], 
        k=6
    )
    return {"context": context}

# Generate the answer based on the question and context
def generate(state: State):
    messages = prompt.invoke({"question": state["question"], "context": state["context"]})
    response = model.invoke(messages)
    return {"answer": response.content}

# Define application steps
workflow = StateGraph(State).add_sequence([retrieve, generate])
workflow.add_edge(START, "retrieve")
app = workflow.compile()

# Run the application
question = "What is the movie with the pig who wants to be a sheep dog?"
response = app.invoke({"question": question})
print("Answer:", response["answer"])

Run the application, review the results, and experiment with different questions:

What is the movie with the pig who wants to be a sheep dog?
What are some movies about fast cars?
What are 3 movies about aliens coming to earth?
What is a typical budget for a romance movie?

Context

The context is a list of Document objects returned by the similarity_search method. The Document content contains the movie plot and the node’s properties are returned as metadata.

The context is defined in the State class:

python

class State(TypedDict):
    question: str
    context: List[Document]
    answer: str

You can define the context as any data type you want, as long as the retrieve function returns it in the expected format.

It is useful in development to output the context to see what data was used to generate the answer.

Update the code to print the context after the answer is generated:

python

print("Context:", response["context"])

Try running the application again and see how the context relates to the question and generated answer.

When you are ready, continue to the next module.

Lesson Summary

In this lesson, you learned how to integrate a Neo4j vector store with a LangChain agent.

In the next lesson, you will extend the retriever to use data from a graph.

Using Neo4j with LangChain

Neo4j and LangChain

Vectors

Text to Cypher

Vector Retriever

Neo4j connection

Embedding model

Vector store

Retriever

Context

Lesson Summary

Chatbot