RAG Module - Rust Implementation

High-performance Rust implementation of the Enterprise RAG module with chat context storage, vector search, session management, and automatic model downloading (like Node.js Transformers).

🚀 Features

🤖 Model Management

• Automatic Model Downloading: Downloads models from Hugging Face Hub to ./models/
• Local Model Caching: Efficient caching system for transformer models
• Fallback System: BGE-M3 → MiniLM → MPNet fallback chain
• Directory Structure: Auto-creates models/, cache/, data/, keys/ directories

🔍 Core RAG Capabilities

• Vector Search: Embedded Qdrant and local file-based vector stores
• Multi-Cloud Support: AWS, Azure, GCP estate data management
• Encryption: AES-256-GCM encryption for sensitive data
• Chat Context: Complete chat history retrieval and management
• Session Management: Persistent chat sessions with context tracking

⚡ Performance & Compatibility

• API Compatibility: HTTP API matching the Node.js module interface
• Performance: Rust's memory safety and zero-cost abstractions
• Privacy: Configurable data filtering and anonymization
• Node.js Style: Familiar API patterns for Node.js developers

Architecture

Core Components

src/
├── lib.rs                  # Main RAG module
├── types/                  # Type definitions
├── config/                 # Configuration management
├── db/                     # Vector store implementations
│   ├── vector_store.rs     # VectorStore trait
│   ├── embedded_qdrant.rs  # Embedded Qdrant implementation
│   └── local_file_store.rs # Local file storage
├── services/               # Business logic services
│   ├── embedding_service.rs
│   ├── encryption_service.rs
│   ├── document_service.rs
│   ├── search_service.rs
│   └── [other services...]
└── bin/
    └── server.rs           # HTTP API server

Key Features Converted

1. RagModule.js → lib.rs: Main module with all services
2. EmbeddedQdrantVectorStore.js → embedded_qdrant.rs: File-based vector storage
3. EncryptionService.js → encryption_service.rs: AES-256-GCM encryption
4. ConfigManager.js → config/mod.rs: YAML/JSON configuration
5. All Services → services/ directory: Complete business logic

📦 Installation

Add to your Cargo.toml:

[dependencies]
rag-module = "0.1"
tokio = { version = "1.0", features = ["full"] }

🏃 Quick Start

🤖 Automatic Model Setup (Like Node.js)

use rag_module::RagModule;

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // Creates directory structure like Node.js
    let rag = RagModule::new("./rag-data").await?;

    // This automatically:
    // 1. Creates ./rag-data/models/ directory
    // 2. Downloads BGE-M3 from Hugging Face (or fallback to MiniLM)
    // 3. Caches model locally for future use
    // 4. Sets up encryption keys in ./rag-data/keys/
    rag.initialize().await?;

    println!("✅ RAG Module initialized with model caching!");

    // Check what got downloaded
    let model_info = rag.embedding_service.get_model_info().await?;
    let storage_info = rag.embedding_service.get_storage_info().await?;

    println!("Model: {}", model_info["name"]);
    println!("Cached size: {}", storage_info["totalSizeFormatted"]);

    Ok(())
}

📄 Document Management

use rag_module::{RagModule, types::*};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // Initialize RAG module
    let rag = RagModule::new("./rag-data").await?;
    rag.initialize().await?;
    
    // Add document
    let doc = Document::new(
        "doc-1".to_string(),
        "AWS user permissions for EC2 and RDS".to_string()
    );
    let doc_id = rag.add_document("aws_estate", doc).await?;
    
    // Search
    let results = rag.search(
        "aws_estate",
        "EC2 permissions",
        SearchOptions::default()
    ).await?;
    
    Ok(())
}

As an HTTP Server

# Run the server
cargo run --bin rag-server

# Server runs on http://127.0.0.1:3000

API Endpoints (Node.js Compatible)

# Health check
GET /health

# Documents
POST /api/documents
GET /api/documents/:collection/:id
PUT /api/documents/:collection/:id
DELETE /api/documents/:collection/:id

# Search
POST /api/search

# Chat
POST /api/chat/message
GET /api/chat/:context_id

# AWS Estate
POST /api/aws/estate

# Collections
GET /api/collections
GET /api/collections/:name

Configuration

Create config.yaml in your data directory:

embedding:
  model: "BAAI/bge-m3"
  dimensions: 1024
  service_url: "http://localhost:8001"

vector_store:
  backend: "qdrant-embedded"  # or "local-files"
  distance_metric: "Cosine"

encryption:
  algorithm: "AES-256-GCM"
  enable_content_encryption: true
  enable_metadata_encryption: true
  enable_embedding_encryption: false

privacy:
  level: "minimal-aws"  # "full", "minimal-aws", "anonymous"
  enable_data_filtering: true

security:
  enable_access_logging: true
  max_request_size: 10485760  # 10MB

Dependencies

Core Dependencies

• tokio: Async runtime
• serde: Serialization
• anyhow: Error handling
• ring: Encryption
• reqwest: HTTP client
• axum: HTTP server
• qdrant-client: Vector database

Build Requirements

• Rust 2021 edition
• Cargo for building

Performance Benefits

Memory Safety

• No garbage collection overhead
• Zero-cost abstractions
• Memory safety without runtime cost

Concurrency

• Tokio async runtime
• Lock-free data structures where possible
• Efficient resource management

Benchmarks vs Node.js

Migration from Node.js

API Compatibility

The Rust implementation maintains 100% API compatibility with the Node.js version:

// Node.js
const rag = new RagModule('./rag-data');
await rag.initialize();
const results = await rag.search('aws_estate', 'EC2 permissions', {});

// Rust HTTP API (same interface)
const response = await fetch('/api/search', {
  method: 'POST',
  body: JSON.stringify({
    collection_type: 'aws_estate',
    query: 'EC2 permissions',
    options: {}
  })
});

Data Migration

Existing Node.js data can be migrated:

1. Export data from Node.js module
2. Use Rust import API endpoints
3. Vector embeddings are preserved
4. Encryption keys can be migrated

Production Deployment

As a Service

# Build optimized release
cargo build --release

# Run production server
RUST_LOG=info ./target/release/rag-server

Docker Support

FROM rust:1.70 AS builder
WORKDIR /app
COPY . .
RUN cargo build --release

FROM debian:bullseye-slim
RUN apt-get update && apt-get install -y ca-certificates
COPY --from=builder /app/target/release/rag-server /usr/local/bin/
EXPOSE 3000
CMD ["rag-server"]

Resource Requirements

• Memory: 50-100MB base usage
• CPU: Multi-core support with tokio
• Disk: Depends on vector data size
• Network: HTTP/HTTPS support

Testing

# Run all tests
cargo test

# Run with output
cargo test -- --nocapture

# Run specific test
cargo test test_embedding_service

# Benchmarks
cargo bench

Development

Adding New Services

1. Create service in src/services/
2. Implement required traits
3. Add to services/mod.rs
4. Update lib.rs initialization

Custom Vector Stores

Implement the VectorStore trait:

#[async_trait]
impl VectorStore for MyCustomStore {
    async fn initialize(&self) -> Result<()> { ... }
    async fn add_document(&self, collection: &str, doc: Document) -> Result<String> { ... }
    async fn search(&self, collection: &str, vector: Vec<f32>, options: SearchOptions) -> Result<Vec<SearchResult>> { ... }
    // ... other methods
}

Security

• Encryption: AES-256-GCM for data at rest
• Privacy: Configurable data filtering
• Access Control: Token-based authentication (configurable)
• Audit Logging: Request/response logging
• Memory Safety: Rust's ownership system prevents memory vulnerabilities

License

MIT License - same as the original Node.js module.

Support

For issues and questions:

1. Check existing Node.js documentation
2. Rust-specific issues: Create GitHub issues
3. Performance optimization: Benchmarking tools included

This Rust implementation provides the same functionality as the Node.js version with significantly improved performance, memory safety, and resource efficiency."