Search lesson content

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:

  1. Create a VectorCypherRetriever retriever that uses the Chunk vector index.

  2. Define a new tool function that uses this retriever to search for lesson content.

  3. 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:

  1. Add an embedding model for the retriever to convert the user query into a vector.

    python
    embedder
    from neo4j_graphrag.embeddings.openai import OpenAIEmbeddings
    
    # Create embedder
    embedder = OpenAIEmbeddings(model="text-embedding-ada-002")
  2. Create a retrieval_query that the retriever will use to add additional context to the vector search results.

    python
    retrieval_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
    """
  3. Create a VectorCypherRetriever using the chunkEmbedding index, Neo4j driver, and embedder.

    python
    retriever
    from neo4j_graphrag.retrievers import VectorCypherRetriever
    
    vector_retriever = VectorCypherRetriever(
        driver,
        neo4j_database=os.getenv("NEO4J_DATABASE"),
        index_name="chunkEmbedding",
        embedder=embedder,
        retrieval_query=retrieval_query,
    )
  4. Define a tool function to search for lesson content using the retriever.

    python
    Search-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 context

    Tool description

    The tool name search-lesson-content and docstring Search for lesson content related to the query. help the agent decide when to use this tool.

  5. Update the tools list to include the new lesson search tool.

    python
    tools
    tools = [get_schema, search_lessons]
  6. Modify the query variable to test the new lesson search tool.

    python
    query
    query = "What are the benefits of using GraphRAG"
Reveal the complete code
python
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.

Chatbot

How can I help you today?