poimen-memory/crates/mem-llm/src/vllm.rs

// M5.4 — vLLM LoRA Serving Client
//
// Client for vLLM with LoRA adapter support.
// Communicates over OpenAI-compatible API endpoint.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// vLLM chat completion request
#[derive(Debug, Clone, Serialize)]
pub struct VllmCompletionRequest {
    /// Model name (base or adapter)
    pub model: String,
    
    /// Messages (OpenAI format)
    pub messages: Vec<ChatMessage>,
    
    /// Temperature for sampling
    pub temperature: Option<f32>,
    
    /// Max tokens to generate
    pub max_tokens: Option<usize>,
    
    /// Optional seed for reproducibility
    pub seed: Option<u64>,
}

/// Chat message (OpenAI format)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
    pub role: String,  // "user", "assistant", "system"
    pub content: String,
}

/// vLLM chat completion response
#[derive(Debug, Clone, Deserialize)]
pub struct VllmCompletionResponse {
    pub choices: Vec<Choice>,
    pub usage: Usage,
}

#[derive(Debug, Clone, Deserialize)]
pub struct Choice {
    pub message: ChatMessage,
    pub finish_reason: Option<String>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct Usage {
    pub prompt_tokens: usize,
    pub completion_tokens: usize,
    pub total_tokens: usize,
}

/// vLLM model info response
#[derive(Debug, Clone, Deserialize)]
pub struct VllmModelsResponse {
    pub object: String,
    pub data: Vec<Model>,
}

#[derive(Debug, Clone, Deserialize)]
pub struct Model {
    pub id: String,
    pub object: String,
    pub owned_by: String,
}

/// vLLM health check response
#[derive(Debug, Clone, Deserialize)]
pub struct VllmHealthResponse {
    pub status: String,
}

/// vLLM configuration for LoRA serving
#[derive(Debug, Clone)]
pub struct VllmConfig {
    /// Base model (e.g., "qwen2.5-3b-instruct")
    pub base_model: String,
    
    /// Served model name for API
    pub served_model_name: String,
    
    /// Max LoRA rank
    pub max_lora_rank: usize,
    
    /// Max model context length
    pub max_model_len: usize,
    
    /// LoRA adapters: name → path mapping
    pub lora_modules: HashMap<String, String>,
    
    /// Endpoint URL
    pub endpoint: String,
    
    /// API key (optional)
    pub api_key: Option<String>,
}

impl Default for VllmConfig {
    fn default() -> Self {
        Self {
            base_model: "qwen2.5-3b-instruct".to_string(),
            served_model_name: "memory".to_string(),
            max_lora_rank: 32,
            max_model_len: 32768,
            lora_modules: HashMap::new(),
            endpoint: "http://localhost:8000/v1".to_string(),
            api_key: None,
        }
    }
}

impl VllmConfig {
    /// Add a LoRA adapter module
    pub fn add_adapter(&mut self, name: String, path: String) {
        self.lora_modules.insert(name, path);
    }
    
    /// Generate K8s container args for vLLM
    pub fn to_container_args(&self) -> Vec<String> {
        let mut args = vec![
            "python".to_string(),
            "-m".to_string(),
            "vllm.entrypoints.openai.api_server".to_string(),
            "--model".to_string(),
            self.base_model.clone(),
            "--served-model-name".to_string(),
            self.served_model_name.clone(),
            "--enable-lora".to_string(),
            "--max-lora-rank".to_string(),
            self.max_lora_rank.to_string(),
            "--max-model-len".to_string(),
            self.max_model_len.to_string(),
        ];
        
        // Add LoRA modules
        for (name, path) in &self.lora_modules {
            args.push("--lora-modules".to_string());
            args.push(format!("{}={}", name, path));
        }
        
        args
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_vllm_config_default() {
        let config = VllmConfig::default();
        assert_eq!(config.base_model, "qwen2.5-3b-instruct");
        assert_eq!(config.served_model_name, "memory");
        assert_eq!(config.max_lora_rank, 32);
        assert_eq!(config.max_model_len, 32768);
    }
    
    #[test]
    fn test_add_adapter() {
        let mut config = VllmConfig::default();
        config.add_adapter("memory-v1".to_string(), "/mnt/adapters/memory-v1".to_string());
        
        assert_eq!(config.lora_modules.len(), 1);
        assert_eq!(config.lora_modules.get("memory-v1"), Some(&"/mnt/adapters/memory-v1".to_string()));
    }
    
    #[test]
    fn test_container_args_includes_lora() {
        let mut config = VllmConfig::default();
        config.add_adapter("memory-v1".to_string(), "/mnt/adapters/memory-v1".to_string());
        
        let args = config.to_container_args();
        assert!(args.contains(&"--enable-lora".to_string()));
        assert!(args.contains(&"--max-lora-rank".to_string()));
        assert!(args.contains(&"32".to_string()));
    }
    
    #[test]
    fn test_completion_request_serde() {
        let req = VllmCompletionRequest {
            model: "memory-v1".to_string(),
            messages: vec![ChatMessage {
                role: "user".to_string(),
                content: "Hello".to_string(),
            }],
            temperature: Some(0.7),
            max_tokens: Some(100),
            seed: None,
        };
        
        let json = serde_json::to_string(&req).expect("Should serialize");
        assert!(json.contains("memory-v1"));
        assert!(json.contains("user"));
    }
    
    #[test]
    fn test_health_response_serde() {
        let json = r#"{"status": "healthy"}"#;
        let response: VllmHealthResponse = serde_json::from_str(json)
            .expect("Should deserialize");
        assert_eq!(response.status, "healthy");
    }
}