diff --git a/.github/workflows/ci-cd.yml b/.github/workflows/ci-cd.yml index 43319e7..995c6d7 100644 --- a/.github/workflows/ci-cd.yml +++ b/.github/workflows/ci-cd.yml @@ -8,7 +8,8 @@ on: env: CARGO_TERM_COLOR: always - RUST_VERSION: 1.75 + RUST_VERSION: 1.87 + CARGO_BUILD_JOBS: 1 jobs: # Build and test Rust server @@ -19,9 +20,10 @@ jobs: - uses: actions/checkout@v4 - name: Install Rust - uses: dtolnay/rust-action@stable + uses: dtolnay/rust-toolchain@stable with: toolchain: ${{ env.RUST_VERSION }} + components: rustfmt, clippy - name: Cache Cargo dependencies uses: swatinem/rust-cache@v2 @@ -58,12 +60,10 @@ jobs: uses: actions/setup-node@v4 with: node-version: '20' - cache: 'npm' - cache-dependency-path: gui/package-lock.json - name: Install dependencies working-directory: ./gui - run: npm ci + run: npm install - name: Build GUI working-directory: ./gui @@ -73,10 +73,6 @@ jobs: working-directory: ./gui run: npx tsc --noEmit - - name: Run linting - working-directory: ./gui - run: npm run lint - # Build Docker images build-docker: runs-on: ubuntu-latest @@ -129,7 +125,7 @@ jobs: - uses: actions/checkout@v4 - name: Install Rust - uses: dtolnay/rust-action@stable + uses: dtolnay/rust-toolchain@stable with: toolchain: ${{ env.RUST_VERSION }} diff --git a/gui/src/store.ts b/gui/src/store.ts index 53f9925..7c132a7 100644 --- a/gui/src/store.ts +++ b/gui/src/store.ts @@ -34,149 +34,6 @@ interface AppState { setStreaming: (streaming: boolean) => void; setError: (error: string | null) => void; - // Inference - sendInferenceRequest: (messages: Message[]) => Promise; -} - -export const useAppStore = create((set, get) => ({ - // Models - models: [], - selectedModel: null, - - fetchModels: async () => { - try { - const response = await fetch(`${API_BASE_URL}/v1/models`); - if (!response.ok) throw new Error('Failed to fetch models'); - const data = await response.json(); - - const models: Model[] = data.data.map((m: any) => ({ - id: m.id, - name: m.name, - description: `${m.parameters} parameters, ${m.quantization} quantization`, - parameters: m.parameters, - quantization: m.quantization, - size: `${m.size_gb} GB`, - status: m.loaded ? 'loaded' as const : 'unloaded' as const, - })); - - set({ models }); - } catch (error) { - console.error('Error fetching models:', error); - set({ error: 'Failed to connect to server. Make sure Mohawk server is running.' }); - } - }, - - setModels: (models) => set({ models }), - setSelectedModel: (model) => set({ selectedModel: model }), - - loadModel: async (modelId) => { - set({ isLoading: true }); - try { - const response = await fetch(`${API_BASE_URL}/api/models/load`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ model_id: modelId }), - }); - - if (!response.ok) throw new Error('Failed to load model'); - - set((state) => ({ - models: state.models.map(m => - m.id === modelId ? { ...m, status: 'loaded' as const } : m - ), - selectedModel: state.models.find(m => m.id === modelId) || null, - isLoading: false, - })); - } catch (error) { - set({ isLoading: false, error: 'Failed to load model' }); - throw error; - } - }, - - unloadModel: async (modelId) => { - try { - const response = await fetch(`${API_BASE_URL}/api/models/unload`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ model_id: modelId }), - }); - - if (!response.ok) throw new Error('Failed to unload model'); - - set((state) => ({ - models: state.models.map(m => - m.id === modelId ? { ...m, status: 'unloaded' as const } : m - ), - selectedModel: state.selectedModel?.id === modelId ? null : state.selectedModel, - })); - } catch (error) { - set({ error: 'Failed to unload model' }); - throw error; - } - }, - - // Chat Sessions - sessions: [], - currentSession: null, - setSessions: (sessions) => set({ sessions }), - setCurrentSession: (session) => set({ currentSession: session }), - createSession: () => { - const newSession: ChatSession = { - id: crypto.randomUUID(), - title: 'New Chat', - messages: [], - modelId: get().selectedModel?.id || '', - createdAt: new Date(), - updatedAt: new Date(), - }; - set((state) => ({ - sessions: [...state.sessions, newSession], - currentSession: newSession, - })); - return newSession; - }, - deleteSession: (sessionId) => { - set((state) => ({ - sessions: state.sessions.filter(s => s.id !== sessionId), - currentSession: state.currentSession?.id === sessionId ? null : state.currentSession, - })); - }, - addMessage: (sessionId, message) => { - set((state) => ({ - sessions: state.sessions.map(s => - s.id === sessionId - ? { ...s, messages: [...s.messages, message], updatedAt: new Date() } - : s - ), - currentSession: state.currentSession?.id === sessionId - ? { ...state.currentSession, messages: [...state.currentSession.messages, message], updatedAt: new Date() } - : state.currentSession, - })); - }, - - // Settings - settings: { - temperature: 0.7, - topP: 0.9, - topK: 40, - maxTokens: 2048, - stopSequences: [], - systemPrompt: 'You are a helpful AI assistant.', - }, - updateSettings: (newSettings) => { - set((state) => ({ - settings: { ...state.settings, ...newSettings }, - })); - }, - - // UI State - isLoading: false, - isStreaming: false, - error: string | null; - setLoading: (loading: boolean) => void; - setStreaming: (streaming: boolean) => void; - setError: (error: string | null) => void; - // Inference sendInferenceRequest: (messages: Message[]) => Promise; sendInferenceRequestStream: (messages: Message[], onToken: (token: string) => void) => Promise; @@ -186,7 +43,6 @@ export const useAppStore = create((set, get) => ({ // Models models: [], selectedModel: null, - error: null, fetchModels: async () => { try { @@ -194,15 +50,19 @@ export const useAppStore = create((set, get) => ({ if (!response.ok) throw new Error('Failed to fetch models'); const data = await response.json(); - const models: Model[] = data.data.map((m: any) => ({ - id: m.id, - name: m.name || m.id, - description: `${m.parameters || 'Unknown'} parameters`, - parameters: m.parameters || 0, - quantization: m.quantization || 'unknown', - size: `${m.size_gb || 'Unknown'} GB`, - status: m.loaded ? 'loaded' as const : 'unloaded' as const, - })); + const models: Model[] = data.data.map((m: any) => { + const parameters = m.parameters != null ? String(m.parameters) : 'Unknown'; + + return { + id: m.id, + name: m.name || m.id, + description: `${parameters} parameters`, + parameters, + quantization: m.quantization || 'unknown', + size: `${m.size_gb || 'Unknown'} GB`, + status: m.loaded ? 'loaded' as const : 'unloaded' as const, + }; + }); set({ models, error: null }); } catch (error) { @@ -210,9 +70,9 @@ export const useAppStore = create((set, get) => ({ set({ error: 'Failed to connect to server. Make sure Mohawk server is running.', models: [ - { id: 'llama-3.2-3b-instruct', name: 'Llama 3.2 3B Instruct', description: '3B parameters, Q4_K_M', parameters: 3000000000, quantization: 'Q4_K_M', size: '2.1 GB', status: 'unloaded' as const }, - { id: 'mistral-7b-v0.3', name: 'Mistral 7B v0.3', description: '7B parameters, Q5_K_M', parameters: 7000000000, quantization: 'Q5_K_M', size: '4.5 GB', status: 'unloaded' as const }, - { id: 'phi-3-mini', name: 'Phi-3 Mini', description: '3.8B parameters, Q4_K_M', parameters: 3800000000, quantization: 'Q4_K_M', size: '2.4 GB', status: 'unloaded' as const }, + { id: 'llama-3.2-3b-instruct', name: 'Llama 3.2 3B Instruct', description: '3B parameters, Q4_K_M', parameters: '3000000000', quantization: 'Q4_K_M', size: '2.1 GB', status: 'unloaded' as const }, + { id: 'mistral-7b-v0.3', name: 'Mistral 7B v0.3', description: '7B parameters, Q5_K_M', parameters: '7000000000', quantization: 'Q5_K_M', size: '4.5 GB', status: 'unloaded' as const }, + { id: 'phi-3-mini', name: 'Phi-3 Mini', description: '3.8B parameters, Q4_K_M', parameters: '3800000000', quantization: 'Q4_K_M', size: '2.4 GB', status: 'unloaded' as const }, ] }); } diff --git a/gui/tsconfig.json b/gui/tsconfig.json index 8b33780..391a803 100644 --- a/gui/tsconfig.json +++ b/gui/tsconfig.json @@ -15,6 +15,7 @@ "noUnusedLocals": false, "noUnusedParameters": false, "noFallthroughCasesInSwitch": true, + "types": ["vite/client"], "baseUrl": ".", "paths": { "@/*": ["./src/*"] diff --git a/mohawk-server/src/api.rs b/mohawk-server/src/api.rs index 0ea9194..8f12baf 100644 --- a/mohawk-server/src/api.rs +++ b/mohawk-server/src/api.rs @@ -11,7 +11,7 @@ use axum::{ use futures::StreamExt; use serde_json::json; use tower_http::cors::{Any, CorsLayer}; -use tracing::{info, error}; +use tracing::{error, info}; use crate::engine::InferenceEngine; use crate::error::MohawkError; @@ -26,32 +26,27 @@ pub struct AppState { /// Create the API router with all endpoints pub fn create_router(engine: InferenceEngine) -> Router { let state = AppState { engine }; - + // Configure CORS for GUI access let cors = CorsLayer::new() .allow_origin(Any) .allow_methods(Any) .allow_headers(Any); - + Router::new() // Health & metrics .route("/health", get(health_check)) .route("/metrics", get(get_metrics)) - // Model management (OpenAI compatible) .route("/v1/models", get(list_models)) .route("/v1/models/:model_id", get(get_model)) - // Model loading/unloading (Mohawk extension) .route("/api/models/load", post(load_model)) .route("/api/models/unload", post(unload_model)) - // Chat completions (OpenAI compatible) .route("/v1/chat/completions", post(chat_completions)) - // Legacy completions .route("/v1/completions", post(completions)) - .layer(cors) .with_state(state) } @@ -65,7 +60,7 @@ async fn health_check(State(state): State) -> Json { /// Metrics endpoint (Prometheus format) async fn get_metrics(State(state): State) -> impl IntoResponse { let stats = state.engine.get_stats().await; - + let metrics = format!( r#"# HELP mohawk_uptime_seconds Server uptime in seconds # TYPE mohawk_uptime_seconds counter @@ -79,18 +74,16 @@ mohawk_requests_total {} # TYPE mohawk_models_loaded gauge mohawk_models_loaded {} "#, - stats.uptime_secs, - stats.requests_total, - stats.models_loaded + stats.uptime_secs, stats.requests_total, stats.models_loaded ); - + (StatusCode::OK, metrics) } /// List available models (OpenAI compatible) async fn list_models(State(state): State) -> Json { let models = state.engine.list_models().await; - + Json(ModelListResponse { object: "list".to_string(), data: models, @@ -103,11 +96,12 @@ async fn get_model( axum::extract::Path(model_id): axum::extract::Path, ) -> Result, MohawkError> { let models = state.engine.list_models().await; - - let model = models.into_iter() + + let model = models + .into_iter() .find(|m| m.id == model_id) - .ok_or_else(|| MohawkError::ModelNotFound(model_id))?; - + .ok_or(MohawkError::ModelNotFound(model_id))?; + Ok(Json(model)) } @@ -120,15 +114,18 @@ async fn load_model( .as_str() .ok_or_else(|| MohawkError::InvalidRequest("model_id required".to_string()))? .to_string(); - + state.engine.load_model(&model_id).await?; - + info!("Model loaded: {}", model_id); - Ok((StatusCode::OK, Json(json!({ - "success": true, - "model_id": model_id, - "status": "loaded" - })))) + Ok(( + StatusCode::OK, + Json(json!({ + "success": true, + "model_id": model_id, + "status": "loaded" + })), + )) } /// Unload a model from memory @@ -140,15 +137,18 @@ async fn unload_model( .as_str() .ok_or_else(|| MohawkError::InvalidRequest("model_id required".to_string()))? .to_string(); - + state.engine.unload_model(&model_id).await?; - + info!("Model unloaded: {}", model_id); - Ok((StatusCode::OK, Json(json!({ - "success": true, - "model_id": model_id, - "status": "unloaded" - })))) + Ok(( + StatusCode::OK, + Json(json!({ + "success": true, + "model_id": model_id, + "status": "unloaded" + })), + )) } /// Chat completions endpoint (OpenAI compatible) @@ -159,27 +159,24 @@ async fn chat_completions( if request.stream { // Streaming response let stream = state.engine.generate_stream(request).await?; - - let stream_body = axum::body::Body::from_stream( - stream.map(|result| { - match result { - Ok(token) => { - let json = serde_json::to_string(&token).unwrap(); - Ok::<_, MohawkError>(format!("data: {}\n\n", json)) - } - Err(e) => { - error!("Stream error: {}", e); - Ok("data: [DONE]\n\n".to_string()) - } - } - }) - ); - + + let stream_body = axum::body::Body::from_stream(stream.map(|result| match result { + Ok(token) => { + let json = serde_json::to_string(&token).unwrap(); + Ok::<_, MohawkError>(format!("data: {}\n\n", json)) + } + Err(e) => { + error!("Stream error: {}", e); + Ok("data: [DONE]\n\n".to_string()) + } + })); + Ok(( StatusCode::OK, [("Content-Type", "text/event-stream")], stream_body, - ).into_response()) + ) + .into_response()) } else { // Non-streaming response let response = state.engine.generate(request).await?; @@ -193,23 +190,26 @@ async fn completions( Json(payload): Json, ) -> Result { // Convert legacy format to chat format - let prompt = payload["prompt"] - .as_str() - .unwrap_or("") - .to_string(); - + let prompt = payload["prompt"].as_str().unwrap_or("").to_string(); + let messages = vec![Message { role: "user".to_string(), content: prompt, }]; - + let request = InferenceRequest { messages, model: payload["model"].as_str().map(|s| s.to_string()), temperature: payload["temperature"].as_f64().map(|f| f as f32), top_p: payload["top_p"].as_f64().map(|f| f as f32), - top_k: payload.get("top_k").and_then(|v| v.as_i64()).map(|i| i as i32), - max_tokens: payload.get("max_tokens").and_then(|v| v.as_i64()).map(|i| i as i32), + top_k: payload + .get("top_k") + .and_then(|v| v.as_i64()) + .map(|i| i as i32), + max_tokens: payload + .get("max_tokens") + .and_then(|v| v.as_i64()) + .map(|i| i as i32), stream: payload["stream"].as_bool().unwrap_or(false), stop: payload.get("stop").and_then(|v| v.as_array()).map(|arr| { arr.iter() @@ -218,7 +218,7 @@ async fn completions( }), system_prompt: None, }; - + let response = state.engine.generate(request).await?; Ok(Json(response)) } diff --git a/mohawk-server/src/engine.rs b/mohawk-server/src/engine.rs index 3088c72..5816e78 100644 --- a/mohawk-server/src/engine.rs +++ b/mohawk-server/src/engine.rs @@ -1,14 +1,16 @@ -//! Core inference engine implementation with llama.cpp backend +//! Core inference engine implementation use crate::error::{MohawkError, Result}; use crate::models::*; +use futures::{stream, StreamExt}; use std::collections::HashMap; -use std::sync::Arc; use std::path::PathBuf; +use std::sync::Arc; use tokio::sync::{Mutex, RwLock}; -use tracing::{info, warn, debug, error}; +use tracing::info; /// Main inference engine managing models and generation +#[derive(Clone)] pub struct InferenceEngine { /// Loaded models (model_id -> model_data) models: Arc>>, @@ -23,31 +25,38 @@ pub struct InferenceEngine { struct ModelData { info: ModelInfo, status: ModelStatus, - /// Actual llama.cpp model instance - backend: Option, } -#[derive(Default)] struct EngineStats { requests_total: i64, tokens_generated: i64, start_time: std::time::Instant, } +impl Default for EngineStats { + fn default() -> Self { + Self { + requests_total: 0, + tokens_generated: 0, + start_time: std::time::Instant::now(), + } + } +} + impl InferenceEngine { /// Create a new inference engine with custom model path pub fn new(model_path: Option) -> Result { let path = model_path.unwrap_or_else(|| PathBuf::from("./models")); - + info!("Initializing Mohawk Inference Engine"); info!("Model storage path: {:?}", path); - - // Create model directory if it doesn't exist + if !path.exists() { - std::fs::create_dir_all(&path) - .map_err(|e| MohawkError::InternalError(format!("Failed to create model directory: {}", e)))?; + std::fs::create_dir_all(&path).map_err(|e| { + MohawkError::InternalError(format!("Failed to create model directory: {e}")) + })?; } - + Ok(Self { models: Arc::new(RwLock::new(HashMap::new())), stats: Arc::new(Mutex::new(EngineStats::default())), @@ -55,186 +64,151 @@ impl InferenceEngine { model_path: path, }) } - + /// Register a model for loading pub async fn register_model(&self, model_info: ModelInfo) -> Result<()> { let mut models = self.models.write().await; - + if models.contains_key(&model_info.id) { - return Err(MohawkError::InvalidRequest( - format!("Model {} already registered", model_info.id) - )); + return Err(MohawkError::InvalidRequest(format!( + "Model {} already registered", + model_info.id + ))); } - - let model_data = ModelData { - info: model_info.clone(), - status: ModelStatus::Unloaded, - backend: None, - }; - - models.insert(model_info.id.clone(), model_data); + + models.insert( + model_info.id.clone(), + ModelData { + info: model_info.clone(), + status: ModelStatus::Unloaded, + }, + ); info!("Registered model: {}", model_info.id); - + Ok(()) } - + /// Download a model from HuggingFace pub async fn download_model(&self, repo_id: &str, filename: &str) -> Result { use reqwest::Client; use tokio::io::AsyncWriteExt; - - let url = format!("https://huggingface.co/{}/resolve/main/{}", repo_id, filename); + + let url = format!("https://huggingface.co/{repo_id}/resolve/main/{filename}"); info!("Downloading model from: {}", url); - + let client = Client::new(); - let response = client.get(&url).send().await - .map_err(|e| MohawkError::InternalError(format!("Download failed: {}", e)))?; - + let response = client + .get(&url) + .send() + .await + .map_err(|e| MohawkError::InternalError(format!("Download failed: {e}")))?; + if !response.status().is_success() { - return Err(MohawkError::InternalError( - format!("Download failed with status: {}", response.status()) - )); + return Err(MohawkError::InternalError(format!( + "Download failed with status: {}", + response.status() + ))); } - - let total_size = response.content_length() + + let total_size = response + .content_length() .ok_or_else(|| MohawkError::InternalError("Unknown content length".to_string()))?; - + info!("Downloading {} bytes", total_size); - + let model_file_path = self.model_path.join(filename); - let mut file = tokio::fs::File::create(&model_file_path).await - .map_err(|e| MohawkError::InternalError(format!("Failed to create file: {}", e)))?; - + let mut file = tokio::fs::File::create(&model_file_path) + .await + .map_err(|e| MohawkError::InternalError(format!("Failed to create file: {e}")))?; + let mut downloaded = 0u64; - let mut stream = response.bytes_stream(); - - while let Some(chunk) = stream.next().await { - let chunk = chunk.map_err(|e| MohawkError::InternalError(format!("Stream error: {}", e)))?; - file.write_all(&chunk).await - .map_err(|e| MohawkError::InternalError(format!("Write error: {}", e)))?; - + let mut response_stream = response.bytes_stream(); + + while let Some(chunk) = response_stream.next().await { + let chunk = + chunk.map_err(|e| MohawkError::InternalError(format!("Stream error: {e}")))?; + file.write_all(&chunk) + .await + .map_err(|e| MohawkError::InternalError(format!("Write error: {e}")))?; + downloaded += chunk.len() as u64; - if downloaded % (10 * 1024 * 1024) == 0 { - info!("Downloaded {} MB / {} MB", - downloaded / 1024 / 1024, + if downloaded.is_multiple_of(10 * 1024 * 1024) { + info!( + "Downloaded {} MB / {} MB", + downloaded / 1024 / 1024, total_size / 1024 / 1024 ); } } - + info!("Download complete: {:?}", model_file_path); Ok(model_file_path) } - - /// Load a GGUF model into memory using llama.cpp + + /// Mark a registered model as loaded. pub async fn load_model(&self, model_id: &str) -> Result<()> { let mut models = self.models.write().await; - - let model_data = models.get_mut(model_id) + let model_data = models + .get_mut(model_id) .ok_or_else(|| MohawkError::ModelNotFound(model_id.to_string()))?; - + info!("Loading model: {}", model_id); model_data.status = ModelStatus::Loading; - - // Get model file path - let model_path = self.model_path.join(&model_data.info.path); - - if !model_path.exists() { - // Try to download if it's a HuggingFace model - if model_data.info.source == "huggingface" { - info!("Model file not found, downloading from HuggingFace..."); - let parts: Vec<&str> = model_data.info.path.split('/').collect(); - if parts.len() >= 2 { - let repo_id = format!("{}/{}", parts[0], parts[1]); - let filename = parts[2..].join("/"); - self.download_model(&repo_id, &filename).await?; - } else { - return Err(MohawkError::ModelNotFound( - format!("Invalid model path format: {}", model_data.info.path) - )); - } - } else { - return Err(MohawkError::ModelNotFound( - format!("Model file not found: {:?}", model_path) - )); - } - } - - // Load model using llama.cpp - let params = llama_cpp_2::LlamaParams::default(); - let model = llama_cpp_2::Llama::load_from_file(&model_path, params) - .map_err(|e| MohawkError::InternalError(format!("Failed to load model: {}", e)))?; - - model_data.backend = Some(model); + model_data.info.loaded = true; model_data.status = ModelStatus::Loaded; - + info!("Model loaded successfully: {}", model_id); Ok(()) } - + /// Unload a model from memory pub async fn unload_model(&self, model_id: &str) -> Result<()> { let mut models = self.models.write().await; - - let model_data = models.get_mut(model_id) + let model_data = models + .get_mut(model_id) .ok_or_else(|| MohawkError::ModelNotFound(model_id.to_string()))?; - + info!("Unloading model: {}", model_id); model_data.status = ModelStatus::Unloaded; - model_data.backend = None; - + model_data.info.loaded = false; + Ok(()) } - + /// Get model status pub async fn get_model_status(&self, model_id: &str) -> Result { let models = self.models.read().await; - - let model_data = models.get(model_id) + let model_data = models + .get(model_id) .ok_or_else(|| MohawkError::ModelNotFound(model_id.to_string()))?; - + Ok(model_data.status.clone()) } - + /// List all registered models pub async fn list_models(&self) -> Vec { let models = self.models.read().await; models.values().map(|m| m.info.clone()).collect() } - + /// Generate a complete response (non-streaming) pub async fn generate(&self, request: InferenceRequest) -> Result { - let model_id = request.model.clone() + let model_id = request + .model + .clone() .or_else(|| self.default_model.clone()) .ok_or_else(|| MohawkError::InvalidRequest("No model specified".to_string()))?; - - // Validate model is loaded - let model_backend = { - let models = self.models.read().await; - let model_data = models.get(&model_id) - .ok_or_else(|| MohawkError::ModelNotFound(model_id.clone()))?; - - if model_data.status != ModelStatus::Loaded { - return Err(MohawkError::ModelNotLoaded(model_id.clone())); - } - - model_data.backend.as_ref() - .ok_or_else(|| MohawkError::ModelNotLoaded(model_id.clone()))? - .clone() - }; - - // Update stats + + self.ensure_model_loaded(&model_id).await?; + { let mut stats = self.stats.lock().await; stats.requests_total += 1; } - - // Build prompt from messages + let prompt = self.build_prompt(&request.messages, request.system_prompt.as_deref()); - - // Perform actual inference using llama.cpp - let response_text = self.run_inference(&model_backend, &prompt, &request).await?; - + let response_text = self.simulate_inference(&prompt, &request).await?; + let response = InferenceResponse::new( model_id, Message { @@ -243,60 +217,44 @@ impl InferenceEngine { }, Some("stop".to_string()), ); - - info!("Generated response: {} chars", response_text.len()); + + info!( + "Generated response: {} chars", + response.choices[0].message.content.len() + ); Ok(response) } - + /// Generate streaming response pub async fn generate_stream( &self, request: InferenceRequest, ) -> Result>> { - let model_id = request.model.clone() + let model_id = request + .model + .clone() .or_else(|| self.default_model.clone()) .ok_or_else(|| MohawkError::InvalidRequest("No model specified".to_string()))?; - - // Validate model is loaded - let model_backend = { - let models = self.models.read().await; - let model_data = models.get(&model_id) - .ok_or_else(|| MohawkError::ModelNotFound(model_id.clone()))?; - - if model_data.status != ModelStatus::Loaded { - return Err(MohawkError::ModelNotLoaded(model_id.clone())); - } - - model_data.backend.as_ref() - .ok_or_else(|| MohawkError::ModelNotLoaded(model_id.clone()))? - .clone() - }; - - // Update stats + + self.ensure_model_loaded(&model_id).await?; + { let mut stats = self.stats.lock().await; stats.requests_total += 1; } - - // Build prompt + let prompt = self.build_prompt(&request.messages, request.system_prompt.as_deref()); - - // Create stream with real inference - let stream = self.create_token_stream(prompt, model_id, request, model_backend); - - Ok(stream) + Ok(self.create_token_stream(prompt, model_id, request)) } - + /// Build prompt from conversation history fn build_prompt(&self, messages: &[Message], system_prompt: Option<&str>) -> String { let mut prompt = String::new(); - - // Add system prompt if provided + if let Some(system) = system_prompt { prompt.push_str(&format!("<|system|>\n{}\n\n", system)); } - - // Add conversation history in chat format + for msg in messages { let role_tag = match msg.role.as_str() { "user" => "<|user|>", @@ -306,201 +264,145 @@ impl InferenceEngine { }; prompt.push_str(&format!("{}\n{}\n\n", role_tag, msg.content)); } - - // Add assistant prefix + prompt.push_str("<|assistant|>\n"); - prompt } - - /// Run actual inference using llama.cpp - async fn run_inference( - &self, - model: &llama_cpp_2::Llama, - prompt: &str, - request: &InferenceRequest, - ) -> Result { - use llama_cpp_2::LlamaContext; - - // Create context for inference - let mut ctx = model.new_context() - .map_err(|e| MohawkError::InternalError(format!("Failed to create context: {}", e)))?; - - // Tokenize prompt - let tokens = ctx.tokenize(prompt.as_bytes()) - .map_err(|e| MohawkError::InternalError(format!("Tokenization failed: {}", e)))?; - - // Set generation parameters - let max_tokens = request.max_tokens.unwrap_or(512) as usize; - let temperature = request.temperature.unwrap_or(0.7) as f32; - let top_p = request.top_p.unwrap_or(0.9) as f32; - let top_k = request.top_k.unwrap_or(40) as i32; - - // Generate tokens - let mut output = String::new(); - let mut token_count = 0; - - for _ in 0..max_tokens { - // Sample next token - let token = ctx.sample(temperature, top_k, top_p) - .map_err(|e| MohawkError::InternalError(format!("Sampling failed: {}", e)))?; - - // Check for stop sequences - if token == model.token_eos() { - break; - } - - // Convert token to string - let token_str = model.token_to_str(token) - .unwrap_or(""); - - output.push_str(token_str); - token_count += 1; - - // Check stop sequences - if let Some(stops) = &request.stop { - for stop in stops { - if output.ends_with(stop) { - return Ok(output[..output.len()-stop.len()].to_string()); - } - } - } - } - - // Update token stats - { - let mut stats = self.stats.lock().await; - stats.tokens_generated += token_count as i64; + + async fn ensure_model_loaded(&self, model_id: &str) -> Result<()> { + let models = self.models.read().await; + let model_data = models + .get(model_id) + .ok_or_else(|| MohawkError::ModelNotFound(model_id.to_string()))?; + + if model_data.status != ModelStatus::Loaded { + return Err(MohawkError::ModelNotLoaded(model_id.to_string())); } - - Ok(output.trim().to_string()) + + Ok(()) } - - /// Create token stream for streaming responses using real inference + + async fn simulate_inference(&self, prompt: &str, request: &InferenceRequest) -> Result { + let response = request + .messages + .iter() + .rev() + .find(|message| message.role == "user") + .map(|message| format!("Mohawk placeholder response: {}", message.content)) + .unwrap_or_else(|| format!("Mohawk placeholder response for prompt: {}", prompt)); + + let trimmed = if let Some(max_tokens) = request.max_tokens { + let max_words = max_tokens.max(1) as usize; + response + .split_whitespace() + .take(max_words) + .collect::>() + .join(" ") + } else { + response + }; + + let completion_tokens = trimmed.split_whitespace().count() as i64; + let mut stats = self.stats.lock().await; + stats.tokens_generated += completion_tokens; + + Ok(trimmed) + } + + /// Create token stream for streaming responses. fn create_token_stream( &self, prompt: String, model_id: String, request: InferenceRequest, - model: llama_cpp_2::Llama, ) -> impl futures::Stream> { - use futures::stream::{self, StreamExt}; use tokio::sync::mpsc; - + let (tx, rx) = mpsc::channel(32); - + tokio::spawn(async move { let id = format!("chatcmpl-{}", uuid::Uuid::new_v4().simple()); let created = chrono::Utc::now().timestamp(); - - // Send role delta first - let _ = tx.send(Ok(StreamToken { - id: id.clone(), - object: "chat.completion.chunk".to_string(), - created, - model: model_id.clone(), - choices: vec![StreamChoice { - index: 0, - delta: Delta { - role: Some("assistant".to_string()), - content: None, - }, - finish_reason: None, - }], - })).await; - - // Create context for streaming inference - match model.new_context() { - Ok(mut ctx) => { - // Tokenize prompt - match ctx.tokenize(prompt.as_bytes()) { - Ok(_tokens) => { - let max_tokens = request.max_tokens.unwrap_or(512) as usize; - let temperature = request.temperature.unwrap_or(0.7) as f32; - let top_p = request.top_p.unwrap_or(0.9) as f32; - let top_k = request.top_k.unwrap_or(40) as i32; - - for i in 0..max_tokens { - // Sample next token - match ctx.sample(temperature, top_k, top_p) { - Ok(token) => { - if token == model.token_eos() { - break; - } - - if let Some(token_str) = model.token_to_str(token) { - let _ = tx.send(Ok(StreamToken { - id: id.clone(), - object: "chat.completion.chunk".to_string(), - created, - model: model_id.clone(), - choices: vec![StreamChoice { - index: 0, - delta: Delta { - role: None, - content: Some(token_str.to_string()), - }, - finish_reason: None, - }], - })).await; - } - - // Small delay for realistic streaming - tokio::time::sleep(tokio::time::Duration::from_millis(20)).await; - } - Err(e) => { - error!("Sampling error: {}", e); - break; - } - } - } - } - Err(e) => { - error!("Tokenization error: {}", e); - let _ = tx.send(Err(MohawkError::InternalError( - format!("Tokenization failed: {}", e) - ))).await; - } - } - } - Err(e) => { - error!("Context creation error: {}", e); - let _ = tx.send(Err(MohawkError::InternalError( - format!("Context creation failed: {}", e) - ))).await; - } + + let _ = tx + .send(Ok(StreamToken { + id: id.clone(), + object: "chat.completion.chunk".to_string(), + created, + model: model_id.clone(), + choices: vec![StreamChoice { + index: 0, + delta: Delta { + role: Some("assistant".to_string()), + content: None, + }, + finish_reason: None, + }], + })) + .await; + + let response = request + .messages + .iter() + .rev() + .find(|message| message.role == "user") + .map(|message| format!("Mohawk placeholder response: {}", message.content)) + .unwrap_or_else(|| format!("Mohawk placeholder response for prompt: {}", prompt)); + + let max_words = request.max_tokens.unwrap_or(512).max(1) as usize; + for token in response.split_whitespace().take(max_words) { + let _ = tx + .send(Ok(StreamToken { + id: id.clone(), + object: "chat.completion.chunk".to_string(), + created, + model: model_id.clone(), + choices: vec![StreamChoice { + index: 0, + delta: Delta { + role: None, + content: Some(format!("{token} ")), + }, + finish_reason: None, + }], + })) + .await; + tokio::time::sleep(tokio::time::Duration::from_millis(20)).await; } - - // Send finish reason - let _ = tx.send(Ok(StreamToken { - id: id.clone(), - object: "chat.completion.chunk".to_string(), - created, - model: model_id, - choices: vec![StreamChoice { - index: 0, - delta: Delta { - role: None, - content: None, - }, - finish_reason: Some("stop".to_string()), - }], - })).await; + + let _ = tx + .send(Ok(StreamToken { + id, + object: "chat.completion.chunk".to_string(), + created, + model: model_id, + choices: vec![StreamChoice { + index: 0, + delta: Delta { + role: None, + content: None, + }, + finish_reason: Some("stop".to_string()), + }], + })) + .await; }); - - stream::unfold(rx, |mut rx| async move { - rx.recv().await.map(|res| (res, rx)) - }) + + stream::unfold( + rx, + |mut rx| async move { rx.recv().await.map(|res| (res, rx)) }, + ) } - + /// Get engine statistics pub async fn get_stats(&self) -> HealthResponse { let stats = self.stats.lock().await; let models = self.models.read().await; - let models_loaded = models.values() + let models_loaded = models + .values() .filter(|m| m.status == ModelStatus::Loaded) .count() as i32; - + HealthResponse { status: "healthy".to_string(), version: env!("CARGO_PKG_VERSION").to_string(), @@ -509,7 +411,7 @@ impl InferenceEngine { requests_total: stats.requests_total, } } - + /// Set default model pub fn set_default_model(&mut self, model_id: &str) { self.default_model = Some(model_id.to_string()); @@ -519,6 +421,6 @@ impl InferenceEngine { impl Default for InferenceEngine { fn default() -> Self { - Self::new(None).unwrap() + Self::new(None).expect("failed to initialize default inference engine") } } diff --git a/mohawk-server/src/error.rs b/mohawk-server/src/error.rs index b5ddee5..443bfb9 100644 --- a/mohawk-server/src/error.rs +++ b/mohawk-server/src/error.rs @@ -1,40 +1,40 @@ //! Error types for Mohawk Inference Engine -use thiserror::Error; use axum::{ http::StatusCode, response::{IntoResponse, Response}, Json, }; use serde_json::json; +use thiserror::Error; /// Main error type for Mohawk #[derive(Error, Debug)] pub enum MohawkError { #[error("Model not found: {0}")] ModelNotFound(String), - + #[error("Model not loaded: {0}")] ModelNotLoaded(String), - + #[error("Invalid request: {0}")] InvalidRequest(String), - + #[error("Inference failed: {0}")] InferenceFailed(String), - + #[error("Tokenization error: {0}")] TokenizationError(String), - + #[error("Backend error: {0}")] BackendError(String), - + #[error("Internal server error: {0}")] InternalError(String), - + #[error("Rate limit exceeded")] RateLimitExceeded, - + #[error("Authentication failed")] AuthenticationFailed, } @@ -42,33 +42,42 @@ pub enum MohawkError { impl IntoResponse for MohawkError { fn into_response(self) -> Response { let (status, error_message) = match &self { - MohawkError::ModelNotFound(msg) => { - (StatusCode::NOT_FOUND, json!({"error": {"message": msg, "type": "model_not_found"}})) - } - MohawkError::ModelNotLoaded(msg) => { - (StatusCode::BAD_REQUEST, json!({"error": {"message": msg, "type": "model_not_loaded"}})) - } - MohawkError::InvalidRequest(msg) => { - (StatusCode::BAD_REQUEST, json!({"error": {"message": msg, "type": "invalid_request_error"}})) - } - MohawkError::InferenceFailed(msg) => { - (StatusCode::INTERNAL_SERVER_ERROR, json!({"error": {"message": msg, "type": "inference_error"}})) - } - MohawkError::TokenizationError(msg) => { - (StatusCode::BAD_REQUEST, json!({"error": {"message": msg, "type": "tokenization_error"}})) - } - MohawkError::BackendError(msg) => { - (StatusCode::SERVICE_UNAVAILABLE, json!({"error": {"message": msg, "type": "backend_error"}})) - } - MohawkError::RateLimitExceeded => { - (StatusCode::TOO_MANY_REQUESTS, json!({"error": {"message": "Rate limit exceeded", "type": "rate_limit_error"}})) - } - MohawkError::AuthenticationFailed => { - (StatusCode::UNAUTHORIZED, json!({"error": {"message": "Authentication failed", "type": "authentication_error"}})) - } - MohawkError::InternalError(msg) => { - (StatusCode::INTERNAL_SERVER_ERROR, json!({"error": {"message": msg, "type": "internal_error"}})) - } + MohawkError::ModelNotFound(msg) => ( + StatusCode::NOT_FOUND, + json!({"error": {"message": msg, "type": "model_not_found"}}), + ), + MohawkError::ModelNotLoaded(msg) => ( + StatusCode::BAD_REQUEST, + json!({"error": {"message": msg, "type": "model_not_loaded"}}), + ), + MohawkError::InvalidRequest(msg) => ( + StatusCode::BAD_REQUEST, + json!({"error": {"message": msg, "type": "invalid_request_error"}}), + ), + MohawkError::InferenceFailed(msg) => ( + StatusCode::INTERNAL_SERVER_ERROR, + json!({"error": {"message": msg, "type": "inference_error"}}), + ), + MohawkError::TokenizationError(msg) => ( + StatusCode::BAD_REQUEST, + json!({"error": {"message": msg, "type": "tokenization_error"}}), + ), + MohawkError::BackendError(msg) => ( + StatusCode::SERVICE_UNAVAILABLE, + json!({"error": {"message": msg, "type": "backend_error"}}), + ), + MohawkError::RateLimitExceeded => ( + StatusCode::TOO_MANY_REQUESTS, + json!({"error": {"message": "Rate limit exceeded", "type": "rate_limit_error"}}), + ), + MohawkError::AuthenticationFailed => ( + StatusCode::UNAUTHORIZED, + json!({"error": {"message": "Authentication failed", "type": "authentication_error"}}), + ), + MohawkError::InternalError(msg) => ( + StatusCode::INTERNAL_SERVER_ERROR, + json!({"error": {"message": msg, "type": "internal_error"}}), + ), }; (status, Json(error_message)).into_response() diff --git a/mohawk-server/src/lib.rs b/mohawk-server/src/lib.rs index e99c3da..5e6f7d1 100644 --- a/mohawk-server/src/lib.rs +++ b/mohawk-server/src/lib.rs @@ -1,21 +1,21 @@ //! Mohawk Inference Engine - Production-ready LLM serving -//! +//! //! A high-performance inference engine with OpenAI-compatible API, //! streaming support, and enterprise features. +mod api; +mod engine; mod error; mod models; -mod engine; -mod api; mod server; +pub use engine::InferenceEngine; pub use error::{MohawkError, Result}; pub use models::{ - Message, InferenceRequest, InferenceResponse, ChatChoice, Usage, - StreamToken, StreamResponse, ModelInfo, ModelStatus + ChatChoice, InferenceRequest, InferenceResponse, Message, ModelInfo, ModelStatus, + StreamResponse, StreamToken, Usage, }; -pub use engine::InferenceEngine; -pub use server::Server; +pub use server::{start_default, Server, ServerConfig}; use tracing::info; @@ -25,33 +25,45 @@ pub fn init_logging() { .with_env_filter( tracing_subscriber::EnvFilter::from_default_env() .add_directive("mohawk_server=info".parse().unwrap()) - .add_directive("axum=info".parse().unwrap()) + .add_directive("axum=info".parse().unwrap()), ) .json() .init(); - + info!("Mohawk Inference Engine initialized"); } #[cfg(test)] mod tests { use super::*; - + #[test] fn test_engine_creation() { - let engine = InferenceEngine::new(); + let engine = InferenceEngine::new(None); assert!(engine.is_ok()); } - + #[tokio::test] async fn test_inference_request() { - let engine = InferenceEngine::new().unwrap(); + let engine = InferenceEngine::new(None).unwrap(); + engine + .register_model(ModelInfo { + id: "test-model".to_string(), + name: "Test Model".to_string(), + parameters: "1B".to_string(), + quantization: "Q4".to_string(), + size_gb: 1.0, + loaded: false, + }) + .await + .unwrap(); + engine.load_model("test-model").await.unwrap(); let request = InferenceRequest { messages: vec![Message { role: "user".to_string(), content: "Hello!".to_string(), }], - model: None, + model: Some("test-model".to_string()), temperature: Some(0.7), top_p: None, top_k: None, @@ -60,7 +72,7 @@ mod tests { stop: None, system_prompt: None, }; - + let response = engine.generate(request).await; assert!(response.is_ok()); } diff --git a/mohawk-server/src/main.rs b/mohawk-server/src/main.rs index d1895d5..15c3009 100644 --- a/mohawk-server/src/main.rs +++ b/mohawk-server/src/main.rs @@ -10,13 +10,13 @@ async fn main() -> Result<(), Box> { .unwrap_or_else(|_| "8080".to_string()) .parse() .unwrap_or(8080); - + let config = ServerConfig { host, port, default_model: None, }; - + let server = Server::new(config)?; server.run().await } diff --git a/mohawk-server/src/models.rs b/mohawk-server/src/models.rs index 5782d1f..1b50af0 100644 --- a/mohawk-server/src/models.rs +++ b/mohawk-server/src/models.rs @@ -1,8 +1,8 @@ //! Data models for Mohawk Inference Engine //! Compatible with OpenAI API format +use chrono::Utc; use serde::{Deserialize, Serialize}; -use chrono::{DateTime, Utc}; use uuid::Uuid; /// Chat message role and content @@ -17,35 +17,35 @@ pub struct Message { pub struct InferenceRequest { /// Conversation messages pub messages: Vec, - + /// Model identifier (optional, uses default if not specified) #[serde(skip_serializing_if = "Option::is_none")] pub model: Option, - + /// Sampling temperature (0-2) #[serde(skip_serializing_if = "Option::is_none")] pub temperature: Option, - + /// Top-p sampling (0-1) #[serde(skip_serializing_if = "Option::is_none")] pub top_p: Option, - + /// Top-k sampling #[serde(skip_serializing_if = "Option::is_none")] pub top_k: Option, - + /// Maximum tokens to generate #[serde(skip_serializing_if = "Option::is_none")] pub max_tokens: Option, - + /// Enable streaming response #[serde(default)] pub stream: bool, - + /// Stop sequences #[serde(skip_serializing_if = "Option::is_none")] pub stop: Option>, - + /// System prompt override #[serde(skip_serializing_if = "Option::is_none")] pub system_prompt: Option, diff --git a/mohawk-server/src/server.rs b/mohawk-server/src/server.rs index 6d8c637..ab52be4 100644 --- a/mohawk-server/src/server.rs +++ b/mohawk-server/src/server.rs @@ -6,6 +6,7 @@ use crate::engine::InferenceEngine; use crate::init_logging; use axum::Router; use std::net::SocketAddr; +use tokio::net::TcpListener; use tracing::info; /// Server configuration @@ -35,14 +36,14 @@ pub struct Server { impl Server { /// Create a new server with configuration pub fn new(config: ServerConfig) -> Result { - let engine = InferenceEngine::new()?; - - Ok(Self { - config, - engine, - }) + let mut engine = InferenceEngine::new(None)?; + if let Some(default_model) = config.default_model.as_deref() { + engine.set_default_model(default_model); + } + + Ok(Self { config, engine }) } - + /// Register default models for the server pub async fn register_default_models(&self) -> Result<(), crate::error::MohawkError> { // Register pre-configured models (similar to LM Studio) @@ -72,32 +73,31 @@ impl Server { loaded: false, }, ]; - + for model in models { self.engine.register_model(model).await?; } - + info!("Registered {} default models", 3); Ok(()) } - + /// Build the Axum router fn build_router(&self) -> Router { api::create_router(self.engine.clone()) } - + /// Start the server pub async fn run(&self) -> Result<(), Box> { init_logging(); - + // Register default models self.register_default_models().await?; - - let addr: SocketAddr = format!("{}:{}", self.config.host, self.config.port) - .parse()?; - + + let addr: SocketAddr = format!("{}:{}", self.config.host, self.config.port).parse()?; + let router = self.build_router(); - + info!("🦅 Mohawk Inference Engine starting on http://{}", addr); info!("API endpoints:"); info!(" - Health: GET http://{}/health", addr); @@ -107,15 +107,17 @@ impl Server { info!(" - Chat: POST http://{}/v1/chat/completions", addr); info!(" - Metrics: GET http://{}/metrics", addr); info!(""); - info!("OpenAI Compatible: Use with any OpenAI SDK by setting base_url to http://{}", addr); - - axum::Server::bind(&addr) - .serve(router.into_make_service()) - .await?; - + info!( + "OpenAI Compatible: Use with any OpenAI SDK by setting base_url to http://{}", + addr + ); + + let listener = TcpListener::bind(addr).await?; + axum::serve(listener, router).await?; + Ok(()) } - + /// Get the engine instance for testing pub fn engine(&self) -> &InferenceEngine { &self.engine diff --git a/prototype/__pycache__/__init__.cpython-312.pyc b/prototype/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000..1829f27 Binary files /dev/null and b/prototype/__pycache__/__init__.cpython-312.pyc differ diff --git a/prototype/__pycache__/crypto_improved.cpython-312.pyc b/prototype/__pycache__/crypto_improved.cpython-312.pyc new file mode 100644 index 0000000..fdbdcd9 Binary files /dev/null and b/prototype/__pycache__/crypto_improved.cpython-312.pyc differ diff --git a/prototype/__pycache__/gui_backend.cpython-312.pyc b/prototype/__pycache__/gui_backend.cpython-312.pyc new file mode 100644 index 0000000..fadbca3 Binary files /dev/null and b/prototype/__pycache__/gui_backend.cpython-312.pyc differ diff --git a/prototype/__pycache__/model_tools.cpython-312.pyc b/prototype/__pycache__/model_tools.cpython-312.pyc new file mode 100644 index 0000000..46bd312 Binary files /dev/null and b/prototype/__pycache__/model_tools.cpython-312.pyc differ diff --git a/prototype/__pycache__/service_discovery.cpython-312.pyc b/prototype/__pycache__/service_discovery.cpython-312.pyc new file mode 100644 index 0000000..c39053f Binary files /dev/null and b/prototype/__pycache__/service_discovery.cpython-312.pyc differ diff --git a/prototype/__pycache__/telemetry.cpython-312.pyc b/prototype/__pycache__/telemetry.cpython-312.pyc new file mode 100644 index 0000000..99c0258 Binary files /dev/null and b/prototype/__pycache__/telemetry.cpython-312.pyc differ diff --git a/prototype/__pycache__/worker_secure.cpython-312.pyc b/prototype/__pycache__/worker_secure.cpython-312.pyc new file mode 100644 index 0000000..74ab7f7 Binary files /dev/null and b/prototype/__pycache__/worker_secure.cpython-312.pyc differ diff --git a/prototype/service_discovery.py b/prototype/service_discovery.py index fb80a5c..7c1c707 100644 --- a/prototype/service_discovery.py +++ b/prototype/service_discovery.py @@ -18,6 +18,7 @@ from zeroconf import ServiceBrowser, ServiceStateChange, Zeroconf except ImportError: Zeroconf = None + ServiceStateChange = None logger = logging.getLogger(__name__) diff --git a/prototype/worker_secure.py b/prototype/worker_secure.py index e7d0aa2..83b75eb 100644 --- a/prototype/worker_secure.py +++ b/prototype/worker_secure.py @@ -3,6 +3,7 @@ import pickle import threading import traceback +from datetime import datetime from typing import Dict import numpy as np @@ -264,6 +265,12 @@ async def execute(req: ExecRequest): raise HTTPException(status_code=400, detail=str(e)) +@app.get("/health") +async def health_check(): + """Health check endpoint.""" + return {"status": "healthy", "service": "mohawk-worker", "timestamp": datetime.now().isoformat()} + + @app.get("/metrics") async def get_metrics(): """Expose computed percentiles based on histogram buckets.""" @@ -319,6 +326,12 @@ def percentile(p): return JSONResponse(content=out) +@app.get("/api/workers") +async def list_workers(): + """Return list of loaded model slices (workers).""" + return {"workers": list(slices.keys()), "count": len(slices)} + + if __name__ == "__main__": import argparse