Skip to content

Commit

Permalink
feat: Add mode for fully customizable system prompt template (#1249)
Browse files Browse the repository at this point in the history
  • Loading branch information
jtorreggiani authored Feb 20, 2025
1 parent 0c53411 commit 9d20d88
Show file tree
Hide file tree
Showing 5 changed files with 36 additions and 2 deletions.
8 changes: 8 additions & 0 deletions crates/goose-cli/src/session/builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,14 @@ pub async fn build_session(
.extend_system_prompt(super::prompt::get_cli_prompt())
.await;

// Only override system prompt if a system override exists
let system_prompt_file: Option<String> = config.get("GOOSE_SYSTEM_PROMPT_FILE_PATH").ok();
if let Some(ref path) = system_prompt_file {
let override_prompt =
std::fs::read_to_string(path).expect("Failed to read system prompt file");
session.agent.override_system_prompt(override_prompt).await;
}

output::display_session_info(resume, &provider_name, &model, &session_file);
session
}
Expand Down
3 changes: 3 additions & 0 deletions crates/goose/src/agents/agent.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,4 +31,7 @@ pub trait Agent: Send + Sync {

/// Add custom text to be included in the system prompt
async fn extend_system_prompt(&mut self, extension: String);

/// Override the system prompt with custom text
async fn override_system_prompt(&mut self, template: String);
}
17 changes: 15 additions & 2 deletions crates/goose/src/agents/capabilities.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use tokio::sync::Mutex;
use tracing::{debug, instrument};

use super::extension::{ExtensionConfig, ExtensionError, ExtensionInfo, ExtensionResult};
use crate::prompt_template::load_prompt_file;
use crate::prompt_template::{load_prompt, load_prompt_file};
use crate::providers::base::{Provider, ProviderUsage};
use mcp_client::client::{ClientCapabilities, ClientInfo, McpClient, McpClientTrait};
use mcp_client::transport::{SseTransport, StdioTransport, Transport};
Expand All @@ -30,6 +30,7 @@ pub struct Capabilities {
resource_capable_extensions: HashSet<String>,
provider: Box<dyn Provider>,
provider_usage: Mutex<Vec<ProviderUsage>>,
system_prompt_override: Option<String>,
system_prompt_extensions: Vec<String>,
}

Expand Down Expand Up @@ -89,6 +90,7 @@ impl Capabilities {
resource_capable_extensions: HashSet::new(),
provider,
provider_usage: Mutex::new(Vec::new()),
system_prompt_override: None,
system_prompt_extensions: Vec::new(),
}
}
Expand Down Expand Up @@ -171,6 +173,11 @@ impl Capabilities {
self.system_prompt_extensions.push(extension);
}

/// Override the system prompt with custom text
pub fn set_system_prompt_override(&mut self, template: String) {
self.system_prompt_override = Some(template);
}

/// Get a reference to the provider
pub fn provider(&self) -> &dyn Provider {
&*self.provider
Expand Down Expand Up @@ -310,7 +317,13 @@ impl Capabilities {
context.insert("extensions", serde_json::to_value(extensions_info).unwrap());
context.insert("current_date_time", Value::String(current_date_time));

let base_prompt = load_prompt_file("system.md", &context).expect("Prompt should render");
// Conditionally load the override prompt or the default system prompt
// and set the base prompt to the context
let base_prompt = if let Some(override_prompt) = &self.system_prompt_override {
load_prompt(override_prompt, &context).expect("Prompt should render")
} else {
load_prompt_file("system.md", &context).expect("Prompt should render")
};

if self.system_prompt_extensions.is_empty() {
base_prompt
Expand Down
5 changes: 5 additions & 0 deletions crates/goose/src/agents/reference.rs
Original file line number Diff line number Diff line change
Expand Up @@ -189,6 +189,11 @@ impl Agent for ReferenceAgent {
let mut capabilities = self.capabilities.lock().await;
capabilities.add_system_prompt_extension(extension);
}

async fn override_system_prompt(&mut self, template: String) {
let mut capabilities = self.capabilities.lock().await;
capabilities.set_system_prompt_override(template);
}
}

register_agent!("reference", ReferenceAgent);
5 changes: 5 additions & 0 deletions crates/goose/src/agents/truncate.rs
Original file line number Diff line number Diff line change
Expand Up @@ -297,6 +297,11 @@ impl Agent for TruncateAgent {
let mut capabilities = self.capabilities.lock().await;
capabilities.add_system_prompt_extension(extension);
}

async fn override_system_prompt(&mut self, template: String) {
let mut capabilities = self.capabilities.lock().await;
capabilities.set_system_prompt_override(template);
}
}

register_agent!("truncate", TruncateAgent);

0 comments on commit 9d20d88

Please sign in to comment.