763 lines
26 KiB
Rust
763 lines
26 KiB
Rust
//! Advanced prompt template system with variable substitution, conditionals, and inheritance
|
|
|
|
use crate::error::{Result, TemplateError};
|
|
use dashmap::DashMap;
|
|
use indexmap::IndexMap;
|
|
use parking_lot::RwLock;
|
|
use regex::Regex;
|
|
use serde_json::Value;
|
|
use sha2::{Digest, Sha256};
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use std::time::Instant;
|
|
use uuid::Uuid;
|
|
|
|
/// Function type for custom template functions
|
|
pub type TemplateFunction = Box<dyn Fn() -> Result<Value> + Send + Sync>;
|
|
|
|
/// Filter type for custom template filters
|
|
pub type TemplateFilter = Box<dyn Fn(&Value, &str) -> Result<Value> + Send + Sync>;
|
|
|
|
/// Template context for variable substitution and function execution
|
|
#[derive(Clone)]
|
|
pub struct TemplateContext {
|
|
variables: IndexMap<String, Value>,
|
|
functions: Arc<RwLock<HashMap<String, TemplateFunction>>>,
|
|
filters: Arc<RwLock<HashMap<String, TemplateFilter>>>,
|
|
}
|
|
|
|
impl std::fmt::Debug for TemplateContext {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("TemplateContext")
|
|
.field("variables", &self.variables)
|
|
.field(
|
|
"functions",
|
|
&format!("{} functions", self.functions.read().len()),
|
|
)
|
|
.field("filters", &format!("{} filters", self.filters.read().len()))
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
impl TemplateContext {
|
|
/// Create a new template context
|
|
pub fn new() -> Self {
|
|
Self {
|
|
variables: IndexMap::new(),
|
|
functions: Arc::new(RwLock::new(HashMap::new())),
|
|
filters: Arc::new(RwLock::new(HashMap::new())),
|
|
}
|
|
}
|
|
|
|
/// Insert a variable into the context
|
|
pub fn insert(&mut self, key: impl Into<String>, value: Value) {
|
|
self.variables.insert(key.into(), value);
|
|
}
|
|
|
|
/// Insert nested JSON structure
|
|
pub fn insert_nested(&mut self, value: Value) {
|
|
if let Value::Object(obj) = value {
|
|
for (key, val) in obj {
|
|
self.insert(key, val);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Get a variable from the context
|
|
pub fn get(&self, key: &str) -> Option<&Value> {
|
|
// Support nested access like "user.profile.name"
|
|
let parts: Vec<&str> = key.split('.').collect();
|
|
if parts.len() == 1 {
|
|
return self.variables.get(key);
|
|
}
|
|
|
|
let mut current = self.variables.get(parts[0])?;
|
|
for &part in &parts[1..] {
|
|
current = current.get(part)?;
|
|
}
|
|
Some(current)
|
|
}
|
|
|
|
/// Register a custom function
|
|
pub fn register_function(&mut self, name: impl Into<String>, func: TemplateFunction) {
|
|
self.functions.write().insert(name.into(), func);
|
|
}
|
|
|
|
/// Register a custom filter
|
|
pub fn register_filter(&mut self, name: impl Into<String>, filter: TemplateFilter) {
|
|
self.filters.write().insert(name.into(), filter);
|
|
}
|
|
|
|
/// Call a registered function
|
|
pub fn call_function(&self, name: &str) -> Result<Value> {
|
|
let functions = self.functions.read();
|
|
let func = functions
|
|
.get(name)
|
|
.ok_or_else(|| TemplateError::FunctionNotFound(name.to_string()))?;
|
|
func()
|
|
}
|
|
|
|
/// Apply a registered filter
|
|
pub fn apply_filter(&self, value: &Value, filter_name: &str, args: &str) -> Result<Value> {
|
|
let filters = self.filters.read();
|
|
let filter = filters
|
|
.get(filter_name)
|
|
.ok_or_else(|| TemplateError::FilterNotFound(filter_name.to_string()))?;
|
|
filter(value, args)
|
|
}
|
|
}
|
|
|
|
impl Default for TemplateContext {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
/// Compiled template node
|
|
#[derive(Debug, Clone)]
|
|
pub enum TemplateNode {
|
|
Text(String),
|
|
Variable {
|
|
name: String,
|
|
filters: Vec<(String, String)>, // (filter_name, args)
|
|
},
|
|
Conditional {
|
|
condition: String,
|
|
if_block: Vec<Self>,
|
|
else_block: Option<Vec<Self>>,
|
|
},
|
|
Loop {
|
|
variable: String,
|
|
iterable: String,
|
|
body: Vec<Self>,
|
|
},
|
|
Function {
|
|
name: String,
|
|
filters: Vec<(String, String)>,
|
|
},
|
|
Block {
|
|
name: String,
|
|
content: Vec<Self>,
|
|
},
|
|
Extends {
|
|
parent: String,
|
|
},
|
|
}
|
|
|
|
/// Prompt template with advanced features
|
|
#[derive(Debug, Clone)]
|
|
pub struct PromptTemplate {
|
|
id: String,
|
|
source: String,
|
|
nodes: Vec<TemplateNode>,
|
|
parent_template: Option<String>,
|
|
blocks: HashMap<String, Vec<TemplateNode>>,
|
|
creation_time: Instant,
|
|
}
|
|
|
|
impl PromptTemplate {
|
|
/// Create template from string
|
|
pub fn from_str(template: &str) -> Result<Self> {
|
|
let id = format!("template_{}", Uuid::new_v4());
|
|
let parser = TemplateParser::new();
|
|
let nodes = parser.parse(template)?;
|
|
|
|
let mut blocks = HashMap::new();
|
|
let mut parent_template = None;
|
|
|
|
// Extract blocks and parent template
|
|
for node in &nodes {
|
|
match node {
|
|
TemplateNode::Block { name, content } => {
|
|
blocks.insert(name.clone(), content.clone());
|
|
}
|
|
TemplateNode::Extends { parent } => {
|
|
parent_template = Some(parent.clone());
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
Ok(Self {
|
|
id,
|
|
source: template.to_string(),
|
|
nodes,
|
|
parent_template,
|
|
blocks,
|
|
creation_time: Instant::now(),
|
|
})
|
|
}
|
|
|
|
/// Get template ID
|
|
pub fn template_id(&self) -> &str {
|
|
&self.id
|
|
}
|
|
|
|
/// Render template with context
|
|
pub async fn render(&self, context: &TemplateContext) -> Result<String> {
|
|
self.render_nodes(&self.nodes, context).await
|
|
}
|
|
|
|
/// Bind partial context to create a new template
|
|
pub async fn bind_partial(&self, partial_context: &TemplateContext) -> Result<Self> {
|
|
// Create a new template with some variables pre-filled
|
|
let mut new_nodes = Vec::new();
|
|
|
|
for node in &self.nodes {
|
|
let processed_node = self.process_partial_binding(node, partial_context).await?;
|
|
new_nodes.push(processed_node);
|
|
}
|
|
|
|
Ok(Self {
|
|
id: format!("{}_partial", self.id),
|
|
source: self.source.clone(),
|
|
nodes: new_nodes,
|
|
parent_template: self.parent_template.clone(),
|
|
blocks: self.blocks.clone(),
|
|
creation_time: Instant::now(),
|
|
})
|
|
}
|
|
|
|
fn render_nodes<'a>(
|
|
&'a self,
|
|
nodes: &'a [TemplateNode],
|
|
context: &'a TemplateContext,
|
|
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<String>> + Send + 'a>> {
|
|
Box::pin(async move {
|
|
let mut result = String::new();
|
|
|
|
for node in nodes {
|
|
match node {
|
|
TemplateNode::Text(text) => {
|
|
result.push_str(text);
|
|
}
|
|
TemplateNode::Variable { name, filters } => {
|
|
let value = context
|
|
.get(name)
|
|
.ok_or_else(|| TemplateError::MissingVariable(name.clone()))?;
|
|
|
|
let mut processed_value = value.clone();
|
|
for (filter_name, args) in filters {
|
|
processed_value =
|
|
context.apply_filter(&processed_value, filter_name, args)?;
|
|
}
|
|
|
|
let text = self.value_to_html_escaped_string(&processed_value);
|
|
result.push_str(&text);
|
|
}
|
|
TemplateNode::Function { name, filters } => {
|
|
let mut value = context.call_function(name)?;
|
|
for (filter_name, args) in filters {
|
|
value = context.apply_filter(&value, filter_name, args)?;
|
|
}
|
|
|
|
let text = self.value_to_string(&value);
|
|
result.push_str(&text);
|
|
}
|
|
TemplateNode::Conditional {
|
|
condition,
|
|
if_block,
|
|
else_block,
|
|
} => {
|
|
let should_render = self.evaluate_condition(condition, context)?;
|
|
if should_render {
|
|
result.push_str(&self.render_nodes(if_block, context).await?);
|
|
} else if let Some(else_block) = else_block {
|
|
result.push_str(&self.render_nodes(else_block, context).await?);
|
|
}
|
|
}
|
|
TemplateNode::Loop {
|
|
variable,
|
|
iterable,
|
|
body,
|
|
} => {
|
|
let items = context
|
|
.get(iterable)
|
|
.ok_or_else(|| TemplateError::MissingVariable(iterable.clone()))?;
|
|
|
|
if let Value::Array(arr) = items {
|
|
for item in arr {
|
|
let mut loop_context = context.clone();
|
|
loop_context.insert(variable, item.clone());
|
|
result.push_str(&self.render_nodes(body, &loop_context).await?);
|
|
}
|
|
}
|
|
}
|
|
TemplateNode::Block { name: _, content } => {
|
|
result.push_str(&self.render_nodes(content, context).await?);
|
|
}
|
|
TemplateNode::Extends { .. } => {
|
|
// Template inheritance is handled at the engine level
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(result)
|
|
})
|
|
}
|
|
|
|
fn evaluate_condition(&self, condition: &str, context: &TemplateContext) -> Result<bool> {
|
|
// Simple boolean evaluation
|
|
if let Some(value) = context.get(condition) {
|
|
return Ok(match value {
|
|
Value::Bool(b) => *b,
|
|
Value::Null => false,
|
|
Value::Number(n) => n.as_f64().unwrap_or(0.0) != 0.0,
|
|
Value::String(s) => !s.is_empty(),
|
|
Value::Array(arr) => !arr.is_empty(),
|
|
Value::Object(obj) => !obj.is_empty(),
|
|
});
|
|
}
|
|
Ok(false)
|
|
}
|
|
|
|
fn value_to_string(&self, value: &Value) -> String {
|
|
match value {
|
|
Value::String(s) => s.clone(),
|
|
Value::Number(n) => n.to_string(),
|
|
Value::Bool(b) => b.to_string(),
|
|
Value::Null => String::new(),
|
|
_ => serde_json::to_string(value).unwrap_or_default(),
|
|
}
|
|
}
|
|
|
|
fn value_to_html_escaped_string(&self, value: &Value) -> String {
|
|
let text = self.value_to_string(value);
|
|
// Basic HTML escaping for security
|
|
text.replace('&', "&")
|
|
.replace('<', "<")
|
|
.replace('>', ">")
|
|
.replace('"', """)
|
|
.replace('\'', "'")
|
|
}
|
|
|
|
async fn process_partial_binding(
|
|
&self,
|
|
node: &TemplateNode,
|
|
context: &TemplateContext,
|
|
) -> Result<TemplateNode> {
|
|
match node {
|
|
TemplateNode::Variable { name, filters: _ } => {
|
|
if let Some(value) = context.get(name) {
|
|
// Replace with text node if variable is bound
|
|
let text = self.value_to_string(value);
|
|
Ok(TemplateNode::Text(text))
|
|
} else {
|
|
// Keep as variable if not bound
|
|
Ok(node.clone())
|
|
}
|
|
}
|
|
_ => Ok(node.clone()), // For simplicity, only handle variables in partial binding
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Template parser for converting template strings into AST
|
|
struct TemplateParser {
|
|
variable_regex: Regex,
|
|
conditional_regex: Regex,
|
|
loop_regex: Regex,
|
|
function_regex: Regex,
|
|
block_regex: Regex,
|
|
}
|
|
|
|
impl TemplateParser {
|
|
fn new() -> Self {
|
|
Self {
|
|
variable_regex: Regex::new(r"\{([^}]+)\}").unwrap(),
|
|
conditional_regex: Regex::new(r"\{%-?\s*if\s+([^%}]+)\s*-?%\}").unwrap(),
|
|
loop_regex: Regex::new(r"\{%-?\s*for\s+(\w+)\s+in\s+(\w+)\s*-?%\}").unwrap(),
|
|
function_regex: Regex::new(r"\{([^}]+\(\)(?:\s*\|\s*[^}]+)*)\}").unwrap(),
|
|
block_regex: Regex::new(r"\{%-?\s*block\s+(\w+)\s*-?%\}").unwrap(),
|
|
}
|
|
}
|
|
|
|
fn parse(&self, template: &str) -> Result<Vec<TemplateNode>> {
|
|
let mut nodes = Vec::new();
|
|
let mut remaining = template;
|
|
|
|
while !remaining.is_empty() {
|
|
if let Some(next_node) = self.parse_next_node(&mut remaining)? {
|
|
nodes.push(next_node);
|
|
} else {
|
|
// If no special syntax found, treat rest as text
|
|
if !remaining.is_empty() {
|
|
nodes.push(TemplateNode::Text(remaining.to_string()));
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(nodes)
|
|
}
|
|
|
|
fn parse_next_node(&self, remaining: &mut &str) -> Result<Option<TemplateNode>> {
|
|
if remaining.is_empty() {
|
|
return Ok(None);
|
|
}
|
|
|
|
// Find the earliest position of any special syntax
|
|
let variable_pos = self.variable_regex.find(remaining).map(|m| m.start());
|
|
let conditional_pos = self.conditional_regex.find(remaining).map(|m| m.start());
|
|
let loop_pos = self.loop_regex.find(remaining).map(|m| m.start());
|
|
let block_pos = self.block_regex.find(remaining).map(|m| m.start());
|
|
let function_pos = self.function_regex.find(remaining).map(|m| m.start());
|
|
let extends_pos = if remaining.starts_with("{%- extends") {
|
|
Some(0)
|
|
} else {
|
|
None
|
|
};
|
|
|
|
// Find the earliest special syntax
|
|
let earliest_pos = [
|
|
variable_pos,
|
|
conditional_pos,
|
|
loop_pos,
|
|
block_pos,
|
|
function_pos,
|
|
extends_pos,
|
|
]
|
|
.iter()
|
|
.filter_map(|&pos| pos)
|
|
.min();
|
|
|
|
if let Some(pos) = earliest_pos {
|
|
// If there's text before the special syntax, return that first
|
|
if pos > 0 {
|
|
let text = remaining[..pos].to_string();
|
|
*remaining = &remaining[pos..];
|
|
return Ok(Some(TemplateNode::Text(text)));
|
|
}
|
|
|
|
// Otherwise, parse the special syntax
|
|
if extends_pos == Some(0) {
|
|
return self.parse_extends(remaining);
|
|
} else if conditional_pos == Some(0) {
|
|
return self.parse_conditional(remaining);
|
|
} else if loop_pos == Some(0) {
|
|
return self.parse_loop(remaining);
|
|
} else if block_pos == Some(0) {
|
|
return self.parse_block(remaining);
|
|
} else if function_pos == Some(0) {
|
|
return self.parse_function(remaining);
|
|
} else if variable_pos == Some(0) {
|
|
return self.parse_variable(remaining);
|
|
}
|
|
} else {
|
|
// No special syntax found, treat remaining as text
|
|
if !remaining.is_empty() {
|
|
let text = remaining.to_string();
|
|
*remaining = "";
|
|
return Ok(Some(TemplateNode::Text(text)));
|
|
}
|
|
}
|
|
|
|
Ok(None)
|
|
}
|
|
|
|
fn parse_variable(&self, remaining: &mut &str) -> Result<Option<TemplateNode>> {
|
|
if let Some(caps) = self.variable_regex.captures(remaining) {
|
|
let full_match = caps.get(0).unwrap();
|
|
let variable_content = caps.get(1).unwrap().as_str().trim();
|
|
|
|
// Parse filters
|
|
let parts: Vec<&str> = variable_content.split('|').collect();
|
|
let name = parts[0].trim().to_string();
|
|
let mut filters = Vec::new();
|
|
|
|
for filter_part in parts.iter().skip(1) {
|
|
let filter_parts: Vec<&str> = filter_part.trim().splitn(2, '(').collect();
|
|
let filter_name = filter_parts[0].trim().to_string();
|
|
let args = if filter_parts.len() > 1 {
|
|
filter_parts[1].trim_end_matches(')').to_string()
|
|
} else {
|
|
String::new()
|
|
};
|
|
filters.push((filter_name, args));
|
|
}
|
|
|
|
*remaining = &remaining[full_match.end()..];
|
|
return Ok(Some(TemplateNode::Variable { name, filters }));
|
|
}
|
|
Ok(None)
|
|
}
|
|
|
|
fn parse_conditional(&self, remaining: &mut &str) -> Result<Option<TemplateNode>> {
|
|
if let Some(caps) = self.conditional_regex.captures(remaining) {
|
|
let condition = caps.get(1).unwrap().as_str().trim().to_string();
|
|
let if_start = caps.get(0).unwrap().end();
|
|
|
|
// Find endif or else
|
|
let if_content_start = if_start;
|
|
let (else_pos, endif_pos) = self.find_conditional_boundaries(&remaining[if_start..])?;
|
|
|
|
let if_content =
|
|
&remaining[if_content_start..if_content_start + else_pos.unwrap_or(endif_pos)];
|
|
let if_block = self.parse(if_content)?;
|
|
|
|
let else_block = if let Some(else_start) = else_pos {
|
|
let else_content =
|
|
&remaining[if_content_start + else_start..if_content_start + endif_pos];
|
|
// Skip the else tag
|
|
let else_tag_end = else_content.find("%}").unwrap_or(0) + 2;
|
|
Some(self.parse(&else_content[else_tag_end..])?)
|
|
} else {
|
|
None
|
|
};
|
|
|
|
*remaining = &remaining[if_content_start + endif_pos + self.find_endif_tag_length()..];
|
|
|
|
return Ok(Some(TemplateNode::Conditional {
|
|
condition,
|
|
if_block,
|
|
else_block,
|
|
}));
|
|
}
|
|
Ok(None)
|
|
}
|
|
|
|
fn parse_loop(&self, remaining: &mut &str) -> Result<Option<TemplateNode>> {
|
|
if let Some(caps) = self.loop_regex.captures(remaining) {
|
|
let variable = caps.get(1).unwrap().as_str().to_string();
|
|
let iterable = caps.get(2).unwrap().as_str().to_string();
|
|
let loop_start = caps.get(0).unwrap().end();
|
|
|
|
let endfor_pos = self.find_endfor_position(&remaining[loop_start..])?;
|
|
let loop_content = &remaining[loop_start..loop_start + endfor_pos];
|
|
let body = self.parse(loop_content)?;
|
|
|
|
*remaining = &remaining[loop_start + endfor_pos + self.find_endfor_tag_length()..];
|
|
|
|
return Ok(Some(TemplateNode::Loop {
|
|
variable,
|
|
iterable,
|
|
body,
|
|
}));
|
|
}
|
|
Ok(None)
|
|
}
|
|
|
|
fn parse_block(&self, remaining: &mut &str) -> Result<Option<TemplateNode>> {
|
|
if let Some(caps) = self.block_regex.captures(remaining) {
|
|
let name = caps.get(1).unwrap().as_str().to_string();
|
|
let block_start = caps.get(0).unwrap().end();
|
|
|
|
let endblock_pos = self.find_endblock_position(&remaining[block_start..])?;
|
|
let block_content = &remaining[block_start..block_start + endblock_pos];
|
|
let content = self.parse(block_content)?;
|
|
|
|
*remaining = &remaining[block_start + endblock_pos + self.find_endblock_tag_length()..];
|
|
|
|
return Ok(Some(TemplateNode::Block { name, content }));
|
|
}
|
|
Ok(None)
|
|
}
|
|
|
|
fn parse_function(&self, remaining: &mut &str) -> Result<Option<TemplateNode>> {
|
|
if let Some(caps) = self.function_regex.captures(remaining) {
|
|
let function_content = caps.get(1).unwrap().as_str().trim();
|
|
|
|
// Parse function name and filters
|
|
let parts: Vec<&str> = function_content.split('|').collect();
|
|
let func_call = parts[0].trim();
|
|
let name = func_call.trim_end_matches("()").to_string();
|
|
|
|
let mut filters = Vec::new();
|
|
for filter_part in parts.iter().skip(1) {
|
|
let filter_parts: Vec<&str> = filter_part.trim().splitn(2, '(').collect();
|
|
let filter_name = filter_parts[0].trim().to_string();
|
|
let args = if filter_parts.len() > 1 {
|
|
filter_parts[1]
|
|
.trim_end_matches(')')
|
|
.trim_matches('\'')
|
|
.to_string()
|
|
} else {
|
|
String::new()
|
|
};
|
|
filters.push((filter_name, args));
|
|
}
|
|
|
|
*remaining = &remaining[caps.get(0).unwrap().end()..];
|
|
return Ok(Some(TemplateNode::Function { name, filters }));
|
|
}
|
|
Ok(None)
|
|
}
|
|
|
|
fn parse_extends(&self, remaining: &mut &str) -> Result<Option<TemplateNode>> {
|
|
if let Some(end) = remaining.find("%}") {
|
|
let extends_content = &remaining[..end];
|
|
if let Some(parent_start) = extends_content.rfind(' ') {
|
|
let parent = extends_content[parent_start + 1..].trim().to_string();
|
|
*remaining = &remaining[end + 2..];
|
|
return Ok(Some(TemplateNode::Extends { parent }));
|
|
}
|
|
}
|
|
Err(TemplateError::SyntaxError("Invalid extends syntax".to_string()).into())
|
|
}
|
|
|
|
fn find_conditional_boundaries(&self, content: &str) -> Result<(Option<usize>, usize)> {
|
|
let else_regex = Regex::new(r"\{%-?\s*else\s*-?%\}").unwrap();
|
|
let endif_regex = Regex::new(r"\{%-?\s*endif\s*-?%\}").unwrap();
|
|
|
|
let else_pos = else_regex.find(content).map(|m| m.start());
|
|
let endif_pos = endif_regex
|
|
.find(content)
|
|
.ok_or_else(|| TemplateError::SyntaxError("Missing endif".to_string()))?
|
|
.start();
|
|
|
|
Ok((else_pos, endif_pos))
|
|
}
|
|
|
|
fn find_endfor_position(&self, content: &str) -> Result<usize> {
|
|
let endfor_regex = Regex::new(r"\{%-?\s*endfor\s*-?%\}").unwrap();
|
|
endfor_regex
|
|
.find(content)
|
|
.map(|m| m.start())
|
|
.ok_or_else(|| TemplateError::SyntaxError("Missing endfor".to_string()).into())
|
|
}
|
|
|
|
fn find_endblock_position(&self, content: &str) -> Result<usize> {
|
|
let endblock_regex = Regex::new(r"\{%-?\s*endblock\s*-?%\}").unwrap();
|
|
endblock_regex
|
|
.find(content)
|
|
.map(|m| m.start())
|
|
.ok_or_else(|| TemplateError::SyntaxError("Missing endblock".to_string()).into())
|
|
}
|
|
|
|
fn find_endif_tag_length(&self) -> usize {
|
|
"{%- endif -%}".len() // Approximate, should parse actual tag
|
|
}
|
|
|
|
fn find_endfor_tag_length(&self) -> usize {
|
|
"{%- endfor -%}".len()
|
|
}
|
|
|
|
fn find_endblock_tag_length(&self) -> usize {
|
|
"{%- endblock -%}".len()
|
|
}
|
|
}
|
|
|
|
/// Template engine with caching and inheritance support
|
|
pub struct TemplateEngine {
|
|
templates: Arc<DashMap<String, PromptTemplate>>,
|
|
compiled_cache: Arc<DashMap<String, PromptTemplate>>,
|
|
cache_enabled: bool,
|
|
}
|
|
|
|
impl TemplateEngine {
|
|
/// Create a new template engine
|
|
pub fn new() -> Self {
|
|
Self {
|
|
templates: Arc::new(DashMap::new()),
|
|
compiled_cache: Arc::new(DashMap::new()),
|
|
cache_enabled: false,
|
|
}
|
|
}
|
|
|
|
/// Enable or disable template caching
|
|
pub fn enable_caching(&mut self, enabled: bool) {
|
|
self.cache_enabled = enabled;
|
|
}
|
|
|
|
/// Register a template with name
|
|
pub async fn register_template(&mut self, name: &str, template: &str) -> Result<()> {
|
|
let compiled_template = PromptTemplate::from_str(template)?;
|
|
self.templates.insert(name.to_string(), compiled_template);
|
|
Ok(())
|
|
}
|
|
|
|
/// Compile a template with inheritance support
|
|
pub async fn compile_template(&self, template: &str) -> Result<PromptTemplate> {
|
|
// Check cache first
|
|
if self.cache_enabled {
|
|
let hash = self.compute_template_hash(template);
|
|
if let Some(cached) = self.compiled_cache.get(&hash) {
|
|
return Ok(cached.clone());
|
|
}
|
|
}
|
|
|
|
let mut compiled = PromptTemplate::from_str(template)?;
|
|
|
|
// Handle template inheritance
|
|
if let Some(parent_name) = &compiled.parent_template
|
|
&& let Some(parent_template) = self.templates.get(parent_name)
|
|
{
|
|
compiled = self.merge_with_parent(&compiled, &parent_template)?;
|
|
}
|
|
|
|
// Cache the compiled template
|
|
if self.cache_enabled {
|
|
let hash = self.compute_template_hash(template);
|
|
self.compiled_cache.insert(hash, compiled.clone());
|
|
}
|
|
|
|
Ok(compiled)
|
|
}
|
|
|
|
fn merge_with_parent(
|
|
&self,
|
|
child: &PromptTemplate,
|
|
parent: &PromptTemplate,
|
|
) -> Result<PromptTemplate> {
|
|
let mut merged_nodes = parent.nodes.clone();
|
|
|
|
// Replace blocks with child's version
|
|
for (block_name, child_block) in &child.blocks {
|
|
// Find and replace the corresponding block in parent
|
|
self.replace_block_in_nodes(&mut merged_nodes, block_name, child_block);
|
|
}
|
|
|
|
Ok(PromptTemplate {
|
|
id: format!("{}_extended", child.id),
|
|
source: format!("{}\n{}", parent.source, child.source),
|
|
nodes: merged_nodes,
|
|
parent_template: None,
|
|
blocks: child.blocks.clone(),
|
|
creation_time: Instant::now(),
|
|
})
|
|
}
|
|
|
|
fn replace_block_in_nodes(
|
|
&self,
|
|
nodes: &mut Vec<TemplateNode>,
|
|
block_name: &str,
|
|
replacement: &[TemplateNode],
|
|
) {
|
|
for node in nodes.iter_mut() {
|
|
if let TemplateNode::Block { name, content } = node
|
|
&& name == block_name
|
|
{
|
|
*content = replacement.to_vec();
|
|
}
|
|
}
|
|
}
|
|
|
|
fn compute_template_hash(&self, template: &str) -> String {
|
|
let mut hasher = Sha256::new();
|
|
hasher.update(template.as_bytes());
|
|
format!("{:x}", hasher.finalize())
|
|
}
|
|
}
|
|
|
|
impl Default for TemplateEngine {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use tokio_test;
|
|
|
|
#[tokio::test]
|
|
async fn test_basic_template_parsing() {
|
|
let template = PromptTemplate::from_str("Hello {name}!").unwrap();
|
|
let mut context = TemplateContext::new();
|
|
context.insert("name", Value::String("World".to_string()));
|
|
|
|
let result = template.render(&context).await.unwrap();
|
|
assert_eq!(result, "Hello World!");
|
|
}
|
|
}
|