Building High-Performance AI Infrastructure with Rust: Cortex Memory's Technology Selection and Engineering Experience

// Rust's borrow checker prevents memory errors at compile time
pub struct Memory {
    pub id: String,
    pub content: String,
    pub embedding: Vec<f32>,  // Automatic memory management
}

// Compile-time checks: no dangling pointers, double frees, etc.
pub fn process_memory(memory: &Memory) -> Vec<f32> {
    let embedding = memory.embedding.clone();  // Explicit ownership transfer
    // ... process embedding
    embedding  // Return, ownership transfers to caller
}

// High-level abstractions don't incur runtime overhead
pub trait VectorStore: Send + Sync {
    async fn insert(&self, memory: &Memory) -> Result<()>;
    async fn search(&self, query: &[f32], limit: usize) -> Result<Vec<ScoredMemory>>;
}

// Compiles to machine code equivalent to hand-written C code
impl VectorStore for QdrantStore {
    async fn insert(&self, memory: &Memory) -> Result<()> {
        // Directly call Qdrant gRPC API
        self.client.upsert_point(...).await?;
        Ok(())
    }
}

// Tokio async runtime provides efficient concurrent processing
pub async fn handle_concurrent_requests(
    requests: Vec<Request>,
    manager: Arc<MemoryManager>,
) -> Vec<Response> {
    let tasks: Vec<_> = requests
        .into_iter()
        .map(|req| {
            let manager = manager.clone();
            tokio::spawn(async move {
                manager.process_request(req).await
            })
        })
        .collect();

    // Execute all requests concurrently
    let results = futures::future::join_all(tasks).await;

    results
        .into_iter()
        .filter_map(|r| r.ok())
        .collect()
}

# Cargo.toml - Workspace configuration
[workspace]
resolver = "2"
members = [
    "cortex-mem-core",      # Core business logic
    "cortex-mem-service",   # REST API service
    "cortex-mem-cli",       # Command line tool
    "cortex-mem-mcp",       # MCP protocol adapter
    "cortex-mem-rig",       # AI framework integration
    "cortex-mem-config",    # Configuration management
    "cortex-mem-tools",     # Tool library
]

[workspace.dependencies]
# Unified dependency version management
tokio = { version = "1.48", features = ["full"] }
serde = { version = "1.0", features = ["derive"] }
axum = { version = "0.8", features = ["json"] }
qdrant-client = "1.11"

// Define core Trait
#[async_trait]
pub trait LLMClient: Send + Sync {
    async fn complete(&self, prompt: &str) -> Result<String>;
    async fn embed(&self, text: &str) -> Result<Vec<f32>>;
    async fn health_check(&self) -> Result<bool>;
}

// OpenAI implementation
pub struct OpenAILLMClient {
    client: Client,
    completion_model: Agent<CompletionModel>,
    embedding_model: OpenAIEmbeddingModel,
}

impl LLMClient for OpenAILLMClient {
    async fn complete(&self, prompt: &str) -> Result<String> {
        self.completion_model
            .prompt(prompt)
            .await
            .map_err(|e| MemoryError::LLM(e.to_string()))
    }

    async fn embed(&self, text: &str) -> Result<Vec<f32>> {
        let builder = EmbeddingsBuilder::new(self.embedding_model.clone())
            .document(text)
            .map_err(|e| MemoryError::LLM(e.to_string()))?;

        let embeddings = builder.build().await
            .map_err(|e| MemoryError::LLM(e.to_string()))?;

        embeddings.first()
            .map(|(_, emb)| emb.first().vec.iter().map(|&x| x as f32).collect())
            .ok_or_else(|| MemoryError::LLM("No embedding generated".to_string()))
    }
}

// Use dependency injection
pub struct MemoryManager {
    llm_client: Box<dyn LLMClient>,
    vector_store: Box<dyn VectorStore>,
    // ...
}

impl MemoryManager {
    pub fn new(
        llm_client: Box<dyn LLMClient>,
        vector_store: Box<dyn VectorStore>,
    ) -> Self {
        Self { llm_client, vector_store }
    }
}

use thiserror::Error;

#[derive(Error, Debug)]
pub enum MemoryError {
    #[error("LLM error: {0}")]
    LLM(String),

    #[error("Vector store error: {0}")]
    VectorStore(String),

    #[error("Memory not found: {id}")]
    NotFound { id: String },

    #[error("Validation error: {0}")]
    Validation(String),

    #[error("Configuration error: {0}")]
    Config(String),

    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),

    #[error("Serialization error: {0}")]
    Serialization(#[from] serde_json::Error),
}

pub type Result<T> = std::result::Result<T, MemoryError>;

// Use Tokio async runtime
#[tokio::main]
async fn main() -> Result<()> {
    // Initialize logging
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    // Load configuration
    let config = load_config("config.toml").await?;

    // Create dependencies
    let llm_client = create_llm_client(&config)?;
    let vector_store = create_vector_store(&config).await?;

    // Create MemoryManager
    let manager = Arc::new(MemoryManager::new(
        llm_client,
        vector_store,
        config.memory,
    ));

    // Start HTTP service
    let app = create_router(manager.clone());
    let listener = tokio::net::TcpListener::bind(&config.server.address).await?;

    info!("Server listening on {}", config.server.address);
    axum::serve(listener, app).await?;

    Ok(())
}

use axum::{
    extract::{Path, State},
    http::StatusCode,
    response::Json,
    routing::{get, post},
    Router,
};

pub fn create_router(manager: Arc<MemoryManager>) -> Router {
    Router::new()
        .route("/health", get(health_check))
        .route("/memories", post(create_memory).get(list_memories))
        .route("/memories/:id", get(get_memory).put(update_memory).delete(delete_memory))
        .route("/memories/search", post(search_memories))
        .route("/optimization", post(start_optimization))
        .with_state(manager)
}

// Create memory
pub async fn create_memory(
    State(manager): State<Arc<MemoryManager>>,
    Json(request): Json<CreateMemoryRequest>,
) -> Result<Json<MemoryResponse>, AppError> {
    let memory = manager
        .create_memory(request.content, request.metadata)
        .await?;

    Ok(Json(MemoryResponse::from(memory)))
}

// Search memories
pub async fn search_memories(
    State(manager): State<Arc<MemoryManager>>,
    Json(request): Json<SearchRequest>,
) -> Result<Json<SearchResponse>, AppError> {
    let results = manager
        .search(&request.query, &request.filters, request.limit)
        .await?;

    Ok(Json(SearchResponse {
        results: results.into_iter().map(Into::into).collect(),
        total: results.len(),
    }))
}

// Error handling
pub struct AppError(MemoryError);

impl IntoResponse for AppError {
    fn into_response(self) -> Response {
        let (status, message) = match self.0 {
            MemoryError::NotFound { id } => (StatusCode::NOT_FOUND, format!("Memory not found: {}", id)),
            MemoryError::Validation(msg) => (StatusCode::BAD_REQUEST, msg),
            MemoryError::LLM(msg) => (StatusCode::SERVICE_UNAVAILABLE, format!("LLM error: {}", msg)),
            _ => (StatusCode::INTERNAL_SERVER_ERROR, "Internal server error".to_string()),
        };

        (status, Json(json!({ "error": message }))).into_response()
    }
}

pub struct VectorUtils;

impl VectorUtils {
    /// Calculate cosine similarity (optimized version)
    #[inline]
    pub fn cosine_similarity(vec1: &[f32], vec2: &[f32]) -> f32 {
        // Use SIMD instructions for acceleration
        let dot_product = Self::dot_product_simd(vec1, vec2);
        let norm1 = Self::norm_simd(vec1);
        let norm2 = Self::norm_simd(vec2);

        if norm1 == 0.0 || norm2 == 0.0 {
            return 0.0;
        }

        dot_product / (norm1 * norm2)
    }

    /// Dot product calculation with SIMD acceleration
    #[inline]
    fn dot_product_simd(vec1: &[f32], vec2: &[f32]) -> f32 {
        // Check if SIMD can be used
        if vec1.len() != vec2.len() || vec1.is_empty() {
            return 0.0;
        }

        // Use standard library iterators (compiler will automatically optimize to SIMD)
        vec1.iter()
            .zip(vec2.iter())
            .map(|(a, b)| a * b)
            .sum()
    }

    /// Norm calculation with SIMD acceleration
    #[inline]
    fn norm_simd(vec: &[f32]) -> f32 {
        vec.iter()
            .map(|x| x * x)
            .sum::<f32>()
            .sqrt()
    }

    /// Batch cosine similarity calculation
    pub fn batch_cosine_similarity(
        query: &[f32],
        vectors: &[Vec<f32>],
    ) -> Vec<f32> {
        vectors
            .iter()
            .map(|vec| Self::cosine_similarity(query, vec))
            .collect()
    }
}

use bb8::{Pool, PooledConnection};
use bb8_qdrant::QdrantConnectionManager;

pub struct QdrantPool {
    pool: Pool<QdrantConnectionManager>,
}

impl QdrantPool {
    pub async fn new(url: &str) -> Result<Self> {
        let manager = QdrantConnectionManager::new(url);
        let pool = Pool::builder()
            .max_size(10)
            .min_idle(Some(2))
            .build(manager)
            .await?;

        Ok(Self { pool })
    }

    pub async fn get(&self) -> Result<PooledConnection<QdrantConnectionManager>> {
        self.pool.get().await
            .map_err(|e| MemoryError::VectorStore(e.to_string()))
    }
}

// Use connection pool
impl QdrantStore {
    pub async fn insert(&self, memory: &Memory) -> Result<()> {
        let mut conn = self.pool.get().await?;

        conn.upsert_point(PointStruct::new(
            memory.id.parse()?,
            memory.embedding.clone(),
            self.build_payload(memory),
        )).await?;

        Ok(())
    }
}

use serde::{Deserialize, Serialize};
use config::{Config, ConfigError, Environment, File};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AppConfig {
    pub server: ServerConfig,
    pub qdrant: QdrantConfig,
    pub llm: LLMConfig,
    pub embedding: EmbeddingConfig,
    pub memory: MemoryConfig,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ServerConfig {
    pub host: String,
    pub port: u16,
    pub cors_origins: Vec<String>,
}

impl AppConfig {
    pub fn from_file(path: &str) -> Result<Self, ConfigError> {
        let config = Config::builder()
            .add_source(File::with_name(path))
            .add_source(Environment::with_prefix("CORTEX_MEM"))
            .build()?;

        config.try_deserialize()
    }
}

// Use configuration
pub async fn load_config(path: &str) -> Result<AppConfig> {
    let config = AppConfig::from_file(path)
        .map_err(|e| MemoryError::Config(e.to_string()))?;

    // Validate configuration
    config.validate()?;

    Ok(config)
}

impl AppConfig {
    fn validate(&self) -> Result<()> {
        if self.llm.api_key.is_empty() {
            return Err(MemoryError::Validation("LLM API key is required".to_string()));
        }

        if self.qdrant.url.is_empty() {
            return Err(MemoryError::Validation("Qdrant URL is required".to_string()));
        }

        Ok(())
    }
}

// Use Arc to share large objects
pub struct MemoryManager {
    llm_client: Arc<dyn LLMClient>,
    vector_store: Arc<dyn VectorStore>,
    config: MemoryConfig,
}

// Avoid unnecessary cloning
pub async fn search_with_shared(
    &self,
    query: &str,
    filters: &Filters,
) -> Result<Vec<ScoredMemory>> {
    let query_embedding = self.llm_client.embed(query).await?;

    // Use references instead of cloning
    self.vector_store
        .search(&query_embedding, filters, 10)
        .await
}

// Use Cow to avoid unnecessary allocations
use std::borrow::Cow;

pub fn process_text(text: &str) -> Cow<str> {
    if text.contains("  ") {
        // Needs modification, return owned string
        Cow::Owned(text.replace("  ", " "))
    } else {
        // No modification needed, return borrowed
        Cow::Borrowed(text)
    }
}

// Use Tokio task pool
pub async fn batch_process(
    &self,
    items: Vec<String>,
) -> Result<Vec<Memory>> {
    let semaphore = Arc::new(Semaphore::new(10)); // Limit concurrency

    let tasks: Vec<_> = items
        .into_iter()
        .map(|item| {
            let semaphore = semaphore.clone();
            let manager = self.clone();

            tokio::spawn(async move {
                let _permit = semaphore.acquire().await.unwrap();
                manager.process_item(item).await
            })
        })
        .collect();

    let results = futures::future::join_all(tasks).await;

    results
        .into_iter()
        .filter_map(|r| r.ok())
        .collect()
}

// Use channels for task distribution
pub async fn process_with_channel(
    &self,
    items: Vec<String>,
) -> Result<Vec<Memory>> {
    let (tx, mut rx) = mpsc::channel(100);

    // Send tasks
    for item in items {
        tx.send(item).await?;
    }
    drop(tx); // Close sender

    // Start worker threads
    let mut handles = vec![];
    for _ in 0..4 {
        let mut rx = rx.clone();
        let manager = self.clone();

        let handle = tokio::spawn(async move {
            let mut results = Vec::new();
            while let Some(item) = rx.recv().await {
                if let Ok(memory) = manager.process_item(item).await {
                    results.push(memory);
                }
            }
            results
        });

        handles.push(handle);
    }

    // Collect results
    let mut all_results = Vec::new();
    for handle in handles {
        all_results.extend(handle.await?);
    }

    Ok(all_results)
}

use moka::future::Cache;

pub struct EmbeddingCache {
    cache: Cache<String, Vec<f32>>,
}

impl EmbeddingCache {
    pub fn new(capacity: u64, ttl: Duration) -> Self {
        Self {
            cache: Cache::builder()
                .max_capacity(capacity)
                .time_to_live(ttl)
                .build(),
        }
    }

    pub async fn get_or_compute<F, Fut>(
        &self,
        key: &str,
        compute: F,
    ) -> Result<Vec<f32>>
    where
        F: FnOnce() -> Fut,
        Fut: Future<Output = Result<Vec<f32>>>,
    {
        // Try to get from cache
        if let Some(embedding) = self.cache.get(key).await {
            return Ok(embedding);
        }

        // Compute new value
        let embedding = compute().await?;

        // Store in cache
        self.cache.insert(key.to_string(), embedding.clone()).await;

        Ok(embedding)
    }
}

// Use cache
pub async fn embed_with_cache(
    &self,
    text: &str,
) -> Result<Vec<f32>> {
    let cache_key = self.hash_content(text);

    self.embedding_cache
        .get_or_compute(&cache_key, || {
            self.llm_client.embed(text)
        })
        .await
}

use serde::{Deserialize, Serialize};
use serde_json::Value;

// Use enums to reduce memory usage
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum MemoryContent {
    Short(String),           // Short content stored directly
    Long {                   // Long content uses reference
        id: String,
        content: String,
    },
}

// Use compact numeric types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryMetadata {
    pub importance_score: f32,  // Use f32 instead of f64
    pub created_at: i64,        // Use Unix timestamp instead of DateTime
    pub memory_type: u8,        // Use enum index instead of string
}

// Custom serialization
impl Serialize for Memory {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        // Only serialize necessary fields
        #[derive(Serialize)]
        struct CompactMemory<'a> {
            id: &'a str,
            c: &'a str,           // Shorten field names
            e: &'a [f32],
            m: &'a MemoryMetadata,
        }

        CompactMemory {
            id: &self.id,
            c: &self.content,
            e: &self.embedding,
            m: &self.metadata,
        }.serialize(serializer)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_create_memory() {
        let config = create_test_config();
        let manager = create_test_manager(config).await;

        let memory = manager
            .create_memory("Test content".to_string(), MemoryMetadata::default())
            .await
            .unwrap();

        assert!(!memory.id.is_empty());
        assert_eq!(memory.content, "Test content");
        assert!(!memory.embedding.is_empty());
    }

    #[tokio::test]
    async fn test_search_memories() {
        let config = create_test_config();
        let manager = create_test_manager(config).await;

        // Create test data
        manager.create_memory("I like programming".to_string(), metadata()).await.unwrap();
        manager.create_memory("I love coding".to_string(), metadata()).await.unwrap();

        // Search
        let results = manager
            .search("My hobbies", &Filters::default(), 10)
            .await
            .unwrap();

        assert!(results.len() > 0);
        assert!(results[0].score > 0.5);
    }

    #[test]
    fn test_cosine_similarity() {
        let vec1 = vec![1.0, 0.0, 0.0];
        let vec2 = vec![1.0, 0.0, 0.0];

        let similarity = VectorUtils::cosine_similarity(&vec1, &vec2);

        assert!((similarity - 1.0).abs() < 0.001);
    }
}

#[tokio::test]
async fn test_full_workflow() {
    // Start test server
    let config = create_test_config();
    let manager = Arc::new(create_test_manager(config).await);
    let app = create_router(manager.clone());

    // Create test client
    let client = reqwest::Client::new();
    let base_url = "http://127.0.0.1:8080";

    // Create memory
    let create_resp = client
        .post(&format!("{}/memories", base_url))
        .json(&json!({
            "content": "Test memory",
            "metadata": {
                "user_id": "test_user"
            }
        }))
        .send()
        .await
        .unwrap();

    assert_eq!(create_resp.status(), 201);

    let memory: Memory = create_resp.json().await.unwrap();
    let memory_id = memory.id;

    // Search memory
    let search_resp = client
        .post(&format!("{}/memories/search", base_url))
        .json(&json!({
            "query": "test",
            "filters": {
                "user_id": "test_user"
            }
        }))
        .send()
        .await
        .unwrap();

    assert_eq!(search_resp.status(), 200);

    let search_results: SearchResponse = search_resp.json().await.unwrap();
    assert!(search_results.total > 0);

    // Delete memory
    let delete_resp = client
        .delete(&format!("{}/memories/{}", base_url, memory_id))
        .send()
        .await
        .unwrap();

    assert_eq!(delete_resp.status(), 204);
}

use criterion::{black_box, criterion_group, criterion_main, Criterion, BenchmarkId};

fn bench_cosine_similarity(c: &mut Criterion) {
    let mut group = c.benchmark_group("cosine_similarity");

    for size in [128, 512, 1024, 1536].iter() {
        let vec1: Vec<f32> = (0..*size).map(|_| rand::random()).collect();
        let vec2: Vec<f32> = (0..*size).map(|_| rand::random()).collect();

        group.bench_with_input(BenchmarkId::from_parameter(size), size, |b, _| {
            b.iter(|| {
                black_box(VectorUtils::cosine_similarity(&vec1, &vec2))
            });
        });
    }

    group.finish();
}

criterion_group! {
    name = benches;
    config = Criterion::default().sample_size(100);
    targets = bench_cosine_similarity
}

criterion_main!(benches);

# Cargo.toml - Release configuration
[profile.release]
opt-level = 3           # Highest optimization level
lto = true              # Link-time optimization
codegen-units = 1       # Single code generation unit (better optimization)
strip = true            # Remove symbol table
panic = "abort"         # Reduce binary size

# Compile optimized version
cargo build --release

# Generated binary can run directly
./target/release/cortex-mem-service --config config.toml

# Dockerfile
FROM rust:1.75 as builder

WORKDIR /app
COPY . .

# Compile
RUN cargo build --release

# Runtime image
FROM debian:bookworm-slim

RUN apt-get update && apt-get install -y \
    ca-certificates \
    && rm -rf /var/lib/apt/lists/*

COPY --from=builder /app/target/release/cortex-mem-service /usr/local/bin/

EXPOSE 8000

CMD ["cortex-mem-service", "--config", "/config/config.toml"]

# Build image
docker build -t cortex-mem:latest .

# Run container
docker run -d \
  -p 8000:8000 \
  -v $(pwd)/config.toml:/config/config.toml \
  cortex-mem:latest

use tracing::{info, warn, error};
use tracing_subscriber::{fmt, EnvFilter};

// Initialize logging
pub fn init_logging() {
    tracing_subscriber::fmt()
        .with_env_filter(
            EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| EnvFilter::new("info"))
        )
        .with_target(false)
        .with_thread_ids(true)
        .init();
}

// Use logging
pub async fn process_memory(&self, memory: Memory) -> Result<()> {
    info!("Processing memory: {}", memory.id);

    match self.store_memory(&memory).await {
        Ok(_) => {
            info!("Memory stored successfully: {}", memory.id);
            Ok(())
        }
        Err(e) => {
            error!("Failed to store memory {}: {}", memory.id, e);
            Err(e)
        }
    }
}

// Compile to WebAssembly
#[cfg(target_arch = "wasm32")]
use wasm_bindgen::prelude::*;

#[cfg(target_arch = "wasm32")]
#[wasm_bindgen]
pub async fn search_memories(query: &str) -> JsValue {
    let results = manager.search(query, &Filters::default(), 10).await.unwrap();
    serde_wasm_bindgen::to_value(&results).unwrap()
}

use burn::tensor::Tensor;

pub fn cosine_similarity_gpu(
    vec1: Tensor<Backend, 1>,
    vec2: Tensor<Backend, 1>,
) -> Tensor<Backend, 1> {
    let dot = vec1.clone().matmul(vec2.clone());
    let norm1 = vec1.clone().powf_scalar(2.0).sum().sqrt();
    let norm2 = vec2.powf_scalar(2.0).sum().sqrt();

    dot / (norm1 * norm2)
}

use tonic::transport::Server;

pub struct DistributedMemoryManager {
    local_manager: Arc<MemoryManager>,
    remote_nodes: Vec<Channel>,
}

impl DistributedMemoryManager {
    pub async fn search_distributed(
        &self,
        query: &str,
    ) -> Result<Vec<ScoredMemory>> {
        let mut tasks = Vec::new();

        // Local search
        tasks.push(self.local_manager.search(query, &Filters::default(), 10));

        // Remote search
        for node in &self.remote_nodes {
            let mut client = MemoryServiceClient::new(node.clone());
            let query = query.to_string();
            tasks.push(async move {
                client.search(SearchRequest { query, limit: 10 }).await
            });
        }

        // Merge results
        let results = futures::future::join_all(tasks).await;

        // ... Merge and sort logic

        Ok(final_results)
    }
}