rust-ai-framework

Rust AI Framework

A comprehensive Rust framework for AI/ML applications with support for Large Language Models (LLMs), embeddings, RAG (Retrieval-Augmented Generation), reranking, function calling, workflow management, and Model Context Protocol (MCP).

Features

• LLM Integration: Support for OpenAI, Anthropic, and local models
• Embedding Management: Vector embeddings for semantic search
• RAG System: Retrieval-Augmented Generation with document storage and search
• Reranking: Result reranking for improved search quality
• Function Calling: Dynamic function execution framework
• Workflow Management: Multi-step workflow orchestration
• MCP Support: Model Context Protocol server, client, and proxy
• Configuration Management: Flexible configuration system
• Error Handling: Comprehensive error types and handling

Project Structure

rust-ai/
├── src/
│   ├── core.rs              # Core types and traits
│   ├── error.rs             # Error definitions
│   ├── config.rs            # Configuration management
│   ├── utils.rs             # Utility functions
│   ├── lib.rs               # Library entry point
│   ├── llm/                 # Large Language Model providers
│   │   ├── mod.rs
│   │   ├── openai.rs        # OpenAI integration
│   │   ├── anthropic.rs     # Anthropic integration
│   │   └── local.rs         # Local model support
│   ├── embedding/           # Embedding providers
│   │   ├── mod.rs
│   │   ├── openai.rs        # OpenAI embeddings
│   │   └── local.rs         # Local embedding models
│   ├── rag/                 # RAG system
│   │   └── mod.rs
│   ├── reranking/           # Result reranking
│   │   └── mod.rs
│   ├── function_call/       # Function calling framework
│   │   └── mod.rs
│   ├── flow/                # Workflow management
│   │   └── mod.rs
│   ├── mcp/                 # Model Context Protocol
│   │   ├── mod.rs
│   │   ├── protocol.rs      # MCP protocol definitions
│   │   ├── server.rs        # MCP server implementation
│   │   ├── client.rs        # MCP client implementation
│   │   └── proxy.rs         # MCP proxy implementation
│   └── bin/                 # Binary executables
│       ├── mcp_server.rs    # MCP server binary
│       ├── mcp_client.rs    # MCP client binary
│       └── mcp_proxy.rs     # MCP proxy binary
├── Cargo.toml               # Project dependencies
├── example_config.json      # Example configuration
└── README.md               # This file

Installation

1. Clone the repository:

git clone <repository-url>
cd rust-ai

2. Build the project:

cargo build

3. Run tests:

cargo test

Configuration

Create a configuration file (e.g., config.json) with your API keys and settings:

{
  "llm": {
    "openai": {
      "api_key": "your-openai-api-key",
      "base_url": "https://api.openai.com/v1",
      "model": "gpt-4"
    },
    "anthropic": {
      "api_key": "your-anthropic-api-key",
      "model": "claude-3-sonnet-20240229"
    }
  },
  "embedding": {
    "openai": {
      "api_key": "your-openai-api-key",
      "model": "text-embedding-3-small"
    }
  },
  "mcp": {
    "server": {
      "host": "127.0.0.1",
      "port": 8080
    }
  }
}

Usage

Basic LLM Usage

use rust_ai_framework::{LlmProvider, LlmRequest, LlmResponse};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Initialize with configuration
    let config = rust_ai_framework::config::Config::from_file("config.json")?;
    rust_ai_framework::init_with_config(&config)?;
    
    // Create LLM provider
    let provider = rust_ai_framework::llm::OpenAiProvider::new(config.llm.openai);
    
    // Generate text
    let request = LlmRequest {
        prompt: "Hello, how are you?".to_string(),
        max_tokens: 100,
        temperature: 0.7,
        ..Default::default()
    };
    
    let response = provider.generate(request).await?;
    println!("Response: {}", response.text);
    
    Ok(())
}

RAG System Usage

use rust_ai_framework::{rag::RagSystem, Document};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let rag = RagSystem::new();
    
    // Add documents
    let document = Document {
        id: uuid::Uuid::new_v4(),
        content: "This is a sample document about AI.".to_string(),
        metadata: Default::default(),
    };
    
    rag.add_document(document).await?;
    
    // Search documents
    let results = rag.search("AI", 5).await?;
    for result in results {
        println!("Found: {}", result.document.content);
    }
    
    Ok(())
}

MCP Server Usage

Run the MCP server:

cargo run --bin mcp-server -- --host 127.0.0.1 --port 8080

Run the MCP client:

cargo run --bin mcp-client -- --server 127.0.0.1:8080

Run the MCP proxy:

cargo run --bin mcp-proxy -- --upstream 127.0.0.1:8080 --port 8081

API Reference

Core Types

• LlmRequest: Request structure for LLM generation
• LlmResponse: Response structure from LLM generation
• Document: Document structure for RAG system
• SearchResult: Search result structure
• FunctionCall: Function call definition
• WorkflowStep: Workflow step definition

Traits

• LlmProvider: Trait for LLM providers
• EmbeddingProvider: Trait for embedding providers
• RagProvider: Trait for RAG system providers
• RerankingProvider: Trait for reranking providers
• FunctionCallProvider: Trait for function calling providers
• WorkflowProvider: Trait for workflow management providers

Error Handling

The framework provides comprehensive error handling with specific error types:

• Error::Config: Configuration-related errors
• Error::Llm: LLM-related errors
• Error::Embedding: Embedding-related errors
• Error::Rag: RAG system errors
• Error::Mcp: MCP protocol errors
• Error::Io: I/O errors

Contributing

1. Fork the repository
2. Create a feature branch
3. Make your changes
4. Add tests for new functionality
5. Run the test suite
6. Submit a pull request

License

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

Dependencies

• tokio: Async runtime
• serde: Serialization/deserialization
• serde_json: JSON handling
• uuid: UUID generation
• chrono: Date/time handling
• anyhow: Error handling
• thiserror: Error type definitions

Roadmap

• Add support for more LLM providers (Cohere, Hugging Face, etc.)
• Implement vector database integration (Pinecone, Weaviate, etc.)
• Add streaming support for LLM responses
• Implement caching layer
• Add monitoring and metrics
• Create web UI for configuration and management
• Add support for fine-tuning workflows
• Implement plugin system for custom providers