rustytorch/crates/models/rtx-llm-tools/tests/prompt_template_tests.rs

use rtx_llm_tools::templates::{PromptTemplate, TemplateContext, TemplateEngine};
use serde_json::Value;

#[tokio::test]
async fn test_simple_variable_substitution() {
    let template_str = "Hello {name}, welcome to {system}!";
    let template = PromptTemplate::from_str(template_str).unwrap();

    let mut context = TemplateContext::new();
    context.insert("name", Value::String("Alice".to_string()));
    context.insert("system", Value::String("RustyTorch".to_string()));

    let result = template.render(&context).await;
    assert!(result.is_ok());
    assert_eq!(result.unwrap(), "Hello Alice, welcome to RustyTorch!");
}

#[tokio::test]
async fn test_conditional_rendering() {
    let template_str = r#"
{%- if is_admin -%}
Welcome admin {name}! You have full access.
{%- else -%}
Hello {name}, you have limited access.
{%- endif -%}
"#;

    let template = PromptTemplate::from_str(template_str).unwrap();

    let mut admin_context = TemplateContext::new();
    admin_context.insert("name", Value::String("Bob".to_string()));
    admin_context.insert("is_admin", Value::Bool(true));

    let result = template.render(&admin_context).await.unwrap();
    assert!(result.contains("Welcome admin Bob! You have full access."));

    let mut user_context = TemplateContext::new();
    user_context.insert("name", Value::String("Charlie".to_string()));
    user_context.insert("is_admin", Value::Bool(false));

    let result = template.render(&user_context).await.unwrap();
    assert!(result.contains("Hello Charlie, you have limited access."));
}

#[tokio::test]
async fn test_loop_iteration() {
    let template_str = r#"
Items:
{%- for item in items %}
- {item.name}: {item.price}
{%- endfor -%}
"#;

    let template = PromptTemplate::from_str(template_str).unwrap();

    let items = vec![
        serde_json::json!({"name": "Apple", "price": "$1.50"}),
        serde_json::json!({"name": "Banana", "price": "$0.75"}),
    ];

    let mut context = TemplateContext::new();
    context.insert("items", Value::Array(items));

    let result = template.render(&context).await.unwrap();
    assert!(result.contains("- Apple: $1.50"));
    assert!(result.contains("- Banana: $0.75"));
}

#[tokio::test]
async fn test_nested_object_access() {
    let template_str = "User: {user.profile.name}, Age: {user.profile.age}";
    let template = PromptTemplate::from_str(template_str).unwrap();

    let user_data = serde_json::json!({
        "user": {
            "profile": {
                "name": "David",
                "age": 30
            }
        }
    });

    let mut context = TemplateContext::new();
    context.insert_nested(user_data);

    let result = template.render(&context).await.unwrap();
    assert_eq!(result, "User: David, Age: 30");
}

#[tokio::test]
async fn test_template_inheritance() {
    let base_template = r#"
System: {system_name}
{%- block content -%}
Default content
{%- endblock -%}
"#;

    let child_template = r#"
{%- extends base -%}
{%- block content -%}
Task: {task_description}
Parameters: {params}
{%- endblock -%}
"#;

    let mut engine = TemplateEngine::new();
    engine
        .register_template("base", base_template)
        .await
        .unwrap();

    let template = engine.compile_template(child_template).await.unwrap();

    let mut context = TemplateContext::new();
    context.insert("system_name", Value::String("RustyTorch".to_string()));
    context.insert(
        "task_description",
        Value::String("Generate text".to_string()),
    );
    context.insert("params", Value::String("{\"max_tokens\": 100}".to_string()));

    let result = template.render(&context).await.unwrap();
    assert!(result.contains("System: RustyTorch"));
    assert!(result.contains("Task: Generate text"));
    assert!(result.contains("Parameters: {\"max_tokens\": 100}"));
}

#[tokio::test]
async fn test_custom_functions() {
    let template_str = "Current time: {now() | format_time('%Y-%m-%d %H:%M')}";
    let template = PromptTemplate::from_str(template_str).unwrap();

    let mut context = TemplateContext::new();
    context.register_function(
        "now",
        Box::new(|| Ok(Value::String(chrono::Utc::now().to_rfc3339()))),
    );

    context.register_filter(
        "format_time",
        Box::new(|_value: &Value, _format: &str| {
            // Mock implementation for test
            Ok(Value::String("2024-01-15 10:30".to_string()))
        }),
    );

    let result = template.render(&context).await.unwrap();
    assert!(result.contains("Current time: 2024-01-15 10:30"));
}

#[tokio::test]
async fn test_error_handling_missing_variable() {
    let template_str = "Hello {missing_name}!";
    let template = PromptTemplate::from_str(template_str).unwrap();

    let context = TemplateContext::new();
    let result = template.render(&context).await;

    assert!(result.is_err());
    let error = result.unwrap_err();
    assert!(
        error.to_string().contains("missing_name"),
        "Expected MissingVariable error but got: {}",
        error
    );
}

#[tokio::test]
async fn test_template_security_injection_prevention() {
    let malicious_input = "'; DROP TABLE users; --";
    let template_str = "Query: {user_input}";
    let template = PromptTemplate::from_str(template_str).unwrap();

    let mut context = TemplateContext::new();
    context.insert("user_input", Value::String(malicious_input.to_string()));

    let result = template.render(&context).await.unwrap();
    // Should escape dangerous characters
    assert!(result.contains("Query: '; DROP TABLE users; --"));
}

#[tokio::test]
async fn test_template_caching() {
    let template_str = "Hello {name}!";

    let mut engine = TemplateEngine::new();
    engine.enable_caching(true);

    // First compilation should cache the template
    let template1 = engine.compile_template(template_str).await.unwrap();
    let template2 = engine.compile_template(template_str).await.unwrap();

    // Should use cached version (test by checking compilation time or cache stats)
    assert_eq!(template1.template_id(), template2.template_id());
}

#[tokio::test]
async fn test_template_validation() {
    let invalid_template = "Hello {unclosed_brace";
    let result = PromptTemplate::from_str(invalid_template);

    assert!(result.is_err());
    let error = result.unwrap_err();
    assert!(
        error.to_string().contains("unclosed") || error.to_string().contains("syntax"),
        "Expected SyntaxError but got: {}",
        error
    );
}

#[tokio::test]
async fn test_partial_template_application() {
    let template_str = "System: {system}, Task: {task}, User: {user}";
    let template = PromptTemplate::from_str(template_str).unwrap();

    let mut partial_context = TemplateContext::new();
    partial_context.insert("system", Value::String("RustyTorch".to_string()));
    partial_context.insert("task", Value::String("Generate".to_string()));

    let partial_template = template.bind_partial(&partial_context).await.unwrap();

    let mut final_context = TemplateContext::new();
    final_context.insert("user", Value::String("Alice".to_string()));

    let result = partial_template.render(&final_context).await.unwrap();
    assert_eq!(result, "System: RustyTorch, Task: Generate, User: Alice");
}