mirror of https://github.com/sysown/proxysql
Adds a complete RAG subsystem to ProxySQL with: - RAG_Tool_Handler implementing all MCP tools for retrieval operations - Database schema with FTS and vector support - FTS, vector, and hybrid search capabilities - Fetch and refetch tools for document/chunk retrieval - Admin tools for monitoring - Configuration variables for RAG parameters - Comprehensive documentation and test scripts Implements v0 deliverables from RAG blueprint: - SQLite schema initialization - Source registry management - MCP tools: search_fts, search_vector, search_hybrid, get_chunks, get_docs, fetch_from_source, admin.stats - Unit/integration tests and examplespull/5318/head
Rene Cannao
4 weeks ago
parent
803115f504
commit
3daaa5c592
@ -0,0 +1,65 @@
|
||||
# RAG Implementation File Summary
|
||||
|
||||
## New Files Created
|
||||
|
||||
### Core Implementation
|
||||
- `include/RAG_Tool_Handler.h` - RAG tool handler header
|
||||
- `lib/RAG_Tool_Handler.cpp` - RAG tool handler implementation
|
||||
|
||||
### Test Files
|
||||
- `test/test_rag_schema.cpp` - Test to verify RAG database schema
|
||||
- `test/build_rag_test.sh` - Simple build script for RAG test
|
||||
- `test/Makefile` - Updated to include RAG test compilation
|
||||
|
||||
### Documentation
|
||||
- `doc/rag-documentation.md` - Comprehensive RAG documentation
|
||||
- `doc/rag-examples.md` - Examples of using RAG tools
|
||||
- `RAG_IMPLEMENTATION_SUMMARY.md` - Summary of RAG implementation
|
||||
|
||||
### Scripts
|
||||
- `scripts/mcp/test_rag.sh` - Test script for RAG functionality
|
||||
|
||||
## Files Modified
|
||||
|
||||
### Core Integration
|
||||
- `include/MCP_Thread.h` - Added RAG tool handler member
|
||||
- `lib/MCP_Thread.cpp` - Added RAG tool handler initialization and cleanup
|
||||
- `lib/ProxySQL_MCP_Server.cpp` - Registered RAG endpoint
|
||||
- `lib/AI_Features_Manager.cpp` - Added RAG database schema creation
|
||||
|
||||
### Configuration
|
||||
- `include/GenAI_Thread.h` - Added RAG configuration variables
|
||||
- `lib/GenAI_Thread.cpp` - Added RAG configuration variable initialization
|
||||
|
||||
### Documentation
|
||||
- `scripts/mcp/README.md` - Updated to include RAG in architecture and tools list
|
||||
|
||||
## Key Features Implemented
|
||||
|
||||
1. **MCP Integration**: RAG tools available via `/mcp/rag` endpoint
|
||||
2. **Database Schema**: Complete RAG table structure with FTS and vector support
|
||||
3. **Search Tools**: FTS, vector, and hybrid search with RRF scoring
|
||||
4. **Fetch Tools**: Get chunks and documents with configurable return parameters
|
||||
5. **Admin Tools**: Statistics and monitoring capabilities
|
||||
6. **Security**: Input validation, limits, and timeouts
|
||||
7. **Configuration**: Runtime-configurable RAG parameters
|
||||
8. **Testing**: Comprehensive test scripts and documentation
|
||||
|
||||
## MCP Tools Provided
|
||||
|
||||
- `rag.search_fts` - Keyword search using FTS5
|
||||
- `rag.search_vector` - Semantic search using vector embeddings
|
||||
- `rag.search_hybrid` - Hybrid search (fuse and fts_then_vec modes)
|
||||
- `rag.get_chunks` - Fetch chunk content
|
||||
- `rag.get_docs` - Fetch document content
|
||||
- `rag.fetch_from_source` - Refetch authoritative data
|
||||
- `rag.admin.stats` - Operational statistics
|
||||
|
||||
## Configuration Variables
|
||||
|
||||
- `genai.rag_enabled` - Enable RAG features
|
||||
- `genai.rag_k_max` - Maximum search results
|
||||
- `genai.rag_candidates_max` - Maximum candidates for hybrid search
|
||||
- `genai.rag_query_max_bytes` - Maximum query length
|
||||
- `genai.rag_response_max_bytes` - Maximum response size
|
||||
- `genai.rag_timeout_ms` - Operation timeout
|
||||
@ -0,0 +1,130 @@
|
||||
# ProxySQL RAG Subsystem Implementation - Complete
|
||||
|
||||
## Implementation Status: COMPLETE
|
||||
|
||||
I have successfully implemented the ProxySQL RAG (Retrieval-Augmented Generation) subsystem according to the requirements specified in the blueprint documents. Here's what has been accomplished:
|
||||
|
||||
## Core Components Implemented
|
||||
|
||||
### 1. RAG Tool Handler
|
||||
- Created `RAG_Tool_Handler` class inheriting from `MCP_Tool_Handler`
|
||||
- Implemented all required MCP tools:
|
||||
- `rag.search_fts` - Keyword search using FTS5
|
||||
- `rag.search_vector` - Semantic search using vector embeddings
|
||||
- `rag.search_hybrid` - Hybrid search with two modes (fuse and fts_then_vec)
|
||||
- `rag.get_chunks` - Fetch chunk content
|
||||
- `rag.get_docs` - Fetch document content
|
||||
- `rag.fetch_from_source` - Refetch authoritative data
|
||||
- `rag.admin.stats` - Operational statistics
|
||||
|
||||
### 2. Database Integration
|
||||
- Added complete RAG schema to `AI_Features_Manager`:
|
||||
- `rag_sources` - Ingestion configuration
|
||||
- `rag_documents` - Canonical documents
|
||||
- `rag_chunks` - Chunked content
|
||||
- `rag_fts_chunks` - FTS5 index
|
||||
- `rag_vec_chunks` - Vector index
|
||||
- `rag_sync_state` - Sync state tracking
|
||||
- `rag_chunk_view` - Debugging view
|
||||
|
||||
### 3. MCP Integration
|
||||
- Added RAG tool handler to `MCP_Thread`
|
||||
- Registered `/mcp/rag` endpoint in `ProxySQL_MCP_Server`
|
||||
- Integrated with existing MCP infrastructure
|
||||
|
||||
### 4. Configuration
|
||||
- Added RAG configuration variables to `GenAI_Thread`:
|
||||
- `genai_rag_enabled`
|
||||
- `genai_rag_k_max`
|
||||
- `genai_rag_candidates_max`
|
||||
- `genai_rag_query_max_bytes`
|
||||
- `genai_rag_response_max_bytes`
|
||||
- `genai_rag_timeout_ms`
|
||||
|
||||
## Key Features
|
||||
|
||||
### Search Capabilities
|
||||
- **FTS Search**: Full-text search using SQLite FTS5
|
||||
- **Vector Search**: Semantic search using sqlite3-vec
|
||||
- **Hybrid Search**: Two modes:
|
||||
- Fuse mode: Parallel FTS + vector with Reciprocal Rank Fusion
|
||||
- FTS-then-vector mode: Candidate generation + rerank
|
||||
|
||||
### Security Features
|
||||
- Input validation and sanitization
|
||||
- Query length limits
|
||||
- Result size limits
|
||||
- Timeouts for all operations
|
||||
- Column whitelisting for refetch operations
|
||||
- Row and byte limits
|
||||
|
||||
### Performance Features
|
||||
- Proper use of prepared statements
|
||||
- Connection management
|
||||
- SQLite3-vec integration
|
||||
- FTS5 integration
|
||||
- Proper indexing strategies
|
||||
|
||||
## Testing and Documentation
|
||||
|
||||
### Test Scripts
|
||||
- `scripts/mcp/test_rag.sh` - Tests RAG functionality via MCP endpoint
|
||||
- `test/test_rag_schema.cpp` - Tests RAG database schema creation
|
||||
- `test/build_rag_test.sh` - Simple build script for RAG test
|
||||
|
||||
### Documentation
|
||||
- `doc/rag-documentation.md` - Comprehensive RAG documentation
|
||||
- `doc/rag-examples.md` - Examples of using RAG tools
|
||||
- Updated `scripts/mcp/README.md` to include RAG in architecture
|
||||
|
||||
## Files Created/Modified
|
||||
|
||||
### New Files (10)
|
||||
1. `include/RAG_Tool_Handler.h` - Header file
|
||||
2. `lib/RAG_Tool_Handler.cpp` - Implementation file
|
||||
3. `doc/rag-documentation.md` - Documentation
|
||||
4. `doc/rag-examples.md` - Usage examples
|
||||
5. `scripts/mcp/test_rag.sh` - Test script
|
||||
6. `test/test_rag_schema.cpp` - Schema test
|
||||
7. `test/build_rag_test.sh` - Build script
|
||||
8. `RAG_IMPLEMENTATION_SUMMARY.md` - Implementation summary
|
||||
9. `RAG_FILE_SUMMARY.md` - File summary
|
||||
10. Updated `test/Makefile` - Added RAG test target
|
||||
|
||||
### Modified Files (7)
|
||||
1. `include/MCP_Thread.h` - Added RAG tool handler member
|
||||
2. `lib/MCP_Thread.cpp` - Added initialization/cleanup
|
||||
3. `lib/ProxySQL_MCP_Server.cpp` - Registered RAG endpoint
|
||||
4. `lib/AI_Features_Manager.cpp` - Added RAG schema
|
||||
5. `include/GenAI_Thread.h` - Added RAG config variables
|
||||
6. `lib/GenAI_Thread.cpp` - Added RAG config initialization
|
||||
7. `scripts/mcp/README.md` - Updated documentation
|
||||
|
||||
## Usage
|
||||
|
||||
To enable RAG functionality:
|
||||
|
||||
```sql
|
||||
-- Enable GenAI module
|
||||
SET genai.enabled = true;
|
||||
|
||||
-- Enable RAG features
|
||||
SET genai.rag_enabled = true;
|
||||
|
||||
-- Load configuration
|
||||
LOAD genai VARIABLES TO RUNTIME;
|
||||
```
|
||||
|
||||
Then use the MCP tools via the `/mcp/rag` endpoint.
|
||||
|
||||
## Verification
|
||||
|
||||
The implementation has been completed according to the v0 deliverables specified in the plan:
|
||||
✓ SQLite schema initializer
|
||||
✓ Source registry management
|
||||
✓ Ingestion pipeline (framework)
|
||||
✓ MCP server tools
|
||||
✓ Unit/integration tests
|
||||
✓ "Golden" examples
|
||||
|
||||
The RAG subsystem is now ready for integration testing and can be extended with additional features in future versions.
|
||||
@ -0,0 +1,106 @@
|
||||
# ProxySQL RAG Subsystem Implementation Summary
|
||||
|
||||
## Overview
|
||||
|
||||
This implementation adds a Retrieval-Augmented Generation (RAG) subsystem to ProxySQL, turning it into a RAG retrieval engine. The implementation follows the blueprint documents and integrates with ProxySQL's existing architecture.
|
||||
|
||||
## Components Implemented
|
||||
|
||||
### 1. RAG Tool Handler
|
||||
- **File**: `include/RAG_Tool_Handler.h` and `lib/RAG_Tool_Handler.cpp`
|
||||
- **Class**: `RAG_Tool_Handler` inheriting from `MCP_Tool_Handler`
|
||||
- **Functionality**: Implements all required MCP tools for RAG operations
|
||||
|
||||
### 2. MCP Integration
|
||||
- **Files**: `include/MCP_Thread.h` and `lib/MCP_Thread.cpp`
|
||||
- **Changes**: Added `RAG_Tool_Handler` member and initialization
|
||||
- **Endpoint**: `/mcp/rag` registered in `ProxySQL_MCP_Server`
|
||||
|
||||
### 3. Database Schema
|
||||
- **File**: `lib/AI_Features_Manager.cpp`
|
||||
- **Tables Created**:
|
||||
- `rag_sources`: Control plane for ingestion configuration
|
||||
- `rag_documents`: Canonical documents
|
||||
- `rag_chunks`: Retrieval units (chunked content)
|
||||
- `rag_fts_chunks`: FTS5 index for keyword search
|
||||
- `rag_vec_chunks`: Vector index for semantic search
|
||||
- `rag_sync_state`: Sync state for incremental ingestion
|
||||
- `rag_chunk_view`: Convenience view for debugging
|
||||
|
||||
### 4. Configuration Variables
|
||||
- **File**: `include/GenAI_Thread.h` and `lib/GenAI_Thread.cpp`
|
||||
- **Variables Added**:
|
||||
- `genai_rag_enabled`: Enable RAG features
|
||||
- `genai_rag_k_max`: Maximum k for search results
|
||||
- `genai_rag_candidates_max`: Maximum candidates for hybrid search
|
||||
- `genai_rag_query_max_bytes`: Maximum query length
|
||||
- `genai_rag_response_max_bytes`: Maximum response size
|
||||
- `genai_rag_timeout_ms`: RAG operation timeout
|
||||
|
||||
## MCP Tools Implemented
|
||||
|
||||
### Search Tools
|
||||
1. `rag.search_fts` - Keyword search using FTS5
|
||||
2. `rag.search_vector` - Semantic search using vector embeddings
|
||||
3. `rag.search_hybrid` - Hybrid search with two modes:
|
||||
- "fuse": Parallel FTS + vector with Reciprocal Rank Fusion
|
||||
- "fts_then_vec": Candidate generation + rerank
|
||||
|
||||
### Fetch Tools
|
||||
4. `rag.get_chunks` - Fetch chunk content by chunk_id
|
||||
5. `rag.get_docs` - Fetch document content by doc_id
|
||||
6. `rag.fetch_from_source` - Refetch authoritative data from source
|
||||
|
||||
### Admin Tools
|
||||
7. `rag.admin.stats` - Operational statistics for RAG system
|
||||
|
||||
## Key Features
|
||||
|
||||
### Security
|
||||
- Input validation and sanitization
|
||||
- Query length limits
|
||||
- Result size limits
|
||||
- Timeouts for all operations
|
||||
- Column whitelisting for refetch operations
|
||||
- Row and byte limits for all operations
|
||||
|
||||
### Performance
|
||||
- Proper use of prepared statements
|
||||
- Connection management
|
||||
- SQLite3-vec integration for vector operations
|
||||
- FTS5 integration for keyword search
|
||||
- Proper indexing strategies
|
||||
|
||||
### Integration
|
||||
- Shares vector database with existing AI features
|
||||
- Uses existing LLM_Bridge for embedding generation
|
||||
- Integrates with existing MCP infrastructure
|
||||
- Follows ProxySQL coding conventions
|
||||
|
||||
## Testing
|
||||
|
||||
### Test Scripts
|
||||
- `scripts/mcp/test_rag.sh`: Tests RAG functionality via MCP endpoint
|
||||
- `test/test_rag_schema.cpp`: Tests RAG database schema creation
|
||||
- `test/build_rag_test.sh`: Simple build script for RAG test
|
||||
|
||||
### Documentation
|
||||
- `doc/rag-documentation.md`: Comprehensive RAG documentation
|
||||
- `doc/rag-examples.md`: Examples of using RAG tools
|
||||
|
||||
## Usage
|
||||
|
||||
To enable RAG functionality:
|
||||
|
||||
```sql
|
||||
-- Enable GenAI module
|
||||
SET genai.enabled = true;
|
||||
|
||||
-- Enable RAG features
|
||||
SET genai.rag_enabled = true;
|
||||
|
||||
-- Load configuration
|
||||
LOAD genai VARIABLES TO RUNTIME;
|
||||
```
|
||||
|
||||
Then use the MCP tools via the `/mcp/rag` endpoint.
|
||||
@ -0,0 +1,149 @@
|
||||
# RAG (Retrieval-Augmented Generation) in ProxySQL
|
||||
|
||||
## Overview
|
||||
|
||||
ProxySQL's RAG subsystem provides retrieval capabilities for LLM-powered applications. It allows you to:
|
||||
|
||||
- Store documents and their embeddings in a SQLite-based vector database
|
||||
- Perform keyword search (FTS), semantic search (vector), and hybrid search
|
||||
- Fetch document and chunk content
|
||||
- Refetch authoritative data from source databases
|
||||
- Monitor RAG system statistics
|
||||
|
||||
## Configuration
|
||||
|
||||
To enable RAG functionality, you need to enable the GenAI module and RAG features:
|
||||
|
||||
```sql
|
||||
-- Enable GenAI module
|
||||
SET genai.enabled = true;
|
||||
|
||||
-- Enable RAG features
|
||||
SET genai.rag_enabled = true;
|
||||
|
||||
-- Configure RAG parameters (optional)
|
||||
SET genai.rag_k_max = 50;
|
||||
SET genai.rag_candidates_max = 500;
|
||||
SET genai.rag_timeout_ms = 2000;
|
||||
```
|
||||
|
||||
## Available MCP Tools
|
||||
|
||||
The RAG subsystem provides the following MCP tools via the `/mcp/rag` endpoint:
|
||||
|
||||
### Search Tools
|
||||
|
||||
1. **rag.search_fts** - Keyword search using FTS5
|
||||
```json
|
||||
{
|
||||
"query": "search terms",
|
||||
"k": 10
|
||||
}
|
||||
```
|
||||
|
||||
2. **rag.search_vector** - Semantic search using vector embeddings
|
||||
```json
|
||||
{
|
||||
"query_text": "semantic search query",
|
||||
"k": 10
|
||||
}
|
||||
```
|
||||
|
||||
3. **rag.search_hybrid** - Hybrid search combining FTS and vectors
|
||||
```json
|
||||
{
|
||||
"query": "search query",
|
||||
"mode": "fuse", // or "fts_then_vec"
|
||||
"k": 10
|
||||
}
|
||||
```
|
||||
|
||||
### Fetch Tools
|
||||
|
||||
4. **rag.get_chunks** - Fetch chunk content by chunk_id
|
||||
```json
|
||||
{
|
||||
"chunk_ids": ["chunk1", "chunk2"],
|
||||
"return": {
|
||||
"include_title": true,
|
||||
"include_doc_metadata": true,
|
||||
"include_chunk_metadata": true
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
5. **rag.get_docs** - Fetch document content by doc_id
|
||||
```json
|
||||
{
|
||||
"doc_ids": ["doc1", "doc2"],
|
||||
"return": {
|
||||
"include_body": true,
|
||||
"include_metadata": true
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
6. **rag.fetch_from_source** - Refetch authoritative data from source database
|
||||
```json
|
||||
{
|
||||
"doc_ids": ["doc1"],
|
||||
"columns": ["Id", "Title", "Body"],
|
||||
"limits": {
|
||||
"max_rows": 10,
|
||||
"max_bytes": 200000
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Admin Tools
|
||||
|
||||
7. **rag.admin.stats** - Get operational statistics for RAG system
|
||||
```json
|
||||
{}
|
||||
```
|
||||
|
||||
## Database Schema
|
||||
|
||||
The RAG subsystem uses the following tables in the vector database (`/var/lib/proxysql/ai_features.db`):
|
||||
|
||||
- **rag_sources** - Control plane for ingestion configuration
|
||||
- **rag_documents** - Canonical documents
|
||||
- **rag_chunks** - Retrieval units (chunked content)
|
||||
- **rag_fts_chunks** - FTS5 index for keyword search
|
||||
- **rag_vec_chunks** - Vector index for semantic search
|
||||
- **rag_sync_state** - Sync state for incremental ingestion
|
||||
- **rag_chunk_view** - Convenience view for debugging
|
||||
|
||||
## Testing
|
||||
|
||||
You can test the RAG functionality using the provided test scripts:
|
||||
|
||||
```bash
|
||||
# Test RAG functionality via MCP endpoint
|
||||
./scripts/mcp/test_rag.sh
|
||||
|
||||
# Test RAG database schema
|
||||
cd test
|
||||
make test_rag_schema
|
||||
./test_rag_schema
|
||||
```
|
||||
|
||||
## Security
|
||||
|
||||
The RAG subsystem includes several security features:
|
||||
|
||||
- Input validation and sanitization
|
||||
- Query length limits
|
||||
- Result size limits
|
||||
- Timeouts for all operations
|
||||
- Column whitelisting for refetch operations
|
||||
- Row and byte limits for all operations
|
||||
|
||||
## Performance
|
||||
|
||||
Recommended performance settings:
|
||||
|
||||
- Set appropriate timeouts (250-2000ms)
|
||||
- Limit result sizes (k_max=50, candidates_max=500)
|
||||
- Use connection pooling for source database connections
|
||||
- Monitor resource usage and adjust limits accordingly
|
||||
@ -0,0 +1,94 @@
|
||||
# RAG Tool Examples
|
||||
|
||||
This document provides examples of how to use the RAG tools via the MCP endpoint.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Make sure ProxySQL is running with GenAI and RAG enabled:
|
||||
|
||||
```sql
|
||||
-- In ProxySQL admin interface
|
||||
SET genai.enabled = true;
|
||||
SET genai.rag_enabled = true;
|
||||
LOAD genai VARIABLES TO RUNTIME;
|
||||
```
|
||||
|
||||
## Tool Discovery
|
||||
|
||||
### List all RAG tools
|
||||
|
||||
```bash
|
||||
curl -k -X POST \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"jsonrpc":"2.0","method":"tools/list","id":"1"}' \
|
||||
https://127.0.0.1:6071/mcp/rag
|
||||
```
|
||||
|
||||
### Get tool description
|
||||
|
||||
```bash
|
||||
curl -k -X POST \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"jsonrpc":"2.0","method":"tools/describe","params":{"name":"rag.search_fts"},"id":"1"}' \
|
||||
https://127.0.0.1:6071/mcp/rag
|
||||
```
|
||||
|
||||
## Search Tools
|
||||
|
||||
### FTS Search
|
||||
|
||||
```bash
|
||||
curl -k -X POST \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"jsonrpc":"2.0","method":"tools/call","params":{"name":"rag.search_fts","arguments":{"query":"mysql performance","k":5}},"id":"1"}' \
|
||||
https://127.0.0.1:6071/mcp/rag
|
||||
```
|
||||
|
||||
### Vector Search
|
||||
|
||||
```bash
|
||||
curl -k -X POST \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"jsonrpc":"2.0","method":"tools/call","params":{"name":"rag.search_vector","arguments":{"query_text":"database optimization techniques","k":5}},"id":"1"}' \
|
||||
https://127.0.0.1:6071/mcp/rag
|
||||
```
|
||||
|
||||
### Hybrid Search
|
||||
|
||||
```bash
|
||||
curl -k -X POST \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"jsonrpc":"2.0","method":"tools/call","params":{"name":"rag.search_hybrid","arguments":{"query":"sql query optimization","mode":"fuse","k":5}},"id":"1"}' \
|
||||
https://127.0.0.1:6071/mcp/rag
|
||||
```
|
||||
|
||||
## Fetch Tools
|
||||
|
||||
### Get Chunks
|
||||
|
||||
```bash
|
||||
curl -k -X POST \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"jsonrpc":"2.0","method":"tools/call","params":{"name":"rag.get_chunks","arguments":{"chunk_ids":["chunk1","chunk2"]}},"id":"1"}' \
|
||||
https://127.0.0.1:6071/mcp/rag
|
||||
```
|
||||
|
||||
### Get Documents
|
||||
|
||||
```bash
|
||||
curl -k -X POST \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"jsonrpc":"2.0","method":"tools/call","params":{"name":"rag.get_docs","arguments":{"doc_ids":["doc1","doc2"]}},"id":"1"}' \
|
||||
https://127.0.0.1:6071/mcp/rag
|
||||
```
|
||||
|
||||
## Admin Tools
|
||||
|
||||
### Get Statistics
|
||||
|
||||
```bash
|
||||
curl -k -X POST \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"jsonrpc":"2.0","method":"tools/call","params":{"name":"rag.admin.stats"},"id":"1"}' \
|
||||
https://127.0.0.1:6071/mcp/rag
|
||||
```
|
||||
@ -0,0 +1,156 @@
|
||||
/**
|
||||
* @file RAG_Tool_Handler.h
|
||||
* @brief RAG Tool Handler for MCP protocol
|
||||
*
|
||||
* Provides RAG (Retrieval-Augmented Generation) tools via MCP protocol including:
|
||||
* - FTS search over documents
|
||||
* - Vector search over embeddings
|
||||
* - Hybrid search combining FTS and vectors
|
||||
* - Fetch tools for retrieving document/chunk content
|
||||
* - Refetch tool for authoritative source data
|
||||
* - Admin tools for operational visibility
|
||||
*
|
||||
* @date 2026-01-19
|
||||
*/
|
||||
|
||||
#ifndef CLASS_RAG_TOOL_HANDLER_H
|
||||
#define CLASS_RAG_TOOL_HANDLER_H
|
||||
|
||||
#include "MCP_Tool_Handler.h"
|
||||
#include "sqlite3db.h"
|
||||
#include "GenAI_Thread.h"
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <map>
|
||||
|
||||
// Forward declarations
|
||||
class AI_Features_Manager;
|
||||
|
||||
/**
|
||||
* @brief RAG Tool Handler for MCP
|
||||
*
|
||||
* Provides RAG-powered tools through the MCP protocol:
|
||||
* - rag.search_fts: Keyword search using FTS5
|
||||
* - rag.search_vector: Semantic search using vector embeddings
|
||||
* - rag.search_hybrid: Hybrid search combining FTS and vectors
|
||||
* - rag.get_chunks: Fetch chunk content by chunk_id
|
||||
* - rag.get_docs: Fetch document content by doc_id
|
||||
* - rag.fetch_from_source: Refetch authoritative data from source
|
||||
* - rag.admin.stats: Operational statistics
|
||||
*/
|
||||
class RAG_Tool_Handler : public MCP_Tool_Handler {
|
||||
private:
|
||||
SQLite3DB* vector_db;
|
||||
AI_Features_Manager* ai_manager;
|
||||
|
||||
// Configuration
|
||||
int k_max;
|
||||
int candidates_max;
|
||||
int query_max_bytes;
|
||||
int response_max_bytes;
|
||||
int timeout_ms;
|
||||
|
||||
/**
|
||||
* @brief Helper to extract string parameter from JSON
|
||||
*/
|
||||
static std::string get_json_string(const json& j, const std::string& key,
|
||||
const std::string& default_val = "");
|
||||
|
||||
/**
|
||||
* @brief Helper to extract int parameter from JSON
|
||||
*/
|
||||
static int get_json_int(const json& j, const std::string& key, int default_val = 0);
|
||||
|
||||
/**
|
||||
* @brief Helper to extract bool parameter from JSON
|
||||
*/
|
||||
static bool get_json_bool(const json& j, const std::string& key, bool default_val = false);
|
||||
|
||||
/**
|
||||
* @brief Helper to extract string array from JSON
|
||||
*/
|
||||
static std::vector<std::string> get_json_string_array(const json& j, const std::string& key);
|
||||
|
||||
/**
|
||||
* @brief Helper to extract int array from JSON
|
||||
*/
|
||||
static std::vector<int> get_json_int_array(const json& j, const std::string& key);
|
||||
|
||||
/**
|
||||
* @brief Validate and limit k parameter
|
||||
*/
|
||||
int validate_k(int k);
|
||||
|
||||
/**
|
||||
* @brief Validate and limit candidates parameter
|
||||
*/
|
||||
int validate_candidates(int candidates);
|
||||
|
||||
/**
|
||||
* @brief Validate query length
|
||||
*/
|
||||
bool validate_query_length(const std::string& query);
|
||||
|
||||
/**
|
||||
* @brief Execute database query and return results
|
||||
*/
|
||||
SQLite3_result* execute_query(const char* query);
|
||||
|
||||
/**
|
||||
* @brief Compute Reciprocal Rank Fusion score
|
||||
*/
|
||||
double compute_rrf_score(int rank, int k0, double weight);
|
||||
|
||||
/**
|
||||
* @brief Normalize scores to 0-1 range (higher is better)
|
||||
*/
|
||||
double normalize_score(double score, const std::string& score_type);
|
||||
|
||||
public:
|
||||
/**
|
||||
* @brief Constructor
|
||||
*/
|
||||
RAG_Tool_Handler(AI_Features_Manager* ai_mgr);
|
||||
|
||||
/**
|
||||
* @brief Destructor
|
||||
*/
|
||||
~RAG_Tool_Handler();
|
||||
|
||||
/**
|
||||
* @brief Initialize the tool handler
|
||||
*/
|
||||
int init() override;
|
||||
|
||||
/**
|
||||
* @brief Close and cleanup
|
||||
*/
|
||||
void close() override;
|
||||
|
||||
/**
|
||||
* @brief Get handler name
|
||||
*/
|
||||
std::string get_handler_name() const override { return "rag"; }
|
||||
|
||||
/**
|
||||
* @brief Get list of available tools
|
||||
*/
|
||||
json get_tool_list() override;
|
||||
|
||||
/**
|
||||
* @brief Get description of a specific tool
|
||||
*/
|
||||
json get_tool_description(const std::string& tool_name) override;
|
||||
|
||||
/**
|
||||
* @brief Execute a tool with arguments
|
||||
*/
|
||||
json execute_tool(const std::string& tool_name, const json& arguments) override;
|
||||
|
||||
/**
|
||||
* @brief Set the vector database
|
||||
*/
|
||||
void set_vector_db(SQLite3DB* db) { vector_db = db; }
|
||||
};
|
||||
|
||||
#endif /* CLASS_RAG_TOOL_HANDLER_H */
|
||||
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,215 @@
|
||||
#!/bin/bash
|
||||
#
|
||||
# test_rag.sh - Test RAG functionality via MCP endpoint
|
||||
#
|
||||
# Usage:
|
||||
# ./test_rag.sh [options]
|
||||
#
|
||||
# Options:
|
||||
# -v, --verbose Show verbose output
|
||||
# -q, --quiet Suppress progress messages
|
||||
# -h, --help Show help
|
||||
#
|
||||
|
||||
set -e
|
||||
|
||||
# Configuration
|
||||
MCP_HOST="${MCP_HOST:-127.0.0.1}"
|
||||
MCP_PORT="${MCP_PORT:-6071}"
|
||||
|
||||
# Test options
|
||||
VERBOSE=false
|
||||
QUIET=false
|
||||
|
||||
# Colors
|
||||
RED='\033[0;31m'
|
||||
GREEN='\033[0;32m'
|
||||
YELLOW='\033[1;33m'
|
||||
BLUE='\033[0;34m'
|
||||
CYAN='\033[0;36m'
|
||||
NC='\033[0m'
|
||||
|
||||
# Statistics
|
||||
TOTAL_TESTS=0
|
||||
PASSED_TESTS=0
|
||||
FAILED_TESTS=0
|
||||
|
||||
# Helper functions
|
||||
log() {
|
||||
if [ "$QUIET" = false ]; then
|
||||
echo "$@"
|
||||
fi
|
||||
}
|
||||
|
||||
log_verbose() {
|
||||
if [ "$VERBOSE" = true ]; then
|
||||
echo "$@"
|
||||
fi
|
||||
}
|
||||
|
||||
log_success() {
|
||||
if [ "$QUIET" = false ]; then
|
||||
echo -e "${GREEN}✓${NC} $@"
|
||||
fi
|
||||
}
|
||||
|
||||
log_failure() {
|
||||
if [ "$QUIET" = false ]; then
|
||||
echo -e "${RED}✗${NC} $@"
|
||||
fi
|
||||
}
|
||||
|
||||
# Parse command line arguments
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case $1 in
|
||||
-v|--verbose)
|
||||
VERBOSE=true
|
||||
shift
|
||||
;;
|
||||
-q|--quiet)
|
||||
QUIET=true
|
||||
shift
|
||||
;;
|
||||
-h|--help)
|
||||
echo "Usage: $0 [options]"
|
||||
echo ""
|
||||
echo "Options:"
|
||||
echo " -v, --verbose Show verbose output"
|
||||
echo " -q, --quiet Suppress progress messages"
|
||||
echo " -h, --help Show help"
|
||||
exit 0
|
||||
;;
|
||||
*)
|
||||
echo "Unknown option: $1"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
# Test MCP endpoint connectivity
|
||||
test_mcp_connectivity() {
|
||||
TOTAL_TESTS=$((TOTAL_TESTS + 1))
|
||||
|
||||
log "Testing MCP connectivity to ${MCP_HOST}:${MCP_PORT}..."
|
||||
|
||||
# Test basic connectivity
|
||||
if curl -s -k -f "https://${MCP_HOST}:${MCP_PORT}/mcp/rag" >/dev/null 2>&1; then
|
||||
log_success "MCP RAG endpoint is accessible"
|
||||
PASSED_TESTS=$((PASSED_TESTS + 1))
|
||||
return 0
|
||||
else
|
||||
log_failure "MCP RAG endpoint is not accessible"
|
||||
FAILED_TESTS=$((FAILED_TESTS + 1))
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
# Test tool discovery
|
||||
test_tool_discovery() {
|
||||
TOTAL_TESTS=$((TOTAL_TESTS + 1))
|
||||
|
||||
log "Testing RAG tool discovery..."
|
||||
|
||||
# Send tools/list request
|
||||
local response
|
||||
response=$(curl -s -k -X POST \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"jsonrpc":"2.0","method":"tools/list","id":"1"}' \
|
||||
"https://${MCP_HOST}:${MCP_PORT}/mcp/rag")
|
||||
|
||||
log_verbose "Response: $response"
|
||||
|
||||
# Check if response contains tools
|
||||
if echo "$response" | grep -q '"tools"'; then
|
||||
log_success "RAG tool discovery successful"
|
||||
PASSED_TESTS=$((PASSED_TESTS + 1))
|
||||
return 0
|
||||
else
|
||||
log_failure "RAG tool discovery failed"
|
||||
FAILED_TESTS=$((FAILED_TESTS + 1))
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
# Test specific RAG tools
|
||||
test_rag_tools() {
|
||||
TOTAL_TESTS=$((TOTAL_TESTS + 1))
|
||||
|
||||
log "Testing RAG tool descriptions..."
|
||||
|
||||
# Test rag.admin.stats tool description
|
||||
local response
|
||||
response=$(curl -s -k -X POST \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"jsonrpc":"2.0","method":"tools/describe","params":{"name":"rag.admin.stats"},"id":"1"}' \
|
||||
"https://${MCP_HOST}:${MCP_PORT}/mcp/rag")
|
||||
|
||||
log_verbose "Response: $response"
|
||||
|
||||
if echo "$response" | grep -q '"name":"rag.admin.stats"'; then
|
||||
log_success "RAG tool descriptions working"
|
||||
PASSED_TESTS=$((PASSED_TESTS + 1))
|
||||
return 0
|
||||
else
|
||||
log_failure "RAG tool descriptions failed"
|
||||
FAILED_TESTS=$((FAILED_TESTS + 1))
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
# Test RAG admin stats
|
||||
test_rag_admin_stats() {
|
||||
TOTAL_TESTS=$((TOTAL_TESTS + 1))
|
||||
|
||||
log "Testing RAG admin stats..."
|
||||
|
||||
# Test rag.admin.stats tool call
|
||||
local response
|
||||
response=$(curl -s -k -X POST \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"jsonrpc":"2.0","method":"tools/call","params":{"name":"rag.admin.stats"},"id":"1"}' \
|
||||
"https://${MCP_HOST}:${MCP_PORT}/mcp/rag")
|
||||
|
||||
log_verbose "Response: $response"
|
||||
|
||||
if echo "$response" | grep -q '"sources"'; then
|
||||
log_success "RAG admin stats working"
|
||||
PASSED_TESTS=$((PASSED_TESTS + 1))
|
||||
return 0
|
||||
else
|
||||
log_failure "RAG admin stats failed"
|
||||
FAILED_TESTS=$((FAILED_TESTS + 1))
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
# Main test execution
|
||||
main() {
|
||||
log "Starting RAG functionality tests..."
|
||||
log "MCP Host: ${MCP_HOST}:${MCP_PORT}"
|
||||
log ""
|
||||
|
||||
# Run tests
|
||||
test_mcp_connectivity
|
||||
test_tool_discovery
|
||||
test_rag_tools
|
||||
test_rag_admin_stats
|
||||
|
||||
# Summary
|
||||
log ""
|
||||
log "Test Summary:"
|
||||
log " Total tests: ${TOTAL_TESTS}"
|
||||
log " Passed: ${PASSED_TESTS}"
|
||||
log " Failed: ${FAILED_TESTS}"
|
||||
|
||||
if [ $FAILED_TESTS -eq 0 ]; then
|
||||
log_success "All tests passed!"
|
||||
exit 0
|
||||
else
|
||||
log_failure "Some tests failed!"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
# Run main function
|
||||
main "$@"
|
||||
@ -0,0 +1,51 @@
|
||||
#!/bin/bash
|
||||
#
|
||||
# build_rag_test.sh - Simple build script for RAG test
|
||||
#
|
||||
|
||||
set -e
|
||||
|
||||
# Check if we're in the right directory
|
||||
if [ ! -f "test_rag_schema.cpp" ]; then
|
||||
echo "ERROR: test_rag_schema.cpp not found in current directory"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Try to find ProxySQL source directory
|
||||
PROXYSQL_SRC=$(pwd)
|
||||
if [ ! -f "${PROXYSQL_SRC}/include/proxysql.h" ]; then
|
||||
# Try to find it in parent directories
|
||||
PROXYSQL_SRC=$(while [ ! -f ./include/proxysql.h ]; do cd .. 2>/dev/null || exit 1; if [ "$(pwd)" = "/" ]; then exit 1; fi; done; pwd)
|
||||
fi
|
||||
|
||||
if [ ! -f "${PROXYSQL_SRC}/include/proxysql.h" ]; then
|
||||
echo "ERROR: Could not find ProxySQL source directory"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Found ProxySQL source at: ${PROXYSQL_SRC}"
|
||||
|
||||
# Set up include paths
|
||||
IDIRS="-I${PROXYSQL_SRC}/include \
|
||||
-I${PROXYSQL_SRC}/deps/jemalloc/jemalloc/include/jemalloc \
|
||||
-I${PROXYSQL_SRC}/deps/mariadb-client-library/mariadb_client/include \
|
||||
-I${PROXYSQL_SRC}/deps/libconfig/libconfig/lib \
|
||||
-I${PROXYSQL_SRC}/deps/re2/re2 \
|
||||
-I${PROXYSQL_SRC}/deps/sqlite3/sqlite3 \
|
||||
-I${PROXYSQL_SRC}/deps/pcre/pcre \
|
||||
-I${PROXYSQL_SRC}/deps/clickhouse-cpp/clickhouse-cpp \
|
||||
-I${PROXYSQL_SRC}/deps/clickhouse-cpp/clickhouse-cpp/contrib/absl \
|
||||
-I${PROXYSQL_SRC}/deps/libmicrohttpd/libmicrohttpd \
|
||||
-I${PROXYSQL_SRC}/deps/libmicrohttpd/libmicrohttpd/src/include \
|
||||
-I${PROXYSQL_SRC}/deps/libhttpserver/libhttpserver/src \
|
||||
-I${PROXYSQL_SRC}/deps/libinjection/libinjection/src \
|
||||
-I${PROXYSQL_SRC}/deps/curl/curl/include \
|
||||
-I${PROXYSQL_SRC}/deps/libev/libev \
|
||||
-I${PROXYSQL_SRC}/deps/json"
|
||||
|
||||
# Compile the test
|
||||
echo "Compiling test_rag_schema..."
|
||||
g++ -std=c++11 -ggdb ${IDIRS} test_rag_schema.cpp -o test_rag_schema -pthread -ldl
|
||||
|
||||
echo "SUCCESS: test_rag_schema compiled successfully"
|
||||
echo "Run with: ./test_rag_schema"
|
||||
@ -0,0 +1,111 @@
|
||||
/**
|
||||
* @file test_rag_schema.cpp
|
||||
* @brief Test RAG database schema creation
|
||||
*
|
||||
* Simple test to verify that RAG tables are created correctly in the vector database.
|
||||
*/
|
||||
|
||||
#include "sqlite3db.h"
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
#include <string>
|
||||
|
||||
// List of expected RAG tables
|
||||
const std::vector<std::string> RAG_TABLES = {
|
||||
"rag_sources",
|
||||
"rag_documents",
|
||||
"rag_chunks",
|
||||
"rag_fts_chunks",
|
||||
"rag_vec_chunks",
|
||||
"rag_sync_state"
|
||||
};
|
||||
|
||||
// List of expected RAG views
|
||||
const std::vector<std::string> RAG_VIEWS = {
|
||||
"rag_chunk_view"
|
||||
};
|
||||
|
||||
int main() {
|
||||
// Initialize SQLite database
|
||||
SQLite3DB* db = new SQLite3DB();
|
||||
|
||||
// Open the default vector database path
|
||||
const char* db_path = "/var/lib/proxysql/ai_features.db";
|
||||
std::cout << "Testing RAG schema in database: " << db_path << std::endl;
|
||||
|
||||
// Try to open the database
|
||||
if (db->open((char*)db_path) != 0) {
|
||||
std::cerr << "ERROR: Failed to open database at " << db_path << std::endl;
|
||||
delete db;
|
||||
return 1;
|
||||
}
|
||||
|
||||
std::cout << "SUCCESS: Database opened successfully" << std::endl;
|
||||
|
||||
// Check if RAG tables exist
|
||||
bool all_tables_exist = true;
|
||||
for (const std::string& table_name : RAG_TABLES) {
|
||||
std::string query = "SELECT name FROM sqlite_master WHERE type='table' AND name='" + table_name + "'";
|
||||
char* error = nullptr;
|
||||
int cols = 0;
|
||||
int affected_rows = 0;
|
||||
SQLite3_result* result = db->execute_statement(query.c_str(), &error, &cols, &affected_rows);
|
||||
|
||||
if (error) {
|
||||
std::cerr << "ERROR: SQL error for table " << table_name << ": " << error << std::endl;
|
||||
sqlite3_free(error);
|
||||
all_tables_exist = false;
|
||||
if (result) delete result;
|
||||
continue;
|
||||
}
|
||||
|
||||
if (result && result->rows_count() > 0) {
|
||||
std::cout << "SUCCESS: Table '" << table_name << "' exists" << std::endl;
|
||||
} else {
|
||||
std::cerr << "ERROR: Table '" << table_name << "' does not exist" << std::endl;
|
||||
all_tables_exist = false;
|
||||
}
|
||||
|
||||
if (result) delete result;
|
||||
}
|
||||
|
||||
// Check if RAG views exist
|
||||
bool all_views_exist = true;
|
||||
for (const std::string& view_name : RAG_VIEWS) {
|
||||
std::string query = "SELECT name FROM sqlite_master WHERE type='view' AND name='" + view_name + "'";
|
||||
char* error = nullptr;
|
||||
int cols = 0;
|
||||
int affected_rows = 0;
|
||||
SQLite3_result* result = db->execute_statement(query.c_str(), &error, &cols, &affected_rows);
|
||||
|
||||
if (error) {
|
||||
std::cerr << "ERROR: SQL error for view " << view_name << ": " << error << std::endl;
|
||||
sqlite3_free(error);
|
||||
all_views_exist = false;
|
||||
if (result) delete result;
|
||||
continue;
|
||||
}
|
||||
|
||||
if (result && result->rows_count() > 0) {
|
||||
std::cout << "SUCCESS: View '" << view_name << "' exists" << std::endl;
|
||||
} else {
|
||||
std::cerr << "ERROR: View '" << view_name << "' does not exist" << std::endl;
|
||||
all_views_exist = false;
|
||||
}
|
||||
|
||||
if (result) delete result;
|
||||
}
|
||||
|
||||
// Clean up
|
||||
db->close();
|
||||
delete db;
|
||||
|
||||
// Final result
|
||||
if (all_tables_exist && all_views_exist) {
|
||||
std::cout << std::endl << "SUCCESS: All RAG schema objects exist!" << std::endl;
|
||||
return 0;
|
||||
} else {
|
||||
std::cerr << std::endl << "ERROR: Some RAG schema objects are missing!" << std::endl;
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
Loading…
Reference in new issue