Newer
Older
TelosDB / src-tauri / src / lib.rs
pub mod db;
pub mod entities;
pub mod llama;
pub mod mcp;

use crate::llama::LlamaClient;
use dotenvy::dotenv;
use sea_orm::DatabaseConnection;
use std::env;
use std::path::PathBuf;
use std::sync::Arc;
use tauri::menu::{Menu, MenuItem};
use tauri::tray::{TrayIconBuilder, TrayIconEvent};
use tauri::Manager;

pub struct AppState {
    pub db: DatabaseConnection,
    pub llama: Arc<LlamaClient>,
}

#[tauri::command]
fn greet(name: &str) -> String {
    format!("Hello, {}! You've been greeted from Rust!", name)
}

#[cfg_attr(mobile, tauri::mobile_entry_point)]
pub fn run() {
    tauri::Builder::default()
        .invoke_handler(tauri::generate_handler![greet])
        .on_window_event(|window, event| {
            if let tauri::WindowEvent::CloseRequested { api, .. } = event {
                api.prevent_close();
                let _ = window.hide();
            }
        })
        .setup(|app| {
            dotenv().ok();

            let app_handle = app.handle().clone();

            // Initialize Shell Plugin
            app.handle().plugin(tauri_plugin_shell::init())?;

            // Spawn llama-server sidecar
            let _model_path = env::var("LLAMA_CPP_MODEL_PATH").ok();
            // サイドカーの起動 (std::process::Command を使用してDLL問題を確実に回避)
            let mut sidecar_path = std::env::current_dir().unwrap_or_default();
            println!("Current working directory: {:?}", sidecar_path);

            // src-tauri フォルダの中にいる場合は1つ上がる
            if sidecar_path.ends_with("src-tauri") {
                sidecar_path.pop();
            }

            let project_root = sidecar_path.clone();
            sidecar_path.push("src-tauri");
            sidecar_path.push("bin");
            let sidecar_exe = sidecar_path.join("llama-server-x86_64-pc-windows-msvc.exe");

            println!("Calculated sidecar path: {:?}", sidecar_exe);

            // モデルパスもルートからの絶対パスに変換
            let model_rel_path = env::var("LLAMA_CPP_MODEL_PATH").unwrap_or_default();
            let model_abs_path = project_root.join(&model_rel_path);
            let args = vec![
                "--model".to_string(),
                model_abs_path.to_string_lossy().to_string(),
                "--port".to_string(),
                "8080".to_string(),
                "--embedding".to_string(),
                "--host".to_string(),
                "127.0.0.1".to_string(),
            ];

            let mut cmd = std::process::Command::new(&sidecar_exe);
            cmd.args(&args);
            cmd.current_dir(&sidecar_path); // DLLのあるディレクトリをカレントにする

            // PATHにも追加
            let mut current_path = env::var("PATH").unwrap_or_default();
            current_path = format!("{};{}", sidecar_path.display(), current_path);
            cmd.env("PATH", current_path);

            match cmd.spawn() {
                Ok(child) => {
                    println!("llama-server started with PID: {}", child.id());
                    // アプリ終了時にプロセスを殺すためのハンドルを保持(簡易実装)
                    let pid = child.id();
                    std::thread::spawn(move || {
                        // 子プロセスの終了を待機
                        let _ = child.wait_with_output();
                        println!("llama-server (PID {}) exited", pid);
                    });
                }
                Err(e) => {
                    eprintln!("Failed to spawn llama-server: {}", e);
                }
            }

            tauri::async_runtime::block_on(async move {
                let db_path = "vector.db";

                // 拡張機能 (vec0.dll) のパスを動的に解決する
                let exe_dir = env::current_exe()
                    .map(|p| p.parent().unwrap().to_path_buf())
                    .unwrap_or_else(|_| env::current_dir().unwrap());

                let candidates = [
                    exe_dir.join("vec0.dll"), // 実行ファイルと同階層
                    exe_dir.join("../node_modules/sqlite-vec-windows-x64/vec0.dll"), // Tauri dev (target/debug/..)
                    exe_dir.join("../../node_modules/sqlite-vec-windows-x64/vec0.dll"), // target/debug/../..
                    exe_dir.join("../../../node_modules/sqlite-vec-windows-x64/vec0.dll"), // root
                    PathBuf::from("../node_modules/sqlite-vec-windows-x64/vec0.dll"), // 従来の相対パス
                ];

                let mut extension_path = candidates[0].to_str().unwrap().to_string();
                for cand in &candidates {
                    if cand.exists() {
                        extension_path = cand.to_str().unwrap().to_string();
                        println!("Found sqlite-vec at: {}", extension_path);
                        break;
                    }
                }

                let conn = db::init_db(db_path, &extension_path)
                    .await
                    .expect("Failed to init database");
                let llama_base_url = env::var("LLAMA_CPP_BASE_URL")
                    .unwrap_or_else(|_| "http://localhost:8080".to_string());
                let embedding_model = env::var("LLAMA_CPP_EMBEDDING_MODEL")
                    .unwrap_or_else(|_| "nomic-embed-text".to_string());
                let completion_model =
                    env::var("LLAMA_CPP_MODEL").unwrap_or_else(|_| "mistral".to_string());

                let llama = LlamaClient::new(llama_base_url, embedding_model, completion_model);

                let state = Arc::new(AppState {
                    db: conn,
                    llama: Arc::new(llama),
                });

                app_handle.manage(state.clone());

                // Start MCP Server (SSE)
                let mcp_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, mcp_port).await;
                });
            });

            // Logging implementation
            let mut log_builder = tauri_plugin_log::Builder::default()
                .targets([
                    tauri_plugin_log::Target::new(tauri_plugin_log::TargetKind::Stdout),
                    tauri_plugin_log::Target::new(tauri_plugin_log::TargetKind::LogDir { file_name: None }),
                ])
                .max_file_size(10 * 1024 * 1024)
                .level(log::LevelFilter::Info);

            if cfg!(debug_assertions) {
                // デバッグ時はプロジェクトルートの logs フォルダにも出力
                log_builder = log_builder.target(tauri_plugin_log::Target::new(
                    tauri_plugin_log::TargetKind::Folder {
                        path: std::path::PathBuf::from("logs"),
                        file_name: None,
                    },
                ));
            }

            app.handle().plugin(log_builder.build())?;

            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 _tray = TrayIconBuilder::new()
                .icon(app.default_window_icon().unwrap().clone())
                .menu(&menu)
                .on_menu_event(|app, event| match event.id.as_ref() {
                    "quit" => {
                        app.exit(0);
                    }
                    "show" => {
                        if let Some(window) = app.get_webview_window("main") {
                            let _ = window.show();
                            let _ = window.set_focus();
                        }
                    }
                    _ => {}
                })
                .on_tray_icon_event(|tray, event| {
                    if let TrayIconEvent::DoubleClick { .. } = event {
                        let app = tray.app_handle();
                        if let Some(window) = app.get_webview_window("main") {
                            let _ = window.show();
                            let _ = window.set_focus();
                        }
                    }
                })
                .build(app)?;

            Ok(())
        })
        .run(tauri::generate_context!())
        .expect("error while running tauri application");
}