Skip to content

Commit

Permalink
feat: add experiment manager to control whether we enable a feature (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
yingjiehe-xyz authored Feb 21, 2025
1 parent 7355afe commit 5d4f824
Show file tree
Hide file tree
Showing 3 changed files with 121 additions and 5 deletions.
63 changes: 58 additions & 5 deletions crates/goose-cli/src/commands/configure.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use cliclack::spinner;
use console::style;
use goose::agents::{extension::Envs, ExtensionConfig};
use goose::config::{Config, ConfigError, ExtensionEntry, ExtensionManager};
use goose::config::{Config, ConfigError, ExperimentManager, ExtensionEntry, ExtensionManager};
use goose::message::Message;
use goose::providers::{create, providers};
use mcp_core::Tool;
Expand Down Expand Up @@ -153,7 +153,7 @@ pub async fn handle_configure() -> Result<(), Box<dyn Error>> {
.item(
"settings",
"Goose Settings",
"Set the Goose Mode, Tool Output, and more",
"Set the Goose Mode, Tool Output, Experiment and more",
)
.interact()?;

Expand Down Expand Up @@ -622,14 +622,24 @@ pub fn remove_extension_dialog() -> Result<(), Box<dyn Error>> {
}

pub fn configure_settings_dialog() -> Result<(), Box<dyn Error>> {
let setting_type = cliclack::select("What setting would you like to configure?")
let mut setting_select_builder = cliclack::select("What setting would you like to configure?")
.item("goose_mode", "Goose Mode", "Configure Goose mode")
.item(
"tool_output",
"Tool Output",
"Show more or less tool output",
)
.interact()?;
);

// Conditionally add the "Toggle Experiment" option
if ExperimentManager::is_enabled("EXPERIMENT_CONFIG")? {
setting_select_builder = setting_select_builder.item(
"experiment",
"Toggle Experiment",
"Enable or disable an experiment feature",
);
}

let setting_type = setting_select_builder.interact()?;

match setting_type {
"goose_mode" => {
Expand All @@ -638,6 +648,9 @@ pub fn configure_settings_dialog() -> Result<(), Box<dyn Error>> {
"tool_output" => {
configure_tool_output_dialog()?;
}
"experiment" => {
toggle_experiments_dialog()?;
}
_ => unreachable!(),
};

Expand Down Expand Up @@ -718,3 +731,43 @@ pub fn configure_tool_output_dialog() -> Result<(), Box<dyn Error>> {

Ok(())
}

/// Configure experiment features that can be used with goose
/// Dialog for toggling which experiments are enabled/disabled
pub fn toggle_experiments_dialog() -> Result<(), Box<dyn Error>> {
let experiments = ExperimentManager::get_all()?;

if experiments.is_empty() {
cliclack::outro("No experiments supported yet.")?;
return Ok(());
}

// Get currently enabled experiments for the selection
let enabled_experiments: Vec<&String> = experiments
.iter()
.filter(|(_, enabled)| *enabled)
.map(|(name, _)| name)
.collect();

// Let user toggle experiments
let selected = cliclack::multiselect(
"enable experiments: (use \"space\" to toggle and \"enter\" to submit)",
)
.required(false)
.items(
&experiments
.iter()
.map(|(name, _)| (name, name.as_str(), ""))
.collect::<Vec<_>>(),
)
.initial_values(enabled_experiments)
.interact()?;

// Update enabled status for each experiments
for name in experiments.iter().map(|(name, _)| name) {
ExperimentManager::set_enabled(name, selected.iter().any(|&s| s.as_str() == name))?;
}

cliclack::outro("Experiments settings updated successfully")?;
Ok(())
}
61 changes: 61 additions & 0 deletions crates/goose/src/config/experiments.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
use super::base::Config;
use anyhow::Result;
use std::collections::HashMap;

/// It is the ground truth for init experiments. The experiment names in users' experiment list but not
/// in the list will be remove from user list; The experiment names in the ground-truth list but not
/// in users' experiment list will be added to user list with default value false;
const ALL_EXPERIMENTS: &[(&str, bool)] = &[("EXPERIMENT_CONFIG", false)];

/// Experiment configuration management
pub struct ExperimentManager;

impl ExperimentManager {
/// Get all experiments and their configurations
///
/// - Ensures the user's experiment list is synchronized with `ALL_EXPERIMENTS`.
/// - Adds missing experiments from `ALL_EXPERIMENTS` with the default value.
/// - Removes experiments not in `ALL_EXPERIMENTS`.
pub fn get_all() -> Result<Vec<(String, bool)>> {
let config = Config::global();
let mut experiments: HashMap<String, bool> = config.get("experiments").unwrap_or_default();

// Synchronize the user's experiments with the ground truth (`ALL_EXPERIMENTS`)
for &(key, default_value) in ALL_EXPERIMENTS {
experiments.entry(key.to_string()).or_insert(default_value);
}

// Remove experiments not in `ALL_EXPERIMENTS`
experiments.retain(|key, _| ALL_EXPERIMENTS.iter().any(|(k, _)| k == key));

Ok(experiments.into_iter().collect())
}

/// Enable or disable an experiment
pub fn set_enabled(name: &str, enabled: bool) -> Result<()> {
let config = Config::global();

// Load existing experiments or initialize a new map
let mut experiments: HashMap<String, bool> =
config.get("experiments").unwrap_or_else(|_| HashMap::new());

// Update the status of the experiment
experiments.insert(name.to_string(), enabled);

// Save the updated experiments map
config.set("experiments", serde_json::to_value(experiments)?)?;
Ok(())
}

/// Check if an experiment is enabled
pub fn is_enabled(name: &str) -> Result<bool> {
let config = Config::global();

// Load existing experiments or initialize a new map
let experiments: HashMap<String, bool> =
config.get("experiments").unwrap_or_else(|_| HashMap::new());

// Return whether the experiment is enabled, defaulting to false
Ok(*experiments.get(name).unwrap_or(&false))
}
}
2 changes: 2 additions & 0 deletions crates/goose/src/config/mod.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
mod base;
mod experiments;
mod extensions;

pub use crate::agents::ExtensionConfig;
pub use base::{Config, ConfigError, APP_STRATEGY};
pub use experiments::ExperimentManager;
pub use extensions::{ExtensionEntry, ExtensionManager};

0 comments on commit 5d4f824

Please sign in to comment.