feat: load global prompts at startup, add minijinja crate (#1467)

This commit is contained in:
Salman Mohammed
2025-03-04 13:44:37 -05:00
committed by GitHub
parent 3c1b4f627d
commit f26808b7ee
4 changed files with 342 additions and 422 deletions
Generated
+147 -327
View File
File diff suppressed because it is too large Load Diff
+1 -1
View File
@@ -37,7 +37,7 @@ uuid = { version = "1.0", features = ["v4"] }
regex = "1.11.1" regex = "1.11.1"
async-trait = "0.1" async-trait = "0.1"
async-stream = "0.3" async-stream = "0.3"
tera = "1.20.0" minijinja = "2.8.0"
tokenizers = "0.20.3" tokenizers = "0.20.3"
include_dir = "0.7.4" include_dir = "0.7.4"
chrono = { version = "0.4.38", features = ["serde"] } chrono = { version = "0.4.38", features = ["serde"] }
+6 -5
View File
@@ -11,7 +11,7 @@ use tokio::sync::Mutex;
use tracing::{debug, instrument}; use tracing::{debug, instrument};
use super::extension::{ExtensionConfig, ExtensionError, ExtensionInfo, ExtensionResult}; use super::extension::{ExtensionConfig, ExtensionError, ExtensionInfo, ExtensionResult};
use crate::prompt_template::{load_prompt, load_prompt_file}; use crate::prompt_template;
use crate::providers::base::{Provider, ProviderUsage}; use crate::providers::base::{Provider, ProviderUsage};
use mcp_client::client::{ClientCapabilities, ClientInfo, McpClient, McpClientTrait}; use mcp_client::client::{ClientCapabilities, ClientInfo, McpClient, McpClientTrait};
use mcp_client::transport::{SseTransport, StdioTransport, Transport}; use mcp_client::transport::{SseTransport, StdioTransport, Transport};
@@ -340,12 +340,13 @@ impl Capabilities {
context.insert("extensions", serde_json::to_value(extensions_info).unwrap()); context.insert("extensions", serde_json::to_value(extensions_info).unwrap());
context.insert("current_date_time", Value::String(current_date_time)); context.insert("current_date_time", Value::String(current_date_time));
// Conditionally load the override prompt or the default system prompt // Conditionally load the override prompt or the global system prompt
// and set the base prompt to the context
let base_prompt = if let Some(override_prompt) = &self.system_prompt_override { let base_prompt = if let Some(override_prompt) = &self.system_prompt_override {
load_prompt(override_prompt, &context).expect("Prompt should render") prompt_template::render_inline_once(override_prompt, &context)
.expect("Prompt should render")
} else { } else {
load_prompt_file("system.md", &context).expect("Prompt should render") prompt_template::render_global_file("system.md", &context)
.expect("Prompt should render")
}; };
if self.system_prompt_extensions.is_empty() { if self.system_prompt_extensions.is_empty() {
+188 -89
View File
@@ -1,78 +1,162 @@
use include_dir::{include_dir, Dir}; use include_dir::{include_dir, Dir};
use minijinja::{Environment, Error as MiniJinjaError, Value as MJValue};
use once_cell::sync::Lazy;
use serde::Serialize; use serde::Serialize;
use std::path::PathBuf; use std::path::PathBuf;
use tera::{Context, Error as TeraError, Tera}; use std::sync::{Arc, RwLock};
// The prompts directory needs to be embedded in the binary (so it works when distributed) /// This directory will be embedded into the final binary.
static PROMPTS_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/src/prompts"); /// Typically used to store "core" or "system" prompts.
static CORE_PROMPTS_DIR: Dir = include_dir!("$CARGO_MANIFEST_DIR/src/prompts");
pub fn load_prompt<T: Serialize>(template: &str, context_data: &T) -> Result<String, TeraError> { /// A global MiniJinja environment storing the "core" prompts.
let mut tera = Tera::default(); ///
tera.add_raw_template("inline_template", template)?; /// - Loaded at startup from the `CORE_PROMPTS_DIR`.
let context = Context::from_serialize(context_data)?; /// - Ideal for "system" templates that don't change often.
let rendered = tera.render("inline_template", &context)?; /// - *Not* used for extension prompts (which are ephemeral).
static GLOBAL_ENV: Lazy<Arc<RwLock<Environment<'static>>>> = Lazy::new(|| {
let mut env = Environment::new();
// Pre-load all core templates from the embedded dir.
for file in CORE_PROMPTS_DIR.files() {
let name = file.path().to_string_lossy().to_string();
let source = String::from_utf8_lossy(file.contents()).to_string();
// Since we're using 'static lifetime for the Environment, we need to ensure
// the strings we add as templates live for the entire program duration.
// We can achieve this by leaking the strings (acceptable for initialization).
let static_name: &'static str = Box::leak(name.into_boxed_str());
let static_source: &'static str = Box::leak(source.into_boxed_str());
if let Err(e) = env.add_template(static_name, static_source) {
tracing::error!("Failed to add template {}: {}", static_name, e);
}
}
Arc::new(RwLock::new(env))
});
/// Renders a prompt from the global environment by name.
///
/// # Arguments
/// * `template_name` - The name of the template (usually the file path or a custom ID).
/// * `context_data` - Data to be inserted into the template (must be `Serialize`).
pub fn render_global_template<T: Serialize>(
template_name: &str,
context_data: &T,
) -> Result<String, MiniJinjaError> {
let env = GLOBAL_ENV.read().expect("GLOBAL_ENV lock poisoned");
let tmpl = env.get_template(template_name)?;
let ctx = MJValue::from_serialize(context_data);
let rendered = tmpl.render(ctx)?;
Ok(rendered.trim().to_string()) Ok(rendered.trim().to_string())
} }
pub fn load_prompt_file<T: Serialize>( /// Renders a file from `CORE_PROMPTS_DIR` within the global environment.
///
/// # Arguments
/// * `template_file` - The file path within the embedded directory (e.g. "system.md").
/// * `context_data` - Data to be inserted into the template (must be `Serialize`).
///
/// This function **assumes** the file is already in `CORE_PROMPTS_DIR`. If it wasn't
/// added to the global environment at startup (due to parse errors, etc.), this will error out.
pub fn render_global_file<T: Serialize>(
template_file: impl Into<PathBuf>, template_file: impl Into<PathBuf>,
context_data: &T, context_data: &T,
) -> Result<String, TeraError> { ) -> Result<String, MiniJinjaError> {
let template_path = template_file.into(); let file_path = template_file.into();
let template_name = file_path.to_string_lossy().to_string();
// Get the file content from the embedded directory render_global_template(&template_name, context_data)
let template_content = if let Some(file) = PROMPTS_DIR.get_file(template_path.to_str().unwrap()) }
{
String::from_utf8_lossy(file.contents()).into_owned()
} else {
return Err(TeraError::chain(
"Failed to find template file",
std::io::Error::new(
std::io::ErrorKind::NotFound,
"Template file not found in embedded directory",
),
));
};
load_prompt(&template_content, context_data) /// Alias for render_global_file for backward compatibility
pub fn render_global_from_file<T: Serialize>(
template_file: impl Into<PathBuf>,
context_data: &T,
) -> Result<String, MiniJinjaError> {
render_global_file(template_file, context_data)
}
/// Renders a **one-off ephemeral** template (inline string).
///
/// This does *not* store anything in the global environment and is best for
/// extension prompts or user-supplied templates that are used infrequently.
///
/// # Arguments
/// * `template_str` - The raw template string.
/// * `context_data` - Data to be inserted into the template (must be `Serialize`).
pub fn render_inline_once<T: Serialize>(
template_str: &str,
context_data: &T,
) -> Result<String, MiniJinjaError> {
let mut env = Environment::new();
env.add_template("inline_ephemeral", template_str)?;
let tmpl = env.get_template("inline_ephemeral")?;
let ctx = MJValue::from_serialize(context_data);
let rendered = tmpl.render(ctx)?;
Ok(rendered.trim().to_string())
} }
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use mcp_core::tool::Tool;
use serde_json::json; use serde_json::json;
use std::collections::HashMap; use std::collections::HashMap;
#[test] /// For convenience in tests, define a small struct or use a HashMap to provide context.
fn test_load_prompt() { #[derive(Serialize)]
let template = "Hello, {{ name }}! You are {{ age }} years old."; struct TestContext {
let mut context = HashMap::new(); name: String,
context.insert("name".to_string(), "Alice".to_string()); age: u32,
context.insert("age".to_string(), 30.to_string()); }
let result = load_prompt(template, &context).unwrap(); // A simple function to help us test missing or partial data
fn build_context(name: Option<&str>, age: Option<u32>) -> HashMap<String, serde_json::Value> {
let mut ctx = HashMap::new();
if let Some(n) = name {
ctx.insert("name".to_string(), json!(n));
}
if let Some(a) = age {
ctx.insert("age".to_string(), json!(a));
}
ctx
}
#[test]
fn test_render_inline_once_basic() {
let template_str = "Hello, {{ name }}! You are {{ age }} years old.";
let context = TestContext {
name: "Alice".to_string(),
age: 30,
};
let result = render_inline_once(template_str, &context).unwrap();
assert_eq!(result, "Hello, Alice! You are 30 years old."); assert_eq!(result, "Hello, Alice! You are 30 years old.");
} }
#[test] #[test]
fn test_load_prompt_missing_variable() { fn test_render_inline_missing_variable() {
let template = "Hello, {{ name }}! You are {{ age }} years old."; let template_str = "Hello, {{ name }}! You are {{ age }} years old.";
let mut context = HashMap::new(); let context = build_context(Some("Alice"), None);
context.insert("name".to_string(), "Alice".to_string()); // MiniJinja doesn't fail on missing variables, it renders them as empty strings
// 'age' is missing from context // So we should check that it renders successfully but with missing data
let result = load_prompt(template, &context); let result = render_inline_once(template_str, &context).unwrap();
assert!(result.is_err()); assert!(result.contains("Hello, Alice! You are years old."));
} }
#[test] #[test]
fn test_load_prompt_file() { fn test_global_file_render() {
// since we are embedding the prompts directory, the file path needs to be relative to the prompts directory // "mock.md" should exist in the embedded CORE_PROMPTS_DIR
let file_path = PathBuf::from("mock.md"); // and have placeholders for `name` and `age`.
let mut context = HashMap::new(); let context = TestContext {
context.insert("name".to_string(), "Alice".to_string()); name: "Alice".to_string(),
context.insert("age".to_string(), 30.to_string()); age: 30,
let result = load_prompt_file(file_path, &context).unwrap(); };
let result = render_global_file("mock.md", &context).unwrap();
// Assume mock.md content is something like:
// "This prompt is only used for testing.\n\nHello, {{ name }}! You are {{ age }} years old."
assert_eq!( assert_eq!(
result, result,
"This prompt is only used for testing.\n\nHello, Alice! You are 30 years old." "This prompt is only used for testing.\n\nHello, Alice! You are 30 years old."
@@ -80,60 +164,75 @@ mod tests {
} }
#[test] #[test]
fn test_load_prompt_file_missing_file() { fn test_global_file_not_found() {
let file_path = PathBuf::from("non_existent_template.txt"); let context = TestContext {
let context: HashMap<String, String> = HashMap::new(); // Add type annotation here name: "Unused".to_string(),
age: 99,
};
let result = load_prompt_file(file_path, &context); let result = render_global_file("non_existent.md", &context);
assert!(result.is_err()); assert!(result.is_err(), "Should fail because file is missing");
} }
#[test] #[test]
fn test_load_prompt_with_tools() { fn test_inline_complex_object() {
let template = "### Tool Descriptions\n{% for tool in tools %}\n{{tool.name}}: {{tool.description}}{% endfor %}"; // Example with more complex data.
#[derive(Serialize)]
struct Tool {
name: String,
description: String,
}
let tools = vec![ #[derive(Serialize)]
Tool::new( struct ToolsContext {
"calculator", tools: Vec<Tool>,
"Performs basic math operations", }
json!({
"type": "object",
"properties": {
"operation": {"type": "string"},
"numbers": {"type": "array"}
}
}),
),
Tool::new(
"weather",
"Gets weather information",
json!({
"type": "object",
"properties": {
"location": {"type": "string"}
}
}),
),
];
let mut context = HashMap::new(); let template_str = "\
context.insert("tools".to_string(), tools); ### Tool Descriptions
{% for tool in tools %}
- {{ tool.name }}: {{ tool.description }}
{% endfor %}";
let result = load_prompt(template, &context).unwrap(); let context = ToolsContext {
let expected = "### Tool Descriptions\n\ncalculator: Performs basic math operations\nweather: Gets weather information"; tools: vec![
assert_eq!(result, expected); Tool {
name: "calculator".to_string(),
description: "Performs basic math operations".to_string(),
},
Tool {
name: "weather".to_string(),
description: "Gets weather information".to_string(),
},
],
};
let rendered = render_inline_once(template_str, &context).unwrap();
let expected = "\
### Tool Descriptions
- calculator: Performs basic math operations
- weather: Gets weather information";
assert_eq!(rendered, expected);
} }
#[test] #[test]
fn test_load_prompt_with_empty_tools() { fn test_inline_with_empty_list() {
let template = "### Tool Descriptions\n{% for tool in tools %}\n{{tool.name}}: {{tool.description}}{% endfor %}"; let template_str = "\
### Tool Descriptions
{% for tool in tools %}
- {{ tool.name }}: {{ tool.description }}
{% endfor %}";
let tools: Vec<Tool> = vec![]; #[derive(Serialize)]
let mut context = HashMap::new(); struct ToolsContext {
context.insert("tools".to_string(), tools); tools: Vec<String>, // or a struct if needed
}
let result = load_prompt(template, &context).unwrap(); let context = ToolsContext { tools: vec![] };
let rendered = render_inline_once(template_str, &context).unwrap();
let expected = "### Tool Descriptions"; let expected = "### Tool Descriptions";
assert_eq!(result, expected); assert_eq!(rendered, expected);
} }
} }