pub mod db;
pub mod entities;
pub mod llama;
pub mod mcp;
#[cfg(test)]
mod tests;
use crate::llama::LlamaClient;
use dotenvy::dotenv;
use sea_orm::DatabaseConnection;
use std::env;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tauri::menu::{Menu, MenuItem};
use tauri::tray::{TrayIconBuilder, TrayIconEvent};
use tauri::Manager;
use sea_orm::{ConnectionTrait, Statement, DatabaseBackend};
use std::sync::atomic::AtomicUsize;
use tauri::Emitter;
pub struct AppState {
pub db: DatabaseConnection,
pub llama: Arc<LlamaClient>,
pub mcp_tx: tokio::sync::broadcast::Sender<serde_json::Value>,
pub connection_count: Arc<AtomicUsize>,
pub app_handle: tauri::AppHandle,
}
// --- Path Resolution Helpers ---
pub fn find_build_assets_dir_logic(res_dir: Option<PathBuf>, exe_dir: Option<PathBuf>) -> Option<PathBuf> {
// 1. Check Resource Dir (Release / Installed) - prioritizes nested structure
if let Some(res_dir) = res_dir {
let candidate = res_dir.join("build_assets");
if candidate.exists() {
log::info!("Found build_assets in resource_dir: {:?}", candidate);
return Some(candidate);
}
}
// 2. Check relative to Executable (Debug / Dev) - walks up directory tree
if let Some(exe_dir) = exe_dir {
let mut p = exe_dir;
for _ in 0..10 { // Increased depth to check up to 9 levels up
let candidate = p.join("build_assets");
if candidate.exists() {
log::info!("Found build_assets relative to exe: {:?}", candidate);
return Some(candidate);
}
if !p.pop() { break; }
}
}
None
}
pub fn find_build_assets_dir(app_handle: &tauri::AppHandle) -> Option<PathBuf> {
let res_dir = app_handle.path().resource_dir().ok();
let exe_dir = env::current_exe().ok().and_then(|p| p.parent().map(|p| p.to_path_buf()));
find_build_assets_dir_logic(res_dir, exe_dir)
}
pub fn resolve_resource_path(app_handle: &tauri::AppHandle, relative_path: &str) -> Option<PathBuf> {
find_build_assets_dir(app_handle).map(|dir| dir.join(relative_path))
}
fn diagnose_environment(app_handle: &tauri::AppHandle) {
log::info!("--- [DIAGNOSIS] Environment Check Start ---");
// 1. Build Assets Dir
match find_build_assets_dir(app_handle) {
Some(path) => log::info!("✅ build_assets found at: {:?}", path),
None => log::error!("❌ build_assets NOT FOUND. Sidecar and resources will fail."),
}
// 2. Critical Files Check
let critical_files = ["mcp.json", "config.json", "vector.dll", "llama.dll"];
for file in critical_files {
match resolve_resource_path(app_handle, file) {
Some(p) if p.exists() => log::info!("✅ found {}: {:?}", file, p),
Some(p) => log::error!("❌ missing {}: path resolved to {:?} but file does not exist", file, p),
None => log::error!("❌ could not resolve path for {}", file),
}
}
// 3. Model file check (sample)
if let Some(base) = find_build_assets_dir(app_handle) {
let model_dir = base.join("models");
if model_dir.exists() {
log::info!("✅ models directory exists at: {:?}", model_dir);
} else {
log::warn!("⚠️ models directory NOT found at: {:?}", model_dir);
}
}
log::info!("--- [DIAGNOSIS] Environment Check End ---");
}
#[tauri::command]
fn get_mcp_info(app_handle: tauri::AppHandle) -> Result<serde_json::Value, String> {
let mcp_path = resolve_resource_path(&app_handle, "mcp.json")
.ok_or_else(|| "Failed to resolve mcp.json path".to_string())?;
if !mcp_path.exists() {
return Err(format!("mcp.json not found at {:?}", mcp_path));
}
let content = std::fs::read_to_string(&mcp_path).map_err(|e| e.to_string())?;
let mcp_data: serde_json::Value = serde_json::from_str(&content).map_err(|e| e.to_string())?;
Ok(mcp_data)
}
#[tauri::command]
async fn get_db_stats(state: tauri::State<'_, Arc<AppState>>) -> Result<serde_json::Value, String> {
use crate::entities::items;
use sea_orm::{EntityTrait, PaginatorTrait};
let count = items::Entity::find().count(&state.db).await.map_err(|e| e.to_string())?;
Ok(serde_json::json!({ "itemCount": count }))
}
#[tauri::command]
async fn get_sidecar_status(state: tauri::State<'_, Arc<AppState>>) -> Result<bool, String> {
Ok(state.llama.check_health().await)
}
// --- DB Browser Commands ---
fn validate_table_name(table_name: &str) -> Result<(), String> {
if table_name.chars().all(|c| c.is_alphanumeric() || c == '_' || c == '-') {
Ok(())
} else {
Err("Invalid table name".to_string())
}
}
#[tauri::command]
async fn get_table_list(state: tauri::State<'_, Arc<AppState>>) -> Result<Vec<String>, String> {
let results = state.db.query_all(Statement::from_string(
DatabaseBackend::Sqlite,
"SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%';".to_string(),
)).await.map_err(|e| e.to_string())?;
Ok(results.into_iter().map(|r| r.try_get::<String>("", "name").unwrap_or_default()).collect())
}
#[tauri::command]
async fn get_table_schema(state: tauri::State<'_, Arc<AppState>>, table: String) -> Result<serde_json::Value, String> {
validate_table_name(&table)?;
let query = format!("PRAGMA table_info({});", table);
let results = state.db.query_all(Statement::from_string(
DatabaseBackend::Sqlite,
query,
)).await.map_err(|e| e.to_string())?;
let mut schema = Vec::new();
for row in results {
schema.push(serde_json::json!({
"name": row.try_get::<String>("", "name").unwrap_or_default(),
"type": row.try_get::<String>("", "type").unwrap_or_default(),
"notnull": row.try_get::<i64>("", "notnull").unwrap_or(0),
"pk": row.try_get::<i64>("", "pk").unwrap_or(0),
}));
}
Ok(serde_json::json!(schema))
}
#[tauri::command]
async fn get_table_data(state: tauri::State<'_, Arc<AppState>>, table: String, limit: u64, offset: u64) -> Result<serde_json::Value, String> {
validate_table_name(&table)?;
let query = format!("SELECT * FROM {} LIMIT {} OFFSET {};", table, limit, offset);
let results = state.db.query_all(Statement::from_string(
DatabaseBackend::Sqlite,
query,
)).await.map_err(|e| e.to_string())?;
let mut data = Vec::new();
for row in results {
let mut item = serde_json::Map::new();
// Since we don't know column names dynamically easily with sea-orm query_all without knowing schema,
// we might just return keys from the first result if we had them.
// For now, let's keep it simple or use a specialized query.
item.insert("id".to_string(), serde_json::json!(row.try_get::<i64>("", "id").unwrap_or(0)));
if let Ok(content) = row.try_get::<String>("", "content") {
item.insert("content".to_string(), serde_json::json!(content));
}
data.push(serde_json::Value::Object(item));
}
Ok(serde_json::json!(data))
}
// --- Configuration ---
#[derive(serde::Deserialize, Clone, Debug)]
pub struct AppConfig {
pub database: DatabasePathConfig,
pub model: ModelConfig,
pub llama_server: LlamaServerConfig,
}
#[derive(serde::Deserialize, Clone, Debug)]
pub struct DatabasePathConfig {
pub path: String,
}
#[derive(serde::Deserialize, Clone, Debug)]
pub struct ModelConfig {
pub path: String,
}
#[derive(serde::Deserialize, Clone, Debug)]
pub struct LlamaServerConfig {
pub port: u16,
}
impl Default for AppConfig {
fn default() -> Self {
Self {
database: DatabasePathConfig { path: "data/vector.db".to_string() },
model: ModelConfig { path: "models/embeddinggemma-300m-q4_0.gguf".to_string() },
llama_server: LlamaServerConfig { port: 8080 },
}
}
}
pub fn get_config(app_handle: &tauri::AppHandle) -> AppConfig {
let config_path = resolve_resource_path(app_handle, "config.json");
if let Some(path) = config_path {
if path.exists() {
if let Ok(content) = std::fs::read_to_string(path) {
if let Ok(config) = serde_json::from_str(&content) {
return config;
}
}
}
}
AppConfig::default()
}
pub fn resolve_db_path(app_handle: &tauri::AppHandle, config: &AppConfig) -> String {
let base_dir = app_handle.path().app_data_dir().unwrap_or_else(|_| PathBuf::from("."));
if !base_dir.exists() {
let _ = std::fs::create_dir_all(&base_dir);
}
let db_path = base_dir.join(&config.database.path);
if let Some(parent) = db_path.parent() {
if !parent.exists() {
let _ = std::fs::create_dir_all(parent);
}
}
db_path.to_string_lossy().to_string()
}
pub fn resolve_extension_path(app_handle: &tauri::AppHandle) -> String {
let ext_path = resolve_resource_path(app_handle, "vector.dll");
match ext_path {
Some(p) => p.to_string_lossy().to_string(),
None => "vector".to_string(),
}
}
// --- Sidecar ---
pub fn spawn_llama_server(app_handle: &tauri::AppHandle, config: &AppConfig) {
use tauri_plugin_shell::ShellExt;
let base_dir = find_build_assets_dir(app_handle).unwrap_or_else(|| {
log::error!("Failed to find build_assets for sidecar");
PathBuf::from(".")
});
let model_path = base_dir.join(&config.model.path);
log::info!("Starting llama-server with model: {:?}", model_path);
if !model_path.exists() {
log::error!("❌ Model file does NOT exist: {:?}", model_path);
return;
}
let port_str = config.llama_server.port.to_string();
// Add build_assets to PATH so llama-server can find its DLLs (ggml.dll, etc.)
let old_path = std::env::var("PATH").unwrap_or_default();
let new_path = format!("{};{}", base_dir.display(), old_path).replace("\\\\?\\", "");
match app_handle.shell().sidecar("llama-server") {
Ok(sidecar) => {
let sidecar = sidecar
.args(["--model", model_path.to_str().unwrap_or(""), "--port", &port_str, "--embedding", "--host", "127.0.0.1", "-c", "8192", "-b", "8192", "-ub", "8192", "--parallel", "1"])
.env("PATH", &new_path);
let (mut rx, _child) = sidecar.spawn().expect("Failed to spawn llama-server sidecar");
tauri::async_runtime::spawn(async move {
while let Some(event) = rx.recv().await {
match event {
tauri_plugin_shell::process::CommandEvent::Stdout(line) => {
log::debug!("llama-server stdout: {}", String::from_utf8_lossy(&line));
}
tauri_plugin_shell::process::CommandEvent::Stderr(line) => {
log::debug!("llama-server stderr: {}", String::from_utf8_lossy(&line));
}
_ => {}
}
}
});
}
Err(e) => {
log::error!("Failed to initialize llama-server sidecar: {}", e);
}
}
}
// --- Run ---
pub fn run() {
tauri::Builder::default()
.plugin(tauri_plugin_shell::init())
.plugin(tauri_plugin_log::Builder::new().target(tauri_plugin_log::Target::new(tauri_plugin_log::TargetKind::LogDir {
file_name: Some("TelosDB".to_string()),
})).build())
.setup(|app| {
log::info!("Application starting (Tauri 2)...");
// 2. DIAGNOSIS
diagnose_environment(app.handle());
dotenv().ok();
let app_handle = app.handle().clone();
// 2.5 Update process PATH to include build_assets directory so DLLs can find each other
if let Some(assets_dir) = find_build_assets_dir(&app_handle) {
if assets_dir.exists() {
let old_path = env::var("PATH").unwrap_or_default();
let new_path = format!("{};{}", assets_dir.display(), old_path).replace("\\\\?\\", "");
env::set_var("PATH", &new_path);
log::info!("Updated process PATH with build_assets: {}", new_path);
}
}
// 3. Start Sidecar
let config = get_config(&app_handle);
spawn_llama_server(&app_handle, &config);
// 4. Database and MCP initialization
let mcp_tx = tokio::sync::broadcast::channel(100).0;
let connection_count = Arc::new(AtomicUsize::new(0));
let llama = Arc::new(LlamaClient::new(
env::var("LLAMA_CPP_BASE_URL").unwrap_or_else(|_| "http://localhost:8080".to_string()),
env::var("LLAMA_CPP_EMBEDDING_MODEL").unwrap_or_else(|_| "nomic-embed-text".to_string()),
env::var("LLAMA_CPP_MODEL").unwrap_or_else(|_| "mistral".to_string()),
));
let app_handle_for_db = app.handle().clone();
let mcp_tx_for_db = mcp_tx.clone();
let connection_count_for_db = connection_count.clone();
let llama_for_db = llama.clone();
tauri::async_runtime::block_on(async move {
let db_path = resolve_db_path(&app_handle_for_db, &config);
let ext_path = resolve_extension_path(&app_handle_for_db);
log::info!("DB Path: {}, Ext Path: {}", db_path, ext_path);
let vec_dim = env::var("VEC_DIM").unwrap_or_else(|_| "768".to_string()).parse::<usize>().unwrap_or(768);
match db::init_db(&db_path, &ext_path, vec_dim).await {
Ok(conn) => {
log::info!("✅ Database initialized successfully");
let state = Arc::new(AppState {
db: conn,
llama: llama_for_db,
mcp_tx: mcp_tx_for_db,
connection_count: connection_count_for_db,
app_handle: app_handle_for_db.clone(),
});
app_handle_for_db.manage(state.clone());
let port = env::var("MCP_PORT").unwrap_or_else(|_| "3000".to_string()).parse::<u16>().unwrap_or(3000);
tokio::spawn(async move { mcp::start_mcp_server(state, port).await; });
}
Err(e) => {
log::error!("CRITICAL: Failed to initialize database: {}", e);
log::error!("This is likely due to SQLite extension load failure. Check vector.dll and its dependencies.");
panic!("Failed to init db: {}", e);
}
}
});
// 5. Tray setup
let quit_i = MenuItem::with_id(app, "quit", "Quit", true, None::<&str>)?;
let show_i = MenuItem::with_id(app, "show", "Show", true, None::<&str>)?;
let menu = Menu::with_items(app, &[&show_i, &quit_i])?;
let mut tray_builder = TrayIconBuilder::new();
if let Some(icon) = app.default_window_icon() {
tray_builder = tray_builder.icon(icon.clone());
} else {
log::warn!("Default window icon not found, tray icon might be empty.");
}
let _tray = tray_builder
.menu(&menu)
.on_menu_event(|app, event| match event.id.as_ref() {
"quit" => app.exit(0),
"show" => if let Some(w) = app.get_webview_window("main") { let _ = w.show(); let _ = w.set_focus(); }
_ => {}
})
.on_tray_icon_event(|tray, event| if let TrayIconEvent::DoubleClick { .. } = event {
let app = tray.app_handle();
if let Some(w) = app.get_webview_window("main") { let _ = w.show(); let _ = w.set_focus(); }
})
.build(app)?;
log::info!("Tauri setup completed");
Ok(())
})
.invoke_handler(tauri::generate_handler![
get_mcp_info,
get_db_stats,
get_sidecar_status,
get_table_list,
get_table_schema,
get_table_data
])
.run(tauri::generate_context!())
.expect("error while running tauri application");
}
