From 9d20d882f321102b19c5366c4a24731d50688f48 Mon Sep 17 00:00:00 2001
From: Joe Torreggiani <joetorreggiani@gmail.com>
Date: Thu, 20 Feb 2025 11:19:44 -0500
Subject: [PATCH] feat: Add mode for fully customizable system prompt template
 (#1249)

---
 crates/goose-cli/src/session/builder.rs |  8 ++++++++
 crates/goose/src/agents/agent.rs        |  3 +++
 crates/goose/src/agents/capabilities.rs | 17 +++++++++++++++--
 crates/goose/src/agents/reference.rs    |  5 +++++
 crates/goose/src/agents/truncate.rs     |  5 +++++
 5 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/crates/goose-cli/src/session/builder.rs b/crates/goose-cli/src/session/builder.rs
index 1e61f1ce46..642bc26d95 100644
--- a/crates/goose-cli/src/session/builder.rs
+++ b/crates/goose-cli/src/session/builder.rs
@@ -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
 }
diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs
index ca681de4dc..4500f95d14 100644
--- a/crates/goose/src/agents/agent.rs
+++ b/crates/goose/src/agents/agent.rs
@@ -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);
 }
diff --git a/crates/goose/src/agents/capabilities.rs b/crates/goose/src/agents/capabilities.rs
index e9a8764681..783f15def7 100644
--- a/crates/goose/src/agents/capabilities.rs
+++ b/crates/goose/src/agents/capabilities.rs
@@ -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};
@@ -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>,
 }
 
@@ -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(),
         }
     }
@@ -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
@@ -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
diff --git a/crates/goose/src/agents/reference.rs b/crates/goose/src/agents/reference.rs
index 0cb2f67e9b..6c30435d95 100644
--- a/crates/goose/src/agents/reference.rs
+++ b/crates/goose/src/agents/reference.rs
@@ -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);
diff --git a/crates/goose/src/agents/truncate.rs b/crates/goose/src/agents/truncate.rs
index d6df60cb72..685524d4bf 100644
--- a/crates/goose/src/agents/truncate.rs
+++ b/crates/goose/src/agents/truncate.rs
@@ -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);