From bb9118d477b7550ce9f73ed83ec0e1a600aaa75c Mon Sep 17 00:00:00 2001 From: Logan King Date: Sat, 19 Jul 2025 18:15:36 -0700 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20(assistant=5Fv2):=20add=20process?= =?UTF-8?q?=20management?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- assistant_v2/Cargo.lock | 63 +++++++++++ assistant_v2/Cargo.toml | 1 + assistant_v2/FEATURE_PROGRESS.md | 2 +- assistant_v2/src/main.rs | 172 +++++++++++++++++++++++++++++++ 4 files changed, 237 insertions(+), 1 deletion(-) diff --git a/assistant_v2/Cargo.lock b/assistant_v2/Cargo.lock index 024fa0f..cd5b06a 100644 --- a/assistant_v2/Cargo.lock +++ b/assistant_v2/Cargo.lock @@ -128,6 +128,7 @@ dependencies = [ "rdev", "serde_json", "speakstream", + "sysinfo", "tempfile", "tokio", "tracing", @@ -682,6 +683,25 @@ dependencies = [ "windows", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.21" @@ -1642,6 +1662,15 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "ntapi" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8a3895c6391c39d7fe7ebc444a87eb2991b2a0bc718fdabd071eec617fc68e4" +dependencies = [ + "winapi", +] + [[package]] name = "nu-ansi-term" version = "0.46.0" @@ -2026,6 +2055,26 @@ dependencies = [ "getrandom 0.3.3", ] +[[package]] +name = "rayon" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b418a60154510ca1a002a752ca9714984e21e4241e804d32555251faf8b78ffa" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1465873a3dfdaa8ae7cb14b4383657caab0b3e8a0aa9ae8e04b044854c8dfce2" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "rdev" version = "0.5.3" @@ -2523,6 +2572,20 @@ dependencies = [ "syn", ] +[[package]] +name = "sysinfo" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c33cd241af0f2e9e3b5c32163b873b29956890b5342e6745b917ce9d490f4af" +dependencies = [ + "core-foundation-sys 0.8.7", + "libc", + "memchr", + "ntapi", + "rayon", + "windows", +] + [[package]] name = "tempfile" version = "3.20.0" diff --git a/assistant_v2/Cargo.toml b/assistant_v2/Cargo.toml index 831c5bb..eb922f0 100644 --- a/assistant_v2/Cargo.toml +++ b/assistant_v2/Cargo.toml @@ -23,3 +23,4 @@ clap = { version = "4.4.6", features = ["derive"] } colored = "2.0.4" clipboard = "0.5.0" open = "5.3.1" +sysinfo = "0.32.1" diff --git a/assistant_v2/FEATURE_PROGRESS.md b/assistant_v2/FEATURE_PROGRESS.md index a85d7d8..3550527 100644 --- a/assistant_v2/FEATURE_PROGRESS.md +++ b/assistant_v2/FEATURE_PROGRESS.md @@ -10,7 +10,7 @@ This document tracks which features from the original assistant have been implem | Launch applications from voice | Pending | | Display log files | Pending | | Get system info | Pending | -| List and kill processes | Pending | +| List and kill processes | Done | | Run internet speed tests | Pending | | Set the clipboard contents | Done | | Timers with alarm sounds | Pending | diff --git a/assistant_v2/src/main.rs b/assistant_v2/src/main.rs index 7912ea3..d8b32c9 100644 --- a/assistant_v2/src/main.rs +++ b/assistant_v2/src/main.rs @@ -10,6 +10,8 @@ use async_openai::{ use clap::Parser; use clipboard::{ClipboardContext, ClipboardProvider}; use colored::Colorize; +use sysinfo::System; +use tracing::debug; use dotenvy::dotenv; use futures::StreamExt; use open; @@ -239,6 +241,32 @@ async fn main() -> Result<(), Box> { strict: None, } .into(), + FunctionObject { + name: "get_system_processes".into(), + description: Some( + "Returns information about running system processes.".into(), + ), + parameters: Some(serde_json::json!({ + "type": "object", + "properties": {}, + "required": [], + })), + strict: None, + } + .into(), + FunctionObject { + name: "kill_processes_with_name".into(), + description: Some( + "Kills all processes with a given name. ALWAYS call \"get_system_processes\" first to get the name of the process you want to kill.".into(), + ), + parameters: Some(serde_json::json!({ + "type": "object", + "properties": {"process_name": {"type": "string"}}, + "required": ["process_name"], + })), + strict: None, + } + .into(), ]) .build()?; @@ -561,6 +589,30 @@ async fn handle_requires_action( output: Some(format!("{}", name).into()), }); } + + if tool.function.name == "get_system_processes" { + let output = get_system_processes(); + tool_outputs.push(ToolsOutputs { + tool_call_id: Some(tool.id.clone()), + output: Some(output.into()), + }); + } + + if tool.function.name == "kill_processes_with_name" { + let name = match serde_json::from_str::(&tool.function.arguments) { + Ok(v) => v["process_name"].as_str().unwrap_or("").to_string(), + Err(_) => String::new(), + }; + let result = kill_processes_with_name(&name); + let msg = match result { + Some(_) => format!("Killed all processes with name: \"{}\"", name), + None => format!("Failed to kill all processes with name: {}", name), + }; + tool_outputs.push(ToolsOutputs { + tool_call_id: Some(tool.id.clone()), + output: Some(msg.into()), + }); + } } if let Err(e) = submit_tool_outputs(client, run_object, tool_outputs, speak_stream).await { @@ -613,6 +665,68 @@ async fn submit_tool_outputs( Ok(()) } +fn get_system_processes() -> String { + let mut info = String::new(); + let mut sys = System::new_all(); + sys.refresh_all(); + info.push_str("=> processes:\n"); + for (pid, process) in sys.processes() { + let path_string = match process.exe() { + Some(path) => path.to_string_lossy().to_string(), + None => "Unknown".to_string(), + }; + info.push_str(&format!( + "[{}] {:?} start_time: {:?} runtime: {} status: {} cpu_usage: {} directory: {}\n", + pid, + process.name(), + process.start_time(), + process.run_time(), + process.status(), + process.cpu_usage(), + path_string, + )); + } + info +} + +fn get_process_names() -> Vec { + let sys = System::new_all(); + let mut process_names = Vec::new(); + for process in sys.processes().values() { + let process_name = process.name().to_string_lossy().to_string(); + if !process_names.contains(&process_name) { + process_names.push(process_name); + } + } + process_names +} + +fn kill_processes_with_name(process_name: &str) -> Option<()> { + let mut sys = System::new_all(); + for (pid, process) in sys.processes() { + if process.name().to_string_lossy() == process_name { + if process.kill() { + debug!( + "Killed process with name: \"{}\" and PID: {}", + process_name, pid + ); + } else { + debug!( + "[Potentially] killed process with name: \"{}\" and PID: {}", + process_name, pid + ); + } + } + } + std::thread::sleep(std::time::Duration::from_secs(1)); + sys.refresh_all(); + if get_process_names().contains(&process_name.to_string()) { + None + } else { + Some(()) + } +} + #[cfg(test)] mod tests { use super::*; @@ -735,4 +849,62 @@ mod tests { _ => false, })); } + + #[test] + fn includes_get_system_processes_function() { + let req = CreateAssistantRequestArgs::default() + .instructions("test") + .model("gpt-4o") + .tools(vec![FunctionObject { + name: "get_system_processes".into(), + description: Some( + "Returns information about running system processes.".into(), + ), + parameters: Some(serde_json::json!({ + "type": "object", + "properties": {}, + "required": [], + })), + strict: None, + } + .into()]) + .build() + .unwrap(); + + let tools = req.tools.unwrap(); + assert!(tools.iter().any(|t| match t { + async_openai::types::AssistantTools::Function(f) => + f.function.name == "get_system_processes", + _ => false, + })); + } + + #[test] + fn includes_kill_processes_with_name_function() { + let req = CreateAssistantRequestArgs::default() + .instructions("test") + .model("gpt-4o") + .tools(vec![FunctionObject { + name: "kill_processes_with_name".into(), + description: Some( + "Kills all processes with a given name. ALWAYS call \"get_system_processes\" first to get the name of the process you want to kill.".into(), + ), + parameters: Some(serde_json::json!({ + "type": "object", + "properties": {"process_name": {"type": "string"}}, + "required": ["process_name"], + })), + strict: None, + } + .into()]) + .build() + .unwrap(); + + let tools = req.tools.unwrap(); + assert!(tools.iter().any(|t| match t { + async_openai::types::AssistantTools::Function(f) => + f.function.name == "kill_processes_with_name", + _ => false, + })); + } }