In recent years, the advancements in AI, particularly in Large Language Models (LLMs) such as GPT-3, have revolutionized various sectors by providing nearly human-like text generation and interaction capabilities. However, when combined with Retrieval-Augmented Generation (RAG) and Knowledge Graphs, these models’ efficiency and accuracy reach unprecedented levels, offering more contextually relevant and precise information.
In this article, we will introduce the basics of RAG and its limitation, and introduce a better RAG solution too.
Overview of LLMs, RAG, and Knowledge Graphs
Brief explanation of Large Language Models (LLMs)
LLMs, such as GPT (Generative Pre-trained Transformer) series, are deep learning algorithms trained on vast amounts of text data. They can predict the next word in a sentence, making them powerful tools for generating human-like text based on the input they’re fed.
Short introduction to RAG
Retrieval-Augmented Generation (RAG) combines the generative power of LLMs with an external knowledge source, allowing the model to “retrieve” relevant information before “generating” a response. This process enhances the model’s ability to provide informed and precise answers.
Short Introduction to Knowledge Graphs
Knowledge Graphs organize and store complex data in an interconnected, graph-like structure. They enable logical reasoning over the stored data, facilitating precise information retrieval and data interconnectivity essential for various AI applications.
Limitation of LLM only
While LLMs have shown remarkable capabilities, their effectiveness is often limited by the data on which they were trained. They can generate outdated or irrelevant information and struggle with queries requiring up-to-date or domain-specific knowledge.
Limitation of LLM + RAG
While RAG models enhance LLMs by integrating external sources, they still face challenges such as integrating these sources smoothly and efficiently extracting the most relevant information from vast databases.
Use Knowledge Graph improve Traditional RAG
Combining LLMs with Knowledge Graphs bridges the gap between generative models’ linguistic capabilities and the structured, detailed data stored in knowledge bases.
Tech design of Knowledge Graph
Go through the code of demo: pingcap/tidb-vector-python/examples/graphrag-demo, we can find the code how we design schema SQL to store knowledge graph:
Prepare Tables
Here, we use TiDB Serverless – MySQL compatible database with built-in vector storage – to store the entity and relationship data:
CREATE TABLE `entities` (
`id` int(11) NOT NULL AUTO_INCREMENT,
`name` varchar(512) DEFAULT NULL,
`description` text DEFAULT NULL,
`description_vec` vector<float> DEFAULT NULL,
PRIMARY KEY (`id`) /*T![clustered_index] CLUSTERED */
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin;
CREATE TABLE `relationships` (
`id` int(11) NOT NULL AUTO_INCREMENT,
`source_entity_id` int(11) DEFAULT NULL,
`target_entity_id` int(11) DEFAULT NULL,
`relationship_desc` text DEFAULT NULL,
PRIMARY KEY (`id`) /*T![clustered_index] CLUSTERED */,
KEY `fk_1` (`source_entity_id`),
KEY `fk_2` (`target_entity_id`),
CONSTRAINT `fk_1` FOREIGN KEY (`source_entity_id`) REFERENCES `entities` (`id`),
CONSTRAINT `fk_2` FOREIGN KEY (`target_entity_id`) REFERENCES `entities` (`id`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin;Insert Entity & Relationship
Sure, it’s a MySQL compatible table schema definition, and the data format would be like:
INSERT INTO entities (id, name, description, description_vec) VALUES (1, 'Elon Reeve Musk', 'Elon Reeve Musk is a businessman and investor, born on June 28, 1971, in Pretoria, South Africa. He is the founder, chairman, CEO, and CTO of SpaceX; angel investor, CEO, product architect, and former chairman of Tesla, Inc.; owner, executive chairman, and CTO of X Corp.; founder of the Boring Company and xAI; co-founder of Neuralink and OpenAI; and president of the Musk Foundation.', [1536 length of vector]);The full demo data is here.
Search knowledge with SQL
After create table and inserting demo data, we can perform a search on top of them with just MySQL skills but with vector search!
from sqlalchemy import create_engine, text
import openai
import getpass
# TiDB Connection String Pattern:
# mysql+pymysql://{TIDB_USER}:{TIDB_PASSWORD}@{TIDB_HOST}:{TIDB_PORT}/{TIDB_DB_NAME}?ssl_verify_cert=True&ssl_verify_identity=True
db_engine = create_engine(getpass.getpass("Input your TIDB connection string:"))
oai_cli = openai.OpenAI(api_key=getpass.getpass("Input your OpenAI API Key:"))
question = input("Enter your question:")
embedding = str(oai_cli.embeddings.create(input=[question], model="text-embedding-3-small").data[0].embedding)
with db_engine.connect() as conn:
result = conn.execute(text("""
WITH initial_entity AS (
SELECT id FROM `entities`
ORDER BY VEC_Cosine_Distance(description_vec, :embedding) LIMIT 1
), entities_ids AS (
SELECT source_entity_id i FROM relationships r INNER JOIN initial_entity i ON r.target_entity_id = i.id
UNION SELECT target_entity_id i FROM relationships r INNER JOIN initial_entity i ON r.source_entity_id = i.id
UNION SELECT initial_entity.id i FROM initial_entity
) SELECT description FROM `entities` WHERE id IN (SELECT i FROM entities_ids);"""), {"embedding": embedding}).fetchall()
print(oai_cli.chat.completions.create(model="gpt-4o", messages=[
{"role": "system", "content": f"Please carefully answer the question by {str(result)}"},
{"role": "user", "content": question}]).choices[0].message.content)It’s clear that we have implemented a RAG but:
- With Knowledge Graph concept in it
- Use MySQL style skills
- With Vector Storage
- With LLM
That’s amazing!
(The full demo code is here.)
How LLMs and Knowledge Graphs complement each other
The integration allows LLMs to leverage structured data from Knowledge Graphs, enriching the context and accuracy of the generated responses. This symbiosis enables AI to not just predict text but to understand and generate content with nuanced, fact-based insights.
Benefits of combining LLMs’ natural language understanding with Knowledge Graphs’ structured data
This combination supports complex query answering, factual consistency, and the exploration of relationships between entities, significantly improving task-oriented and informational dialogues in AI applications.
Workflow Design of GraphRAG
GraphRAG revolutionizes the traditional RAG approach by directly integrating Knowledge Graph data into the indexing and retrieval process, leading to more relevant data extraction and smarter, contextual response generation.
Explain the workflow to use graphrag
- Data Indexing: Knowledge Graph entities and relationships are indexed, enabling efficient traversal and retrieval. Here we use the following tech stacks:
- Knowledge Graphs Extraction: OpenAI GPT-4
- Vector Databases: we use TiDB Serverless to store vectors and entity->relationship graph structures.
- Data Retrieval: Upon receiving a query, GraphRAG retrieves the most pertinent nodes and edges from the Knowledge Graph, considering the query context and available structured data as relevant context.
- Reranking: The retrieved information undergoes a reranking process based on relevance to the query, with the most relevant pieces directed to the generative component for response formulation.
Try TiDB Serverless Vector Storage
Build AI applications confidently with the most advanced MySQL vector solution.
Conclusion
GraphRAG exemplifies the future of RAG by seamlessly integrating the depth of Knowledge Graphs with the generative prowess of LLMs. This pioneering approach significantly refines AI’s ability to interact with and understand the human language within specific contexts, making it a game-changer in fields ranging from customer support to sophisticated data analysis.
As we continue to advance in AI and explore innovative integrations like GraphRAG, the potential for truly intelligent systems seems more attainable than ever. It’s not just about building smarter AI, it’s about creating systems that enhance our decision-making and broaden our understanding of the world around us. Discover more about how TiDB is pioneering the future of data management and AI here.
More Demos
- OpenAI Embedding: use the OpenAI embedding model to generate vectors for text data.
- Image Search: use the OpenAI CLIP model to generate vectors for image and text.
- LlamaIndex RAG with UI: use the LlamaIndex to build an RAG(Retrieval-Augmented Generation) application.
- Chat with URL: use LlamaIndex to build an RAG(Retrieval-Augmented Generation) application that can chat with a URL.
- GraphRAG: 20 lines code of using TiDB Serverless to build a Knowledge Graph based RAG application.
- GraphRAG Step by Step Tutorial: Step by step tutorial to build a Knowledge Graph based RAG application with Colab notebook. In this tutorial, you will learn how to extract knowledge from a text corpus, build a Knowledge Graph, store the Knowledge Graph in TiDB Serverless, and search from the Knowledge Graph.
Embrace the future with us. Explore TiDB’s capabilities and see how it can transform your data management strategies. Join our community for more insights and discussions on the evolving landscape of AI, RAG, GraphRAG: https://discord.gg/XzSW23Jg9p