diff --git a/.gitignore b/.gitignore index 82e032d..89736ed 100644 --- a/.gitignore +++ b/.gitignore @@ -14,3 +14,7 @@ # Added by cargo /target + +# Python +__pycache__/ +*.pyc diff --git a/README.md b/README.md index 88eb929..bf51cc8 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,7 @@ A local knowledge base CLI for [Claude Code](https://docs.anthropic.com/en/docs/ - Self-update from GitHub Releases - Git worktree support — all worktrees share the main worktree's DB, so knowledge is available across worktrees - MCP (Model Context Protocol) server — Claude Code / Claude Desktop can autonomously search, add, and manage knowledge +- Semantic intelligence with Laya (Apple Silicon / MLX) — on-demand background daemon providing semantic duplicate warning, Japanese kanji compound auto-keyword selection, and search result semantic reranking with graceful fallback - Installs Claude Code slash commands for seamless integration ## Installation @@ -576,6 +577,13 @@ gitattributes_generated = true # instead; `lk init` then removes the import line rather than adding it. # `lk init --no-import` sets this for you. claude_md_import = true + +# Laya MLX integration (Apple Silicon) +# Note: Laya is machine-specific and must be enabled globally in ~/.config/lk/config.toml (or LK_LAYA=1). +# A project can opt out by setting enabled = false (overridable via LK_LAYA=1), or configure the duplicate threshold. +# [laya] +# enabled = false # Opt out for this project +# duplicate_threshold = 0.85 # Semantic duplicate threshold (default: 0.85) ``` ### Global config (user scope) @@ -597,6 +605,14 @@ secret_detection = true # Set to false when agents read the instructions from the lk-knowledge MCP server. # `lk init --global --no-import` sets this for you. claude_md_import = true + +# Laya MLX integration (Apple Silicon, opt-in) +# [laya] +# enabled = false # Default: false (set to true to enable) +# idle_timeout = 600 # Seconds before daemon auto-shuts down (default: 600) +# model = "aac6fef/laya-multilingual-mlx" # Hugging Face model or local path +# socket_path = "~/.cache/lk/laya.sock" # Optional custom socket path +# duplicate_threshold = 0.85 # Semantic duplicate threshold (default: 0.85) ``` ### Environment variable overrides @@ -607,6 +623,8 @@ Environment variables take precedence over config file values: |----------|--------| | `LK_NO_AUTO_SYNC=1` | Disable auto-sync | | `LK_COMMAND_LOG=1` | Enable command logging | +| `LK_LAYA=1` | Enable Laya semantic features | +| `LK_NO_LAYA=1` | Disable Laya semantic features | ### Auto-sync @@ -631,6 +649,22 @@ lk command-log # Show last 20 entries lk command-log -n 50 # Show last 50 entries ``` +### Semantic intelligence (Laya) + +On Apple Silicon (macOS arm64), `lk` can leverage [laya-mlx](https://github.com/mizorewww/laya-mlx) — a low-latency typed decision model runtime executing on Apple Neural Engine / Metal — to enhance duplicate detection, keyword extraction, and search ranking. + +This feature is **disabled by default (opt-in)**. To enable it, set `[laya] enabled = true` in `~/.config/lk/config.toml`, or set `LK_LAYA=1` in your environment (individual projects can opt out via `.knowledge/config.toml`, which `LK_LAYA=1` overrides): + +- **Semantic Duplicate Warning**: Identifies paraphrased entries that share identical meaning during `lk add`, flagging them in `possibly_related` with match reason `semantic-duplicate` without blocking the add. +- **Smart Auto-Keywords**: When adding an entry without explicit keywords, `lk` extracts candidates including Japanese kanji compound words (e.g. `排他制御`, `有効期限`) and uses Laya to filter out uninformative terms, retaining high-quality domain keywords. +- **Search Reranking**: Re-scores the top BM25 search results with query-content semantic relevance while strictly preserving current-project (`mine`) priority. + +#### Daemon Lifecycle + +- **On-Demand & Singleton**: A single Python daemon process is spawned on the first command requiring semantic decisions, communicating via a Unix Domain Socket (`~/.cache/lk/laya.sock`). All CLI and MCP sessions share this single daemon without redundant memory usage or model loading latency. +- **Auto-Shutdown**: If no requests are received for 10 minutes (configurable via `idle_timeout`), the daemon shuts down automatically and frees all unified memory. +- **Graceful Fallback**: If Python/MLX or dependencies are unavailable, or on non-macOS platforms, `lk` falls back transparently to trigram FTS5 and frequency-based keyword heuristics without errors. + ## Supported Platforms | Platform | Architecture | Note | diff --git a/scripts/laya_server.py b/scripts/laya_server.py new file mode 100755 index 0000000..11b12b3 --- /dev/null +++ b/scripts/laya_server.py @@ -0,0 +1,360 @@ +#!/usr/bin/env python3 +"""Laya MLX Unix Domain Socket server for local-knowledge-cli (lk). + +Provides fast, low-latency typed decisions (duplicate detection, +keyword filtering, and search reranking) over a local Unix domain socket. +Includes an idle shutdown timer that automatically terminates the daemon +after a period of inactivity to release memory. +""" + +import argparse +import asyncio +import json +import os +import signal +import sys +import time +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from typing import Any, Dict, List, Optional + +try: + import laya_mlx as laya +except ImportError: + print("Error: laya-mlx is required. Run with `uv run --with laya-mlx python3 ...`", file=sys.stderr) + sys.exit(1) + +DEFAULT_MODEL = "aac6fef/laya-multilingual-mlx" +DEFAULT_IDLE_TIMEOUT = 600 # 10 minutes + + +class LayaServer: + def __init__( + self, + socket_path: Path, + model_name: str = DEFAULT_MODEL, + idle_timeout: int = DEFAULT_IDLE_TIMEOUT, + ): + self.socket_path = socket_path + self.model_name = model_name + self.idle_timeout = idle_timeout + self.last_active_time = time.time() + self.agent: Optional[Any] = None + self.server: Optional[asyncio.Server] = None + self._running = True + self.executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="laya-worker") + + def load_model(self) -> None: + print(f"Loading Laya model: {self.model_name}...", file=sys.stderr) + start = time.time() + self.agent = laya.load(self.model_name) + elapsed = time.time() - start + print(f"Model loaded in {elapsed:.2f}s", file=sys.stderr) + + def handle_ping(self, _params: Dict[str, Any]) -> Dict[str, Any]: + return { + "status": "ok", + "model": self.model_name, + "idle_timeout": self.idle_timeout, + } + + def handle_duplicate(self, params: Dict[str, Any]) -> Dict[str, Any]: + assert self.agent is not None + entry_a = params.get("entry_a", {}) + entry_b = params.get("entry_b", {}) + + state = { + "entry_a_title": entry_a.get("title", ""), + "entry_a_content": entry_a.get("content", ""), + "entry_b_title": entry_b.get("title", ""), + "entry_b_content": entry_b.get("content", ""), + } + questions = { + "is_duplicate": { + "type": "noul", + "instructions": ( + "Do Entry A and Entry B describe the exact same technical specification, " + "policy, bug solution, or architectural decision?" + ), + } + } + res = self.agent.predict(state, questions) + prob = float(res["answers"]["is_duplicate"]["noul"]) + threshold = float(params.get("threshold", 0.85)) + return { + "is_duplicate": prob >= threshold, + "probability": prob, + } + + def handle_filter_keywords(self, params: Dict[str, Any]) -> Dict[str, Any]: + assert self.agent is not None + title = params.get("title", "") + content = params.get("content", "") + candidates: List[str] = params.get("candidates", []) + + if not candidates: + return {"ranked_keywords": []} + + state = { + "title": title, + "content": content, + } + questions = { + f"kw_{i}": { + "type": "noul", + "instructions": ( + f"Is `{kw}` a core, essential technical keyword representing the main topic " + "in `title` and `content`?" + ), + } + for i, kw in enumerate(candidates) + } + + res = self.agent.predict(state, questions) + ranked: List[Dict[str, Any]] = [] + for i, kw in enumerate(candidates): + prob = float(res["answers"][f"kw_{i}"]["noul"]) + ranked.append({"keyword": kw, "score": prob}) + + ranked.sort(key=lambda x: x["score"], reverse=True) + return {"ranked_keywords": ranked} + + def handle_rerank(self, params: Dict[str, Any]) -> Dict[str, Any]: + assert self.agent is not None + query = params.get("query", "") + candidates: List[Dict[str, Any]] = params.get("candidates", []) + + if not candidates: + return {"scores": []} + + # Evaluate each candidate against the query + scores = [] + for cand in candidates: + cand_id = cand.get("id") + cand_title = cand.get("title", "") + cand_content = cand.get("content", "") + + state = { + "query": query, + "document_title": cand_title, + "document_content": cand_content, + } + questions = { + "relevance": { + "type": "noul", + "instructions": ( + "Does the document described in `document_title` and `document_content` " + "contain relevant information answering or addressing the user search `query`?" + ), + } + } + res = self.agent.predict(state, questions) + score = float(res["answers"]["relevance"]["noul"]) + scores.append({"id": cand_id, "score": score}) + + # Sort descending by score + scores.sort(key=lambda x: x["score"], reverse=True) + return {"scores": scores} + + def handle_categorize(self, params: Dict[str, Any]) -> Dict[str, Any]: + assert self.agent is not None + title = params.get("title", "") + content = params.get("content", "") + + state = { + "title": title, + "content": content, + } + questions = { + "category": { + "type": "choice", + "instructions": "Which category best fits the engineering note described in `title` and `content`?", + "criteria": { + "bugs": "bug fixes, investigation of issues, errors, troubleshooting, regressions", + "architecture": "system architecture, high-level design, database schema, tech stack decisions", + "decisions": "ADR, architectural decision records, trade-off evaluations, rejected alternatives", + "conventions": "coding guidelines, naming rules, workflow rules, formatting", + "features": "new feature design, specifications, requirements, user stories", + "context": "conversation logs, session handoff context, temporary scratchpad notes", + }, + } + } + res = self.agent.predict(state, questions) + ans = res["answers"]["category"] + return { + "category": ans["choice"], + "confidence": float(ans.get("confidence", 0.0)), + "probabilities": {k: float(v) for k, v in ans.get("probabilities", {}).items()}, + } + + def dispatch(self, req: Dict[str, Any]) -> Dict[str, Any]: + req_id = req.get("id") + task = req.get("task", "") + params = req.get("params", {}) + + handlers = { + "ping": self.handle_ping, + "duplicate": self.handle_duplicate, + "filter_keywords": self.handle_filter_keywords, + "rerank": self.handle_rerank, + "categorize": self.handle_categorize, + } + + handler = handlers.get(task) + if not handler: + return {"id": req_id, "error": f"Unknown task: {task}"} + + try: + result = handler(params) + return {"id": req_id, "result": result} + except Exception as e: + return {"id": req_id, "error": str(e)} + + async def handle_client( + self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + while self._running: + line = await reader.readline() + if not line: + break + self.last_active_time = time.time() + text = line.decode("utf-8").strip() + if not text: + continue + + try: + req = json.loads(text) + except json.JSONDecodeError as e: + resp = {"id": None, "error": f"Invalid JSON: {e}"} + writer.write((json.dumps(resp) + "\n").encode("utf-8")) + await writer.drain() + continue + + if req.get("task") == "shutdown": + resp = {"id": req.get("id"), "result": {"status": "shutting_down"}} + writer.write((json.dumps(resp) + "\n").encode("utf-8")) + await writer.drain() + self._running = False + break + + # Handle ping immediately on event loop so liveness checks never wait behind inference + if req.get("task") == "ping": + resp = {"id": req.get("id"), "result": self.handle_ping(req.get("params", {}))} + else: + loop = asyncio.get_running_loop() + resp = await loop.run_in_executor(self.executor, self.dispatch, req) + + self.last_active_time = time.time() + writer.write((json.dumps(resp) + "\n").encode("utf-8")) + await writer.drain() + + writer.close() + await writer.wait_closed() + + async def idle_checker(self) -> None: + while self._running: + await asyncio.sleep(15) + idle_seconds = time.time() - self.last_active_time + if idle_seconds >= self.idle_timeout: + print( + f"Idle timeout reached ({idle_seconds:.0f}s >= {self.idle_timeout}s). Shutting down...", + file=sys.stderr, + ) + self._running = False + if self.server: + self.server.close() + break + + async def run(self) -> None: + loop = asyncio.get_running_loop() + await loop.run_in_executor(self.executor, self.load_model) + + # Ensure parent directory exists + self.socket_path.parent.mkdir(parents=True, exist_ok=True) + if self.socket_path.exists(): + try: + self.socket_path.unlink() + except OSError: + pass + + pid_file = self.socket_path.with_suffix(".pid") + pid_file.write_text(str(os.getpid())) + + self.server = await asyncio.start_unix_server( + self.handle_client, path=str(self.socket_path) + ) + print(f"Laya server listening on {self.socket_path} (pid: {os.getpid()})", file=sys.stderr) + + idle_task = asyncio.create_task(self.idle_checker()) + + try: + async with self.server: + while self._running: + await asyncio.sleep(1) + finally: + idle_task.cancel() + await asyncio.to_thread( + self.executor.shutdown, + wait=True, + cancel_futures=True, + ) + if self.socket_path.exists(): + try: + self.socket_path.unlink() + except OSError: + pass + if pid_file.exists(): + try: + pid_file.unlink() + except OSError: + pass + print("Laya server shutdown complete.", file=sys.stderr) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Laya MLX Unix Domain Socket server for lk") + parser.add_argument( + "--socket-path", + type=Path, + default=Path.home() / ".cache" / "lk" / "laya.sock", + help="Path to the Unix domain socket", + ) + parser.add_argument( + "--model", + type=str, + default=DEFAULT_MODEL, + help="Laya model name on Hugging Face or local path", + ) + parser.add_argument( + "--idle-timeout", + type=int, + default=DEFAULT_IDLE_TIMEOUT, + help="Seconds of inactivity before auto-shutdown", + ) + args = parser.parse_args() + + server = LayaServer( + socket_path=args.socket_path, + model_name=args.model, + idle_timeout=args.idle_timeout, + ) + + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + def handle_sigterm() -> None: + server._running = False + if server.server: + server.server.close() + + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, handle_sigterm) + + try: + loop.run_until_complete(server.run()) + finally: + loop.close() + + +if __name__ == "__main__": + main() diff --git a/src/cmd/add.rs b/src/cmd/add.rs index bcd9db7..8e789d7 100644 --- a/src/cmd/add.rs +++ b/src/cmd/add.rs @@ -1,6 +1,6 @@ use crate::db; use crate::keywords; -use crate::similarity::Tier; +use crate::similarity::{self, Tier}; use crate::util::open_db_with_migrate; /// User-scope ids collide with project ids, so user entries are referenced by @@ -76,7 +76,32 @@ pub fn cmd_add( super::Scope::Project => open_db_with_migrate()?, super::Scope::User => crate::util::open_or_create_user_db()?, }; - let category = category.unwrap_or(""); + + let laya_config = match scope { + super::Scope::Project => { + crate::config::Config::load(&crate::util::get_knowledge_dir()).laya + } + super::Scope::User => crate::config::GlobalConfig::load().laya, + }; + let mut laya_client = crate::laya::LayaClient::connect(&laya_config); + + // Auto-detect category with Laya if not specified + let auto_category; + let category = match category { + Some(c) if !c.is_empty() => c, + _ => { + if let Some(ref mut client) = laya_client { + auto_category = client + .categorize(title, content.unwrap_or("")) + .map(|r| r.category) + .unwrap_or_default(); + &auto_category + } else { + "" + } + } + }; + // Apply category template if content is not provided or empty let template_content; let content = match content { @@ -136,7 +161,7 @@ pub fn cmd_add( Vec::new() }; if kws.is_empty() { - kws = keywords::extract_keywords(title, content); + kws = keywords::extract_keywords_auto(title, content, laya_client.as_mut()); } kws.sort_by_key(|a| a.to_lowercase()); let mut seen = std::collections::HashSet::new(); @@ -148,8 +173,8 @@ pub fn cmd_add( let result = (|| -> Result<(i64, Vec), Box> { // Duplicate check (skip with --force). Only a Block-tier hit refuses the - // add — a near-identical title. Weaker hits are advisory: they are - // reported *after* the entry is committed, because refusing on them is + // add — a near-identical title or semantic duplicate. Weaker hits are advisory: + // they are reported *after* the entry is committed, because refusing on them is // what made duplicate detection reject almost every add. let similar = if force { Vec::new() @@ -213,8 +238,19 @@ pub fn cmd_add( })(); match result { - Ok((entry_id, similar)) => { + Ok((entry_id, mut similar)) => { conn.execute_batch("COMMIT")?; + + if !force && !similar.is_empty() { + similarity::refine_similar_with_laya( + &mut similar, + title, + content, + laya_config.duplicate_threshold, + laya_client.as_mut(), + ); + } + let uid = db::get_entry(&conn, entry_id) .ok() .flatten() diff --git a/src/cmd/search.rs b/src/cmd/search.rs index 2de9a7f..dd7ed2c 100644 --- a/src/cmd/search.rs +++ b/src/cmd/search.rs @@ -69,13 +69,50 @@ pub fn cmd_search( // the keyword and LIKE fallbacks carry no score and every row ties. Among equals, // knowledge recorded in the repo you are standing in is the better guess; only // then does the per-DB order (updated_at DESC) decide. + let is_mine = |label: &str, e: &db::Entry| { + label == "project" || (here.is_some() && e.project.as_deref() == here.as_deref()) + }; items.sort_by(|a, b| { - let mine = |e: &db::Entry| here.is_some() && e.project.as_deref() == here.as_deref(); a.0.partial_cmp(&b.0) .unwrap_or(std::cmp::Ordering::Equal) - .then_with(|| mine(&b.2).cmp(&mine(&a.2))) + .then_with(|| is_mine(b.1, &b.2).cmp(&is_mine(a.1, &a.2))) .then_with(|| b.2.updated_at.cmp(&a.2.updated_at)) }); + + // Semantic reranking with Laya (if available) + let mut semantic_scores: std::collections::HashMap = + std::collections::HashMap::new(); + if !items.is_empty() { + let mut laya_client = crate::laya::LayaClient::connect(&config.laya); + if laya_client.is_some() { + let cand_count = items.len().min(15); + let candidates: Vec = items[..cand_count] + .iter() + .map(|(_, _, r, _)| crate::laya::RerankCandidate { + id: r.uid.clone(), + title: r.title.clone(), + content: r.content.clone(), + }) + .collect(); + + semantic_scores = + crate::laya::compute_rerank_scores(query, &candidates, laya_client.as_mut()); + + if !semantic_scores.is_empty() { + // Sort by current project priority first, then semantic score descending, falling back to BM25 order + items.sort_by(|a, b| { + let sa = semantic_scores.get(&a.2.uid).copied().unwrap_or(0.0); + let sb = semantic_scores.get(&b.2.uid).copied().unwrap_or(0.0); + is_mine(b.1, &b.2) + .cmp(&is_mine(a.1, &a.2)) + .then_with(|| sb.partial_cmp(&sa).unwrap_or(std::cmp::Ordering::Equal)) + .then_with(|| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal)) + .then_with(|| b.2.updated_at.cmp(&a.2.updated_at)) + }); + } + } + } + items.truncate(limit); let result_count = items.len().to_string(); @@ -101,6 +138,9 @@ pub fn cmd_search( "status": r.status, "stale": stale, }); + if let Some(sem) = semantic_scores.get(&r.uid) { + obj["semantic_score"] = serde_json::json!(crate::util::round2(*sem)); + } if let Some(ref project) = r.project { obj["project"] = serde_json::json!(project); } diff --git a/src/config.rs b/src/config.rs index 0b8271a..a860025 100644 --- a/src/config.rs +++ b/src/config.rs @@ -37,6 +37,35 @@ pub struct Config { /// Set to false when agents receive the instructions over the lk-knowledge MCP /// server instead, which serves the same content in its `initialize` response. pub claude_md_import: bool, + /// Laya MLX integration settings. + pub laya: LayaConfig, +} + +/// Configuration for Laya MLX integration (typed decisions on Apple Silicon). +#[derive(Debug, Clone)] +pub struct LayaConfig { + /// Whether Laya integration is enabled (default: false, opt-in). + pub enabled: bool, + /// Path to the Unix domain socket. Defaults to `~/.cache/lk/laya.sock`. + pub socket_path: Option, + /// Seconds of inactivity before the background daemon auto-shuts down (default: 600). + pub idle_timeout: u64, + /// Hugging Face model repository or local path. + pub model: String, + /// Semantic duplicate detection probability threshold (default: 0.85). + pub duplicate_threshold: f64, +} + +impl Default for LayaConfig { + fn default() -> Self { + Self { + enabled: false, + socket_path: None, + idle_timeout: 600, + model: "aac6fef/laya-multilingual-mlx".to_string(), + duplicate_threshold: 0.85, + } + } } impl Default for Config { @@ -50,6 +79,7 @@ impl Default for Config { command_log: false, gitattributes_generated: true, claude_md_import: true, + laya: LayaConfig::default(), } } } @@ -69,20 +99,39 @@ impl Config { /// - `LK_NO_AUTO_SYNC=1` → auto_sync = false /// - `LK_COMMAND_LOG=1` or `LK_SEARCH_LOG=1` → command_log = true pub fn load(knowledge_dir: &Path) -> Self { - let mut config = Self::default(); + let global = GlobalConfig::load(); + Self::load_with_global(knowledge_dir, &global) + } + + /// Load config from `.knowledge/config.toml`, inheriting `laya` from `global` as default. + pub fn load_with_global(knowledge_dir: &Path, global: &GlobalConfig) -> Self { + let mut config = Self { + laya: global.laya.clone(), + ..Self::default() + }; let config_path = knowledge_dir.join("config.toml"); if let Ok(content) = std::fs::read_to_string(&config_path) { + let mut section = String::new(); for line in content.lines() { let line = line.trim(); // Skip comments and empty lines if line.is_empty() || line.starts_with('#') { continue; } + if line.starts_with('[') && line.ends_with(']') { + section = line[1..line.len() - 1].trim().to_ascii_lowercase(); + continue; + } if let Some((key, value)) = line.split_once('=') { let key = key.trim(); - let value = value.trim(); - match key { + let value = value.trim().trim_matches('"').trim_matches('\'').trim(); + let full_key = if section.is_empty() { + key.to_string() + } else { + format!("{section}.{key}") + }; + match full_key.as_str() { "stale_threshold_days" => { if let Ok(v) = value.parse::() && v > 0 @@ -122,6 +171,20 @@ impl Config { config.claude_md_import = parse_bool(value, key, config.claude_md_import); } + "laya.enabled" | "laya_enabled" => { + // Project config may only disable Laya (opt-out); it cannot + // enable Laya if it is disabled in the global user configuration. + if !parse_bool(value, key, true) { + config.laya.enabled = false; + } + } + "laya.duplicate_threshold" | "laya_duplicate_threshold" => { + if let Ok(v) = value.parse::() + && (0.0..=1.0).contains(&v) + { + config.laya.duplicate_threshold = v; + } + } _ => {} // Ignore unknown keys } } @@ -137,6 +200,13 @@ impl Config { { config.command_log = true; } + if std::env::var("LK_NO_LAYA").unwrap_or_default() == "1" + || std::env::var("LK_LAYA").unwrap_or_default() == "0" + { + config.laya.enabled = false; + } else if std::env::var("LK_LAYA").unwrap_or_default() == "1" { + config.laya.enabled = true; + } config } @@ -196,6 +266,8 @@ pub struct GlobalConfig { /// Add the lk-instructions import to `~/.claude/CLAUDE.md` (default: true). /// The user-scope counterpart of [`Config::claude_md_import`]. pub claude_md_import: bool, + /// Laya MLX integration settings. + pub laya: LayaConfig, } impl GlobalConfig { @@ -204,6 +276,7 @@ impl GlobalConfig { user_knowledge_dir: home.join(".config").join("lk").join("knowledge"), secret_detection: true, claude_md_import: true, + laya: LayaConfig::default(), } } @@ -220,16 +293,26 @@ impl GlobalConfig { let mut config = Self::default_with_home(home); if let Ok(content) = std::fs::read_to_string(config_path) { + let mut section = String::new(); for line in content.lines() { let line = line.trim(); if line.is_empty() || line.starts_with('#') { continue; } + if line.starts_with('[') && line.ends_with(']') { + section = line[1..line.len() - 1].trim().to_ascii_lowercase(); + continue; + } if let Some((key, value)) = line.split_once('=') { let key = key.trim(); // Strip optional surrounding quotes from the value. let value = value.trim().trim_matches('"').trim_matches('\'').trim(); - match key { + let full_key = if section.is_empty() { + key.to_string() + } else { + format!("{section}.{key}") + }; + match full_key.as_str() { "user_knowledge_dir" => { if !value.is_empty() { config.user_knowledge_dir = resolve_path(value, home); @@ -243,12 +326,47 @@ impl GlobalConfig { config.claude_md_import = parse_bool(value, key, config.claude_md_import); } + "laya.enabled" | "laya_enabled" => { + config.laya.enabled = parse_bool(value, key, config.laya.enabled); + } + "laya.idle_timeout" | "laya_idle_timeout" => { + if let Ok(v) = value.parse::() + && v > 0 + { + config.laya.idle_timeout = v; + } + } + "laya.model" | "laya_model" => { + if !value.is_empty() { + config.laya.model = value.to_string(); + } + } + "laya.socket_path" | "laya_socket_path" => { + if !value.is_empty() { + config.laya.socket_path = Some(resolve_path(value, home)); + } + } + "laya.duplicate_threshold" | "laya_duplicate_threshold" => { + if let Ok(v) = value.parse::() + && (0.0..=1.0).contains(&v) + { + config.laya.duplicate_threshold = v; + } + } _ => {} // Ignore unknown keys } } } } + if std::env::var("LK_NO_LAYA").unwrap_or_default() == "1" + || std::env::var("LK_LAYA").unwrap_or_default() == "0" + { + config.laya.enabled = false; + } else if std::env::var("LK_LAYA").unwrap_or_default() == "1" { + config.laya.enabled = true; + } + config } @@ -273,7 +391,7 @@ fn resolve_path(value: &str, home: &Path) -> PathBuf { .any(|c| matches!(c, std::path::Component::ParentDir)) { eprintln!( - "Warning: `user_knowledge_dir` contains `..` ({value:?}); \ + "Warning: config path contains `..` ({value:?}); \ prefer an absolute path or one under `~/` to avoid writing outside your home." ); } @@ -411,6 +529,70 @@ mod tests { assert!(!GlobalConfig::load_from(&config_path, home.path()).claude_md_import); } + #[test] + fn test_config_inherits_global_laya_config() { + let home = TempDir::new().unwrap(); + let config_path = home.path().join("config.toml"); + std::fs::write(&config_path, "[laya]\nenabled = true\n").unwrap(); + let global = GlobalConfig::load_from(&config_path, home.path()); + assert!(global.laya.enabled); + + let project_dir = TempDir::new().unwrap(); + // Project without laya config inherits global laya + let config = Config::load_with_global(project_dir.path(), &global); + assert!(config.laya.enabled); + + // Project with laya.enabled = false overrides global laya + std::fs::write( + project_dir.path().join("config.toml"), + "[laya]\nenabled = false\n", + ) + .unwrap(); + let config = Config::load_with_global(project_dir.path(), &global); + assert!(!config.laya.enabled); + + // Project config cannot force-enable laya if disabled globally, + // and cannot override socket_path or model + let disabled_global = + GlobalConfig::load_from(&home.path().join("nonexistent.toml"), home.path()); + assert!(!disabled_global.laya.enabled); + let untrusted_dir = TempDir::new().unwrap(); + std::fs::write( + untrusted_dir.path().join("config.toml"), + "[laya]\nenabled = true\nsocket_path = \"/tmp/attacker.sock\"\nmodel = \"attacker/model\"\n", + ) + .unwrap(); + let config = Config::load_with_global(untrusted_dir.path(), &disabled_global); + assert!( + !config.laya.enabled, + "project cannot force-enable Laya when globally disabled" + ); + assert_eq!( + config.laya.socket_path, None, + "project cannot override socket_path" + ); + assert_eq!( + config.laya.model, disabled_global.laya.model, + "project cannot override model" + ); + } + + #[test] + fn test_global_config_laya_socket_path_expansion() { + let home = TempDir::new().unwrap(); + let config_path = home.path().join("config.toml"); + std::fs::write( + &config_path, + "[laya]\nsocket_path = \"~/.cache/custom.sock\"\n", + ) + .unwrap(); + let global = GlobalConfig::load_from(&config_path, home.path()); + assert_eq!( + global.laya.socket_path, + Some(home.path().join(".cache").join("custom.sock")) + ); + } + #[test] fn test_set_config_bool_rewrites_existing_key() { let dir = TempDir::new().unwrap(); diff --git a/src/keywords.rs b/src/keywords.rs index 93767c0..154c6ae 100644 --- a/src/keywords.rs +++ b/src/keywords.rs @@ -20,6 +20,13 @@ const STOP_WORDS: &[&str] = &[ /// search and duplicate detection. pub const MAX_AUTO_KEYWORDS: usize = 15; +/// Japanese general stopwords for candidate filtering. +const JP_STOP_WORDS: &[&str] = &[ + "場合", "設定", "追加", "確認", "対応", "実行", "作成", "更新", "取得", "処理", "使用", "問題", + "修正", "機能", "変更", "必要", "可能", "利用", "対象", "詳細", "理由", "手順", "記述", "指定", + "管理", "表示", "発生", "関係", "存在", "関連", "完了", "状況", "現在", +]; + /// A term occurring in the title counts this many times a content occurrence. const TITLE_WEIGHT: u32 = 5; @@ -27,6 +34,57 @@ const TITLE_WEIGHT: u32 = 5; /// high-signal identifiers (module/file names) worth keeping over prose words. const PATH_WEIGHT: u32 = 3; +/// Extract keywords with Laya if available, falling back to heuristic frequency extraction. +pub fn extract_keywords_auto( + title: &str, + content: &str, + mut laya: Option<&mut crate::laya::LayaClient>, +) -> Vec { + if let Some(ref mut client) = laya { + let candidates = extract_candidates_for_laya(title, content); + if !candidates.is_empty() + && let Ok(ranked) = client.filter_keywords(title, content, &candidates) + { + // Keep terms with score >= 0.50, capped at MAX_AUTO_KEYWORDS + let mut selected: Vec = ranked + .iter() + .filter(|(_, score)| *score >= 0.50) + .map(|(kw, _)| kw.clone()) + .take(MAX_AUTO_KEYWORDS) + .collect(); + + // If threshold filtered everything out, take top 5 from Laya's ranked results + if selected.is_empty() && !ranked.is_empty() { + selected = ranked.into_iter().take(5).map(|(kw, _)| kw).collect(); + } else if selected.is_empty() && !candidates.is_empty() { + selected = candidates.into_iter().take(5).collect(); + } + + if !selected.is_empty() { + selected.sort(); + return selected; + } + } + } + + extract_keywords(title, content) +} + +/// Extract candidate keywords for Laya ranking, including ASCII words, katakana, +/// and Japanese kanji compounds (2-6 chars). +pub fn extract_candidates_for_laya(title: &str, content: &str) -> Vec { + let mut scores: HashMap = HashMap::new(); + score_text_with_kanji(title, TITLE_WEIGHT, &mut scores); + score_text_with_kanji(content, 1, &mut scores); + + let mut ranked: Vec<(String, u32)> = scores.into_iter().collect(); + ranked.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0))); + // Offer up to 25 candidates to Laya + ranked.truncate(25); + + ranked.into_iter().map(|(kw, _)| kw).collect() +} + /// Extract keywords from title and content, ranked by weighted frequency and /// capped at `MAX_AUTO_KEYWORDS`. Title occurrences and file-path segments are /// weighted higher than plain content words. @@ -50,8 +108,39 @@ pub fn extract_keywords(title: &str, content: &str) -> Vec { static PATH_RE: LazyLock = LazyLock::new(|| Regex::new(r"[\w./\\-]+\.[\w]+").unwrap()); static WORD_RE: LazyLock = LazyLock::new(|| Regex::new(r"[A-Za-z_][A-Za-z0-9_]*").unwrap()); static KATAKANA_RE: LazyLock = LazyLock::new(|| Regex::new(r"[\u30A0-\u30FF]{4,}").unwrap()); +static KANJI_RE: LazyLock = LazyLock::new(|| Regex::new(r"[\u4E00-\u9FFF]{2,6}").unwrap()); static CAMEL_RE: LazyLock = LazyLock::new(|| Regex::new(r"([a-z])([A-Z])").unwrap()); +fn score_text_with_kanji(text: &str, weight: u32, scores: &mut HashMap) { + score_text(text, weight, scores); + + // Kanji compounds (2-6 chars) + for mat in KANJI_RE.find_iter(text) { + let s = mat.as_str(); + if !JP_STOP_WORDS.contains(&s) { + *scores.entry(s.to_string()).or_insert(0) += weight; + } + + let chars: Vec = s.chars().collect(); + let len = chars.len(); + if len >= 4 { + // Sub-compound candidates: 2-char and 4-char prefixes/suffixes + for w in [2, 4] { + if w < len { + let prefix: String = chars[..w].iter().collect(); + if !JP_STOP_WORDS.contains(&prefix.as_str()) { + *scores.entry(prefix).or_insert(0) += weight; + } + let suffix: String = chars[len - w..].iter().collect(); + if !JP_STOP_WORDS.contains(&suffix.as_str()) { + *scores.entry(suffix).or_insert(0) += weight; + } + } + } + } + } +} + fn score_text(text: &str, weight: u32, scores: &mut HashMap) { // File path segments. Path ranges are masked out of the word scan below so // path-derived tokens are scored exactly once, at PATH_WEIGHT. @@ -215,4 +304,33 @@ mod tests { let kws = extract_keywords("", &content); assert!(kws.contains(&"webhook".to_string())); } + + #[test] + fn test_candidates_for_laya_extracts_kanji() { + let candidates = extract_candidates_for_laya( + "OAuth2トークン失効時の自動リフレッシュ処理", + "アクセストークンの有効期限切れにより401 Unauthorizedが返却された場合の排他制御手順。", + ); + assert!( + candidates.contains(&"oauth2".to_string()) + || candidates.contains(&"OAuth2".to_string()) + ); + assert!(candidates.contains(&"トークン".to_string())); + assert!( + candidates.contains(&"失効時".to_string()) + || candidates.contains(&"失効".to_string()) + || candidates.contains(&"排他制御".to_string()) + ); + assert!(candidates.contains(&"排他制御".to_string())); + // Stopwords like "処理" and "場合" should be excluded + assert!(!candidates.contains(&"処理".to_string())); + assert!(!candidates.contains(&"場合".to_string())); + } + + #[test] + fn test_extract_keywords_auto_none_laya() { + let kws = extract_keywords_auto("Session Management", "Token expiration handling", None); + assert!(!kws.is_empty()); + assert!(kws.contains(&"session".to_string())); + } } diff --git a/src/laya.rs b/src/laya.rs new file mode 100644 index 0000000..a1ac590 --- /dev/null +++ b/src/laya.rs @@ -0,0 +1,475 @@ +//! Client for the Laya MLX typed decision sidecar server. +//! +//! Connects via Unix domain socket (default: `~/.cache/lk/laya.sock`). +//! If the daemon is not running and Laya is enabled, it automatically spawns +//! the Python server (`scripts/laya_server.py`) using `uv run`. +//! If the server cannot be reached, times out, or errors, operations fail open +//! (returning `None` or failing over to standard heuristics) so core commands +//! never crash. + +use std::collections::HashMap; +use std::io::{BufRead, BufReader, Write}; +#[cfg(unix)] +use std::os::unix::net::UnixStream; +use std::path::{Path, PathBuf}; +use std::process::Command; +use std::time::Duration; +#[cfg(all(target_os = "macos", target_arch = "aarch64"))] +use std::time::Instant; + +use serde::{Deserialize, Serialize}; +use serde_json::json; + +use crate::config::LayaConfig; + +#[allow(dead_code)] +pub const DEFAULT_SOCKET_NAME: &str = "laya.sock"; + +/// Embedded server script content so `lk` can unpack and run it even when +/// installed without the repository source directory. +#[allow(dead_code)] +pub const EMBEDDED_SERVER_SCRIPT: &str = include_str!("../scripts/laya_server.py"); + +#[derive(Debug, Clone, Deserialize)] +pub struct DuplicateResult { + pub is_duplicate: bool, + #[allow(dead_code)] + pub probability: f64, +} + +#[derive(Debug, Clone, Serialize)] +pub struct RerankCandidate { + pub id: String, + pub title: String, + pub content: String, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct RerankScore { + pub id: String, + pub score: f64, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct CategorizeResult { + pub category: String, + #[allow(dead_code)] + pub confidence: f64, + #[allow(dead_code)] + pub probabilities: HashMap, +} + +#[derive(Debug, Deserialize)] +struct KeywordItem { + keyword: String, + score: f64, +} + +#[derive(Debug, Deserialize)] +struct FilterKeywordsResponse { + ranked_keywords: Vec, +} + +#[derive(Debug, Deserialize)] +struct RerankResponse { + scores: Vec, +} + +pub struct LayaClient { + #[cfg(unix)] + writer: UnixStream, + #[cfg(unix)] + reader: BufReader, + #[allow(dead_code)] + socket_path: PathBuf, + #[allow(dead_code)] + broken: bool, +} + +#[allow(dead_code)] +fn is_process_alive(pid: u32) -> bool { + #[cfg(unix)] + { + Command::new("kill") + .args(["-0", &pid.to_string()]) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status() + .map(|s| s.success()) + .unwrap_or(false) + } + #[cfg(not(unix))] + { + let _ = pid; + false + } +} + +#[allow(dead_code)] +impl LayaClient { + /// Connect to an existing Laya daemon, or spawn one if enabled and missing. + /// Returns `None` if disabled, unsupported, or if connection/spawning fails + /// (providing transparent graceful degradation). + #[cfg(all(target_os = "macos", target_arch = "aarch64"))] + pub fn connect(config: &LayaConfig) -> Option { + if !config.enabled { + return None; + } + + let socket_path = resolve_socket_path(config); + + // 1. Try connecting to an already-running daemon. + if let Some((mut client, model_matches)) = Self::try_connect(&socket_path, &config.model) { + if model_matches { + return Some(client); + } + // Model mismatch: shut down the outdated daemon + let _ = client.send_request("shutdown", json!({})); + let pid_path = socket_path.with_extension("pid"); + if let Ok(content) = std::fs::read_to_string(&pid_path) + && let Ok(pid) = content.trim().parse::() + { + let wait_start = Instant::now(); + while is_process_alive(pid) && wait_start.elapsed() < Duration::from_millis(1000) { + std::thread::sleep(Duration::from_millis(50)); + } + } + } + + // 2. Not running (or shut down): attempt to spawn the daemon in the background. + if Self::spawn_daemon(config, &socket_path) { + // Poll for socket readiness up to 5 seconds. + let start = Instant::now(); + while start.elapsed() < Duration::from_secs(5) { + std::thread::sleep(Duration::from_millis(150)); + if let Some((client, model_matches)) = + Self::try_connect(&socket_path, &config.model) + && model_matches + { + return Some(client); + } + } + } + + None + } + + /// On platforms other than macOS Apple Silicon, Laya is unsupported. + #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))] + pub fn connect(_config: &LayaConfig) -> Option { + None + } + + /// Try connecting to the socket and verifying with a quick ping and model check. + #[allow(dead_code)] + fn try_connect(path: &Path, expected_model: &str) -> Option<(Self, bool)> { + #[cfg(unix)] + { + let stream = UnixStream::connect(path).ok()?; + let _ = stream.set_read_timeout(Some(Duration::from_millis(3000))); + let _ = stream.set_write_timeout(Some(Duration::from_millis(3000))); + let reader_stream = stream.try_clone().ok()?; + let mut client = Self { + writer: stream, + reader: BufReader::new(reader_stream), + socket_path: path.to_path_buf(), + broken: false, + }; + let ping_res = client.ping().ok()?; + let model_matches = ping_res + .get("model") + .and_then(|m| m.as_str()) + .map(|m| m == expected_model) + .unwrap_or(true); + Some((client, model_matches)) + } + #[cfg(not(unix))] + { + let _ = (path, expected_model); + None + } + } + + /// Spawn the Python daemon process using `uv run`. + #[allow(dead_code)] + fn spawn_daemon(config: &LayaConfig, socket_path: &Path) -> bool { + let pid_path = socket_path.with_extension("pid"); + if pid_path.exists() { + if let Ok(content) = std::fs::read_to_string(&pid_path) + && let Ok(pid) = content.trim().parse::() + && is_process_alive(pid) + { + // Daemon process is already alive; don't remove files or re-spawn + return true; + } + let _ = std::fs::remove_file(&pid_path); + } + if socket_path.exists() { + let _ = std::fs::remove_file(socket_path); + } + + let script_path = match locate_or_extract_server_script() { + Ok(p) => p, + Err(_) => return false, + }; + + // Check if `uv` command is available + if Command::new("uv") + .arg("--version") + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status() + .map(|s| s.success()) + .unwrap_or(false) + { + let child = Command::new("uv") + .args([ + "run", + "--with", + "laya-mlx", + "python3", + script_path.to_str().unwrap_or(""), + "--socket-path", + socket_path.to_str().unwrap_or(""), + "--model", + &config.model, + "--idle-timeout", + &config.idle_timeout.to_string(), + ]) + .stdin(std::process::Stdio::null()) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .spawn(); + + child.is_ok() + } else { + false + } + } + + #[cfg(unix)] + fn send_request( + &mut self, + task: &str, + params: serde_json::Value, + ) -> Result { + if self.broken { + return Err("Laya socket is in a broken state from a previous error".to_string()); + } + + let req = json!({ + "id": 1, + "task": task, + "params": params, + }); + let mut line = serde_json::to_string(&req).map_err(|e| e.to_string())?; + line.push('\n'); + + if let Err(e) = self.writer.write_all(line.as_bytes()) { + self.broken = true; + return Err(format!("Failed to send to Laya socket: {e}")); + } + if let Err(e) = self.writer.flush() { + self.broken = true; + return Err(format!("Failed to flush Laya socket: {e}")); + } + + let mut resp_line = String::new(); + if let Err(e) = self.reader.read_line(&mut resp_line) { + self.broken = true; + return Err(format!("Failed to read from Laya socket: {e}")); + } + if resp_line.is_empty() { + self.broken = true; + return Err("Laya socket closed by peer".to_string()); + } + + let resp: serde_json::Value = match serde_json::from_str(resp_line.trim()) { + Ok(v) => v, + Err(e) => { + self.broken = true; + return Err(format!("Failed to parse Laya response: {e}")); + } + }; + + if let Some(err) = resp.get("error") { + let msg = err.as_str().unwrap_or("Unknown server error"); + return Err(format!("Laya daemon error: {msg}")); + } + + resp.get("result") + .cloned() + .ok_or_else(|| "Missing result in Laya response".to_string()) + } + + #[cfg(not(unix))] + fn send_request( + &mut self, + _task: &str, + _params: serde_json::Value, + ) -> Result { + Err("Laya is only supported on Unix systems".to_string()) + } + + pub fn ping(&mut self) -> Result { + self.send_request("ping", json!({})) + } + + pub fn check_duplicate( + &mut self, + entry_a_title: &str, + entry_a_content: &str, + entry_b_title: &str, + entry_b_content: &str, + threshold: f64, + ) -> Result { + let params = json!({ + "entry_a": { + "title": entry_a_title, + "content": entry_a_content, + }, + "entry_b": { + "title": entry_b_title, + "content": entry_b_content, + }, + "threshold": threshold, + }); + let res = self.send_request("duplicate", params)?; + serde_json::from_value(res).map_err(|e| e.to_string()) + } + + pub fn filter_keywords( + &mut self, + title: &str, + content: &str, + candidates: &[String], + ) -> Result, String> { + let params = json!({ + "title": title, + "content": content, + "candidates": candidates, + }); + let res = self.send_request("filter_keywords", params)?; + let parsed: FilterKeywordsResponse = + serde_json::from_value(res).map_err(|e| e.to_string())?; + Ok(parsed + .ranked_keywords + .into_iter() + .map(|item| (item.keyword, item.score)) + .collect()) + } + + pub fn rerank( + &mut self, + query: &str, + candidates: &[RerankCandidate], + ) -> Result, String> { + let params = json!({ + "query": query, + "candidates": candidates, + }); + let res = self.send_request("rerank", params)?; + let parsed: RerankResponse = serde_json::from_value(res).map_err(|e| e.to_string())?; + Ok(parsed.scores) + } + + pub fn categorize(&mut self, title: &str, content: &str) -> Result { + let params = json!({ + "title": title, + "content": content, + }); + let res = self.send_request("categorize", params)?; + serde_json::from_value(res).map_err(|e| e.to_string()) + } +} + +/// Compute semantic rerank scores for candidates using Laya, if available. +/// Returns a map of candidate id -> semantic score (0.0 to 1.0). +pub fn compute_rerank_scores( + query: &str, + candidates: &[RerankCandidate], + mut laya: Option<&mut LayaClient>, +) -> HashMap { + if let Some(ref mut client) = laya + && let Ok(scores) = client.rerank(query, candidates) + { + return scores.into_iter().map(|s| (s.id, s.score)).collect(); + } + HashMap::new() +} + +/// Resolve the socket path to use, defaulting to `~/.cache/lk/laya.sock`. +#[allow(dead_code)] +pub fn resolve_socket_path(config: &LayaConfig) -> PathBuf { + if let Some(ref p) = config.socket_path { + p.clone() + } else { + crate::util::home_dir() + .join(".cache") + .join("lk") + .join(DEFAULT_SOCKET_NAME) + } +} + +/// Locate server script via explicit environment override (`LK_LAYA_SERVER_SCRIPT`), or write +/// the embedded script to `~/.cache/lk/laya_server.py`. +#[allow(dead_code)] +fn locate_or_extract_server_script() -> Result> { + // 1. Check explicit override via environment variable (e.g. for development) + if let Ok(env_path) = std::env::var("LK_LAYA_SERVER_SCRIPT") { + let p = PathBuf::from(env_path); + if p.exists() { + return Ok(p); + } + } + + // 2. Extract embedded script into cache directory + let cache_dir = crate::util::home_dir().join(".cache").join("lk"); + std::fs::create_dir_all(&cache_dir)?; + let target = cache_dir.join("laya_server.py"); + let need_write = match std::fs::read_to_string(&target) { + Ok(existing) => existing != EMBEDDED_SERVER_SCRIPT, + Err(_) => true, + }; + if need_write { + std::fs::write(&target, EMBEDDED_SERVER_SCRIPT)?; + } + Ok(target) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + #[cfg(unix)] + fn test_is_process_alive() { + let my_pid = std::process::id(); + assert!(is_process_alive(my_pid)); + assert!(!is_process_alive(4_000_000)); + } + + #[test] + #[cfg(not(unix))] + fn test_is_process_alive_non_unix() { + assert!(!is_process_alive(std::process::id())); + } + + #[test] + fn test_locate_or_extract_server_script() { + let script = locate_or_extract_server_script().expect("script extracted"); + assert!(script.exists()); + let content = std::fs::read_to_string(&script).expect("readable script"); + assert_eq!(content, EMBEDDED_SERVER_SCRIPT); + } + + #[test] + fn test_resolve_socket_path() { + let mut cfg = LayaConfig::default(); + let default_sock = resolve_socket_path(&cfg); + assert!(default_sock.ends_with(DEFAULT_SOCKET_NAME)); + + cfg.socket_path = Some(PathBuf::from("/tmp/custom.sock")); + assert_eq!(resolve_socket_path(&cfg), PathBuf::from("/tmp/custom.sock")); + } +} diff --git a/src/main.rs b/src/main.rs index a20d440..3a42423 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,6 +2,7 @@ mod cmd; mod config; mod db; mod keywords; +mod laya; mod markdown; mod mcp; mod prompts; diff --git a/src/mcp.rs b/src/mcp.rs index 1ce7dc4..4000d68 100644 --- a/src/mcp.rs +++ b/src/mcp.rs @@ -942,14 +942,56 @@ fn call_tool(name: &str, params: &Value, registry: &ProjectRegistry) -> Result = + std::collections::HashMap::new(); + if !items.is_empty() { + let mut laya_client = crate::laya::LayaClient::connect(&config.laya); + if laya_client.is_some() { + let cand_count = items.len().min(15); + let candidates: Vec = items[..cand_count] + .iter() + .map(|(_, _, r, _)| crate::laya::RerankCandidate { + id: r.uid.clone(), + title: r.title.clone(), + content: r.content.clone(), + }) + .collect(); + + semantic_scores = crate::laya::compute_rerank_scores( + query, + &candidates, + laya_client.as_mut(), + ); + + if !semantic_scores.is_empty() { + items.sort_by(|a, b| { + let sa = semantic_scores.get(&a.2.uid).copied().unwrap_or(0.0); + let sb = semantic_scores.get(&b.2.uid).copied().unwrap_or(0.0); + is_mine(b.1, &b.2) + .cmp(&is_mine(a.1, &a.2)) + .then_with(|| { + sb.partial_cmp(&sa).unwrap_or(std::cmp::Ordering::Equal) + }) + .then_with(|| { + a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal) + }) + .then_with(|| b.2.updated_at.cmp(&a.2.updated_at)) + }); + } + } + } + items.truncate(limit); let results: Vec = items @@ -957,6 +999,9 @@ fn call_tool(name: &str, params: &Value, registry: &ProjectRegistry) -> Result Result crate::config::GlobalConfig::load().laya, + _ => config.laya.clone(), + }; + let mut laya_client = crate::laya::LayaClient::connect(&laya_config); + + let mut final_keywords = keywords; + if final_keywords.is_empty() { + final_keywords = crate::keywords::extract_keywords_auto( + title, + effective_content, + laya_client.as_mut(), + ); + } + + // Duplicate check. Only a Block-tier hit (a near-identical title or + // semantic duplicate) refuses the add; weaker hits are reported + // alongside a successful add as `possibly_related` further down. + let mut similar = if force { Vec::new() } else { - db::find_similar_entries(&conn, title, &keywords, category) + db::find_similar_entries(&conn, title, &final_keywords, category) .map_err(|e| format!("duplicate check error: {e}"))? }; + // Shared with `lk add --json` so a hit looks identical on both surfaces. let describe = |s: &db::SimilarEntry| -> Value { util::similar_entry_json(&conn, s, effective_scope) @@ -1142,7 +1203,7 @@ fn call_tool(name: &str, params: &Value, registry: &ProjectRegistry) -> Result Result "same-title", Reason::SimilarTitle => "similar-title", Reason::SimilarKeywords => "similar-keywords", + Reason::SemanticDuplicate => "semantic-duplicate", } } } @@ -278,6 +281,38 @@ pub fn reason(title_sim: f64, kw_sim: f64, title_exact: bool) -> Reason { } } +/// Refine candidate similar entries with Laya. If Laya finds that a candidate +/// is a semantic duplicate, ensure its tier is at least Tier::Warn and record +/// Reason::SemanticDuplicate. +pub fn refine_similar_with_laya( + similar: &mut [crate::db::SimilarEntry], + incoming_title: &str, + incoming_content: &str, + threshold: f64, + mut laya: Option<&mut crate::laya::LayaClient>, +) { + if let Some(ref mut client) = laya { + for s in similar.iter_mut() { + // Only examine candidates that aren't already hard-blocked by title equality. + if s.tier != Tier::Block + && let Ok(res) = client.check_duplicate( + incoming_title, + incoming_content, + &s.entry.title, + &s.entry.content, + threshold, + ) + && res.is_duplicate + { + if s.tier == Tier::None { + s.tier = Tier::Warn; + } + s.reason = Reason::SemanticDuplicate; + } + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/util.rs b/src/util.rs index f801c75..7e07e44 100644 --- a/src/util.rs +++ b/src/util.rs @@ -703,7 +703,7 @@ pub fn similar_entry_json( /// Two decimal places — the scores are advisory, and full float noise in the /// output only invites agents to treat them as precise. -fn round2(v: f64) -> f64 { +pub fn round2(v: f64) -> f64 { (v * 100.0).round() / 100.0 } diff --git a/tests/integration_test.rs b/tests/integration_test.rs index b0a2eec..4a932ff 100644 --- a/tests/integration_test.rs +++ b/tests/integration_test.rs @@ -2,7 +2,14 @@ use std::path::Path; use std::process::Command; fn lk_bin() -> Command { - Command::new(env!("CARGO_BIN_EXE_lk")) + let mut cmd = Command::new(env!("CARGO_BIN_EXE_lk")); + // By default, disable Laya MLX daemon during integration tests to prevent + // parallel test workers from spawning dozens of concurrent Python/MLX processes + // in temp directories. Run with LK_LAYA=1 to test with Laya active. + if std::env::var("LK_LAYA").is_err() { + cmd.env("LK_NO_LAYA", "1"); + } + cmd } /// `lk` run in `dir`, with `HOME` pointed at it too.