1use rpi_ai::providers::llama_cpp::{LlamaCppConfig, LlamaCppModelManager};
6use std::path::PathBuf;
7
8#[derive(Debug)]
10pub enum LlamaCommand {
11 Start {
13 model: PathBuf,
14 port: u16,
15 ctx_size: usize,
16 threads: Option<usize>,
17 },
18 Stop,
20 List { server_url: String },
22 Status { server_url: String },
24 Download {
26 model: String,
27 output: Option<PathBuf>,
28 },
29}
30
31pub 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
178pub 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 if !model.exists() {
211 return Err(format!("Model file not found: {}", model.display()));
212 }
213
214 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 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 tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
237
238 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 #[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 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 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 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 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 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}