Agentic RAG Knowledge Graph

A sophisticated AI agent system that combines traditional RAG (Retrieval-Augmented Generation) with knowledge graphs to provide comprehensive insights using vector similarity and knowledge graph traversal.

🚀 Quick Start

1. Prerequisites

• Python 3.11 or higher
• PostgreSQL with pgvector extension
• Neo4j database
• LLM Provider API key (OpenAI, Anthropic, Gemini, etc.)

2. Installation

# Clone and enter the project
cd agentic-rag-knowledge-graph

# Create virtual environment
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install dependencies
pip install -r requirements.txt

3. Database Setup

PostgreSQL with pgvector:

# Execute the schema (adjust embedding dimensions if needed)
psql -d your_database -f sql/schema.sql

Neo4j:

• Install Neo4j Desktop or use a cloud instance
• Create a new database and note the connection details

4. Configuration

# Copy example environment file
cp .env.example .env

# Edit .env with your actual configuration
nano .env

Key configuration:

• DATABASE_URL: PostgreSQL connection string
• NEO4J_*: Neo4j connection details
• LLM_PROVIDER & LLM_API_KEY: Your preferred LLM provider
• EMBEDDING_MODEL: Embedding model (ensure dimensions match schema)

5. Document Ingestion

# Add documents to the documents/ folder
mkdir -p documents
# Copy your markdown documents here

# Run ingestion (this may take a while for knowledge graph extraction)
python -m ingestion.ingest

# Options:
python -m ingestion.ingest --clean  # Clean database first
python -m ingestion.ingest --no-semantic  # Skip semantic chunking
python -m ingestion.ingest --no-entities  # Skip entity extraction

6. Start the API Server

# Start the FastAPI server
python -m agent.api

# Server will run on http://localhost:8058

7. Use the CLI

# In a new terminal, start the interactive CLI
python cli.py

# The CLI will connect to your API server automatically

📚 Usage Examples

CLI Interaction

You: What are Microsoft's main AI initiatives?

🤖 Assistant:
Microsoft has several major AI initiatives including Azure OpenAI Service, 
Copilot integration across Office 365, and significant investments in AI 
research through Microsoft Research...

🛠 Tools Used:
  1. vector_search (query='Microsoft AI initiatives', limit=10)
  → Found 8 relevant chunks about Microsoft's AI research

API Usage

# Health check
curl http://localhost:8058/health

# Chat (non-streaming)
curl -X POST "http://localhost:8058/chat" \
  -H "Content-Type: application/json" \
  -d '{"message": "How are Microsoft and OpenAI connected?"}'

# Streaming chat
curl -X POST "http://localhost:8058/chat/stream" \
  -H "Content-Type: application/json" \
  -d '{"message": "Compare AI strategies of tech giants"}'

🏗️ Architecture

Core Components

1. Agent System (agent/)

   • agent.py: Main Pydantic AI agent with tool orchestration
   • providers.py: Multi-provider LLM abstraction (OpenAI, Anthropic, Gemini, Ollama)
   • models.py: Pydantic models for API and data structures
   • prompts.py: System prompts and templates
   • api.py: FastAPI web server with streaming endpoints

2. Ingestion Pipeline (ingestion/)

   • ingest.py: Main ingestion orchestrator
   • chunker.py: Semantic document chunking using LLM analysis
   • embedder.py: Embedding generation with multiple provider support

3. Database Layer

   • PostgreSQL with pgvector for vector similarity search
   • Neo4j with Graphiti for temporal knowledge graphs
   • Optimized schemas and indexes for performance

Search Capabilities

The agent intelligently chooses between three search strategies:

• Vector Search: Semantic similarity across document chunks
• Graph Search: Entity relationships and knowledge traversal
• Hybrid Search: Combined approach for comprehensive analysis

Tool Selection Logic

• "What is X?" → Vector search for content understanding
• "How are X and Y connected?" → Graph search for relationships
• "Compare X and Y" → Hybrid search for comprehensive analysis

🛠️ Advanced Configuration

Embedding Models

Supported embedding providers and their dimensions:

# OpenAI (1536 dimensions)
EMBEDDING_MODEL=text-embedding-3-small

# Ollama (768 dimensions - update SQL schema!)
EMBEDDING_MODEL=nomic-embed-text

Important: Ensure your SQL schema dimensions match your embedding model.

LLM Providers

# OpenAI
LLM_PROVIDER=openai
LLM_CHOICE=gpt-4o-mini

# Anthropic Claude  
LLM_PROVIDER=anthropic
LLM_CHOICE=claude-3-5-sonnet-20241022

# Local Ollama
LLM_PROVIDER=ollama
LLM_CHOICE=llama3.2:3b

# Google Gemini
LLM_PROVIDER=gemini
LLM_CHOICE=gemini-1.5-flash

# OpenRouter (access to multiple models)
LLM_PROVIDER=openrouter
LLM_CHOICE=anthropic/claude-3.5-sonnet

Performance Tuning

# Chunking parameters
python -m ingestion.ingest --chunk-size 800 --overlap 150

# Embedding batch size
# Modify EmbeddingGenerator.generate_embeddings_batch(batch_size=20)

# Database connection pooling
DB_POOL_SIZE=15
DB_MAX_OVERFLOW=30

📊 Monitoring and Statistics

Health Checks

# API health
curl http://localhost:8058/health

# Database connectivity
curl http://localhost:8058/stats

CLI Commands

health      - Check API server status
stats       - View system statistics  
clear       - Clear conversation session
help        - Show available commands

🧪 Testing

# Run all tests
pytest

# Run with coverage
pytest --cov=agent --cov=ingestion --cov-report=html

# Test specific components
pytest tests/agent/
pytest tests/ingestion/

🔧 Troubleshooting

Common Issues

Database Connection Errors:

# Test PostgreSQL connection
psql -d "$DATABASE_URL" -c "SELECT 1;"

# Verify pgvector extension
psql -d "$DATABASE_URL" -c "SELECT * FROM pg_extension WHERE extname = 'vector';"

Neo4j Connection Issues:

# Test Neo4j connectivity
docker run --rm neo4j:5.11 bolt://localhost:7687

No Search Results:

# Verify documents were ingested
psql -d "$DATABASE_URL" -c "SELECT COUNT(*) FROM documents;"
psql -d "$DATABASE_URL" -c "SELECT COUNT(*) FROM chunks WHERE embedding IS NOT NULL;"

Embedding Dimension Mismatch:

-- Check current schema dimension
\d+ chunks
-- Look for embedding VECTOR(dimension)

-- Update schema if needed (replace 1536 with your model's dimension)
ALTER TABLE chunks ALTER COLUMN embedding TYPE VECTOR(768);
ALTER TABLE relationships ALTER COLUMN embedding TYPE VECTOR(768);

Performance Issues

Slow Ingestion:

• Use --no-semantic for faster chunking
• Use --no-entities to skip knowledge graph extraction
• Reduce chunk size or batch size

Slow Queries:

• Check database indexes: \di in psql
• Monitor query performance: EXPLAIN ANALYZE SELECT ...
• Consider increasing database connection pool size

📁 Project Structure

agentic-rag-knowledge-graph/
├── agent/                  # AI agent and API
│   ├── agent.py           # Main Pydantic AI agent
│   ├── api.py             # FastAPI application  
│   ├── models.py          # Data models
│   ├── providers.py       # LLM provider abstraction
│   └── prompts.py         # System prompts
├── ingestion/             # Document processing
│   ├── ingest.py          # Main ingestion pipeline
│   ├── chunker.py         # Semantic chunking
│   └── embedder.py        # Embedding generation
├── sql/                   # Database schema
│   └── schema.sql         # PostgreSQL schema with pgvector
├── tests/                 # Test suites
├── documents/             # Your documents (create this)
├── cli.py                 # Interactive CLI
├── requirements.txt       # Dependencies
└── .env.example          # Configuration template

🤝 Contributing

1. Fork the repository
2. Create a feature branch: git checkout -b feature-name
3. Make your changes with tests
4. Run the test suite: pytest
5. Submit a pull request

📄 License

This project is licensed under the MIT License - see the LICENSE file for details.

🙏 Acknowledgments

• Pydantic AI for the agent framework
• Graphiti for temporal knowledge graphs
• pgvector for vector similarity search
• FastAPI for the web framework
• Rich for beautiful CLI output

---

Built with ❤️ for the AI community