Build a Graph RAG system with LangChain and GraphRetriever

import json
import os
from langchain_core.documents import Document
from langchain_graph_retriever.transformers.html import HyperlinkTransformer
from langchain_astradb import AstraDBVectorStore
from langchain_openai import OpenAIEmbeddings

endpoint = os.environ.get("API_ENDPOINT")  (1)
application_token = os.environ.get("APPLICATION_TOKEN")
openai_api_key = os.environ.get("OPENAI_API_KEY")

if not endpoint or not application_token or not openai_api_key:
    raise RuntimeError(
        "Environment variables API_ENDPOINT, APPLICATION_TOKEN, OPENAI_API_KEY must be defined."
    )

data_file_path = "PATH_TO_DATA_FILE"  (2)

# Read the JSON file and parse it into a JSON array
with open(data_file_path, "r", encoding="utf8") as file:
    json_data = json.load(file)

# Convert the JSON array into LangChain Documents
documents = [
    Document(page_content=data["html_doc"], metadata={"url": data["url"]})
    for data in json_data
]

# Extract hyperlinks from the HTML and store them in the Document metadata
html_transformer = HyperlinkTransformer()
documents_with_links = html_transformer.transform_documents(documents)

# Creates a vector store that uses
# Astra DB as the backend and
# OpenAI as the embedding service
print("Creating vector store...")
vector_store = AstraDBVectorStore(
    collection_name="graph_rag_tutorial",
    token=application_token,
    api_endpoint=endpoint,
    embedding=OpenAIEmbeddings(api_key=openai_api_key),
)

# In case a collection with this name already existed,
# delete any documents in the collection
vector_store.clear()

# Insert the documents into the collection
print("Inserting documents...")
vector_store.add_documents(documents_with_links)

from pprint import pprint
import os
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from graph_retriever.strategies import Eager
from langchain_graph_retriever import GraphRetriever
from langchain_astradb import AstraDBVectorStore
from langchain_openai import OpenAIEmbeddings

endpoint = os.environ.get("API_ENDPOINT")  (1)
application_token = os.environ.get("APPLICATION_TOKEN")
openai_api_key = os.environ.get("OPENAI_API_KEY")

if not endpoint or not application_token or not openai_api_key:
    raise RuntimeError(
        "Environment variables API_ENDPOINT, APPLICATION_TOKEN, OPENAI_API_KEY must be defined."
    )

# Initialize AstraDBVectorStore
# based on the collection that you created and populated earlier
vector_store = AstraDBVectorStore(
    collection_name="graph_rag_tutorial",
    token=application_token,
    api_endpoint=endpoint,
    embedding=OpenAIEmbeddings(api_key=openai_api_key),
    autodetect_collection=True,
)

# Initialize the LLM
llm = ChatOpenAI(model="gpt-4o", api_key=openai_api_key)

# Define the prompt template
template = """Answer the question based only on the following context:

{context}

Question: {question}
"""
prompt = ChatPromptTemplate.from_template(template)

# Initialize GraphRetriever.
# This retriever first uses vector search to find relevant documents,
# then uses graph traversal to explore their connections.
retriever = GraphRetriever(
    store=vector_store,
    edges=[("hyperlink", "url")],
    strategy=Eager(
        # Number of documents to fetch via vector search for starting the traversal.
        start_k=3,
        # Maximum total documents to retrieve during traversal
        select_k=10,
        # Maximum traversal depth.
        # A value of 0 only performs vector search, but does not do any graph traversal.
        max_depth=1,
    ),
)


# Helper function to format the retrieved documents for LLM context
def format_docs(docs):
    return "\n\n".join([d.page_content for d in docs])


# Build the RAG chain:
# 1. Take the input question, retrieve relevant documents from a vector store,
#   and format the documents into a string.
#   Create a dictionary with the formatted documents and the original question.
# 2. Create the prompt by injecting the dictionary from step 1 into the prompt template.
# 3. Pass the formatted prompt to the LLM.
# 4. Clean and return the LLM output.
chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

# Try these questions to explore the knowledge graph:
QUESTION = "What is close to the Space Needle?"
# Alternative questions:
# - "What is in the Lower Queen Anne neighborhood?"
# - "What is in the same neighborhood as the Space Needle?"
# - "What connects the 1962 World's Fair to modern Seattle?"
# - "Where is Chihuly Garden and Glass?"

print(f"\nQuestion: {QUESTION}\n")

try:
    response = chain.invoke(QUESTION)
    print("Answer:")
    pprint(response)
except Exception as e:
    print(f"Error during RAG query: {e}")