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