Skip to main content

rpi_cli/
llama_command.rs

1//! LLaMA.cpp integration command.
2//!
3//! Provides CLI commands for managing local LLM inference using llama.cpp server.
4
5use std::path::PathBuf;
6use rpi_ai::providers::llama_cpp::{LlamaCppConfig, LlamaCppModelManager};
7
8/// LLaMA subcommand
9#[derive(Debug)]
10pub enum LlamaCommand {
11    /// Start llama.cpp server
12    Start {
13        model: PathBuf,
14        port: u16,
15        ctx_size: usize,
16        threads: Option<usize>,
17    },
18    /// Stop llama.cpp server
19    Stop,
20    /// List available models
21    List {
22        server_url: String,
23    },
24    /// Show server status
25    Status {
26        server_url: String,
27    },
28    /// Download a model
29    Download {
30        model: String,
31        output: Option<PathBuf>,
32    },
33}
34
35/// Parse llama subcommand from args
36pub fn parse_llama_command(args: &[String]) -> Result<LlamaCommand, String> {
37    if args.is_empty() {
38        return Err("Usage: rpi llama <start|stop|list|status|download> [options]".to_string());
39    }
40    
41    let subcommand = &args[0];
42    let sub_args = &args[1..];
43    
44    match subcommand.as_str() {
45        "start" => parse_start_command(sub_args),
46        "stop" => Ok(LlamaCommand::Stop),
47        "list" => parse_list_command(sub_args),
48        "status" => parse_status_command(sub_args),
49        "download" => parse_download_command(sub_args),
50        _ => Err(format!("Unknown llama subcommand: {}", subcommand)),
51    }
52}
53
54fn parse_start_command(args: &[String]) -> Result<LlamaCommand, String> {
55    let mut model = None;
56    let mut port = 8080u16;
57    let mut ctx_size = 2048usize;
58    let mut threads = None;
59    
60    let mut i = 0;
61    while i < args.len() {
62        match args[i].as_str() {
63            "--model" | "-m" => {
64                i += 1;
65                if i >= args.len() {
66                    return Err("--model requires a value".to_string());
67                }
68                model = Some(PathBuf::from(&args[i]));
69            }
70            "--port" | "-p" => {
71                i += 1;
72                if i >= args.len() {
73                    return Err("--port requires a value".to_string());
74                }
75                port = args[i].parse().map_err(|_| "Invalid port number")?;
76            }
77            "--ctx-size" | "-c" => {
78                i += 1;
79                if i >= args.len() {
80                    return Err("--ctx-size requires a value".to_string());
81                }
82                ctx_size = args[i].parse().map_err(|_| "Invalid context size")?;
83            }
84            "--threads" | "-t" => {
85                i += 1;
86                if i >= args.len() {
87                    return Err("--threads requires a value".to_string());
88                }
89                threads = Some(args[i].parse().map_err(|_| "Invalid thread count")?);
90            }
91            _ => {
92                return Err(format!("Unknown option: {}", args[i]));
93            }
94        }
95        i += 1;
96    }
97    
98    let model = model.ok_or("--model is required")?;
99    
100    Ok(LlamaCommand::Start {
101        model,
102        port,
103        ctx_size,
104        threads,
105    })
106}
107
108fn parse_list_command(args: &[String]) -> Result<LlamaCommand, String> {
109    let mut server_url = "http://localhost:8080".to_string();
110    
111    let mut i = 0;
112    while i < args.len() {
113        match args[i].as_str() {
114            "--server" | "-s" => {
115                i += 1;
116                if i >= args.len() {
117                    return Err("--server requires a value".to_string());
118                }
119                server_url = args[i].clone();
120            }
121            _ => {
122                return Err(format!("Unknown option: {}", args[i]));
123            }
124        }
125        i += 1;
126    }
127    
128    Ok(LlamaCommand::List { server_url })
129}
130
131fn parse_status_command(args: &[String]) -> Result<LlamaCommand, String> {
132    let mut server_url = "http://localhost:8080".to_string();
133    
134    let mut i = 0;
135    while i < args.len() {
136        match args[i].as_str() {
137            "--server" | "-s" => {
138                i += 1;
139                if i >= args.len() {
140                    return Err("--server requires a value".to_string());
141                }
142                server_url = args[i].clone();
143            }
144            _ => {
145                return Err(format!("Unknown option: {}", args[i]));
146            }
147        }
148        i += 1;
149    }
150    
151    Ok(LlamaCommand::Status { server_url })
152}
153
154fn parse_download_command(args: &[String]) -> Result<LlamaCommand, String> {
155    if args.is_empty() {
156        return Err("Usage: rpi llama download <model> [--output <path>]".to_string());
157    }
158    
159    let model = args[0].clone();
160    let mut output = None;
161    
162    let mut i = 1;
163    while i < args.len() {
164        match args[i].as_str() {
165            "--output" | "-o" => {
166                i += 1;
167                if i >= args.len() {
168                    return Err("--output requires a value".to_string());
169                }
170                output = Some(PathBuf::from(&args[i]));
171            }
172            _ => {
173                return Err(format!("Unknown option: {}", args[i]));
174            }
175        }
176        i += 1;
177    }
178    
179    Ok(LlamaCommand::Download { model, output })
180}
181
182/// Execute llama command
183pub async fn run_llama_command(cmd: LlamaCommand) -> Result<(), String> {
184    match cmd {
185        LlamaCommand::Start { model, port, ctx_size, threads } => {
186            start_server(model, port, ctx_size, threads).await
187        }
188        LlamaCommand::Stop => {
189            stop_server().await
190        }
191        LlamaCommand::List { server_url } => {
192            list_models(server_url).await
193        }
194        LlamaCommand::Status { server_url } => {
195            show_status(server_url).await
196        }
197        LlamaCommand::Download { model, output } => {
198            download_model(model, output).await
199        }
200    }
201}
202
203async fn start_server(
204    model: PathBuf,
205    port: u16,
206    ctx_size: usize,
207    threads: Option<usize>,
208) -> Result<(), String> {
209    println!("Starting llama.cpp server...");
210    println!("  Model: {}", model.display());
211    println!("  Port: {}", port);
212    println!("  Context size: {}", ctx_size);
213    
214    if let Some(threads) = threads {
215        println!("  Threads: {}", threads);
216    }
217    
218    // Check if model file exists
219    if !model.exists() {
220        return Err(format!("Model file not found: {}", model.display()));
221    }
222    
223    // Build llama-server command
224    let mut cmd = tokio::process::Command::new("llama-server");
225    cmd.arg("--model").arg(&model)
226       .arg("--port").arg(port.to_string())
227       .arg("--ctx-size").arg(ctx_size.to_string());
228    
229    if let Some(threads) = threads {
230        cmd.arg("--threads").arg(threads.to_string());
231    }
232    
233    // Start server in background
234    let child = cmd.spawn().map_err(|e| format!("Failed to start llama-server: {}", e))?;
235    
236    println!("Server started (PID: {:?})", child.id());
237    println!("Server URL: http://localhost:{}", port);
238    
239    // Wait a bit for server to initialize
240    tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
241    
242    // Check if server is responding
243    let config = LlamaCppConfig {
244        server_url: format!("http://localhost:{}", port),
245        timeout: std::time::Duration::from_secs(5),
246        api_key: None,
247    };
248    
249    let manager = LlamaCppModelManager::new(config);
250    match manager.is_healthy().await {
251        Ok(true) => {
252            println!("Server is ready!");
253            Ok(())
254        }
255        Ok(false) => {
256            println!("Warning: Server started but not responding yet");
257            Ok(())
258        }
259        Err(e) => {
260            println!("Warning: Could not verify server status: {}", e);
261            Ok(())
262        }
263    }
264}
265
266async fn stop_server() -> Result<(), String> {
267    println!("Stopping llama.cpp server...");
268    
269    // Try to find and kill llama-server process
270    #[cfg(unix)]
271    {
272        let output = tokio::process::Command::new("pkill")
273            .arg("-f")
274            .arg("llama-server")
275            .output()
276            .await
277            .map_err(|e| format!("Failed to stop server: {}", e))?;
278        
279        if output.status.success() {
280            println!("Server stopped");
281        } else {
282            println!("No running server found");
283        }
284    }
285    
286    #[cfg(windows)]
287    {
288        let output = tokio::process::Command::new("taskkill")
289            .arg("/F")
290            .arg("/IM")
291            .arg("llama-server.exe")
292            .output()
293            .await
294            .map_err(|e| format!("Failed to stop server: {}", e))?;
295        
296        if output.status.success() {
297            println!("Server stopped");
298        } else {
299            println!("No running server found");
300        }
301    }
302    
303    Ok(())
304}
305
306async fn list_models(server_url: String) -> Result<(), String> {
307    let config = LlamaCppConfig {
308        server_url,
309        timeout: std::time::Duration::from_secs(10),
310        api_key: None,
311    };
312    
313    let manager = LlamaCppModelManager::new(config);
314    
315    println!("Fetching model list...");
316    
317    match manager.list_models().await {
318        Ok(models) => {
319            if models.is_empty() {
320                println!("No models available");
321            } else {
322                println!("Available models:");
323                for model in models {
324                    println!("  - {}", model);
325                }
326            }
327            Ok(())
328        }
329        Err(e) => {
330            Err(format!("Failed to list models: {}", e))
331        }
332    }
333}
334
335async fn show_status(server_url: String) -> Result<(), String> {
336    let config = LlamaCppConfig {
337        server_url: server_url.clone(),
338        timeout: std::time::Duration::from_secs(5),
339        api_key: None,
340    };
341    
342    let manager = LlamaCppModelManager::new(config);
343    
344    println!("Checking server status at {}...", server_url);
345    
346    match manager.is_healthy().await {
347        Ok(true) => {
348            println!("✓ Server is running and healthy");
349            
350            // Try to get model info
351            if let Ok(models) = manager.list_models().await {
352                if !models.is_empty() {
353                    println!("\nLoaded models:");
354                    for model in models {
355                        println!("  - {}", model);
356                    }
357                }
358            }
359            
360            Ok(())
361        }
362        Ok(false) => {
363            println!("✗ Server is not responding");
364            Ok(())
365        }
366        Err(e) => {
367            println!("✗ Failed to check server status: {}", e);
368            Ok(())
369        }
370    }
371}
372
373async fn download_model(model: String, output: Option<PathBuf>) -> Result<(), String> {
374    println!("Downloading model: {}", model);
375    
376    // Determine output path
377    let output_path = output.unwrap_or_else(|| {
378        let home = std::env::var("HOME")
379            .or_else(|_| std::env::var("USERPROFILE"))
380            .unwrap_or_else(|_| ".".to_string());
381        PathBuf::from(home).join(".cache").join("llama-cpp").join("models")
382    });
383    
384    // Create output directory
385    std::fs::create_dir_all(&output_path)
386        .map_err(|e| format!("Failed to create output directory: {}", e))?;
387    
388    println!("Output directory: {}", output_path.display());
389    
390    // Check if it's a URL
391    if model.starts_with("http://") || model.starts_with("https://") {
392        println!("Downloading from URL...");
393        
394        let output_file = output_path.join("model.gguf");
395        
396        let response = reqwest::get(&model)
397            .await
398            .map_err(|e| format!("Failed to download model: {}", e))?;
399        
400        if !response.status().is_success() {
401            return Err(format!("Download failed with status: {}", response.status()));
402        }
403        
404        let bytes = response.bytes()
405            .await
406            .map_err(|e| format!("Failed to read response: {}", e))?;
407        
408        std::fs::write(&output_file, bytes)
409            .map_err(|e| format!("Failed to write model file: {}", e))?;
410        
411        println!("Model downloaded to: {}", output_file.display());
412    } else {
413        // Assume it's a Hugging Face model ID
414        println!("Downloading from Hugging Face...");
415        println!("Note: Hugging Face download not yet implemented");
416        println!("Please download manually from: https://huggingface.co/{}", model);
417    }
418    
419    Ok(())
420}