summaryrefslogtreecommitdiffstats
path: root/src/config/mod.rs
diff options
context:
space:
mode:
Diffstat (limited to 'src/config/mod.rs')
-rw-r--r--src/config/mod.rs196
1 files changed, 143 insertions, 53 deletions
diff --git a/src/config/mod.rs b/src/config/mod.rs
index 508b768..72f6d5a 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -5,7 +5,7 @@ mod session;
pub use self::agent::{list_agents, Agent};
pub use self::input::Input;
-pub use self::role::{Role, RoleLike, CODE_ROLE, EXPLAIN_SHELL_ROLE, SHELL_ROLE};
+pub use self::role::{Role, RoleLike, BUILTIN_ROLES, CODE_ROLE, EXPLAIN_SHELL_ROLE, SHELL_ROLE};
use self::session::Session;
use crate::client::{
@@ -19,7 +19,7 @@ use crate::utils::*;
use anyhow::{anyhow, bail, Context, Result};
use indexmap::IndexMap;
-use inquire::{Confirm, Select};
+use inquire::{validator::Validation, Confirm, Select, Text};
use parking_lot::RwLock;
use serde::Deserialize;
use serde_json::json;
@@ -40,7 +40,7 @@ const DARK_THEME: &[u8] = include_bytes!("../../assets/monokai-extended.theme.bi
const LIGHT_THEME: &[u8] = include_bytes!("../../assets/monokai-extended-light.theme.bin");
const CONFIG_FILE_NAME: &str = "config.yaml";
-const ROLES_FILE_NAME: &str = "roles.yaml";
+const ROLES_DIR_NAME: &str = "roles";
const ENV_FILE_NAME: &str = ".env";
const MESSAGES_FILE_NAME: &str = "messages.md";
const SESSIONS_DIR_NAME: &str = "sessions";
@@ -129,8 +129,6 @@ pub struct Config {
pub clients: Vec<ClientConfig>,
#[serde(skip)]
- pub roles: Vec<Role>,
- #[serde(skip)]
pub role: Option<Role>,
#[serde(skip)]
pub session: Option<Session>,
@@ -195,7 +193,6 @@ impl Default for Config {
clients: vec![],
- roles: vec![],
role: None,
session: None,
rag: None,
@@ -236,7 +233,6 @@ impl Config {
}
config.load_functions()?;
- config.load_roles()?;
config.setup_model()?;
config.setup_document_loaders();
@@ -269,13 +265,17 @@ impl Config {
}
}
- pub fn roles_file() -> Result<PathBuf> {
- match env::var(get_env_name("roles_file")) {
+ pub fn roles_dir() -> Result<PathBuf> {
+ match env::var(get_env_name("roles_dir")) {
Ok(value) => Ok(PathBuf::from(value)),
- Err(_) => Self::local_path(ROLES_FILE_NAME),
+ Err(_) => Self::local_path(ROLES_DIR_NAME),
}
}
+ pub fn role_file(name: &str) -> Result<PathBuf> {
+ Ok(Self::roles_dir()?.join(format!("{name}.md")))
+ }
+
pub fn env_file() -> Result<PathBuf> {
match env::var(get_env_name("env_file")) {
Ok(value) => Ok(PathBuf::from(value)),
@@ -487,7 +487,7 @@ impl Config {
} else if let Some(session) = &self.session {
session.export()
} else if let Some(role) = &self.role {
- role.export()
+ Ok(role.export())
} else if let Some(rag) = &self.rag {
rag.export()
} else {
@@ -531,7 +531,7 @@ impl Config {
("highlight", self.highlight.to_string()),
("light_theme", self.light_theme.to_string()),
("config_file", display_path(&Self::config_file()?)),
- ("roles_file", display_path(&Self::roles_file()?)),
+ ("roles_dir", display_path(&Self::roles_dir()?)),
("env_file", display_path(&Self::env_file()?)),
("functions_dir", display_path(&Self::functions_dir()?)),
("rags_dir", display_path(&Self::rags_dir()?)),
@@ -543,9 +543,9 @@ impl Config {
}
let output = items
.iter()
- .map(|(name, value)| format!("{name:<24}{value}"))
+ .map(|(name, value)| format!("{name:<24}{value}\n"))
.collect::<Vec<String>>()
- .join("\n");
+ .join("");
Ok(output)
}
@@ -716,7 +716,7 @@ impl Config {
pub fn role_info(&self) -> Result<String> {
if let Some(role) = &self.role {
- role.export()
+ Ok(role.export())
} else {
bail!("No role")
}
@@ -733,17 +733,17 @@ impl Config {
}
pub fn retrieve_role(&self, name: &str) -> Result<Role> {
- let mut role = self
- .roles
- .iter()
- .find(|v| v.match_name(name))
- .map(|v| {
- let mut role = v.clone();
- role.complete_prompt_args(name);
- role
- })
- .ok_or_else(|| anyhow!("Unknown role `{name}`"))?;
-
+ let mut role = if Self::list_roles(false).contains(&name.to_string()) {
+ let path = Self::role_file(name)?;
+ let content = read_to_string(path)?;
+ Role::new(name, &content)
+ } else {
+ BUILTIN_ROLES
+ .iter()
+ .find(|v| v.name() == name)
+ .cloned()
+ .ok_or_else(|| anyhow!("Unknown role `{name}`"))?
+ };
match role.model_id() {
Some(model_id) => {
if self.model.id() != model_id {
@@ -758,6 +758,110 @@ impl Config {
Ok(role)
}
+ pub fn new_role(&mut self, name: &str) -> Result<()> {
+ let ans = Confirm::new("Create a new role?")
+ .with_default(true)
+ .prompt()?;
+ if ans {
+ self.upsert_role(name)?;
+ }
+ Ok(())
+ }
+
+ pub fn edit_role(&mut self) -> Result<()> {
+ if let Some(name) = self.role.as_ref().map(|v| v.name().to_string()) {
+ self.upsert_role(&name)
+ } else {
+ bail!("No role")
+ }
+ }
+
+ pub fn upsert_role(&mut self, name: &str) -> Result<()> {
+ let role_path = Self::role_file(name)?;
+ ensure_parent_exists(&role_path)?;
+ let editor = self.editor()?;
+ edit_file(&editor, &role_path)?;
+ self.use_role(name)?;
+ Ok(())
+ }
+
+ pub fn save_role(&mut self, name: Option<&str>) -> Result<()> {
+ let mut role_name = match &self.role {
+ Some(role) => match name {
+ Some(v) => v.to_string(),
+ None => role.name().to_string(),
+ },
+ None => bail!("No role"),
+ };
+ if role_name == TEMP_ROLE_NAME {
+ role_name = Text::new("Role name:")
+ .with_validator(|input: &str| {
+ if input.trim().is_empty() {
+ Ok(Validation::Invalid("This field is required".into()))
+ } else {
+ Ok(Validation::Valid)
+ }
+ })
+ .prompt()?;
+ }
+ let role_path = Self::role_file(&role_name)?;
+ if let Some(role) = self.role.as_mut() {
+ let old_name = role.name().to_string();
+ role.save(&role_name, &role_path, self.working_mode.is_repl())?;
+ if old_name != role_name {
+ if let Ok(path) = Self::role_file(&old_name) {
+ let _ = remove_file(&path);
+ }
+ }
+ }
+
+ Ok(())
+ }
+
+ pub fn all_roles() -> Vec<Role> {
+ let mut roles: HashMap<String, Role> = BUILTIN_ROLES
+ .iter()
+ .map(|v| (v.name().to_string(), v.clone()))
+ .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);
+ }
+ }
+ }
+ let mut roles: Vec<_> = roles.into_values().collect();
+ roles.sort_unstable_by(|a, b| a.name().cmp(b.name()));
+ roles
+ }
+
+ pub fn list_roles(with_builtin: bool) -> Vec<String> {
+ let mut names = HashSet::new();
+ if let Some(rd) = Self::roles_dir().ok().and_then(|dir| read_dir(dir).ok()) {
+ for entry in rd.flatten() {
+ if let Some(name) = entry
+ .file_name()
+ .to_str()
+ .and_then(|v| v.strip_suffix(".md"))
+ {
+ names.insert(name.to_string());
+ }
+ }
+ }
+ if with_builtin {
+ names.extend(BUILTIN_ROLES.iter().map(|v| v.name().to_string()));
+ }
+ let mut names: Vec<_> = names.into_iter().collect();
+ names.sort_unstable();
+ names
+ }
+
+ pub fn has_role(name: &str) -> bool {
+ Self::list_roles(true).iter().any(|v| v == name)
+ }
+
pub fn use_session(&mut self, session_name: Option<&str>) -> Result<()> {
if self.session.is_some() {
bail!(
@@ -824,16 +928,22 @@ impl Config {
}
pub fn save_session(&mut self, name: Option<&str>) -> Result<()> {
- let name = match &self.session {
+ let session_name = match &self.session {
Some(session) => match name {
Some(v) => v.to_string(),
None => session.name().to_string(),
},
None => bail!("No session"),
};
- let session_path = self.session_file(&name)?;
+ let session_path = self.session_file(&session_name)?;
if let Some(session) = self.session.as_mut() {
- session.save(&session_path, self.working_mode.is_repl())?;
+ let old_name = session.name().to_string();
+ session.save(&session_name, &session_path, self.working_mode.is_repl())?;
+ if old_name != session_name {
+ if let Ok(path) = self.session_file(&old_name) {
+ let _ = remove_file(&path);
+ }
+ }
}
Ok(())
}
@@ -843,9 +953,9 @@ impl Config {
Some(session) => session.name().to_string(),
None => bail!("No session"),
};
- let editor = self.editor()?;
let session_path = self.session_file(&name)?;
self.save_session(Some(&name))?;
+ let editor = self.editor()?;
edit_file(&editor, &session_path).with_context(|| {
format!(
"Failed to edit '{}' with '{editor}'",
@@ -1201,10 +1311,9 @@ impl Config {
let mut filter = "";
if args.len() == 1 {
values = match cmd {
- ".role" => self
- .roles
- .iter()
- .map(|v| (v.name().to_string(), None))
+ ".role" => Self::list_roles(true)
+ .into_iter()
+ .map(|v| (v, None))
.collect(),
".model" => list_chat_models(self)
.into_iter()
@@ -1702,25 +1811,6 @@ impl Config {
Ok(())
}
- fn load_roles(&mut self) -> Result<()> {
- let path = Self::roles_file()?;
- self.roles = if !path.exists() {
- vec![]
- } else {
- let content = read_to_string(&path)
- .with_context(|| format!("Failed to load roles at {}", path.display()))?;
- serde_yaml::from_str(&content).with_context(|| "Invalid roles config")?
- };
- let exist_roles: HashSet<_> = self.roles.iter().map(|v| v.name().to_string()).collect();
- let builtin_roles = Role::builtin();
- for role in builtin_roles {
- if !exist_roles.contains(role.name()) {
- self.roles.push(role);
- }
- }
- Ok(())
- }
-
fn setup_model(&mut self) -> Result<()> {
let mut model_id = self.model_id.clone();
if model_id.is_empty() {
generated by cgit v1.2.3 (git 2.39.1) at 2026-10-03 05:18:46 +0000