diff --git a/src/ask/chat_agent.rs b/src/ask/chat_agent.rs index 1075e9a5..5ee182af 100644 --- a/src/ask/chat_agent.rs +++ b/src/ask/chat_agent.rs @@ -32,6 +32,7 @@ pub async fn ask_agent( client: &Client, api_model: &str, max_iterations: Option, + workspace_name: Option<&str>, ) -> Result { let max_iterations = max_iterations.unwrap_or(20); let mut result = AskOutput { @@ -88,7 +89,7 @@ pub async fn ask_agent( // Call the appropriate tool let response_content = - call_tool(name, args, &files, model, &mut result).await?; + call_tool(name, args, &files, model, &mut result, workspace_name).await?; // Print summary of the tool response print_tool_summary(&response_content); diff --git a/src/ask/responses_agent.rs b/src/ask/responses_agent.rs index ad59d496..616e2cae 100644 --- a/src/ask/responses_agent.rs +++ b/src/ask/responses_agent.rs @@ -32,6 +32,7 @@ pub async fn ask_agent_responses( client: &Client, api_model: &str, max_iterations: Option, + workspace_name: Option<&str>, ) -> Result { let max_iterations = max_iterations.unwrap_or(20); let mut result = AskOutput { @@ -96,7 +97,8 @@ pub async fn ask_agent_responses( let args = &function_call.arguments; // Call the appropriate tool - let response_content = call_tool(name, args, &files, model, &mut result).await?; + let response_content = + call_tool(name, args, &files, model, &mut result, workspace_name).await?; // Print summary of the tool response print_tool_summary(&response_content); diff --git a/src/ask/tool_calling.rs b/src/ask/tool_calling.rs index b2ca94d4..d3709a62 100644 --- a/src/ask/tool_calling.rs +++ b/src/ask/tool_calling.rs @@ -13,6 +13,7 @@ pub async fn call_tool( files: &[String], model: &StaticModel, cur_output: &mut AskOutput, + workspace_name: Option<&str>, ) -> Result { let function_args: Value = serde_json::from_str(args)?; @@ -96,7 +97,15 @@ pub async fn call_tool( println!(" top_k: {}", top_k); } - SearchTool::search(files, query, model, config, &mut cur_output.files_searched).await + SearchTool::search( + files, + query, + model, + config, + &mut cur_output.files_searched, + workspace_name, + ) + .await } "read" => { let path = function_args["path"] diff --git a/src/ask/tools.rs b/src/ask/tools.rs index 7682aeb1..750e237a 100644 --- a/src/ask/tools.rs +++ b/src/ask/tools.rs @@ -211,6 +211,7 @@ impl SearchTool { model: &StaticModel, config: SearchConfig, files_searched: &mut Vec, + workspace_name: Option<&str>, ) -> Result { let query = if config.ignore_case { query.to_lowercase() @@ -226,9 +227,10 @@ impl SearchTool { // Handle file input with optional workspace integration #[cfg(feature = "workspace")] - if Workspace::active().is_ok() { + if Workspace::active(workspace_name).is_ok() { // Workspace mode: use persisted line embeddings for speed - let ranked_lines = search_with_workspace(files, &query, model, &config).await?; + let ranked_lines = + search_with_workspace(files, &query, model, &config, workspace_name).await?; // Track files that were searched (have results) for ranked_line in &ranked_lines { diff --git a/src/bin/semtools.rs b/src/bin/semtools.rs index e3ea3ec6..831a9263 100644 --- a/src/bin/semtools.rs +++ b/src/bin/semtools.rs @@ -15,9 +15,15 @@ enum WorkspaceCommands { /// Use or create a workspace (prints export command to run) Use { name: String }, /// Show active workspace and basic stats - Status, + Status { + #[clap(default_value = None)] + name: Option, + }, /// Remove stale or missing files from store - Prune {}, + Prune { + #[clap(default_value = None)] + name: Option, + }, } #[derive(Subcommand, Debug)] @@ -70,6 +76,10 @@ enum Commands { /// Output results in JSON format #[clap(short, long)] json: bool, + + /// Use a specific workspace + #[arg(short, long, default_value = None)] + workspace: Option, }, #[cfg(feature = "ask")] /// A CLI tool for document-based question-answering @@ -104,6 +114,10 @@ enum Commands { /// Output results in JSON or text format #[clap(short, long)] json: bool, + + /// Use a specific workspace + #[arg(short, long, default_value = None)] + workspace: Option, }, #[cfg(feature = "workspace")] /// Manage semtools workspaces @@ -130,9 +144,18 @@ async fn main() -> anyhow::Result<()> { model, api_mode, json, + workspace, } => { ask_cmd( - query, files, config, api_key, base_url, model, api_mode, json, + query, + files, + config, + api_key, + base_url, + model, + api_mode, + json, + workspace.as_deref(), ) .await?; } @@ -152,6 +175,7 @@ async fn main() -> anyhow::Result<()> { max_distance, ignore_case, json, + workspace, } => { search_cmd( query, @@ -161,6 +185,7 @@ async fn main() -> anyhow::Result<()> { max_distance, ignore_case, json, + workspace.as_deref(), ) .await?; } @@ -168,11 +193,11 @@ async fn main() -> anyhow::Result<()> { WorkspaceCommands::Use { name } => { workspace_use_cmd(name, json).await?; } - WorkspaceCommands::Prune {} => { - workspace_prune_cmd(json).await?; + WorkspaceCommands::Prune { name } => { + workspace_prune_cmd(json, name.as_deref()).await?; } - WorkspaceCommands::Status => { - workspace_status_cmd(json).await?; + WorkspaceCommands::Status { name } => { + workspace_status_cmd(json, name.as_deref()).await?; } }, } diff --git a/src/cmds/ask.rs b/src/cmds/ask.rs index 1b1c8350..998225c4 100644 --- a/src/cmds/ask.rs +++ b/src/cmds/ask.rs @@ -27,6 +27,7 @@ pub async fn ask_cmd( model: Option, api_mode: Option, json: bool, + workspace_name: Option<&str>, ) -> Result<()> { // Load configuration let config_path = config.unwrap_or_else(SemtoolsConfig::default_config_path); @@ -134,10 +135,28 @@ pub async fn ask_cmd( // Run the appropriate agent based on API mode let output = match api_mode { ApiMode::Chat => { - ask_agent(files, &query, &model, &client, &model_name, max_iterations).await? + ask_agent( + files, + &query, + &model, + &client, + &model_name, + max_iterations, + workspace_name, + ) + .await? } ApiMode::Responses => { - ask_agent_responses(files, &query, &model, &client, &model_name, max_iterations).await? + ask_agent_responses( + files, + &query, + &model, + &client, + &model_name, + max_iterations, + workspace_name, + ) + .await? } }; diff --git a/src/cmds/search.rs b/src/cmds/search.rs index 399d22ea..fb7b0cec 100644 --- a/src/cmds/search.rs +++ b/src/cmds/search.rs @@ -109,6 +109,7 @@ fn print_workspace_search_results(ranked_lines: &[RankedLine], n_lines: usize) { } } +#[allow(clippy::too_many_arguments)] pub async fn search_cmd( query: String, files: Vec, @@ -117,6 +118,7 @@ pub async fn search_cmd( max_distance: Option, ignore_case: bool, json: bool, + workspace_name: Option<&str>, ) -> Result<()> { let model = StaticModel::from_pretrained( MODEL_NAME, // "minishlab/potion-multilingual-128M", @@ -192,7 +194,7 @@ pub async fn search_cmd( // Handle file input with optional workspace integration #[cfg(feature = "workspace")] { - if Workspace::active().is_ok() { + if Workspace::active(workspace_name).is_ok() { // Workspace mode: use persisted line embeddings for speed let config = SearchConfig { n_lines, @@ -200,7 +202,8 @@ pub async fn search_cmd( max_distance, ignore_case, }; - let ranked_lines = search_with_workspace(&files, &query, &model, &config).await?; + let ranked_lines = + search_with_workspace(&files, &query, &model, &config, workspace_name).await?; if json { // Convert workspace results to SearchResultJSON diff --git a/src/cmds/workspace.rs b/src/cmds/workspace.rs index c083eb07..54e16a4c 100644 --- a/src/cmds/workspace.rs +++ b/src/cmds/workspace.rs @@ -46,6 +46,8 @@ pub async fn workspace_use_cmd(name: String, json: bool) -> Result<()> { println!(" export SEMTOOLS_WORKSPACE={name}"); println!(); println!("Or add this to your shell profile (.bashrc, .zshrc, etc.)"); + println!(); + println!("Or use the `--workspace` option on the commands that support it"); } } #[cfg(not(feature = "workspace"))] @@ -64,11 +66,11 @@ pub async fn workspace_use_cmd(name: String, json: bool) -> Result<()> { Ok(()) } -pub async fn workspace_status_cmd(json: bool) -> Result<()> { +pub async fn workspace_status_cmd(json: bool, workspace_name: Option<&str>) -> Result<()> { #[cfg(feature = "workspace")] { - let _name = Workspace::active().context("No active workspace")?; - let ws = Workspace::open()?; + let _name = Workspace::active(workspace_name).context("No active workspace")?; + let ws = Workspace::open(workspace_name)?; // Open store and get stats let store = Store::open(&ws.config.root_dir)?; @@ -110,11 +112,11 @@ pub async fn workspace_status_cmd(json: bool) -> Result<()> { Ok(()) } -pub async fn workspace_prune_cmd(json: bool) -> Result<()> { +pub async fn workspace_prune_cmd(json: bool, workspace_name: Option<&str>) -> Result<()> { #[cfg(feature = "workspace")] { - let _name = Workspace::active().context("No active workspace")?; - let ws = Workspace::open()?; + let _name = Workspace::active(workspace_name).context("No active workspace")?; + let ws = Workspace::open(workspace_name)?; let store = Store::open(&ws.config.root_dir)?; // Get all document paths from the workspace diff --git a/src/search/mod.rs b/src/search/mod.rs index 7d52e437..68a0ef50 100644 --- a/src/search/mod.rs +++ b/src/search/mod.rs @@ -148,9 +148,10 @@ pub async fn search_with_workspace( query: &str, model: &StaticModel, config: &SearchConfig, + workspace_name: Option<&str>, ) -> Result> { let query_embedding = model.encode_single(query); - let ws = Workspace::open()?; + let ws = Workspace::open(workspace_name)?; let store = Store::open(&ws.config.root_dir)?; // Step 1: Analyze document states (changed/new/unchanged) diff --git a/src/workspace/mod.rs b/src/workspace/mod.rs index 30d69b4a..6b24e831 100644 --- a/src/workspace/mod.rs +++ b/src/workspace/mod.rs @@ -30,8 +30,8 @@ pub struct Workspace { } impl Workspace { - pub fn open() -> Result { - let active_workspace = Self::active()?; + pub fn open(workspace_name: Option<&str>) -> Result { + let active_workspace = Self::active(workspace_name)?; let cfg_path = Self::config_path_for(&active_workspace)?; let cfg = std::fs::read_to_string(&cfg_path) .ok() @@ -55,16 +55,22 @@ impl Workspace { Ok(()) } - pub fn active_path() -> Result { - let active = std::env::var("SEMTOOLS_WORKSPACE").unwrap_or_default(); + pub fn active_path(workspace_name: Option<&str>) -> Result { + let active = match workspace_name { + None => std::env::var("SEMTOOLS_WORKSPACE").unwrap_or_default(), + Some(a) => a.to_string(), + }; if active.is_empty() { bail!("No active workspace. Run: workspace use "); } Self::root_path(&active) } - pub fn active() -> Result { - let active = std::env::var("SEMTOOLS_WORKSPACE").unwrap_or_default(); + pub fn active(workspace_name: Option<&str>) -> Result { + let active = match workspace_name { + None => std::env::var("SEMTOOLS_WORKSPACE").unwrap_or_default(), + Some(a) => a.to_string(), + }; if active.is_empty() { bail!("No active workspace. Run: workspace use "); } @@ -135,18 +141,30 @@ mod tests { } #[test] - fn test_workspace_set_and_get_active() { + fn test_workspace_set_and_get_active_from_env() { // Test setting active workspace unsafe { std::env::set_var("SEMTOOLS_WORKSPACE", "test-workspace"); } // Test getting active workspace - let active = Workspace::active().expect("Failed to get active"); + let active = Workspace::active(None).expect("Failed to get active"); + assert_eq!(active, "test-workspace"); + + // Test getting active path + let active_path = Workspace::active_path(None).expect("Failed to get active path"); + assert!(active_path.contains("test-workspace")); + } + + #[test] + fn test_workspace_set_and_get_active_from_name() { + // Test getting active workspace + let active = Workspace::active(Some("test-workspace")).expect("Failed to get active"); assert_eq!(active, "test-workspace"); // Test getting active path - let active_path = Workspace::active_path().expect("Failed to get active path"); + let active_path = + Workspace::active_path(Some("test-workspace")).expect("Failed to get active path"); assert!(active_path.contains("test-workspace")); } @@ -161,7 +179,7 @@ mod tests { } // Should fail when no active workspace - let result = Workspace::active(); + let result = Workspace::active(None); assert!(result.is_err()); // Clear environment variable @@ -170,7 +188,7 @@ mod tests { } // Should fail when no active workspace - let result = Workspace::active_path(); + let result = Workspace::active_path(None); assert!(result.is_err()); // Restore original state @@ -268,7 +286,7 @@ mod tests { } // Since we haven't saved a config file, open should use defaults - let workspace = Workspace::open().expect("Failed to open workspace"); + let workspace = Workspace::open(None).expect("Failed to open workspace"); assert_eq!(workspace.config.name, workspace_name); assert!(!workspace.config.root_dir.is_empty()); @@ -280,4 +298,15 @@ mod tests { } } } + + #[test] + fn test_workspace_open_with_workspace_name_and_defaults() { + let workspace_name = "test-defaults"; + + // Since we haven't saved a config file, open should use defaults + let workspace = Workspace::open(Some(workspace_name)).expect("Failed to open workspace"); + + assert_eq!(workspace.config.name, workspace_name); + assert!(!workspace.config.root_dir.is_empty()); + } }