Initial commit

crates/models/rtx-nlg/src/serving/templates.rs

@@ -0,0 +1,996 @@
+//! Advanced prompt templates and engineering utilities with support for
+//! dynamic templates, chain-of-thought patterns, and complex prompt engineering
+
+use crate::{GenerationConfig, ModelInterface, NlgError, Result};
+use regex::Regex;
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::sync::Arc;
+
+/// Advanced prompt template with support for complex transformations and patterns
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct PromptTemplate {
+    pub name: String,
+    pub template: String,
+    pub variables: Vec<String>,
+    pub description: Option<String>,
+    pub template_type: TemplateType,
+    pub validation_rules: Vec<ValidationRule>,
+    pub transformations: Vec<TemplateTransformation>,
+    pub metadata: TemplateMetadata,
+}
+
+/// Types of prompt templates
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum TemplateType {
+    /// Basic variable substitution template
+    Basic,
+    /// Chain-of-thought reasoning template
+    ChainOfThought,
+    /// Few-shot learning template with examples
+    FewShot { examples: Vec<FewShotExample> },
+    /// Role-based template with system/user/assistant messages
+    RoleBased { roles: Vec<Role> },
+    /// Conditional template with if/else logic
+    Conditional { conditions: Vec<ConditionalBlock> },
+    /// Iterative template with loops
+    Iterative { loop_config: LoopConfig },
+    /// Multi-step template for complex workflows
+    MultiStep { steps: Vec<TemplateStep> },
+}
+
+/// Few-shot example for learning templates
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct FewShotExample {
+    pub input: String,
+    pub output: String,
+    pub explanation: Option<String>,
+}
+
+/// Role in a conversation template
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Role {
+    pub name: String,
+    pub template: String,
+    pub instructions: Option<String>,
+}
+
+/// Conditional block for template logic
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ConditionalBlock {
+    pub condition: String,
+    pub if_template: String,
+    pub else_template: Option<String>,
+}
+
+/// Loop configuration for iterative templates
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct LoopConfig {
+    pub variable: String,
+    pub template: String,
+    pub separator: String,
+    pub max_iterations: usize,
+}
+
+/// Step in a multi-step template
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TemplateStep {
+    pub name: String,
+    pub template: String,
+    pub requires_model_output: bool,
+    pub depends_on: Vec<String>,
+}
+
+/// Validation rule for template inputs
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ValidationRule {
+    pub variable: String,
+    pub rule_type: ValidationType,
+    pub parameters: HashMap<String, String>,
+    pub error_message: String,
+}
+
+/// Types of validation
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum ValidationType {
+    Required,
+    MinLength(usize),
+    MaxLength(usize),
+    Regex(String),
+    OneOf(Vec<String>),
+    Numeric { min: Option<f64>, max: Option<f64> },
+}
+
+/// Template transformation for preprocessing
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TemplateTransformation {
+    pub variable: String,
+    pub transformation: TransformationType,
+}
+
+/// Types of transformations
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum TransformationType {
+    Lowercase,
+    Uppercase,
+    Trim,
+    Truncate(usize),
+    Replace { from: String, to: String },
+    Format(String),
+    Custom(String),
+}
+
+/// Template metadata
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TemplateMetadata {
+    pub author: Option<String>,
+    pub version: String,
+    pub tags: Vec<String>,
+    pub performance_metrics: Option<PerformanceMetrics>,
+    pub usage_stats: Option<UsageStats>,
+}
+
+/// Performance metrics for template evaluation
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct PerformanceMetrics {
+    pub avg_generation_time_ms: f64,
+    pub success_rate: f32,
+    pub quality_score: f32,
+}
+
+/// Usage statistics for templates
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct UsageStats {
+    pub total_uses: u64,
+    pub successful_uses: u64,
+    pub avg_output_length: f32,
+    pub last_used: Option<chrono::DateTime<chrono::Utc>>,
+}
+
+impl Default for TemplateMetadata {
+    fn default() -> Self {
+        Self {
+            author: None,
+            version: "1.0.0".to_string(),
+            tags: Vec::new(),
+            performance_metrics: None,
+            usage_stats: None,
+        }
+    }
+}
+
+impl PromptTemplate {
+    /// Create a new basic template
+    pub fn new(name: String, template: String) -> Self {
+        let variables = Self::extract_variables(&template);
+        Self {
+            name,
+            template,
+            variables,
+            description: None,
+            template_type: TemplateType::Basic,
+            validation_rules: Vec::new(),
+            transformations: Vec::new(),
+            metadata: TemplateMetadata::default(),
+        }
+    }
+
+    /// Create a chain-of-thought template
+    pub fn chain_of_thought(name: String, template: String) -> Self {
+        let variables = Self::extract_variables(&template);
+        Self {
+            name,
+            template,
+            variables,
+            description: Some("Chain-of-thought reasoning template".to_string()),
+            template_type: TemplateType::ChainOfThought,
+            validation_rules: Vec::new(),
+            transformations: Vec::new(),
+            metadata: TemplateMetadata::default(),
+        }
+    }
+
+    /// Create a few-shot learning template
+    pub fn few_shot(name: String, template: String, examples: Vec<FewShotExample>) -> Self {
+        let variables = Self::extract_variables(&template);
+        Self {
+            name,
+            template,
+            variables,
+            description: Some("Few-shot learning template".to_string()),
+            template_type: TemplateType::FewShot { examples },
+            validation_rules: Vec::new(),
+            transformations: Vec::new(),
+            metadata: TemplateMetadata::default(),
+        }
+    }
+
+    /// Create a role-based conversation template
+    pub fn role_based(name: String, roles: Vec<Role>) -> Self {
+        let template = roles
+            .iter()
+            .map(|r| format!("{}: {}", r.name, r.template))
+            .collect::<Vec<_>>()
+            .join("\n");
+        let variables = Self::extract_variables(&template);
+
+        Self {
+            name,
+            template,
+            variables,
+            description: Some("Role-based conversation template".to_string()),
+            template_type: TemplateType::RoleBased { roles },
+            validation_rules: Vec::new(),
+            transformations: Vec::new(),
+            metadata: TemplateMetadata::default(),
+        }
+    }
+
+    /// Add validation rule
+    pub fn with_validation(mut self, rule: ValidationRule) -> Self {
+        self.validation_rules.push(rule);
+        self
+    }
+
+    /// Add transformation
+    pub fn with_transformation(mut self, transformation: TemplateTransformation) -> Self {
+        self.transformations.push(transformation);
+        self
+    }
+
+    /// Render template with advanced features
+    pub fn render(&self, values: &HashMap<String, String>) -> Result<String> {
+        // Validate inputs
+        self.validate_inputs(values)?;
+
+        // Apply transformations
+        let transformed_values = self.apply_transformations(values)?;
+
+        match &self.template_type {
+            TemplateType::Basic => self.render_basic(&transformed_values),
+            TemplateType::ChainOfThought => self.render_chain_of_thought(&transformed_values),
+            TemplateType::FewShot { examples } => {
+                self.render_few_shot(&transformed_values, examples)
+            }
+            TemplateType::RoleBased { roles } => self.render_role_based(&transformed_values, roles),
+            TemplateType::Conditional { conditions } => {
+                self.render_conditional(&transformed_values, conditions)
+            }
+            TemplateType::Iterative { loop_config } => {
+                self.render_iterative(&transformed_values, loop_config)
+            }
+            TemplateType::MultiStep { steps } => self.render_multi_step(&transformed_values, steps),
+        }
+    }
+
+    /// Render with model for multi-step templates
+    pub async fn render_with_model(
+        &self,
+        values: &HashMap<String, String>,
+        model: Option<Arc<dyn ModelInterface>>,
+        config: Option<&GenerationConfig>,
+    ) -> Result<String> {
+        match &self.template_type {
+            TemplateType::MultiStep { steps } => {
+                self.render_multi_step_with_model(values, steps, model, config)
+                    .await
+            }
+            _ => self.render(values),
+        }
+    }
+
+    fn validate_inputs(&self, values: &HashMap<String, String>) -> Result<()> {
+        for rule in &self.validation_rules {
+            let value = values.get(&rule.variable);
+
+            match &rule.rule_type {
+                ValidationType::Required => {
+                    if value.is_none() || value.unwrap().is_empty() {
+                        return Err(NlgError::ValidationFailed {
+                            field: rule.variable.clone(),
+                            message: rule.error_message.clone(),
+                        });
+                    }
+                }
+                ValidationType::MinLength(min_len) => {
+                    if let Some(v) = value
+                        && v.len() < *min_len
+                    {
+                        return Err(NlgError::ValidationFailed {
+                            field: rule.variable.clone(),
+                            message: rule.error_message.clone(),
+                        });
+                    }
+                }
+                ValidationType::MaxLength(max_len) => {
+                    if let Some(v) = value
+                        && v.len() > *max_len
+                    {
+                        return Err(NlgError::ValidationFailed {
+                            field: rule.variable.clone(),
+                            message: rule.error_message.clone(),
+                        });
+                    }
+                }
+                ValidationType::Regex(pattern) => {
+                    if let Some(v) = value {
+                        let regex = Regex::new(pattern).map_err(|e| NlgError::InvalidPattern {
+                            pattern: pattern.clone(),
+                            error: e.to_string(),
+                        })?;
+                        if !regex.is_match(v) {
+                            return Err(NlgError::ValidationFailed {
+                                field: rule.variable.clone(),
+                                message: rule.error_message.clone(),
+                            });
+                        }
+                    }
+                }
+                ValidationType::OneOf(options) => {
+                    if let Some(v) = value
+                        && !options.contains(v)
+                    {
+                        return Err(NlgError::ValidationFailed {
+                            field: rule.variable.clone(),
+                            message: rule.error_message.clone(),
+                        });
+                    }
+                }
+                ValidationType::Numeric { min, max } => {
+                    if let Some(v) = value {
+                        match v.parse::<f64>() {
+                            Ok(num) => {
+                                if let Some(min_val) = min
+                                    && num < *min_val
+                                {
+                                    return Err(NlgError::ValidationFailed {
+                                        field: rule.variable.clone(),
+                                        message: rule.error_message.clone(),
+                                    });
+                                }
+                                if let Some(max_val) = max
+                                    && num > *max_val
+                                {
+                                    return Err(NlgError::ValidationFailed {
+                                        field: rule.variable.clone(),
+                                        message: rule.error_message.clone(),
+                                    });
+                                }
+                            }
+                            Err(_) => {
+                                return Err(NlgError::ValidationFailed {
+                                    field: rule.variable.clone(),
+                                    message: "Value must be numeric".to_string(),
+                                });
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+
+    fn apply_transformations(
+        &self,
+        values: &HashMap<String, String>,
+    ) -> Result<HashMap<String, String>> {
+        let mut result = values.clone();
+
+        for transformation in &self.transformations {
+            if let Some(value) = result.get(&transformation.variable).cloned() {
+                let transformed = match &transformation.transformation {
+                    TransformationType::Lowercase => value.to_lowercase(),
+                    TransformationType::Uppercase => value.to_uppercase(),
+                    TransformationType::Trim => value.trim().to_string(),
+                    TransformationType::Truncate(max_len) => {
+                        if value.len() > *max_len {
+                            value[..*max_len].to_string()
+                        } else {
+                            value
+                        }
+                    }
+                    TransformationType::Replace { from, to } => value.replace(from, to),
+                    TransformationType::Format(format_str) => format_str.replace("{}", &value),
+                    TransformationType::Custom(_) => {
+                        // Would implement custom transformation logic
+                        value
+                    }
+                };
+                result.insert(transformation.variable.clone(), transformed);
+            }
+        }
+
+        Ok(result)
+    }
+
+    fn render_basic(&self, values: &HashMap<String, String>) -> Result<String> {
+        let mut result = self.template.clone();
+
+        for var in &self.variables {
+            let placeholder = format!("{{{{{var}}}}}");
+            let value = values
+                .get(var)
+                .ok_or_else(|| NlgError::missing_parameter(var))?;
+            result = result.replace(&placeholder, value);
+        }
+
+        Ok(result)
+    }
+
+    fn render_chain_of_thought(&self, values: &HashMap<String, String>) -> Result<String> {
+        let mut result = self.template.clone();
+
+        // Add chain-of-thought prompting
+        let cot_suffix = "\n\nLet me think step by step:";
+        if !result.ends_with(cot_suffix) {
+            result.push_str(cot_suffix);
+        }
+
+        for var in &self.variables {
+            let placeholder = format!("{{{{{var}}}}}");
+            let value = values
+                .get(var)
+                .ok_or_else(|| NlgError::missing_parameter(var))?;
+            result = result.replace(&placeholder, value);
+        }
+
+        Ok(result)
+    }
+
+    fn render_few_shot(
+        &self,
+        values: &HashMap<String, String>,
+        examples: &[FewShotExample],
+    ) -> Result<String> {
+        let mut result = String::new();
+
+        // Add examples first
+        for (i, example) in examples.iter().enumerate() {
+            result.push_str(&format!("Example {}:\n", i + 1));
+            result.push_str(&format!("Input: {}\n", example.input));
+            result.push_str(&format!("Output: {}\n", example.output));
+            if let Some(explanation) = &example.explanation {
+                result.push_str(&format!("Explanation: {explanation}\n"));
+            }
+            result.push('\n');
+        }
+
+        // Add the actual template
+        result.push_str("Now solve this:\n");
+        let rendered_template = self.render_basic(values)?;
+        result.push_str(&rendered_template);
+
+        Ok(result)
+    }
+
+    fn render_role_based(
+        &self,
+        values: &HashMap<String, String>,
+        roles: &[Role],
+    ) -> Result<String> {
+        let mut result = String::new();
+
+        for role in roles {
+            let mut role_content = role.template.clone();
+
+            // Replace variables in role content
+            for var in &self.variables {
+                let placeholder = format!("{{{{{var}}}}}");
+                if let Some(value) = values.get(var) {
+                    role_content = role_content.replace(&placeholder, value);
+                }
+            }
+
+            result.push_str(&format!("{}:\n{}\n\n", role.name, role_content));
+        }
+
+        Ok(result.trim_end().to_string())
+    }
+
+    fn render_conditional(
+        &self,
+        values: &HashMap<String, String>,
+        conditions: &[ConditionalBlock],
+    ) -> Result<String> {
+        let mut result = self.render_basic(values)?;
+
+        for condition in conditions {
+            // Simple condition evaluation - would implement proper expression parsing
+            let should_use_if = self.evaluate_condition(&condition.condition, values)?;
+
+            let template_to_use = if should_use_if {
+                &condition.if_template
+            } else if let Some(else_template) = &condition.else_template {
+                else_template
+            } else {
+                continue;
+            };
+
+            // Replace condition placeholder with appropriate template
+            let condition_placeholder = format!("{{{{#{}}}}}", condition.condition);
+            result = result.replace(&condition_placeholder, template_to_use);
+        }
+
+        Ok(result)
+    }
+
+    fn render_iterative(
+        &self,
+        values: &HashMap<String, String>,
+        loop_config: &LoopConfig,
+    ) -> Result<String> {
+        let list_value = values
+            .get(&loop_config.variable)
+            .ok_or_else(|| NlgError::missing_parameter(&loop_config.variable))?;
+
+        // Parse list (simple comma-separated for now)
+        let items: Vec<&str> = list_value.split(',').map(str::trim).collect();
+        let mut results = Vec::new();
+
+        for (i, item) in items.iter().enumerate() {
+            if i >= loop_config.max_iterations {
+                break;
+            }
+
+            let mut item_template = loop_config.template.clone();
+            item_template = item_template.replace("{{item}}", item);
+            item_template = item_template.replace("{{index}}", &i.to_string());
+
+            // Replace other variables
+            for (key, value) in values {
+                if key != &loop_config.variable {
+                    let placeholder = format!("{{{{{key}}}}}");
+                    item_template = item_template.replace(&placeholder, value);
+                }
+            }
+
+            results.push(item_template);
+        }
+
+        Ok(results.join(&loop_config.separator))
+    }
+
+    fn render_multi_step(
+        &self,
+        values: &HashMap<String, String>,
+        steps: &[TemplateStep],
+    ) -> Result<String> {
+        // For basic multi-step without model, just concatenate steps
+        let mut result = String::new();
+        let mut step_outputs: HashMap<String, String> = HashMap::new();
+
+        for step in steps {
+            let mut step_template = step.template.clone();
+
+            // Replace variables from original values
+            for (key, value) in values {
+                let placeholder = format!("{{{{{key}}}}}");
+                step_template = step_template.replace(&placeholder, value);
+            }
+
+            // Replace variables from previous step outputs
+            for (key, value) in &step_outputs {
+                let placeholder = format!("{{{{{key}}}}}");
+                step_template = step_template.replace(&placeholder, value);
+            }
+
+            if step.requires_model_output {
+                step_template.push_str(" [Requires model output]");
+            }
+
+            step_outputs.insert(step.name.clone(), step_template.clone());
+            result.push_str(&format!("Step {}: {}\n\n", step.name, step_template));
+        }
+
+        Ok(result.trim_end().to_string())
+    }
+
+    async fn render_multi_step_with_model(
+        &self,
+        values: &HashMap<String, String>,
+        steps: &[TemplateStep],
+        model: Option<Arc<dyn ModelInterface>>,
+        config: Option<&GenerationConfig>,
+    ) -> Result<String> {
+        let mut result = String::new();
+        let mut step_outputs: HashMap<String, String> = HashMap::new();
+
+        for step in steps {
+            let mut step_template = step.template.clone();
+
+            // Replace variables
+            for (key, value) in values {
+                let placeholder = format!("{{{{{key}}}}}");
+                step_template = step_template.replace(&placeholder, value);
+            }
+
+            for (key, value) in &step_outputs {
+                let placeholder = format!("{{{{{key}}}}}");
+                step_template = step_template.replace(&placeholder, value);
+            }
+
+            let step_output = if step.requires_model_output && model.is_some() && config.is_some() {
+                // Generate output using model
+                let model_ref = model.as_ref().unwrap();
+                let config_ref = config.unwrap();
+
+                let tokenizer = model_ref.tokenizer();
+                let tokens = tokenizer.encode(&step_template)?;
+                let tokens_data: Vec<f32> = tokens.iter().map(|&x| x as f32).collect();
+                let input_tensor = rtx_tensor::Tensor::from_data(
+                    tokens_data,
+                    [1, tokens.len()],
+                    &rtx_tensor::Device::default(),
+                )?;
+
+                let output = model_ref.generate_tokens(&input_tensor, None, config_ref)?;
+
+                tokenizer.decode(&output.sequences[0])?
+            } else {
+                step_template
+            };
+
+            step_outputs.insert(step.name.clone(), step_output.clone());
+            result.push_str(&format!("Step {}: {}\n\n", step.name, step_output));
+        }
+
+        Ok(result.trim_end().to_string())
+    }
+
+    fn evaluate_condition(
+        &self,
+        condition: &str,
+        values: &HashMap<String, String>,
+    ) -> Result<bool> {
+        // Simple condition evaluation - would implement proper expression parser
+        if condition.contains("==") {
+            let parts: Vec<&str> = condition.split("==").collect();
+            if parts.len() == 2 {
+                let left = parts[0].trim();
+                let right = parts[1].trim().trim_matches('"');
+
+                if let Some(value) = values.get(left) {
+                    return Ok(value == right);
+                }
+            }
+        }
+
+        Ok(false)
+    }
+
+    fn extract_variables(template: &str) -> Vec<String> {
+        let mut variables = Vec::new();
+        let mut chars = template.chars().peekable();
+
+        while let Some(ch) = chars.next() {
+            if ch == '{' && chars.peek() == Some(&'{') {
+                chars.next(); // consume second '{'
+                let mut var_name = String::new();
+
+                while let Some(ch) = chars.next() {
+                    if ch == '}' && chars.peek() == Some(&'}') {
+                        chars.next(); // consume second '}'
+                        if !var_name.is_empty() && !var_name.starts_with('#') {
+                            variables.push(var_name.trim().to_string());
+                        }
+                        break;
+                    }
+                    var_name.push(ch);
+                }
+            }
+        }
+
+        variables.sort();
+        variables.dedup();
+        variables
+    }
+}
+
+/// Advanced template manager with categories, optimization, and analytics
+pub struct TemplateManager {
+    templates: HashMap<String, PromptTemplate>,
+    categories: HashMap<String, Vec<String>>,
+    template_cache: lru::LruCache<String, String>,
+    usage_tracker: Arc<std::sync::RwLock<HashMap<String, UsageStats>>>,
+}
+
+impl Default for TemplateManager {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl TemplateManager {
+    /// Create a new template manager
+    pub fn new() -> Self {
+        let mut manager = Self {
+            templates: HashMap::new(),
+            categories: HashMap::new(),
+            template_cache: lru::LruCache::new(std::num::NonZeroUsize::new(1000).unwrap()),
+            usage_tracker: Arc::new(std::sync::RwLock::new(HashMap::new())),
+        };
+        manager.load_default_templates();
+        manager
+    }
+
+    /// Add a template to the manager
+    pub fn add_template(&mut self, template: PromptTemplate) {
+        // Add to categories
+        for tag in &template.metadata.tags {
+            self.categories
+                .entry(tag.clone())
+                .or_default()
+                .push(template.name.clone());
+        }
+
+        self.templates.insert(template.name.clone(), template);
+    }
+
+    /// Get a template by name
+    pub fn get_template(&self, name: &str) -> Option<&PromptTemplate> {
+        self.templates.get(name)
+    }
+
+    /// Get templates by category
+    pub fn get_templates_by_category(&self, category: &str) -> Vec<&PromptTemplate> {
+        if let Some(template_names) = self.categories.get(category) {
+            template_names
+                .iter()
+                .filter_map(|name| self.templates.get(name))
+                .collect()
+        } else {
+            Vec::new()
+        }
+    }
+
+    /// List all template names
+    pub fn list_templates(&self) -> Vec<String> {
+        self.templates.keys().cloned().collect()
+    }
+
+    /// List all categories
+    pub fn list_categories(&self) -> Vec<String> {
+        self.categories.keys().cloned().collect()
+    }
+
+    /// Render a template with caching
+    pub fn render_template(
+        &mut self,
+        name: &str,
+        values: &HashMap<String, String>,
+    ) -> Result<String> {
+        let cache_key = self.create_cache_key(name, values);
+
+        // Check cache first
+        if let Some(cached) = self.template_cache.get(&cache_key) {
+            let result = cached.clone();
+            drop(cached); // Explicitly drop the mutable borrow before calling update_usage_stats
+            self.update_usage_stats(name, true);
+            return Ok(result);
+        }
+
+        // Render template
+        let template = self
+            .get_template(name)
+            .ok_or_else(|| NlgError::missing_parameter(name))?;
+
+        let result = template.render(values)?;
+
+        // Cache the result
+        self.template_cache.put(cache_key, result.clone());
+
+        // Update usage statistics
+        self.update_usage_stats(name, false);
+
+        Ok(result)
+    }
+
+    /// Render template with model support
+    pub async fn render_template_with_model(
+        &mut self,
+        name: &str,
+        values: &HashMap<String, String>,
+        model: Option<Arc<dyn ModelInterface>>,
+        config: Option<&GenerationConfig>,
+    ) -> Result<String> {
+        let template = self
+            .get_template(name)
+            .ok_or_else(|| NlgError::missing_parameter(name))?;
+
+        let result = template.render_with_model(values, model, config).await?;
+        self.update_usage_stats(name, false);
+
+        Ok(result)
+    }
+
+    /// Search templates by content or metadata
+    pub fn search_templates(&self, query: &str) -> Vec<&PromptTemplate> {
+        let query_lower = query.to_lowercase();
+        self.templates
+            .values()
+            .filter(|template| {
+                template.name.to_lowercase().contains(&query_lower)
+                    || template
+                        .description
+                        .as_ref()
+                        .is_some_and(|d| d.to_lowercase().contains(&query_lower))
+                    || template
+                        .metadata
+                        .tags
+                        .iter()
+                        .any(|tag| tag.to_lowercase().contains(&query_lower))
+                    || template.template.to_lowercase().contains(&query_lower)
+            })
+            .collect()
+    }
+
+    /// Get template usage statistics
+    pub fn get_usage_stats(&self, name: &str) -> Option<UsageStats> {
+        let tracker = self.usage_tracker.read().unwrap();
+        tracker.get(name).cloned()
+    }
+
+    /// Export templates to JSON
+    pub fn export_templates(&self) -> Result<String> {
+        let export_data = serde_json::json!({
+            "templates": self.templates,
+            "categories": self.categories,
+        });
+        Ok(serde_json::to_string_pretty(&export_data)?)
+    }
+
+    /// Clear template cache
+    pub fn clear_cache(&mut self) {
+        self.template_cache.clear();
+    }
+
+    /// Get cache statistics
+    pub fn get_cache_stats(&self) -> (usize, usize) {
+        (self.template_cache.len(), self.template_cache.cap().get())
+    }
+
+    // Private helper methods
+
+    fn create_cache_key(&self, name: &str, values: &HashMap<String, String>) -> String {
+        let mut key_parts = vec![name.to_string()];
+        let mut sorted_values: Vec<_> = values.iter().collect();
+        sorted_values.sort_by_key(|(k, _)| *k);
+
+        for (k, v) in sorted_values {
+            key_parts.push(format!("{k}={v}"));
+        }
+
+        key_parts.join("|")
+    }
+
+    fn update_usage_stats(&self, name: &str, cache_hit: bool) {
+        let mut tracker = self.usage_tracker.write().unwrap();
+        let stats = tracker
+            .entry(name.to_string())
+            .or_insert_with(|| UsageStats {
+                total_uses: 0,
+                successful_uses: 0,
+                avg_output_length: 0.0,
+                last_used: None,
+            });
+
+        stats.total_uses += 1;
+        if !cache_hit {
+            stats.successful_uses += 1;
+        }
+        stats.last_used = Some(chrono::Utc::now());
+    }
+
+    fn load_default_templates(&mut self) {
+        // Basic templates with enhanced features
+        let mut summarize_template = PromptTemplate::new(
+            "summarize".to_string(),
+            "Please summarize the following text in {{max_sentences}} sentences:\n\n{{text}}"
+                .to_string(),
+        )
+        .with_validation(ValidationRule {
+            variable: "text".to_string(),
+            rule_type: ValidationType::MinLength(10),
+            parameters: HashMap::new(),
+            error_message: "Text must be at least 10 characters long".to_string(),
+        })
+        .with_transformation(TemplateTransformation {
+            variable: "text".to_string(),
+            transformation: TransformationType::Trim,
+        });
+        summarize_template.metadata.tags = vec!["summarization".to_string(), "nlp".to_string()];
+        self.add_template(summarize_template);
+
+        // Translation template
+        let mut translate_template = PromptTemplate::new(
+            "translate".to_string(),
+            "Translate the following text from {{source_lang}} to {{target_lang}}:\n\n{{text}}"
+                .to_string(),
+        );
+        translate_template.metadata.tags =
+            vec!["translation".to_string(), "multilingual".to_string()];
+        self.add_template(translate_template);
+
+        // Question answering template
+        let mut qa_template = PromptTemplate::new(
+            "qa".to_string(),
+            "Based on the following context, answer the question:\n\nContext: {{context}}\n\nQuestion: {{question}}\n\nAnswer:".to_string(),
+        );
+        qa_template.metadata.tags = vec!["qa".to_string(), "comprehension".to_string()];
+        self.add_template(qa_template);
+
+        // Chain-of-thought template
+        let mut cot_template = PromptTemplate::chain_of_thought(
+            "reasoning".to_string(),
+            "Solve this problem: {{problem}}".to_string(),
+        );
+        cot_template.metadata.tags = vec!["reasoning".to_string(), "chain-of-thought".to_string()];
+        self.add_template(cot_template);
+
+        // Few-shot classification template
+        let few_shot_examples = vec![
+            FewShotExample {
+                input: "I love this product!".to_string(),
+                output: "positive".to_string(),
+                explanation: Some("Expresses satisfaction and positive emotion".to_string()),
+            },
+            FewShotExample {
+                input: "This is terrible".to_string(),
+                output: "negative".to_string(),
+                explanation: Some("Expresses dissatisfaction and negative emotion".to_string()),
+            },
+        ];
+
+        let mut sentiment_template = PromptTemplate::few_shot(
+            "sentiment_classification".to_string(),
+            "Classify the sentiment of: {{text}}".to_string(),
+            few_shot_examples,
+        );
+        sentiment_template.metadata.tags =
+            vec!["classification".to_string(), "sentiment".to_string()];
+        self.add_template(sentiment_template);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_template_variable_extraction() {
+        let template = PromptTemplate::new(
+            "test".to_string(),
+            "Hello {{name}}, you are {{age}} years old.".to_string(),
+        );
+
+        assert_eq!(template.variables, vec!["age", "name"]);
+    }
+
+    #[test]
+    fn test_template_rendering() -> Result<()> {
+        let template = PromptTemplate::new("greeting".to_string(), "Hello {{name}}!".to_string());
+
+        let mut values = HashMap::new();
+        values.insert("name".to_string(), "World".to_string());
+
+        let result = template.render(&values)?;
+        assert_eq!(result, "Hello World!");
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_template_manager() -> Result<()> {
+        let mut manager = TemplateManager::new();
+
+        let mut values = HashMap::new();
+        values.insert("text".to_string(), "This is a test.".to_string());
+        values.insert("max_sentences".to_string(), "1".to_string());
+
+        let result = manager.render_template("summarize", &values)?;
+        assert!(result.contains("This is a test."));
+
+        Ok(())
+    }
+}