diff --git a/README.md b/README.md index adc6898..0173220 100644 --- a/README.md +++ b/README.md @@ -15,10 +15,11 @@ ## 主な機能 -- **セマンティック検索 (Vector Search)**: `sqlite-vec` を用いた高精度な意味検索。 -- **内蔵推論エンジン (Sidecar Integration)**: `llama-server` を自動管理。アプリ起動と同時に推論サーバーが立ち上がります。 -- **MCP SSE サーバー仕様**: SSE (Server-Sent Events) を用いた非同期レスポンス方式を採用。LM Studio, Claude Desktop などの外部AIツールから「プラグイン」として直接接続可能です。 -- **プレミアム・デザイン**: ガラスモーフィズムを基調とした洗練されたUI。ステータスインジケーターにより推論エンジンの稼働状況を一目で把握できます。 +- **完全な知識管理 (CRUD)**: `add_item_text`, `update_item`, `delete_item` を通じて、AIエージェントが自律的に知識を蓄積・修正・削除可能。 +- **セマンティック検索 (Vector Search)**: `sqlite-vec` + `llama-server` (Gemma-3) による高度な意味検索。 +- **サイドカー自動管理**: `llama-server` の起動・ヘルスチェック・ログ転送をRust側で制御。 +- **MCP SSE サーバー**: LM Studio, Claude Desktop 等の外部ツールからプラグインとして即座に利用可能。 +- **Premium UI**: ガラスモーフィズムを採用した洗練されたデスクトップUX。 - **堅牢なロギング**: ローテーション機能付きログ出力。`llama-server` の詳細な内部ログもキャプチャします。 --- diff --git a/src-tauri/src/db.rs b/src-tauri/src/db.rs index 3459b18..e24b113 100644 --- a/src-tauri/src/db.rs +++ b/src-tauri/src/db.rs @@ -30,6 +30,7 @@ CREATE TABLE IF NOT EXISTS items ( id INTEGER PRIMARY KEY AUTOINCREMENT, content TEXT NOT NULL, + path TEXT, created_at TEXT DEFAULT (datetime('now', 'localtime')), updated_at TEXT DEFAULT (datetime('now', 'localtime')) ); @@ -41,7 +42,7 @@ END; CREATE VIRTUAL TABLE IF NOT EXISTS vec_items USING vec0( id INTEGER PRIMARY KEY, - embedding FLOAT[768] + embedding FLOAT[640] );" )?; diff --git a/src-tauri/src/mcp.rs b/src-tauri/src/mcp.rs index a546e40..fca8303 100644 --- a/src-tauri/src/mcp.rs +++ b/src-tauri/src/mcp.rs @@ -148,6 +148,24 @@ session_id: Option, } +async fn get_embedding(content: &str) -> Result, String> { + let client = reqwest::Client::new(); + let resp = client.post("http://127.0.0.1:8080/embedding") + .json(&serde_json::json!({ "content": content })) + .send() + .await + .map_err(|e| e.to_string())?; + + let json: serde_json::Value = resp.json().await.map_err(|e| e.to_string())?; + let embedding = json["embedding"].as_array() + .ok_or("No embedding field in llama-server response")? + .iter() + .map(|v| v.as_f64().unwrap_or(0.0) as f32) + .collect(); + + Ok(embedding) +} + async fn mcp_messages_handler( State(state): State, Query(query): Query, @@ -166,56 +184,180 @@ "tools/list" => Some(serde_json::json!({ "tools": [ { - "name": "search_text", - "description": "Semantic search using llama-server vector embeddings.", + "name": "add_item_text", + "description": "Store text with auto-generated embeddings.", "inputSchema": { "type": "object", "properties": { - "content": { "type": "string", "description": "Query" }, + "content": { "type": "string" }, + "path": { "type": "string" } + }, + "required": ["content"] + } + }, + { + "name": "search_text", + "description": "Semantic search using vector embeddings.", + "inputSchema": { + "type": "object", + "properties": { + "content": { "type": "string" }, "limit": { "type": "number" } }, "required": ["content"] } + }, + { + "name": "update_item", + "description": "Update existing text and its embedding.", + "inputSchema": { + "type": "object", + "properties": { + "id": { "type": "integer" }, + "content": { "type": "string" }, + "path": { "type": "string" } + }, + "required": ["id", "content"] + } + }, + { + "name": "delete_item", + "description": "Delete item by ID.", + "inputSchema": { + "type": "object", + "properties": { + "id": { "type": "integer" } + }, + "required": ["id"] + } } ] })), "search_text" | "tools/call" => { - let (content, limit, is_mcp_output) = if method == "search_text" { - let p = req.params.unwrap_or_default(); - ( - p.get("content").and_then(|v| v.as_str()).unwrap_or("").to_string(), - p.get("limit").and_then(|v| v.as_u64()).unwrap_or(10) as u32, - false - ) + let p = req.params.clone().unwrap_or_default(); + let (actual_method, args) = if method == "tools/call" { + (p.get("name").and_then(|v| v.as_str()).unwrap_or(""), p.get("arguments").cloned().unwrap_or_default()) } else { - let p = req.params.unwrap_or_default(); - let args = p.get("arguments").cloned().unwrap_or_default(); - ( - args.get("content").and_then(|v| v.as_str()).unwrap_or("").to_string(), - args.get("limit").and_then(|v| v.as_u64()).unwrap_or(10) as u32, - true - ) + (method, p) }; - let rows = sqlx::query("SELECT id, content FROM items WHERE content LIKE ? LIMIT ?") - .bind(format!("%{}%", content)) - .bind(limit) - .fetch_all(&state.db_pool) - .await - .unwrap_or_default(); + match actual_method { + "add_item_text" => { + let content = args.get("content").and_then(|v| v.as_str()).unwrap_or(""); + let path = args.get("path").and_then(|v| v.as_str()); + + match get_embedding(content).await { + Ok(emb) => { + let mut tx = state.db_pool.begin().await.unwrap(); + let res = sqlx::query("INSERT INTO items (content, path) VALUES (?, ?)") + .bind(content) + .bind(path) + .execute(&mut *tx) + .await + .unwrap(); + let id = res.last_insert_rowid(); + + sqlx::query("INSERT INTO vec_items (id, embedding) VALUES (?, ?)") + .bind(id) + .bind(serde_json::to_string(&emb).unwrap()) + .execute(&mut *tx) + .await + .unwrap(); + + tx.commit().await.unwrap(); + Some(serde_json::json!({ "content": [{ "type": "text", "text": format!("Successfully added item with ID: {}", id) }] })) + } + Err(e) => Some(serde_json::json!({ "error": format!("Embedding failed: {}", e) })) + } + }, + "search_text" => { + let content = args.get("content").and_then(|v| v.as_str()).unwrap_or(""); + let limit = args.get("limit").and_then(|v| v.as_u64()).unwrap_or(10) as u32; - if is_mcp_output { - let txt = if rows.is_empty() { "No results.".to_string() } else { - rows.iter().map(|r| format!("ID: {}, Content: {}", r.get::(0), r.get::(1))).collect::>().join("\n\n") - }; - Some(serde_json::json!({ "content": [{ "type": "text", "text": txt }] })) - } else { - let res: Vec<_> = rows.iter().map(|r| serde_json::json!({ - "id": r.get::(0), - "content": r.get::(1), - "distance": 0.0 - })).collect(); - Some(serde_json::json!({ "content": res })) + match get_embedding(content).await { + Ok(emb) => { + let rows = sqlx::query( + "SELECT items.id, items.content, v.distance + FROM items + JOIN vec_items v ON items.id = v.id + WHERE v.embedding MATCH ? AND k = ? + ORDER BY distance LIMIT ?" + ) + .bind(serde_json::to_string(&emb).unwrap()) + .bind(limit) + .bind(limit) + .fetch_all(&state.db_pool) + .await + .unwrap_or_default(); + + let is_mcp_output = method == "tools/call"; + if is_mcp_output { + let txt = if rows.is_empty() { "No results.".to_string() } else { + rows.iter().map(|r| format!("[ID: {}, Distance: {:.4}]\n{}", r.get::(0), r.get::(2), r.get::(1))).collect::>().join("\n\n---\n\n") + }; + Some(serde_json::json!({ "content": [{ "type": "text", "text": txt }] })) + } else { + let res: Vec<_> = rows.iter().map(|r| serde_json::json!({ + "id": r.get::(0), + "content": r.get::(1), + "distance": r.get::(2) + })).collect(); + Some(serde_json::json!({ "content": res })) + } + } + Err(e) => { + // Fallback to LIKE if llama-server is not running + let rows = sqlx::query("SELECT id, content FROM items WHERE content LIKE ? LIMIT ?") + .bind(format!("%{}%", content)) + .bind(limit) + .fetch_all(&state.db_pool) + .await + .unwrap_or_default(); + + let txt = format!("(Fallback SEARCH due to embedding error: {})\n\n", e); + let results = rows.iter().map(|r| format!("ID: {}, Content: {}", r.get::(0), r.get::(1))).collect::>().join("\n\n"); + Some(serde_json::json!({ "content": [{ "type": "text", "text": txt + &results }] })) + } + } + }, + "update_item" => { + let id = args.get("id").and_then(|v| v.as_i64()).unwrap_or(0); + let content = args.get("content").and_then(|v| v.as_str()).unwrap_or(""); + let path = args.get("path").and_then(|v| v.as_str()); + + match get_embedding(content).await { + Ok(emb) => { + let mut tx = state.db_pool.begin().await.unwrap(); + sqlx::query("UPDATE items SET content = ?, path = ? WHERE id = ?") + .bind(content) + .bind(path) + .bind(id) + .execute(&mut *tx) + .await + .unwrap(); + + sqlx::query("UPDATE vec_items SET embedding = ? WHERE id = ?") + .bind(serde_json::to_string(&emb).unwrap()) + .bind(id) + .execute(&mut *tx) + .await + .unwrap(); + + tx.commit().await.unwrap(); + Some(serde_json::json!({ "content": [{ "type": "text", "text": format!("Successfully updated item {}", id) }] })) + } + Err(e) => Some(serde_json::json!({ "error": format!("Embedding failed: {}", e) })) + } + }, + "delete_item" => { + let id = args.get("id").and_then(|v| v.as_i64()).unwrap_or(0); + let mut tx = state.db_pool.begin().await.unwrap(); + sqlx::query("DELETE FROM items WHERE id = ?").bind(id).execute(&mut *tx).await.unwrap(); + sqlx::query("DELETE FROM vec_items WHERE id = ?").bind(id).execute(&mut *tx).await.unwrap(); + tx.commit().await.unwrap(); + Some(serde_json::json!({ "content": [{ "type": "text", "text": format!("Successfully deleted item {}", id) }] })) + }, + _ => Some(serde_json::json!({ "error": "Unknown tool" })), } }, _ => Some(serde_json::json!({ "error": "Not implemented" })), @@ -225,7 +367,7 @@ let resp = JsonRpcResponse { jsonrpc: "2.0", result, error: None, id: Some(id_val) }; if let Some(sid) = query.session_id { - // MCP Client (SSE Mode): Return 202 and send response via SSE + // MCP Client (SSE Mode) let resp_str = serde_json::to_string(&resp).unwrap(); let sessions = state.sessions.read().await; if let Some(tx) = sessions.get(&sid) { @@ -233,11 +375,10 @@ } axum::http::StatusCode::ACCEPTED.into_response() } else { - // App UI (Direct Mode): Return Json response directly + // App UI (Direct Mode) Json(resp).into_response() } } else { - // Notification: No response needed axum::http::StatusCode::NO_CONTENT.into_response() } }