150 lines
4.8 KiB
Rust
150 lines
4.8 KiB
Rust
use minijinja::{context, Environment};
|
|
use mlua::prelude::*;
|
|
use serde::{Deserialize, Serialize};
|
|
use std::path::Path;
|
|
use std::sync::{Arc, Mutex};
|
|
|
|
struct State<'a> {
|
|
environment: Mutex<Option<Environment<'a>>>,
|
|
}
|
|
|
|
impl State<'_> {
|
|
fn new() -> Self {
|
|
State {
|
|
environment: Mutex::new(None),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Serialize, Deserialize)]
|
|
struct SelectedCode {
|
|
path: String,
|
|
content: Option<String>,
|
|
file_type: String,
|
|
}
|
|
|
|
#[derive(Debug, Serialize, Deserialize)]
|
|
struct SelectedFile {
|
|
path: String,
|
|
content: Option<String>,
|
|
file_type: String,
|
|
}
|
|
|
|
#[derive(Debug, Serialize, Deserialize)]
|
|
struct TemplateContext {
|
|
ask: bool,
|
|
code_lang: String,
|
|
selected_files: Option<Vec<SelectedFile>>,
|
|
selected_code: Option<SelectedCode>,
|
|
recently_viewed_files: Option<Vec<String>>,
|
|
relevant_files: Option<Vec<String>>,
|
|
project_context: Option<String>,
|
|
diagnostics: Option<String>,
|
|
system_info: Option<String>,
|
|
model_name: Option<String>,
|
|
memory: Option<String>,
|
|
todos: Option<String>,
|
|
enable_fastapply: Option<bool>,
|
|
use_react_prompt: Option<bool>,
|
|
}
|
|
|
|
// Given the file name registered after add, the context table in Lua, resulted in a formatted
|
|
// Lua string
|
|
#[allow(clippy::needless_pass_by_value)]
|
|
fn render(state: &State, template: &str, context: TemplateContext) -> LuaResult<String> {
|
|
let environment = state.environment.lock().unwrap();
|
|
match environment.as_ref() {
|
|
Some(environment) => {
|
|
let jinja_template = environment
|
|
.get_template(template)
|
|
.map_err(LuaError::external)
|
|
.unwrap();
|
|
|
|
Ok(jinja_template
|
|
.render(context! {
|
|
ask => context.ask,
|
|
code_lang => context.code_lang,
|
|
selected_files => context.selected_files,
|
|
selected_code => context.selected_code,
|
|
recently_viewed_files => context.recently_viewed_files,
|
|
relevant_files => context.relevant_files,
|
|
project_context => context.project_context,
|
|
diagnostics => context.diagnostics,
|
|
system_info => context.system_info,
|
|
model_name => context.model_name,
|
|
memory => context.memory,
|
|
todos => context.todos,
|
|
enable_fastapply => context.enable_fastapply,
|
|
use_react_prompt => context.use_react_prompt,
|
|
})
|
|
.map_err(LuaError::external)
|
|
.unwrap())
|
|
}
|
|
None => Err(LuaError::RuntimeError(
|
|
"Environment not initialized".to_string(),
|
|
)),
|
|
}
|
|
}
|
|
|
|
fn initialize(state: &State, cache_directory: String, project_directory: String) {
|
|
let mut environment_mutex = state.environment.lock().unwrap();
|
|
let mut env = Environment::new();
|
|
|
|
// Create a custom loader that searches both cache and project directories
|
|
let cache_dir = cache_directory.clone();
|
|
let project_dir = project_directory.clone();
|
|
|
|
env.set_loader(
|
|
move |name: &str| -> Result<Option<String>, minijinja::Error> {
|
|
// First try the cache directory (for built-in templates)
|
|
let cache_path = Path::new(&cache_dir).join(name);
|
|
if cache_path.exists() {
|
|
match std::fs::read_to_string(&cache_path) {
|
|
Ok(content) => return Ok(Some(content)),
|
|
Err(_) => {} // Continue to try project directory
|
|
}
|
|
}
|
|
|
|
// Then try the project directory (for custom includes)
|
|
let project_path = Path::new(&project_dir).join(name);
|
|
if project_path.exists() {
|
|
match std::fs::read_to_string(&project_path) {
|
|
Ok(content) => return Ok(Some(content)),
|
|
Err(_) => {} // File not found or read error
|
|
}
|
|
}
|
|
|
|
// Template not found in either directory
|
|
Ok(None)
|
|
},
|
|
);
|
|
|
|
*environment_mutex = Some(env);
|
|
}
|
|
|
|
#[mlua::lua_module]
|
|
fn avante_templates(lua: &Lua) -> LuaResult<LuaTable> {
|
|
let core = State::new();
|
|
let state = Arc::new(core);
|
|
let state_clone = Arc::clone(&state);
|
|
|
|
let exports = lua.create_table()?;
|
|
exports.set(
|
|
"initialize",
|
|
lua.create_function(
|
|
move |_, (cache_directory, project_directory): (String, String)| {
|
|
initialize(&state, cache_directory, project_directory);
|
|
Ok(())
|
|
},
|
|
)?,
|
|
)?;
|
|
exports.set(
|
|
"render",
|
|
lua.create_function_mut(move |lua, (template, context): (String, LuaValue)| {
|
|
let ctx = lua.from_value(context)?;
|
|
render(&state_clone, template.as_str(), ctx)
|
|
})?,
|
|
)?;
|
|
Ok(exports)
|
|
}
|