From b2dbd73b8c8938105246f04e6ebcd68221740bd7 Mon Sep 17 00:00:00 2001 From: sigoden Date: Sat, 9 Nov 2024 07:39:21 +0800 Subject: [PATCH] refactor: path related functions (#981) --- src/config/agent.rs | 28 ++++--- src/config/mod.rs | 176 +++++++++++++++++++++----------------------- src/function.rs | 11 ++- src/serve.rs | 8 +- src/utils/path.rs | 8 +- 5 files changed, 106 insertions(+), 125 deletions(-) diff --git a/src/config/agent.rs b/src/config/agent.rs index c80a271d..e16e7a71 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -33,14 +33,14 @@ impl Agent { name: &str, abort_signal: AbortSignal, ) -> Result { - let functions_dir = Config::agent_functions_dir(name)?; + let functions_dir = Config::agent_functions_dir(name); let definition_file_path = functions_dir.join("index.yaml"); if !definition_file_path.exists() { bail!("Unknown agent `{name}`"); } let functions_file_path = functions_dir.join("functions.json"); - let rag_path = Config::agent_rag_file(name, DEFAULT_AGENT_NAME)?; - let config_path = Config::agent_config_file(name)?; + let rag_path = Config::agent_rag_file(name, DEFAULT_AGENT_NAME); + let config_path = Config::agent_config_file(name); let mut agent_config = if config_path.exists() { AgentConfig::load(&config_path)? } else { @@ -172,15 +172,15 @@ impl Agent { if !variables.is_empty() { value["variables"] = serde_json::to_value(variables)?; } - value["functions_dir"] = Config::agent_functions_dir(&self.name)? + value["functions_dir"] = Config::agent_functions_dir(&self.name) .display() .to_string() .into(); - value["data_dir"] = Config::agent_data_dir(&self.name)? + value["data_dir"] = Config::agent_data_dir(&self.name) .display() .to_string() .into(); - value["config_file"] = Config::agent_config_file(&self.name)? + value["config_file"] = Config::agent_config_file(&self.name) .display() .to_string() .into(); @@ -456,13 +456,12 @@ pub struct AgentVariable { } pub fn list_agents() -> Vec { - list_agents_impl().unwrap_or_default() -} - -fn list_agents_impl() -> Result> { - let base_dir = Config::functions_dir()?; - let contents = read_to_string(base_dir.join("agents.txt"))?; - let agents = contents + let agents_file = Config::functions_dir().join("agents.txt"); + let contents = match read_to_string(agents_file) { + Ok(v) => v, + Err(_) => return vec![], + }; + contents .split('\n') .filter_map(|line| { let line = line.trim(); @@ -472,6 +471,5 @@ fn list_agents_impl() -> Result> { Some(line.to_string()) } }) - .collect(); - Ok(agents) + .collect() } diff --git a/src/config/mod.rs b/src/config/mod.rs index 7139bd09..fa0882bf 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -221,7 +221,7 @@ pub type GlobalConfig = Arc>; impl Config { pub fn init(working_mode: WorkingMode) -> Result { - let config_path = Self::config_file()?; + let config_path = Self::config_file(); let mut config = if !config_path.exists() { match env::var(get_env_name("platform")) { Ok(v) => Self::load_dynamic(&v)?, @@ -253,130 +253,126 @@ impl Config { Ok(config) } - pub fn config_dir() -> Result { - let path = if let Ok(v) = env::var(get_env_name("config_dir")) { + pub fn config_dir() -> PathBuf { + if let Ok(v) = env::var(get_env_name("config_dir")) { PathBuf::from(v) } else if let Ok(v) = env::var("XDG_CONFIG_HOME") { PathBuf::from(v).join(env!("CARGO_CRATE_NAME")) } else { - let dir = dirs::config_dir().ok_or_else(|| anyhow!("Not available config dir"))?; + let dir = dirs::config_dir().expect("No user's config directory"); dir.join(env!("CARGO_CRATE_NAME")) - }; - Ok(path) + } } - pub fn local_path(name: &str) -> Result { - let mut path = Self::config_dir()?; - path.push(name); - Ok(path) + pub fn local_path(name: &str) -> PathBuf { + Self::config_dir().join(name) } - pub fn config_file() -> Result { + pub fn config_file() -> PathBuf { match env::var(get_env_name("config_file")) { - Ok(value) => Ok(PathBuf::from(value)), + Ok(value) => PathBuf::from(value), Err(_) => Self::local_path(CONFIG_FILE_NAME), } } - pub fn roles_dir() -> Result { + pub fn roles_dir() -> PathBuf { match env::var(get_env_name("roles_dir")) { - Ok(value) => Ok(PathBuf::from(value)), + Ok(value) => PathBuf::from(value), Err(_) => Self::local_path(ROLES_DIR_NAME), } } - pub fn role_file(name: &str) -> Result { - Ok(Self::roles_dir()?.join(format!("{name}.md"))) + pub fn role_file(name: &str) -> PathBuf { + Self::roles_dir().join(format!("{name}.md")) } - pub fn env_file() -> Result { + pub fn env_file() -> PathBuf { match env::var(get_env_name("env_file")) { - Ok(value) => Ok(PathBuf::from(value)), + Ok(value) => PathBuf::from(value), Err(_) => Self::local_path(ENV_FILE_NAME), } } - pub fn messages_file(&self) -> Result { + pub fn messages_file(&self) -> PathBuf { match &self.agent { None => match env::var(get_env_name("messages_file")) { - Ok(value) => Ok(PathBuf::from(value)), + Ok(value) => PathBuf::from(value), Err(_) => Self::local_path(MESSAGES_FILE_NAME), }, - Some(agent) => Ok(Self::agent_data_dir(agent.name())?.join(MESSAGES_FILE_NAME)), + Some(agent) => Self::agent_data_dir(agent.name()).join(MESSAGES_FILE_NAME), } } - pub fn sessions_dir(&self) -> Result { + pub fn sessions_dir(&self) -> PathBuf { match &self.agent { None => match env::var(get_env_name("sessions_dir")) { - Ok(value) => Ok(PathBuf::from(value)), + Ok(value) => PathBuf::from(value), Err(_) => Self::local_path(SESSIONS_DIR_NAME), }, - Some(agent) => Ok(Self::agent_data_dir(agent.name())?.join(SESSIONS_DIR_NAME)), + Some(agent) => Self::agent_data_dir(agent.name()).join(SESSIONS_DIR_NAME), } } - pub fn rags_dir() -> Result { + pub fn rags_dir() -> PathBuf { match env::var(get_env_name("rags_dir")) { - Ok(value) => Ok(PathBuf::from(value)), + Ok(value) => PathBuf::from(value), Err(_) => Self::local_path(RAGS_DIR_NAME), } } - pub fn functions_dir() -> Result { + pub fn functions_dir() -> PathBuf { match env::var(get_env_name("functions_dir")) { - Ok(value) => Ok(PathBuf::from(value)), + Ok(value) => PathBuf::from(value), Err(_) => Self::local_path(FUNCTIONS_DIR_NAME), } } - pub fn functions_file() -> Result { - Ok(Self::functions_dir()?.join(FUNCTIONS_FILE_NAME)) + pub fn functions_file() -> PathBuf { + Self::functions_dir().join(FUNCTIONS_FILE_NAME) } - pub fn functions_bin_dir() -> Result { - Ok(Self::functions_dir()?.join(FUNCTIONS_BIN_DIR_NAME)) + pub fn functions_bin_dir() -> PathBuf { + Self::functions_dir().join(FUNCTIONS_BIN_DIR_NAME) } - pub fn session_file(&self, name: &str) -> Result { - Ok(self.sessions_dir()?.join(format!("{name}.yaml"))) + pub fn session_file(&self, name: &str) -> PathBuf { + self.sessions_dir().join(format!("{name}.yaml")) } - pub fn rag_file(&self, name: &str) -> Result { - let path = match &self.agent { - Some(agent) => Self::agent_rag_file(agent.name(), name)?, - None => Self::rags_dir()?.join(format!("{name}.yaml")), - }; - Ok(path) + pub fn rag_file(&self, name: &str) -> PathBuf { + match &self.agent { + Some(agent) => Self::agent_rag_file(agent.name(), name), + None => Self::rags_dir().join(format!("{name}.yaml")), + } } - pub fn agents_data_dir() -> Result { + pub fn agents_data_dir() -> PathBuf { Self::local_path(AGENTS_DIR_NAME) } - pub fn agent_data_dir(name: &str) -> Result { + pub fn agent_data_dir(name: &str) -> PathBuf { match env::var(format!("{}_DATA_DIR", normalize_env_name(name))) { - Ok(value) => Ok(PathBuf::from(value)), - Err(_) => Ok(Self::agents_data_dir()?.join(name)), + Ok(value) => PathBuf::from(value), + Err(_) => Self::agents_data_dir().join(name), } } - pub fn agent_config_file(name: &str) -> Result { - Ok(Self::agent_data_dir(name)?.join(CONFIG_FILE_NAME)) + pub fn agent_config_file(name: &str) -> PathBuf { + Self::agent_data_dir(name).join(CONFIG_FILE_NAME) } - pub fn agent_rag_file(agent_name: &str, rag_name: &str) -> Result { - Ok(Self::agent_data_dir(agent_name)?.join(format!("{rag_name}.yaml"))) + pub fn agent_rag_file(agent_name: &str, rag_name: &str) -> PathBuf { + Self::agent_data_dir(agent_name).join(format!("{rag_name}.yaml")) } - pub fn agents_functions_dir() -> Result { - Ok(Self::functions_dir()?.join(AGENTS_DIR_NAME)) + pub fn agents_functions_dir() -> PathBuf { + Self::functions_dir().join(AGENTS_DIR_NAME) } - pub fn agent_functions_dir(name: &str) -> Result { + pub fn agent_functions_dir(name: &str) -> PathBuf { match env::var(format!("{}_FUNCTIONS_DIR", normalize_env_name(name))) { - Ok(value) => Ok(PathBuf::from(value)), - Err(_) => Ok(Self::agents_functions_dir()?.join(name)), + Ok(value) => PathBuf::from(value), + Err(_) => Self::agents_functions_dir().join(name), } } @@ -431,7 +427,7 @@ impl Config { false => Some(Config::local_path(&format!( "{}.log", env!("CARGO_CRATE_NAME") - ))?), + ))), }, }; Ok((log_level, log_path)) @@ -556,13 +552,13 @@ impl Config { ("rag_top_k", rag_top_k.to_string()), ("highlight", self.highlight.to_string()), ("light_theme", self.light_theme.to_string()), - ("env_file", display_path(&Self::env_file()?)), - ("config_file", display_path(&Self::config_file()?)), - ("roles_dir", display_path(&Self::roles_dir()?)), - ("sessions_dir", display_path(&self.sessions_dir()?)), - ("rags_dir", display_path(&Self::rags_dir()?)), - ("functions_dir", display_path(&Self::functions_dir()?)), - ("messages_file", display_path(&self.messages_file()?)), + ("config_file", display_path(&Self::config_file())), + ("env_file", display_path(&Self::env_file())), + ("roles_dir", display_path(&Self::roles_dir())), + ("sessions_dir", display_path(&self.sessions_dir())), + ("rags_dir", display_path(&Self::rags_dir())), + ("functions_dir", display_path(&Self::functions_dir())), + ("messages_file", display_path(&self.messages_file())), ]; if let Ok((_, Some(log_path))) = Self::log_config(self.working_mode.is_serve()) { items.push(("log_path", display_path(&log_path))); @@ -649,10 +645,10 @@ impl Config { pub fn delete(config: &GlobalConfig, kind: &str) -> Result<()> { let (dir, file_ext) = match kind { - "role" => (Self::roles_dir()?, Some(".md")), - "session" => (config.read().sessions_dir()?, Some(".yaml")), - "rag" => (Self::rags_dir()?, Some(".yaml")), - "agent-data" => (Self::agents_data_dir()?, None), + "role" => (Self::roles_dir(), Some(".md")), + "session" => (config.read().sessions_dir(), Some(".yaml")), + "rag" => (Self::rags_dir(), Some(".yaml")), + "agent-data" => (Self::agents_data_dir(), None), _ => bail!("Unknown kind '{kind}'"), }; let names = match read_dir(&dir) { @@ -870,7 +866,7 @@ impl Config { pub fn retrieve_role(&self, name: &str) -> Result { let names = Self::list_roles(false); let mut role = if let Some(role_name) = Role::match_name(&names, name) { - let path = Self::role_file(&role_name)?; + let path = Self::role_file(&role_name); let content = read_to_string(&path)?; Role::new(name, &content) } else { @@ -911,7 +907,7 @@ impl Config { pub fn upsert_role(&mut self, name: &str) -> Result<()> { let names = Self::list_roles(false); let role_name = Role::match_name(&names, name).unwrap_or_else(|| name.to_string()); - let role_path = Self::role_file(&role_name)?; + let role_path = Self::role_file(&role_name); ensure_parent_exists(&role_path)?; let editor = self.editor()?; edit_file(&editor, &role_path)?; @@ -946,7 +942,7 @@ impl Config { }) .prompt()?; } - let role_path = Self::role_file(&role_name)?; + let role_path = Self::role_file(&role_name); if let Some(role) = self.role.as_mut() { role.save(&role_name, &role_path, self.working_mode.is_repl())?; } @@ -961,11 +957,9 @@ impl Config { .collect(); let names = Self::list_roles(false); for name in names { - if let Ok(path) = Self::role_file(&name) { - if let Ok(content) = read_to_string(&path) { - let role = Role::new(&name, &content); - roles.insert(name, role); - } + if let Ok(content) = read_to_string(Self::role_file(&name)) { + let role = Role::new(&name, &content); + roles.insert(name, role); } } let mut roles: Vec<_> = roles.into_values().collect(); @@ -975,7 +969,7 @@ impl Config { pub fn list_roles(with_builtin: bool) -> Vec { let mut names = HashSet::new(); - if let Some(rd) = Self::roles_dir().ok().and_then(|dir| read_dir(dir).ok()) { + if let Ok(rd) = read_dir(Self::roles_dir()) { for entry in rd.flatten() { if let Some(name) = entry .file_name() @@ -1008,7 +1002,7 @@ impl Config { let mut session; match session_name { None | Some(TEMP_SESSION_NAME) => { - let session_file = self.session_file(TEMP_SESSION_NAME)?; + let session_file = self.session_file(TEMP_SESSION_NAME); if session_file.exists() { remove_file(session_file).with_context(|| { format!("Failed to cleanup previous '{TEMP_SESSION_NAME}' session") @@ -1017,7 +1011,7 @@ impl Config { session = Some(Session::new(self, TEMP_SESSION_NAME)); } Some(name) => { - let session_path = self.session_file(name)?; + let session_path = self.session_file(name); if !session_path.exists() { session = Some(Session::new(self, name)); } else { @@ -1058,7 +1052,7 @@ impl Config { pub fn exit_session(&mut self) -> Result<()> { if let Some(mut session) = self.session.take() { - let sessions_dir = self.sessions_dir()?; + let sessions_dir = self.sessions_dir(); session.exit(&sessions_dir, self.working_mode.is_repl())?; self.last_message = None; } @@ -1076,7 +1070,7 @@ impl Config { }, None => bail!("No session"), }; - let session_path = self.session_file(&session_name)?; + let session_path = self.session_file(&session_name); if let Some(session) = self.session.as_mut() { session.save(&session_name, &session_path, self.working_mode.is_repl())?; } @@ -1088,7 +1082,7 @@ impl Config { Some(session) => session.name().to_string(), None => bail!("No session"), }; - let session_path = self.session_file(&name)?; + let session_path = self.session_file(&name); self.save_session(Some(&name))?; let editor = self.editor()?; edit_file(&editor, &session_path).with_context(|| { @@ -1115,7 +1109,7 @@ impl Config { Ok(()) } pub fn list_sessions(&self) -> Vec { - list_file_names(self.sessions_dir().ok(), ".yaml") + list_file_names(self.sessions_dir(), ".yaml") } pub fn should_compress_session(&mut self) -> bool { @@ -1179,7 +1173,7 @@ impl Config { } let rag = match rag { None => { - let rag_path = config.read().rag_file(TEMP_RAG_NAME)?; + let rag_path = config.read().rag_file(TEMP_RAG_NAME); if rag_path.exists() { remove_file(&rag_path).with_context(|| { format!("Failed to cleanup previous '{TEMP_RAG_NAME}' rag") @@ -1188,7 +1182,7 @@ impl Config { Rag::init(config, TEMP_RAG_NAME, &rag_path, &[], abort_signal).await? } Some(name) => { - let rag_path = config.read().rag_file(name)?; + let rag_path = config.read().rag_file(name); if !rag_path.exists() { if config.read().working_mode.is_cmd() { bail!("Unknown RAG '{name}'") @@ -1315,11 +1309,7 @@ impl Config { } pub fn list_rags() -> Vec { - let rags_dir = match Self::rags_dir() { - Ok(dir) => dir, - Err(_) => return vec![], - }; - match read_dir(rags_dir) { + match read_dir(Self::rags_dir()) { Ok(rd) => { let mut names = vec![]; for entry in rd.flatten() { @@ -1633,9 +1623,7 @@ impl Config { values = candidates.into_iter().map(|v| (v, None)).collect(); filter = args[1]; } else if cmd == ".agent" && args.len() >= 2 { - let dir = Self::agent_data_dir(args[0]) - .ok() - .map(|v| v.join(SESSIONS_DIR_NAME)); + let dir = Self::agent_data_dir(args[0]).join(SESSIONS_DIR_NAME); values = list_file_names(dir, ".yaml") .into_iter() .map(|v| (v, None)) @@ -1666,7 +1654,7 @@ impl Config { let theme = if self.highlight { let theme_mode = if self.light_theme { "light" } else { "dark" }; let theme_filename = format!("{theme_mode}.tmTheme"); - let theme_path = Self::local_path(&theme_filename)?; + let theme_path = Self::local_path(&theme_filename); if theme_path.exists() { let theme = ThemeSet::get_theme(&theme_path) .with_context(|| format!("Invalid theme at '{}'", theme_path.display()))?; @@ -1890,7 +1878,7 @@ impl Config { } fn open_message_file(&self) -> Result { - let path = self.messages_file()?; + let path = self.messages_file(); ensure_parent_exists(&path)?; OpenOptions::new() .create(true) @@ -2079,7 +2067,7 @@ impl Config { } fn load_functions(&mut self) -> Result<()> { - self.functions = Functions::init(&Self::functions_file()?)?; + self.functions = Functions::init(&Self::functions_file())?; Ok(()) } @@ -2118,7 +2106,7 @@ impl Config { } pub fn load_env_file() -> Result<()> { - let env_file_path = Config::env_file()?; + let env_file_path = Config::env_file(); let contents = match read_to_string(&env_file_path) { Ok(v) => v, Err(_) => return Ok(()), diff --git a/src/function.rs b/src/function.rs index 9c884e5f..517b9eb8 100644 --- a/src/function.rs +++ b/src/function.rs @@ -221,14 +221,12 @@ impl ToolCall { let mut bin_dirs: Vec = vec![]; if let Some(agent) = config.read().agent.as_ref() { - let dir = Config::agent_functions_dir(agent.name()) - .context("No agent functions dir")? - .join("bin"); + let dir = Config::agent_functions_dir(agent.name()).join("bin"); if dir.exists() { bin_dirs.push(dir); } } - bin_dirs.push(Config::functions_bin_dir().context("No functions bin dir")?); + bin_dirs.push(Config::functions_bin_dir()); let current_path = std::env::var("PATH").context("No PATH environment variable")?; let prepend_path = bin_dirs .iter() @@ -240,6 +238,11 @@ impl ToolCall { let temp_file = temp_file("-eval-", ""); envs.insert("LLM_OUTPUT".into(), temp_file.display().to_string()); + envs.insert( + "LLM_FUNCTIONS_ROOT".into(), + Config::functions_dir().display().to_string(), + ); + #[cfg(windows)] let cmd_name = polyfill_cmd_name(&cmd_name, &bin_dirs); if *IS_STDOUT_TERMINAL { diff --git a/src/serve.rs b/src/serve.rs index ebb79af9..56249030 100644 --- a/src/serve.rs +++ b/src/serve.rs @@ -246,12 +246,8 @@ impl Server { let abort_signal = create_abort_signal(); - let rag = config - .read() - .rag_file(&name) - .ok() - .and_then(|rag_path| Rag::load(&config, &name, &rag_path).ok()) - .ok_or_else(|| anyhow!("Invalid rag"))?; + let rag_path = config.read().rag_file(&name); + let rag = Rag::load(&config, &name, &rag_path)?; let rag_result = Config::search_rag(&config, &rag, &input, abort_signal).await?; diff --git a/src/utils/path.rs b/src/utils/path.rs index b0a5614f..635604ea 100644 --- a/src/utils/path.rs +++ b/src/utils/path.rs @@ -42,12 +42,8 @@ pub async fn expand_glob_paths>(paths: &[T]) -> Result Ok(new_paths) } -pub fn list_file_names>(dir: Option, ext: &str) -> Vec { - let dir = match dir { - Some(v) => v, - None => return vec![], - }; - match std::fs::read_dir(dir) { +pub fn list_file_names>(dir: T, ext: &str) -> Vec { + match std::fs::read_dir(dir.as_ref()) { Ok(rd) => { let mut names = vec![]; for entry in rd.flatten() {