🧠 Building an Advanced RAG System with Knowledge Graphs and Vector Search

Today we open-source our project for MTH088 - Advanced Mathematics for AI, where we built a sophisticated Retrieval-Augmented Generation (RAG) system that combines vector similarity search with knowledge graph extraction. This project showcases how we can leverage advanced AI techniques to create a powerful document understanding system.

🚀 What Makes This RAG System Special?

In the rapidly evolving landscape of AI and natural language processing, Retrieval-Augmented Generation (RAG) systems have become crucial for building intelligent applications that can understand and reason about large document collections. But what if we could go beyond simple vector similarity search and actually understand the relationships between concepts in our documents?

That's exactly what we set out to build! 🎯

Our system doesn't just store and retrieve documents—it understands them by:

• 🕸️ Extracting Knowledge Graphs: Automatically identifying relationships between entities
• 🔍 Performing Vector Search: Finding semantically similar content with high precision
• 🤖 Combining Both Approaches: Leveraging the strengths of structured and unstructured data retrieval

🏗️ Architecture Overview

graph TB
    A[📄 Document Upload] --> B[📖 Docling Processing]
    B --> C[🔗 Knowledge Graph Extraction]
    B --> D[🔢 Vector Embedding]
    C --> E[🗄️ Milvus Storage]
    D --> E
    F[🔍 User Query] --> G[🏷️ Entity Recognition]
    F --> H[🔢 Query Embedding]
    G --> I[🕸️ Graph Search]
    H --> J[📊 Vector Search]
    I --> K[📋 Combined Results]
    J --> K
    K --> L[💬 LLM Response]

Our system follows a sophisticated multi-stage pipeline:

1️⃣ Document Ingestion Phase

When you upload a document (PDF or text), our system:

• 📖 Uses Docling for intelligent chunking and preprocessing
• 🤖 Leverages LLM-powered analysis to extract (subject, relation, object) triplets
• 🔢 Generates high-dimensional embeddings for semantic search
• 💾 Stores everything in Milvus vector database with separate collections for entities and relations

2️⃣ Query Processing Phase

When you ask a question:

• 🏷️ Named Entity Recognition identifies key concepts in your query
• 🔍 Vector similarity search finds semantically related content
• 🕸️ Knowledge graph traversal discovers connected relationships
• 📊 Smart ranking algorithm combines both signals for optimal results

💻 Tech Stack Deep Dive

🐍 Backend Powerhouse

• ⚡ FastAPI (v0.115.5): Lightning-fast async API framework
• 🗄️ Milvus (v2.5.0): High-performance vector database
• 🔴 Redis: Intelligent caching layer for 60-second query caching
• ✅ Pydantic (v2.9.2): Rock-solid data validation
• 📄 Docling: Advanced PDF processing with smart chunking

⚛️ Frontend Excellence

• ⚛️ React (v19.0.0): Modern, responsive UI
• 📘 TypeScript (v4.9.5): Type-safe development
• 🔗 Axios (v1.8.2): Seamless API communication

🐳 Infrastructure Stack

• 🐳 Docker Compose: One-command deployment
• 🪣 MinIO: S3-compatible object storage
• 🔑 etcd: Distributed metadata management

🌟 Key Features That Set Us Apart

🕸️ Intelligent Knowledge Graph Extraction

Our system doesn't just store text—it understands relationships:

# Example extracted triplets from a document about AI
[
    ("Neural Networks", "are used in", "Machine Learning"),
    ("GPT", "is a type of", "Large Language Model"),
    ("Transformers", "revolutionized", "Natural Language Processing")
]

Each relationship is stored with proper entity linking, enabling complex queries like:

• "Show me everything connected to Neural Networks"
• "What are the applications of Transformers?"

⚡ Performance Optimizations

• 🔄 Asynchronous Processing: Non-blocking document ingestion
• 📦 Intelligent Batching: Configurable batch sizes for optimal throughput
• 🚦 Concurrency Limits: Smart rate limiting for external API calls
• ⚡ Redis Caching: Sub-second query responses for repeated searches

🎯 Multi-Modal Search

Our system supports various input types:

• 📄 PDF Documents: Automatic text extraction and processing
• 📝 Raw Text: Direct text input for quick knowledge base building
• 🔗 URL Processing: Fetch and process web content (when integrated)

🚀 Getting Started in Minutes

Want to try it out? Here's how to get running locally:

✅ Prerequisites

• 🐍 Python 3.12+
• 📗 Node.js 16+
• 🐳 Docker & Docker Compose

📦 Quick Setup

# 1️⃣ Clone the repository
git clone https://github.com/hcmus-project-collection/llm-with-knowledge-base
cd llm-with-knowledge-base

# 2️⃣ Start the infrastructure
docker-compose -f milvus-docker-compose.yml up -d

# 3️⃣ Set up the backend
conda create -n llmkb python=3.12 -y
conda activate llmkb
pip install -r requirements.txt
cp .env.template .env
# Edit .env with your configuration
python -O server.py

# 4️⃣ Launch the frontend
cd frontend
npm install
cp .env.template .env
# Edit .env with your settings
npm start

🎯 Real-World Use Cases

📚 Academic Research

• Upload research papers and discover connections between concepts
• Find related work through semantic similarity
• Extract knowledge graphs from literature reviews

🏢 Enterprise Knowledge Management

• Build company-wide knowledge bases from documentation
• Enable semantic search across technical manuals
• Discover hidden relationships in business documents

💻 Development Documentation

• Create searchable code documentation
• Link related functions and classes automatically
• Find usage examples through relationship mapping

🧪 What We Learned Building This

🎓 Mathematical Foundations

This project gave us hands-on experience with:

• Vector Spaces & Similarity Metrics: Understanding cosine similarity, L2 distance, and inner product spaces
• Graph Theory: Implementing knowledge graphs with proper entity linking
• Linear Algebra: Working with high-dimensional embeddings and dimensionality reduction
• Information Retrieval: Combining multiple ranking signals for optimal search results

🛠️ Engineering Challenges

• Scalability: Handling large document collections efficiently
• Latency: Optimizing query response times with smart caching
• Accuracy: Balancing precision and recall in hybrid search
• Robustness: Error handling and graceful degradation

🔮 Future Enhancements

We're excited about potential improvements:

• 🌐 Multi-language Support: Extend to non-English documents
• 📊 Advanced Analytics: Query pattern analysis and optimization
• 🔗 Graph Visualization: Interactive knowledge graph exploration
• 🤖 Auto-categorization: Intelligent document classification
• 📱 Mobile App: Native mobile interface for on-the-go access

📊 Performance Metrics

Our system achieves impressive performance:

🤝 Contributing & Feedback

We'd love to hear from the community! Whether you're:

• 🐛 Finding bugs
• 💡 Suggesting features
• 📚 Improving documentation
• 🔧 Contributing code

Check out our GitHub repository and feel free to open issues or submit pull requests!

🎉 Conclusion

Building this RAG system has been an incredible journey that combines cutting-edge AI research with practical engineering. By integrating knowledge graphs with vector search, we've created a system that doesn't just find similar documents—it understands relationships and provides contextually rich results.

The intersection of mathematics, AI, and software engineering in this project perfectly embodies what modern AI development looks like. We hope this system inspires others to explore the fascinating world of knowledge representation and retrieval!

Ready to dive into the future of intelligent document understanding? 🚀

⭐ Star us on GitHub | 📖 Read the Docs | 🐛 Report Issues

🙏 Acknowledgments

Special thanks to ChatGPT for enhancing this post with suggestions, formatting, and emojis.