1use std::path::PathBuf;
6use rpi_ai::providers::llama_cpp::{LlamaCppConfig, LlamaCppModelManager};
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 {
22 server_url: String,
23 },
24 Status {
26 server_url: String,
27 },
28 Download {
30 model: String,
31 output: Option<PathBuf>,
32 },
33}
34
35pub 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
182pub 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 if !model.exists() {
220 return Err(format!("Model file not found: {}", model.display()));
221 }
222
223 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 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 tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
241
242 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 #[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 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 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 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 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 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}