Retrieval Augemented Generation(RAG) with Knowledge Graphs

Vivek Pandit
3 min readMar 16, 2024

--

What is RAG

Augmented Generation (RAG) models are at the forefront of revolutionizing how machines understand and generate human language. By merging the best of retrieval-based techniques with advanced generative models, RAGs fetch relevant information from vast datasets and then craft responses that are not only accurate but contextually rich. This approach has set new benchmarks in creating AI that can converse, answer, and interact in ways more aligned with human understanding.

https://medium.com/@bijit211987/advanced-rag-for-llms-slms-5bcc6fbba411

Challenges with RAG

While RAG models offer significant advantages, they face key challenges:

While RAG models offer significant advantages, they face key challenges:

  • Data Relevance: Ensuring the retrieved data is contextually appropriate and up-to-date.
  • Handling Complex Queries: Dealing with intricate questions that require a deep understanding of relationships and contexts.
  • Integration of Data Types: Merging structured and unstructured data to provide comprehensive responses.

What is a Knowledge Graph?

A Knowledge Graph is a dynamic, structured representation of knowledge, where entities (nodes) are connected by relationships (edges). This setup not only catalogs information but also the context and the interrelation among data points. For instance, in a medical knowledge graph, nodes could represent symptoms, diseases, and treatments, with edges defining relationships like “symptom of” or “treated by.”

Example: Consider a simplified medical knowledge graph where “Cough” and “Cold” are symptoms connected to “Flu,” a disease, by the relationship “symptom of.” Treatments like “Rest” and “Hydration” are connected to “Flu” with “treated by” relationships.

Addressing RAG Challenges with Knowledge

Graphs Knowledge Graphs can significantly mitigate RAG challenges:

  • Contextual Relevance: The graph’s structure ensures that the information retrieved is not only relevant but also contextual, providing a richer background for response generation.
  • Complex Queries: The interconnected nature of knowledge graphs enables handling complex queries efficiently, as relationships provide a deep understanding of context.
  • Data Integration: Knowledge graphs are adept at integrating various data types, offering a unified view that enriches RAG responses.

Enhancing Search with Embeddings

Embeddings transform text into high-dimensional vectors, capturing semantic meaning. In knowledge graphs, embeddings can improve search by allowing semantic matching, ensuring that retrieved information closely aligns with the query’s context, significantly benefiting RAG models.

Creating Node Embeddings with Neo4J using Langchain and Watsonx

Neo4J, a graph database, facilitates creating node embeddings, enhancing RAG’s retrieval capabilities. Here’s how you can generate embeddings in a LangChain environment:

Example: Building on our medical knowledge graph example, let’s generate embeddings for symptoms and diseases.

  1. Setting Up Neo4J: Ensure your Neo4J instance is running and accessible.
  2. Creating the Graph: Load your entities and relationships into Neo4J, structuring them according to your knowledge graph.
  3. Generating Embeddings: Use the following LangChain code snippet to generate node embeddings:pythonCopy code
from langchain.neural_graphs import Neo4jGraph
from langchain_community.vectorstores import Neo4jVector
from langchain.chains import RetrievalQAWithSourcesChain
from langchain_community.llms import WatsonxLLM
graph = Neo4jGraph()

# Import movie information

movies_query = """
LOAD CSV WITH HEADERS FROM
'https://raw.githubusercontent.com/tomasonjo/blog-datasets/main/movies/movies_small.csv'
AS row
MERGE (m:Movie {id:row.movieId})
SET m.released = date(row.released),
m.title = row.title,
m.imdbRating = toFloat(row.imdbRating)
FOREACH (director in split(row.director, '|') |
MERGE (p:Person {name:trim(director)})
MERGE (p)-[:DIRECTED]->(m))
FOREACH (actor in split(row.actors, '|') |
MERGE (p:Person {name:trim(actor)})
MERGE (p)-[:ACTED_IN]->(m))
FOREACH (genre in split(row.genres, '|') |
MERGE (g:Genre {name:trim(genre)})
MERGE (m)-[:IN_GENRE]->(g))
"""

graph.query(movies_query)

vector_store_extra_text = Neo4jVector.from_existing_index(
embedding=HuggingFaceEmbeddings(),
url=NEO4J_URI,
username=NEO4J_USERNAME,
password=NEO4J_PASSWORD,
database="neo4j",
index_name=VECTOR_INDEX_NAME,
text_node_property=VECTOR_SOURCE_PROPERTY,
retrieval_query=retrieval_query_extra_text, # NEW !!!
)

# Create a retriever from the vector store
retriever_extra_text = vector_store_extra_text.as_retriever()

parameters = {
GenParams.DECODING_METHOD: DecodingMethods.GREEDY,
GenParams.MIN_NEW_TOKENS: 1,
GenParams.MAX_NEW_TOKENS: 200,
GenParams.STOP_SEQUENCES: ["<|endoftext|>","\n\n"]
}

watsonx_model = Model(
model_id=wx_model_id,
credentials=credentials,
project_id=wx_project_id,
params=parameters)
# Create a chatbot Question & Answer chain from the retriever

chain_extra_text = RetrievalQAWithSourcesChain.from_chain_type(
WatsonxLLM(model=watsonx_model),
chain_type="stuff",
retriever=retriever_extra_text
)

This snippet connects to a Neo4J database, generates embeddings for all nodes using a pre-trained model, and retrieves the embedding for a specific node.

By integrating knowledge graphs with RAG models and enhancing them with node embeddings, we can create AI systems that understand and generate human language with unprecedented accuracy and depth. This approach not only addresses the inherent challenges of RAG models but also unlocks new potentials in natural language processing.

--

--