Overview
In this lesson, you will enhance your agent by adding a search lesson tool using a vector + cypher retriever.
The agent will decide which tool is best for each question:
Schema Tool - to understand the database structure:
-
"What entities exist in the graph?"
-
"How are technologies related to concepts?"
Search Lesson Tool - for finding content within the lessons:
-
"What are the benefits of using GraphRAG?"
-
"THow are Knowledge Graphs associated with other technologies?"
Search lessons tool
You will modify the agent.py code to:
-
Create a
VectorCypherRetrieverretriever that uses theChunkvector index. -
Define a new
toolfunction that uses this retriever to search for lesson content. -
Add the new tool to the agent’s list of available tools.
Continue with the lesson to create the text to Cypher retriever.
Update the agent
Open workshop-genai\agent.py and make the following changes:
-
Add an embedding model for the retriever to convert the user query into a vector.
pythonembedderfrom neo4j_graphrag.embeddings.openai import OpenAIEmbeddings # Create embedder embedder = OpenAIEmbeddings(model="text-embedding-ada-002") -
Create a
retrieval_querythat the retriever will use to add additional context to the vector search results.pythonretrieval_query# Define retrieval query retrieval_query = """ MATCH (node)-[:FROM_DOCUMENT]->(d)-[:PDF_OF]->(lesson) RETURN node.text as text, score, lesson.url as lesson_url, collect { MATCH (node)<-[:FROM_CHUNK]-(entity)-[r]->(other)-[:FROM_CHUNK]->() WITH toStringList([ labels(entity)[2], entity.name, entity.type, entity.description, type(r), labels(other)[2], other.name, other.type, other.description ]) as values RETURN reduce(acc = "", item in values | acc || coalesce(item || ' ', '')) } as associated_entities """ -
Create a
VectorCypherRetrieverusing thechunkEmbeddingindex, Neo4jdriver, andembedder.pythonretrieverfrom neo4j_graphrag.retrievers import VectorCypherRetriever vector_retriever = VectorCypherRetriever( driver, neo4j_database=os.getenv("NEO4J_DATABASE"), index_name="chunkEmbedding", embedder=embedder, retrieval_query=retrieval_query, ) -
Define a tool function to search for lesson content using the retriever.
pythonSearch-lesson-content tool@tool("Search-lesson-content") def search_lessons(query: str): """Search for lesson content related to the query.""" # Use the vector to find relevant chunks result = vector_retriever.search( query_text=query, top_k=5 ) context = [item.content for item in result.items] return contextTool description
The tool name
search-lesson-contentand docstringSearch for lesson content related to the query.help the agent decide when to use this tool. -
Update the
toolslist to include the new lesson search tool.pythontoolstools = [get_schema, search_lessons] -
Modify the
queryvariable to test the new lesson search tool.pythonqueryquery = "What are the benefits of using GraphRAG"
Reveal the complete code
import os
from dotenv import load_dotenv
load_dotenv()
from neo4j import GraphDatabase
from neo4j_graphrag.embeddings.openai import OpenAIEmbeddings
from neo4j_graphrag.retrievers import VectorCypherRetriever
from langchain.chat_models import init_chat_model
from langchain.agents import create_agent
from langchain_core.tools import tool
# Initialize the chat model
model = init_chat_model("gpt-4o", model_provider="openai")
# Connect to Neo4j database
driver = GraphDatabase.driver(
os.getenv("NEO4J_URI"),
auth=(
os.getenv("NEO4J_USERNAME"),
os.getenv("NEO4J_PASSWORD")
)
)
# Create embedder
embedder = OpenAIEmbeddings(model="text-embedding-ada-002")
# Define retrieval query
retrieval_query = """
MATCH (node)-[:FROM_DOCUMENT]->(d)-[:PDF_OF]->(lesson)
RETURN
node.text as text, score,
lesson.url as lesson_url,
collect {
MATCH (node)<-[:FROM_CHUNK]-(entity)-[r]->(other)-[:FROM_CHUNK]->()
WITH toStringList([
labels(entity)[2],
entity.name,
entity.type,
entity.description,
type(r),
labels(other)[2],
other.name,
other.type,
other.description
]) as values
RETURN reduce(acc = "", item in values | acc || coalesce(item || ' ', ''))
} as associated_entities
"""
# Create retriever
vector_retriever = VectorCypherRetriever(
driver,
neo4j_database=os.getenv("NEO4J_DATABASE"),
index_name="chunkEmbedding",
embedder=embedder,
retrieval_query=retrieval_query,
)
# Define functions for each tool in the agent
@tool("Get-graph-database-schema")
def get_schema():
"""Get the schema of the graph database."""
results, summary, keys = driver.execute_query(
"CALL db.schema.visualization()",
database_=os.getenv("NEO4J_DATABASE")
)
return results
# Define a tool to retrieve lesson content
@tool("Search-lesson-content")
def search_lessons(query: str):
"""Search for lesson content related to the query."""
# Use the vector to find relevant chunks
result = vector_retriever.search(
query_text=query,
top_k=5
)
context = [item.content for item in result.items]
return context
# Define a list of tools for the agent
tools = [get_schema, search_lessons]
# Create the agent with the model and tools
agent = create_agent(
model,
tools
)
# Run the application
query = "What are the benefits of using GraphRAG"
for step in agent.stream(
{
"messages": [{"role": "user", "content": query}]
},
stream_mode="values",
):
step["messages"][-1].pretty_print()Run the agent. The agent should decide to use the new tool based on the query.
Experiment
Experiment with agent, modify the query to ask different questions, for example:
-
"How are Knowledge Graphs associated with other technologies?" -
"Summarize what concepts are associated with Knowledge Graphs?" -
"How would you minimize hallucinations in LLMs?"
Asking questions related to the graph schema should still use the schema tool, for example:
-
"What entities exist in the graph?"
Lesson Summary
In this lesson, you modified the agent to use an additional tool for searching lesson content using a vector + Cypher retriever.
In the next lesson, you will add another tool that can query the database directly using a text to Cypher retriever.