# RAG Server with MCP Integration

A Retrieval-Augmented Generation (RAG) system with Google AI Studio integration, featuring both a REST API and Model Context Protocol (MCP) server.

## Features

### Core Capabilities
- **Document Storage**: Upload and store text (.txt) and Markdown (.md) documents
- **Hierarchical Chunking**: Structure-aware chunking for markdown that preserves document hierarchy
- **Vector Search**: Efficient similarity search using Qdrant vector database
- **Google AI Integration**: Uses Google AI Studio for embeddings (text-embedding-004) and generation (gemini-1.5-flash)
- **REST API**: FastAPI-based REST API with automatic OpenAPI documentation
- **MCP Server**: Model Context Protocol server for seamless integration with Claude and other MCP clients
- **OpenAI-Compatible API**: Supports OpenAI-compatible chat completions for web UI integration
- **Code Indexing**: Index and search source code repositories with semantic understanding
- **Smart Query Routing**: Automatic query classification and routing to appropriate retrieval methods

### Advanced Features
- **Tag-Based Organization**: Organize documents with multiple tags for easy categorization
- **Section-Aware Retrieval**: Query specific sections of documentation (e.g., "Installation > Prerequisites")
- **Markdown Structure Preservation**: Automatic extraction of heading hierarchy with breadcrumb paths
- **Context-Enhanced Answers**: LLM receives section context for more accurate responses
- **Flexible Filtering**: Filter documents by tags and/or section paths during queries
- **Document Structure API**: Explore table of contents and section organization
- **GitHub Integration**: Parse and extract content from GitHub URLs
- **Reference Following**: Automatically follow documentation references for comprehensive answers
- **Multi-Mode Retrieval**: Choose between standard, enhanced, or smart query modes
- **Rate Limiting**: Built-in rate limiting for API endpoints

## Project Structure

```
mcp-rag-docs/
   config/
      __init__.py
      settings.py                # Configuration and settings
   rag_server/
      __init__.py
      models.py                  # Pydantic models for API
      openai_api.py              # OpenAI-compatible API endpoints
      openai_models.py           # OpenAI API models
      rag_system.py              # Core RAG system logic
      server.py                  # FastAPI server
      smart_query.py             # Smart query routing
   mcp_server/
      __init__.py
      server.py                  # MCP server implementation
   utils/
      __init__.py
      code_indexer.py            # Source code indexing
      code_index_store.py        # Code index storage
      document_processor.py      # Document processing
      embeddings.py              # Google AI embeddings
      frontmatter_parser.py      # YAML frontmatter parsing
      github_parser.py           # GitHub URL parsing
      google_api_client.py       # Google AI API client
      hierarchical_chunker.py    # Hierarchical document chunking
      markdown_parser.py         # Markdown parsing
      query_classifier.py        # Query type classification
      rate_limit_store.py        # Rate limiting
      reference_extractor.py     # Extract doc references
      retrieval_router.py        # Multi-mode retrieval routing
      source_extractor.py        # Extract source code snippets
      text_chunker.py            # Text chunking utility
      vector_store.py            # Qdrant vector store wrapper
   build_code_index.py          # Build code index from repository
   check_github_urls.py         # Validate GitHub URLs
   check_status.py              # System status checker
   example_usage.py             # Example usage scripts
   ingest_docs.py               # Document ingestion utility
   main.py                      # Main entry point
   .env.example                 # Example environment variables
   docker-compose.yml           # Docker setup for Qdrant
   pyproject.toml               # Project dependencies
```

## Installation

### Prerequisites

- Python 3.13 or higher
- Google AI Studio API key ([Get one here](https://makersuite.google.com/app/apikey))

### Setup

1. **Clone or navigate to the project directory**

2. **Install dependencies**

```bash
# Using pip
pip install -e .

# Or using uv (recommended)
uv pip install -e .
```

3. **Configure environment variables**

```bash
# Copy the example env file
cp .env.example .env

# Edit .env and add your Google API key
GOOGLE_API_KEY=your_api_key_here
```

4. **Start Qdrant (optional - using Docker)**

```bash
docker-compose up -d
```

## Usage

### Running the FastAPI Server

Start the REST API server:

```bash
python -m rag_server.server
```

The server will start at `http://localhost:8000`. Visit `http://localhost:8000/docs` for interactive API documentation.

#### API Endpoints

**Core Endpoints:**
- **POST /documents** - Upload a document
- **POST /query** - Query the RAG system (standard mode)
- **POST /query-enhanced** - Query with automatic reference following
- **POST /smart-query** - Smart query with automatic routing
- **GET /documents** - List all documents
- **DELETE /documents/{doc_id}** - Delete a document
- **GET /stats** - Get system statistics
- **GET /health** - Health check
- **GET /tags** - List all available tags
- **GET /documents/{doc_id}/sections** - Get document structure

**OpenAI-Compatible Endpoints:**
- **POST /v1/chat/completions** - OpenAI-compatible chat completions
- **GET /v1/models** - List available models

#### Example Usage with curl

```bash
# Upload a document
curl -X POST "http://localhost:8000/documents" \
  -F "file=@example.txt"

# Upload with tags
curl -X POST "http://localhost:8000/documents" \
  -F "file=@dagster-docs.md" \
  -F "tags=dagster,python,orchestration"

# Query the RAG system
curl -X POST "http://localhost:8000/query" \
  -H "Content-Type: application/json" \
  -d '{"question": "What is the main topic of the documents?", "top_k": 5}'

# Smart query with automatic routing
curl -X POST "http://localhost:8000/smart-query" \
  -H "Content-Type: application/json" \
  -d '{"question": "How do I create a Dagster asset?"}'

# OpenAI-compatible chat completion
curl -X POST "http://localhost:8000/v1/chat/completions" \
  -H "Content-Type: application/json" \
  -d '{
    "model": "rag-smart",
    "messages": [{"role": "user", "content": "What is an asset in Dagster?"}],
    "stream": false
  }'

# List documents
curl "http://localhost:8000/documents"

# Get statistics
curl "http://localhost:8000/stats"
```

### Running the MCP Server

The MCP server allows integration with Claude and other MCP-compatible clients.

```bash
python -m mcp_server.server
```

#### MCP Tools Available

1. **query_rag** - Query the RAG system with a question
2. **query_rag_enhanced** - Query with automatic reference following
3. **smart_query** - Smart query with automatic routing and classification
4. **add_document** - Add a document to the RAG system
5. **list_documents** - List all stored documents
6. **delete_document** - Delete a document by ID
7. **get_rag_stats** - Get system statistics
8. **get_tags** - List all available tags
9. **get_document_structure** - Get document table of contents

#### Using with Claude Desktop

Add to your Claude Desktop configuration (`claude_desktop_config.json`):

```json
{
  "mcpServers": {
    "rag": {
      "command": "uv",
      "args": [
        "--directory",
        "/path/to/mcp-rag-docs",
        "run",
        "python",
        "-m",
        "mcp_server.server"
      ]
    }
  }
}
```

See [QUICK_START.md](QUICK_START.md) for a quick setup guide.

## Configuration

All configuration is managed through environment variables (defined in `.env`):

| Variable | Description | Default |
|----------|-------------|---------|
| `GOOGLE_API_KEY` | Google AI Studio API key | (required) |
| `CHUNK_SIZE` | Size of text chunks in characters | 1000 |
| `CHUNK_OVERLAP` | Overlap between chunks | 200 |
| `TOP_K_RESULTS` | Number of chunks to retrieve | 5 |
| `QDRANT_PATH` | Path to Qdrant storage | ./qdrant_storage |
| `QDRANT_COLLECTION_NAME` | Qdrant collection name | documents |
| `FASTAPI_HOST` | FastAPI server host | 0.0.0.0 |
| `FASTAPI_PORT` | FastAPI server port | 8000 |
| `EMBEDDING_MODEL` | Google embedding model | text-embedding-004 |
| `LLM_MODEL` | Google LLM model | gemini-1.5-flash |

## Architecture

### Document Processing Pipeline

1. **Upload** - User uploads a .txt or .md file
2. **Processing** - Document is read and metadata extracted (including frontmatter)
3. **Chunking** - Text is split using hierarchical chunking for markdown or standard chunking for text
4. **Embedding** - Each chunk is converted to a vector using Google AI embeddings
5. **Storage** - Vectors and metadata are stored in Qdrant

### Query Pipeline

#### Standard Query
1. **Query** - User submits a question
2. **Embedding** - Question is converted to a vector
3. **Retrieval** - Similar chunks are retrieved from Qdrant
4. **Generation** - Context is provided to Google AI Studio model
5. **Response** - Answer is generated and returned with sources

#### Smart Query
1. **Classification** - Query is classified (documentation, code, conceptual, etc.)
2. **Routing** - Automatically selects best retrieval strategy
3. **Multi-Source** - May combine documentation search, code search, and direct answers
4. **Synthesis** - Generates comprehensive answer from multiple sources

### Code Indexing

The system can index source code repositories:

```bash
# Build code index
python build_code_index.py /path/to/repo

# Query code through the API or MCP server
```

Code is indexed with:
- Class and function definitions
- Docstrings and comments
- File structure and imports
- Semantic embeddings for natural language queries

## Development

### Running Tests

```bash
# Install test dependencies
pip install pytest pytest-asyncio httpx

# Run tests
pytest

# Run specific test files
pytest test_openai_api.py
pytest test_mcp_integration.py
```

### Code Style

The project follows Python best practices with type hints and docstrings.

## Troubleshooting

### Common Issues

**Issue**: `GOOGLE_API_KEY not found`
- **Solution**: Ensure you've created a `.env` file and added your Google API key

**Issue**: `Unsupported file type`
- **Solution**: Only .txt and .md files are supported. Convert other formats first.

**Issue**: `Collection already exists` error
- **Solution**: Delete the `qdrant_storage/` directory to reset the database

**Issue**: MCP server not connecting
- **Solution**: Check that the path in your MCP config is correct and the `.env` file is in the project root

## Advanced Usage

### Tag-Based Organization

Organize your documents with tags for easy categorization and filtering:

```bash
# Upload document with tags
curl -X POST "http://localhost:8000/documents" \
  -F "file=@dagster-docs.md" \
  -F "tags=dagster,python,orchestration"

# List all available tags
curl "http://localhost:8000/tags"

# Query only dagster-related documents
curl -X POST "http://localhost:8000/query" \
  -H "Content-Type: application/json" \
  -d '{"question": "How do I create a pipeline?", "tags": ["dagster"]}'

# List documents filtered by tags
curl "http://localhost:8000/documents?tags=dagster,python"
```

### Hierarchical Document Structure

For markdown documents, the system automatically preserves heading hierarchy:

```bash
# Get document structure (table of contents)
curl "http://localhost:8000/documents/{doc_id}/sections"

# Query specific section
curl -X POST "http://localhost:8000/query" \
  -H "Content-Type: application/json" \
  -d '{"question": "What are the prerequisites?", "section_path": "Installation > Prerequisites"}'
```

### Section-Aware Queries

The system includes section context when generating answers:

```python
# Example: Markdown document structure
# Installation
#   Prerequisites
#     Python Version
#   Setup Steps

# When you query about "Python version requirements"
# The system will:
# 1. Retrieve relevant chunks from "Installation > Prerequisites > Python Version"
# 2. Include section path in context sent to LLM
# 3. Cite sources with full section paths
```

### Smart Query Modes

The system supports three query modes:

1. **Standard** (`/query`) - Basic vector search and retrieval
2. **Enhanced** (`/query-enhanced`) - Follows documentation references automatically
3. **Smart** (`/smart-query`) - Automatic classification and routing

Use the OpenAI-compatible API to access different modes:

```bash
# Standard mode
curl -X POST "http://localhost:8000/v1/chat/completions" \
  -H "Content-Type: application/json" \
  -d '{"model": "rag-standard", "messages": [{"role": "user", "content": "What is Dagster?"}]}'

# Enhanced mode with reference following
curl -X POST "http://localhost:8000/v1/chat/completions" \
  -H "Content-Type: application/json" \
  -d '{"model": "rag-enhanced", "messages": [{"role": "user", "content": "What is Dagster?"}]}'

# Smart mode with automatic routing
curl -X POST "http://localhost:8000/v1/chat/completions" \
  -H "Content-Type: application/json" \
  -d '{"model": "rag-smart", "messages": [{"role": "user", "content": "What is Dagster?"}]}'
```

### MCP Tools

The MCP server provides enhanced tools for Claude and other MCP clients:

**query_rag** - Query with optional tags and section filtering
```json
{
  "question": "How do I deploy?",
  "tags": ["dagster"],
  "section_path": "Deployment"
}
```

**smart_query** - Smart query with automatic routing
```json
{
  "question": "What is an asset and how do I use it?"
}
```

**add_document** - Upload with tags
```json
{
  "file_path": "/path/to/doc.md",
  "tags": ["dagster", "docs"]
}
```

**get_tags** - List all tags

**get_document_structure** - Get table of contents
```json
{
  "doc_id": "abc123"
}
```

## API Reference

### Enhanced Endpoints

**POST /documents**
- Body: `file` (multipart), `tags` (comma-separated string)
- Response: Document info with tags and chunk count

**POST /query**
- Body: `{"question": "...", "tags": [...], "section_path": "..."}`
- Response: Answer with section-aware sources

**POST /smart-query**
- Body: `{"question": "..."}`
- Response: Smart answer with automatic routing and classification

**GET /tags**
- Response: `{"tags": [...], "total": N}`

**GET /documents/{doc_id}/sections**
- Response: Document structure with section hierarchy

**GET /documents?tags=tag1,tag2**
- Query filtered by tags
- Response: List of matching documents

**POST /v1/chat/completions**
- OpenAI-compatible chat completion endpoint
- Supports models: `rag-standard`, `rag-enhanced`, `rag-smart`
- Supports streaming with `stream: true`

**GET /v1/models**
- List available RAG models

## Additional Documentation

- [QUICK_START.md](QUICK_START.md) - Quick setup guide for MCP integration
- [MCP_SETUP.md](MCP_SETUP.md) - Detailed MCP server setup
- [OPENAI_API_GUIDE.md](OPENAI_API_GUIDE.md) - OpenAI-compatible API documentation
- [QUERY_ROUTING_GUIDE.md](QUERY_ROUTING_GUIDE.md) - Smart query routing guide
- [MULTI_MODE_RETRIEVAL_GUIDE.md](MULTI_MODE_RETRIEVAL_GUIDE.md) - Multi-mode retrieval documentation
- [CODE_INDEX_GUIDE.md](CODE_INDEX_GUIDE.md) - Code indexing and search guide
- [RATE_LIMITING.md](RATE_LIMITING.md) - Rate limiting configuration
- [TEST_COVERAGE.md](TEST_COVERAGE.md) - Test coverage and testing guide

## License

MIT License

## Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

## Acknowledgments

- Google AI Studio for embeddings and LLM capabilities
- Qdrant for vector database
- FastAPI for the REST API framework
- Anthropic MCP for the Model Context Protocol