diff --git a/_primitives/MANIFEST.toml b/_primitives/MANIFEST.toml index 2a1badb..71e3a55 100644 --- a/_primitives/MANIFEST.toml +++ b/_primitives/MANIFEST.toml @@ -21,7 +21,8 @@ core = ["tomd", "genesis-scan"] frontend = ["mock-render", "visual-diff", "tokens-sync", "design-scrape", "live-preview", "figma-tokens", "frontend-inspect", "screenshot-decode"] ops = ["kei-ledger", "ssh-check", "firewall-diff", "provision-hetzner", "provision-vultr", "harden-base", "metrics-scrape", "log-ship"] dev = ["kei-migrate", "kei-changelog", "kei-ci-lint", "kei-docs-scaffold", "kei-memory", "kei-conflict-scan", "kei-refactor-engine", "kei-graph-check", "kei-store"] -full = ["tomd", "genesis-scan", "kei-ledger", "kei-migrate", "kei-changelog", "ssh-check", "firewall-diff", "mock-render", "visual-diff", "tokens-sync", "design-scrape", "live-preview", "figma-tokens", "frontend-inspect", "screenshot-decode", "provision-hetzner", "provision-vultr", "harden-base", "metrics-scrape", "log-ship", "kei-ci-lint", "kei-docs-scaffold", "kei-memory", "kei-conflict-scan", "kei-refactor-engine", "kei-graph-check", "kei-store"] +mcp = ["kei-router", "kei-sage", "kei-task", "kei-chat-store", "kei-crossdomain", "kei-search-core", "kei-content-store", "kei-social-store", "kei-curator", "kei-auth"] +full = ["tomd", "genesis-scan", "kei-ledger", "kei-migrate", "kei-changelog", "ssh-check", "firewall-diff", "mock-render", "visual-diff", "tokens-sync", "design-scrape", "live-preview", "figma-tokens", "frontend-inspect", "screenshot-decode", "provision-hetzner", "provision-vultr", "harden-base", "metrics-scrape", "log-ship", "kei-ci-lint", "kei-docs-scaffold", "kei-memory", "kei-conflict-scan", "kei-refactor-engine", "kei-graph-check", "kei-store", "kei-router", "kei-sage", "kei-task", "kei-chat-store", "kei-crossdomain", "kei-search-core", "kei-content-store", "kei-social-store", "kei-curator", "kei-auth"] # --- shell primitives (13) ------------------------------------------------- @@ -188,3 +189,65 @@ kind = "rust" crate = "kei-store" deps = ["git2 (vendored libgit2)"] desc = "Memory-repo backend abstraction — GitHub / Forgejo / Gitea / Filesystem / S3 (S3 = MVP stub)" + +# --- v0.14 LBM port (10) --------------------------------------------------- + +[primitive.kei-router] +kind = "rust" +crate = "kei-router" +deps = ["regex"] +desc = "Natural-language query → tool-call router (LBM pkg/keirouter port, no ML)" + +[primitive.kei-sage] +kind = "rust" +crate = "kei-sage" +deps = ["rusqlite bundled (FTS5 enabled)"] +desc = "Obsidian-style knowledge graph with FTS5, BFS, PageRank (LBM internal/sage port)" + +[primitive.kei-task] +kind = "rust" +crate = "kei-task" +deps = ["rusqlite bundled (FTS5 enabled)"] +desc = "Task DAG + deps + milestones (LBM internal/task port)" + +[primitive.kei-chat-store] +kind = "rust" +crate = "kei-chat-store" +deps = ["rusqlite bundled (FTS5 enabled)"] +desc = "Session persistence for Claude chats (LBM internal/chat port)" + +[primitive.kei-crossdomain] +kind = "rust" +crate = "kei-crossdomain" +deps = ["rusqlite bundled"] +desc = "Cross-domain typed-edge store + BFS + auto-link (LBM internal/crossdomain port)" + +[primitive.kei-search-core] +kind = "rust" +crate = "kei-search-core" +deps = ["rusqlite bundled"] +desc = "3-wave research engine with budget cap; fetch interface frozen (LBM internal/search port)" + +[primitive.kei-content-store] +kind = "rust" +crate = "kei-content-store" +deps = ["rusqlite bundled", "sha2"] +desc = "Asset + prompt + campaign registry (LBM internal/content port)" + +[primitive.kei-social-store] +kind = "rust" +crate = "kei-social-store" +deps = ["rusqlite bundled (FTS5 enabled)"] +desc = "People + interaction CRM lite (LBM internal/social port)" + +[primitive.kei-curator] +kind = "rust" +crate = "kei-curator" +deps = ["rusqlite bundled"] +desc = "Edge decay + orphan prune for cross-domain graphs (LBM internal/curator port)" + +[primitive.kei-auth] +kind = "rust" +crate = "kei-auth" +deps = ["rusqlite bundled", "hmac", "sha2"] +desc = "Multi-tenant session tokens with scopes + HMAC-signed expiry (rewrite, not port)" diff --git a/_primitives/_rust/Cargo.lock b/_primitives/_rust/Cargo.lock index eb4bcbb..1d84aa8 100644 --- a/_primitives/_rust/Cargo.lock +++ b/_primitives/_rust/Cargo.lock @@ -909,6 +909,23 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "kei-auth" +version = "0.1.0" +dependencies = [ + "anyhow", + "base64 0.22.1", + "chrono", + "clap", + "hmac", + "rand", + "rusqlite", + "serde", + "serde_json", + "sha2", + "tempfile", +] + [[package]] name = "kei-changelog" version = "0.1.0" @@ -920,6 +937,20 @@ dependencies = [ "regex", ] +[[package]] +name = "kei-chat-store" +version = "0.1.0" +dependencies = [ + "anyhow", + "chrono", + "clap", + "rusqlite", + "serde", + "serde_json", + "tempfile", + "uuid", +] + [[package]] name = "kei-conflict-scan" version = "0.1.0" @@ -933,6 +964,46 @@ dependencies = [ "walkdir", ] +[[package]] +name = "kei-content-store" +version = "0.1.0" +dependencies = [ + "anyhow", + "chrono", + "clap", + "rusqlite", + "serde", + "serde_json", + "sha2", + "tempfile", +] + +[[package]] +name = "kei-crossdomain" +version = "0.1.0" +dependencies = [ + "anyhow", + "chrono", + "clap", + "rusqlite", + "serde", + "serde_json", + "tempfile", +] + +[[package]] +name = "kei-curator" +version = "0.1.0" +dependencies = [ + "anyhow", + "chrono", + "clap", + "rusqlite", + "serde", + "serde_json", + "tempfile", +] + [[package]] name = "kei-graph-check" version = "0.1.0" @@ -996,6 +1067,56 @@ dependencies = [ "tempfile", ] +[[package]] +name = "kei-router" +version = "0.1.0" +dependencies = [ + "anyhow", + "clap", + "regex", + "serde", + "serde_json", +] + +[[package]] +name = "kei-sage" +version = "0.1.0" +dependencies = [ + "anyhow", + "chrono", + "clap", + "rusqlite", + "serde", + "serde_json", + "tempfile", +] + +[[package]] +name = "kei-search-core" +version = "0.1.0" +dependencies = [ + "anyhow", + "chrono", + "clap", + "rusqlite", + "serde", + "serde_json", + "tempfile", +] + +[[package]] +name = "kei-social-store" +version = "0.1.0" +dependencies = [ + "anyhow", + "chrono", + "clap", + "rusqlite", + "serde", + "serde_json", + "tempfile", +] + [[package]] name = "kei-store" version = "0.1.0" @@ -1009,6 +1130,19 @@ dependencies = [ "toml", ] +[[package]] +name = "kei-task" +version = "0.1.0" +dependencies = [ + "anyhow", + "chrono", + "clap", + "rusqlite", + "serde", + "serde_json", + "tempfile", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -2315,6 +2449,17 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "uuid" +version = "1.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd74a9687298c6858e9b88ec8935ec45d22e8fd5e6394fa1bd4e99a87789c76" +dependencies = [ + "getrandom 0.4.2", + "js-sys", + "wasm-bindgen", +] + [[package]] name = "vcpkg" version = "0.2.15" diff --git a/_primitives/_rust/Cargo.toml b/_primitives/_rust/Cargo.toml index 7ef4681..36e041c 100644 --- a/_primitives/_rust/Cargo.toml +++ b/_primitives/_rust/Cargo.toml @@ -15,6 +15,17 @@ members = [ "kei-refactor-engine", "kei-graph-check", "kei-store", + # v0.14 LBM port — 10 new MCP-core primitives + "kei-router", + "kei-sage", + "kei-task", + "kei-chat-store", + "kei-crossdomain", + "kei-search-core", + "kei-content-store", + "kei-social-store", + "kei-curator", + "kei-auth", ] [workspace.package] diff --git a/_primitives/_rust/kei-auth/Cargo.toml b/_primitives/_rust/kei-auth/Cargo.toml new file mode 100644 index 0000000..b42f8a0 --- /dev/null +++ b/_primitives/_rust/kei-auth/Cargo.toml @@ -0,0 +1,29 @@ +[package] +name = "kei-auth" +version = "0.1.0" +edition = "2021" +rust-version = "1.75" +description = "Multi-tenant session tokens with scopes + HMAC-signed expiry (SQLite backend)." + +[[bin]] +name = "kei-auth" +path = "src/main.rs" + +[lib] +name = "kei_auth" +path = "src/lib.rs" + +[dependencies] +rusqlite = { version = "0.31", features = ["bundled"] } +clap = { version = "4", features = ["derive"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +anyhow = "1" +chrono = { version = "0.4", default-features = false, features = ["clock"] } +hmac = "0.12" +sha2 = "0.10" +base64 = "0.22" +rand = "0.8" + +[dev-dependencies] +tempfile = "3" diff --git a/_primitives/_rust/kei-auth/src/hmac.rs b/_primitives/_rust/kei-auth/src/hmac.rs new file mode 100644 index 0000000..be73569 --- /dev/null +++ b/_primitives/_rust/kei-auth/src/hmac.rs @@ -0,0 +1,25 @@ +//! HMAC-SHA256 signer for token bodies. + +use ::hmac::{Hmac, Mac}; +use anyhow::{anyhow, Result}; +use base64::Engine; +use sha2::Sha256; + +type H = Hmac; + +/// Sign `body` with `key`. Returns URL-safe base64 MAC. +pub fn sign(key: &[u8], body: &[u8]) -> String { + let mut mac = ::new_from_slice(key).expect("HMAC accepts any key size"); + mac.update(body); + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(mac.finalize().into_bytes()) +} + +/// Verify `body` against MAC. Returns Err if mismatch. +pub fn verify(key: &[u8], body: &[u8], mac_b64: &str) -> Result<()> { + let mut mac = ::new_from_slice(key).expect("HMAC accepts any key size"); + mac.update(body); + let bytes = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(mac_b64) + .map_err(|e| anyhow!("bad b64 mac: {e}"))?; + mac.verify_slice(&bytes).map_err(|_| anyhow!("hmac mismatch")) +} diff --git a/_primitives/_rust/kei-auth/src/lib.rs b/_primitives/_rust/kei-auth/src/lib.rs new file mode 100644 index 0000000..caa0506 --- /dev/null +++ b/_primitives/_rust/kei-auth/src/lib.rs @@ -0,0 +1,15 @@ +//! kei-auth — multi-tenant token auth. Replaces LBM's single LBM_MCP_TOKEN. +//! +//! Cubes: +//! - [`schema`] — SQLite tables for users + tokens +//! - [`hmac`] — HMAC-SHA256 signing helpers +//! - [`tokens`] — issue / verify / revoke / list +//! - [`scopes`] — read / write / admin enum + checks + +pub mod hmac; +pub mod schema; +pub mod scopes; +pub mod tokens; + +pub use scopes::Scope; +pub use tokens::{issue, revoke, verify, VerifyOutcome}; diff --git a/_primitives/_rust/kei-auth/src/main.rs b/_primitives/_rust/kei-auth/src/main.rs new file mode 100644 index 0000000..f6f0c6b --- /dev/null +++ b/_primitives/_rust/kei-auth/src/main.rs @@ -0,0 +1,70 @@ +//! kei-auth CLI — issue/verify/revoke. + +use clap::{Parser, Subcommand}; +use kei_auth::schema::open; +use kei_auth::scopes::Scope; +use kei_auth::tokens::{issue, revoke, verify}; +use std::path::PathBuf; +use std::process::ExitCode; +use std::str::FromStr; + +#[derive(Parser)] +#[command(name = "kei-auth", version)] +struct Cli { + #[arg(long)] db: Option, + /// HMAC signing key (env KEI_AUTH_KEY fallback). + #[arg(long)] key: Option, + #[command(subcommand)] cmd: Cmd, +} + +#[derive(Subcommand)] +enum Cmd { + Issue { #[arg(long)] user: String, + #[arg(long)] project: String, + #[arg(long, default_value = "read")] scope: String, + #[arg(long, default_value_t = 86400)] ttl: i64 }, + Verify { token: String }, + Revoke { token: String }, +} + +fn db_path(o: Option) -> PathBuf { + if let Some(p) = o { return p; } + if let Ok(e) = std::env::var("KEI_AUTH_DB") { return PathBuf::from(e); } + let home = std::env::var("HOME").unwrap_or_else(|_| ".".into()); + PathBuf::from(home).join(".claude/auth/auth.sqlite") +} + +fn key(cli_key: Option) -> anyhow::Result> { + if let Some(k) = cli_key { return Ok(k.into_bytes()); } + let k = std::env::var("KEI_AUTH_KEY") + .map_err(|_| anyhow::anyhow!("provide --key or set KEI_AUTH_KEY"))?; + Ok(k.into_bytes()) +} + +fn run() -> anyhow::Result<()> { + let cli = Cli::parse(); + let conn = open(&db_path(cli.db))?; + let k = key(cli.key)?; + match cli.cmd { + Cmd::Issue { user, project, scope, ttl } => { + let sc = Scope::from_str(&scope).map_err(|e| anyhow::anyhow!(e))?; + println!("{}", issue(&conn, &user, &project, sc, ttl, &k)?); + } + Cmd::Verify { token } => { + let out = verify(&conn, &token, &k)?; + println!("user={} project={} scope={}", out.user_id, out.project, out.scope); + } + Cmd::Revoke { token } => { + let n = revoke(&conn, &token)?; + println!("revoked {} row(s)", n); + } + } + Ok(()) +} + +fn main() -> ExitCode { + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(e) => { eprintln!("kei-auth: {e:#}"); ExitCode::from(1) } + } +} diff --git a/_primitives/_rust/kei-auth/src/schema.rs b/_primitives/_rust/kei-auth/src/schema.rs new file mode 100644 index 0000000..720750e --- /dev/null +++ b/_primitives/_rust/kei-auth/src/schema.rs @@ -0,0 +1,36 @@ +use anyhow::{Context, Result}; +use rusqlite::Connection; +use std::path::Path; + +pub fn open(path: &Path) -> Result { + if let Some(parent) = path.parent() { + let _ = std::fs::create_dir_all(parent); + } + let conn = Connection::open(path).context("open sqlite")?; + create_schema(&conn)?; + Ok(conn) +} + +pub fn open_memory() -> Result { + let conn = Connection::open_in_memory()?; + create_schema(&conn)?; + Ok(conn) +} + +pub fn create_schema(conn: &Connection) -> Result<()> { + conn.execute_batch(r#" + CREATE TABLE IF NOT EXISTS auth_tokens ( + id INTEGER PRIMARY KEY, + token_hash TEXT NOT NULL UNIQUE, + user_id TEXT NOT NULL, + project TEXT NOT NULL, + scope TEXT NOT NULL CHECK(scope IN ('read','write','admin')), + expires_at INTEGER NOT NULL, + created_at INTEGER NOT NULL, + revoked_at INTEGER DEFAULT 0 + ); + CREATE INDEX IF NOT EXISTS idx_tok_user ON auth_tokens(user_id); + CREATE INDEX IF NOT EXISTS idx_tok_project ON auth_tokens(project); + "#)?; + Ok(()) +} diff --git a/_primitives/_rust/kei-auth/src/scopes.rs b/_primitives/_rust/kei-auth/src/scopes.rs new file mode 100644 index 0000000..85068d5 --- /dev/null +++ b/_primitives/_rust/kei-auth/src/scopes.rs @@ -0,0 +1,45 @@ +use serde::{Deserialize, Serialize}; +use std::fmt; +use std::str::FromStr; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum Scope { + Read, + Write, + Admin, +} + +impl Scope { + pub fn as_str(&self) -> &'static str { + match self { Scope::Read => "read", Scope::Write => "write", Scope::Admin => "admin" } + } + + /// Admin ⊇ Write ⊇ Read. + pub fn allows(&self, required: Scope) -> bool { + use Scope::*; + match (self, required) { + (Admin, _) => true, + (Write, Read) | (Write, Write) => true, + (Read, Read) => true, + _ => false, + } + } +} + +impl fmt::Display for Scope { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + +impl FromStr for Scope { + type Err = String; + fn from_str(s: &str) -> Result { + match s { + "read" => Ok(Scope::Read), + "write" => Ok(Scope::Write), + "admin" => Ok(Scope::Admin), + _ => Err(format!("unknown scope: {s}")), + } + } +} diff --git a/_primitives/_rust/kei-auth/src/tokens.rs b/_primitives/_rust/kei-auth/src/tokens.rs new file mode 100644 index 0000000..7164fda --- /dev/null +++ b/_primitives/_rust/kei-auth/src/tokens.rs @@ -0,0 +1,110 @@ +//! Token issue / verify / revoke. +//! +//! Token layout (URL-safe, no padding): +//! `.` +//! Payload contains {tid, user_id, project, scope, expires_at}. +//! The db keeps sha256(token) to support revocation and lookup. + +use crate::hmac::{sign, verify as verify_mac}; +use crate::scopes::Scope; +use anyhow::{anyhow, Result}; +use base64::Engine; +use chrono::Utc; +use rand::RngCore; +use rusqlite::{params, Connection}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use std::str::FromStr; + +#[derive(Debug, Clone, Serialize, Deserialize)] +struct Payload { + tid: String, + user_id: String, + project: String, + scope: String, + expires_at: i64, +} + +#[derive(Debug)] +pub struct VerifyOutcome { + pub user_id: String, + pub project: String, + pub scope: Scope, +} + +/// Issue a new token. The returned string is the ONLY copy — DB stores only its sha256. +pub fn issue( + conn: &Connection, + user_id: &str, + project: &str, + scope: Scope, + ttl_secs: i64, + key: &[u8], +) -> Result { + let now = Utc::now().timestamp(); + let expires_at = now + ttl_secs; + let mut raw = [0u8; 16]; + rand::thread_rng().fill_bytes(&mut raw); + let tid = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(raw); + let payload = Payload { + tid: tid.clone(), + user_id: user_id.into(), + project: project.into(), + scope: scope.to_string(), + expires_at, + }; + let body = serde_json::to_vec(&payload)?; + let body_b64 = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(&body); + let sig = sign(key, body_b64.as_bytes()); + let token = format!("{}.{}", body_b64, sig); + let hash = sha256_hex(token.as_bytes()); + conn.execute( + "INSERT INTO auth_tokens (token_hash, user_id, project, scope, expires_at, created_at) + VALUES (?1,?2,?3,?4,?5,?6)", + params![hash, user_id, project, scope.as_str(), expires_at, now], + )?; + Ok(token) +} + +/// Verify a token: signature valid, not revoked, not expired, returns identity + scope. +pub fn verify(conn: &Connection, token: &str, key: &[u8]) -> Result { + let (body_b64, sig) = token + .split_once('.') + .ok_or_else(|| anyhow!("malformed token"))?; + verify_mac(key, body_b64.as_bytes(), sig)?; + let body = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(body_b64) + .map_err(|e| anyhow!("bad b64 payload: {e}"))?; + let p: Payload = serde_json::from_slice(&body)?; + if p.expires_at < Utc::now().timestamp() { + return Err(anyhow!("token expired")); + } + let hash = sha256_hex(token.as_bytes()); + let row: Option = conn.query_row( + "SELECT revoked_at FROM auth_tokens WHERE token_hash=?1", + params![hash], |r| r.get(0)).ok(); + match row { + None => Err(anyhow!("token unknown to server")), + Some(rev) if rev > 0 => Err(anyhow!("token revoked")), + _ => Ok(VerifyOutcome { + user_id: p.user_id, + project: p.project, + scope: Scope::from_str(&p.scope).map_err(|e| anyhow!(e))?, + }), + } +} + +/// Mark a token as revoked. Returns number of rows affected (0 = unknown). +pub fn revoke(conn: &Connection, token: &str) -> Result { + let hash = sha256_hex(token.as_bytes()); + let now = Utc::now().timestamp(); + let n = conn.execute( + "UPDATE auth_tokens SET revoked_at=?1 WHERE token_hash=?2 AND revoked_at=0", + params![now, hash], + )?; + Ok(n) +} + +fn sha256_hex(bytes: &[u8]) -> String { + format!("{:x}", Sha256::digest(bytes)) +} diff --git a/_primitives/_rust/kei-auth/tests/integration.rs b/_primitives/_rust/kei-auth/tests/integration.rs new file mode 100644 index 0000000..645d78c --- /dev/null +++ b/_primitives/_rust/kei-auth/tests/integration.rs @@ -0,0 +1,52 @@ +use kei_auth::schema::open_memory; +use kei_auth::scopes::Scope; +use kei_auth::tokens::{issue, revoke, verify}; + +const KEY: &[u8] = b"test-key-must-not-be-used-in-production"; + +#[test] +fn issue_and_verify() { + let conn = open_memory().unwrap(); + let tok = issue(&conn, "alice", "kgl", Scope::Write, 3600, KEY).unwrap(); + let out = verify(&conn, &tok, KEY).unwrap(); + assert_eq!(out.user_id, "alice"); + assert_eq!(out.project, "kgl"); + assert_eq!(out.scope, Scope::Write); +} + +#[test] +fn revoke_blocks_verify() { + let conn = open_memory().unwrap(); + let tok = issue(&conn, "bob", "x", Scope::Read, 3600, KEY).unwrap(); + assert_eq!(revoke(&conn, &tok).unwrap(), 1); + assert!(verify(&conn, &tok, KEY).is_err()); +} + +#[test] +fn expired_token_rejected() { + let conn = open_memory().unwrap(); + let tok = issue(&conn, "carol", "x", Scope::Read, -10, KEY).unwrap(); + let err = verify(&conn, &tok, KEY); + assert!(err.is_err(), "expired must fail"); +} + +#[test] +fn scope_check_admin_implies_write() { + assert!(Scope::Admin.allows(Scope::Write)); + assert!(Scope::Admin.allows(Scope::Read)); + assert!(Scope::Write.allows(Scope::Read)); + assert!(!Scope::Read.allows(Scope::Write)); + assert!(!Scope::Write.allows(Scope::Admin)); +} + +#[test] +fn tampered_token_rejected() { + let conn = open_memory().unwrap(); + let tok = issue(&conn, "dave", "x", Scope::Read, 3600, KEY).unwrap(); + let mut chars: Vec = tok.chars().collect(); + // flip one char in the signature + let last = chars.len() - 1; + chars[last] = if chars[last] == 'A' { 'B' } else { 'A' }; + let tampered: String = chars.into_iter().collect(); + assert!(verify(&conn, &tampered, KEY).is_err()); +} diff --git a/_primitives/_rust/kei-chat-store/Cargo.toml b/_primitives/_rust/kei-chat-store/Cargo.toml new file mode 100644 index 0000000..3b3e03f --- /dev/null +++ b/_primitives/_rust/kei-chat-store/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "kei-chat-store" +version = "0.1.0" +edition = "2021" +rust-version = "1.75" +description = "Session persistence for Claude conversations. Port of LBM internal/chat." + +[[bin]] +name = "kei-chat-store" +path = "src/main.rs" + +[lib] +name = "kei_chat_store" +path = "src/lib.rs" + +[dependencies] +rusqlite = { version = "0.31", features = ["bundled"] } +clap = { version = "4", features = ["derive"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +anyhow = "1" +chrono = { version = "0.4", default-features = false, features = ["clock"] } +uuid = { version = "1", features = ["v4"] } + +[dev-dependencies] +tempfile = "3" diff --git a/_primitives/_rust/kei-chat-store/src/lib.rs b/_primitives/_rust/kei-chat-store/src/lib.rs new file mode 100644 index 0000000..2773de6 --- /dev/null +++ b/_primitives/_rust/kei-chat-store/src/lib.rs @@ -0,0 +1,10 @@ +//! kei-chat-store — SQLite + FTS5 session archive for Claude chats. + +pub mod schema; +pub mod search; +pub mod sessions; +pub mod stats; +pub mod store; + +pub use sessions::{ChatMessage, ChatSession}; +pub use store::Store; diff --git a/_primitives/_rust/kei-chat-store/src/main.rs b/_primitives/_rust/kei-chat-store/src/main.rs new file mode 100644 index 0000000..b300396 --- /dev/null +++ b/_primitives/_rust/kei-chat-store/src/main.rs @@ -0,0 +1,77 @@ +//! kei-chat-store CLI. + +use clap::{Parser, Subcommand}; +use kei_chat_store::search::search; +use kei_chat_store::sessions::{archive_session, save_message, start_session, ChatMessage}; +use kei_chat_store::stats::stats; +use kei_chat_store::Store; +use std::path::PathBuf; +use std::process::ExitCode; + +#[derive(Parser)] +#[command(name = "kei-chat-store", version)] +struct Cli { + #[arg(long)] db: Option, + #[command(subcommand)] cmd: Cmd, +} + +#[derive(Subcommand)] +enum Cmd { + Start { #[arg(long)] project: String, + #[arg(long, default_value = "")] title: String, + #[arg(long, default_value = "")] model: String }, + Save { #[arg(long)] session_id: String, + #[arg(long)] role: String, + content: String, + #[arg(long, default_value_t = 0)] tokens_in: i64, + #[arg(long, default_value_t = 0)] tokens_out: i64, + #[arg(long, default_value_t = 0.0)] cost: f64 }, + Search { query: String, #[arg(long, default_value_t = 20)] limit: i64 }, + Archive { session_id: String }, + Stats, +} + +fn db_path(o: Option) -> PathBuf { + if let Some(p) = o { return p; } + if let Ok(e) = std::env::var("KEI_CHAT_DB") { return PathBuf::from(e); } + let home = std::env::var("HOME").unwrap_or_else(|_| ".".into()); + PathBuf::from(home).join(".claude/chat/chat.sqlite") +} + +fn run() -> anyhow::Result<()> { + let cli = Cli::parse(); + let s = Store::open(&db_path(cli.db))?; + match cli.cmd { + Cmd::Start { project, title, model } => { + println!("{}", start_session(&s, &project, &title, &model)?); + } + Cmd::Save { session_id, role, content, tokens_in, tokens_out, cost } => { + let id = save_message(&s, &ChatMessage { + session_id, role, content, tokens_in, tokens_out, cost, + ..Default::default() + })?; + println!("{}", id); + } + Cmd::Search { query, limit } => { + for m in search(&s, &query, limit)? { + println!("{}\t{}\t{}", m.id, m.role, m.content); + } + } + Cmd::Archive { session_id } => { + archive_session(&s, &session_id)?; + println!("archived {}", session_id); + } + Cmd::Stats => { + let st = stats(&s)?; + println!("{}", serde_json::to_string_pretty(&st)?); + } + } + Ok(()) +} + +fn main() -> ExitCode { + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(e) => { eprintln!("kei-chat-store: {e:#}"); ExitCode::from(1) } + } +} diff --git a/_primitives/_rust/kei-chat-store/src/schema.rs b/_primitives/_rust/kei-chat-store/src/schema.rs new file mode 100644 index 0000000..3d577f0 --- /dev/null +++ b/_primitives/_rust/kei-chat-store/src/schema.rs @@ -0,0 +1,40 @@ +//! Chat SQLite schema. + +use rusqlite::{Connection, Result}; + +pub fn create_schema(conn: &Connection) -> Result<()> { + conn.execute_batch(r#" + CREATE TABLE IF NOT EXISTS chat_sessions ( + id TEXT PRIMARY KEY, + project TEXT NOT NULL, + title TEXT DEFAULT '', + model TEXT DEFAULT '', + status TEXT DEFAULT 'active', + message_count INTEGER DEFAULT 0, + total_tokens INTEGER DEFAULT 0, + total_cost REAL DEFAULT 0.0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_cs_project ON chat_sessions(project); + CREATE INDEX IF NOT EXISTS idx_cs_status ON chat_sessions(status); + + CREATE TABLE IF NOT EXISTS chat_messages ( + id INTEGER PRIMARY KEY, + session_id TEXT NOT NULL REFERENCES chat_sessions(id) ON DELETE CASCADE, + role TEXT NOT NULL, + content TEXT NOT NULL, + tokens_in INTEGER DEFAULT 0, + tokens_out INTEGER DEFAULT 0, + cost REAL DEFAULT 0.0, + created_at INTEGER NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_cm_session ON chat_messages(session_id); + "#)?; + conn.execute_batch(r#" + CREATE VIRTUAL TABLE IF NOT EXISTS fts_chat + USING fts5(message_id UNINDEXED, session_id UNINDEXED, content, + tokenize='porter unicode61'); + "#)?; + Ok(()) +} diff --git a/_primitives/_rust/kei-chat-store/src/search.rs b/_primitives/_rust/kei-chat-store/src/search.rs new file mode 100644 index 0000000..362c619 --- /dev/null +++ b/_primitives/_rust/kei-chat-store/src/search.rs @@ -0,0 +1,26 @@ +//! FTS over messages. + +use crate::sessions::ChatMessage; +use crate::store::Store; +use anyhow::Result; +use rusqlite::params; + +pub fn search(store: &Store, query: &str, limit: i64) -> Result> { + let lim = if limit <= 0 { 20 } else { limit }; + let mut stmt = store.conn().prepare( + "SELECT m.id, m.session_id, m.role, m.content, m.tokens_in, m.tokens_out, + m.cost, m.created_at + FROM fts_chat f + JOIN chat_messages m ON m.id = f.message_id + WHERE fts_chat MATCH ?1 ORDER BY rank LIMIT ?2", + )?; + let rows = stmt.query_map(params![query, lim], |r| { + Ok(ChatMessage { + id: r.get(0)?, session_id: r.get(1)?, role: r.get(2)?, content: r.get(3)?, + tokens_in: r.get(4)?, tokens_out: r.get(5)?, cost: r.get(6)?, created_at: r.get(7)?, + }) + })?; + let mut out = Vec::new(); + for r in rows { out.push(r?); } + Ok(out) +} diff --git a/_primitives/_rust/kei-chat-store/src/sessions.rs b/_primitives/_rust/kei-chat-store/src/sessions.rs new file mode 100644 index 0000000..7f89e6a --- /dev/null +++ b/_primitives/_rust/kei-chat-store/src/sessions.rs @@ -0,0 +1,94 @@ +//! Session + message operations. + +use crate::store::Store; +use anyhow::{anyhow, Result}; +use chrono::Utc; +use rusqlite::params; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct ChatSession { + pub id: String, + pub project: String, + pub title: String, + pub model: String, + pub status: String, + pub message_count: i64, + pub total_tokens: i64, + pub total_cost: f64, + pub created_at: i64, + pub updated_at: i64, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct ChatMessage { + pub id: i64, + pub session_id: String, + pub role: String, + pub content: String, + pub tokens_in: i64, + pub tokens_out: i64, + pub cost: f64, + pub created_at: i64, +} + +pub fn start_session(store: &Store, project: &str, title: &str, model: &str) -> Result { + let id = uuid::Uuid::new_v4().to_string(); + let now = Utc::now().timestamp(); + store.conn().execute( + "INSERT INTO chat_sessions (id, project, title, model, status, created_at, updated_at) + VALUES (?1,?2,?3,?4,'active',?5,?5)", + params![id, project, title, model, now], + )?; + Ok(id) +} + +pub fn save_message(store: &Store, msg: &ChatMessage) -> Result { + let now = Utc::now().timestamp(); + let created = if msg.created_at == 0 { now } else { msg.created_at }; + store.conn().execute( + "INSERT INTO chat_messages (session_id, role, content, tokens_in, tokens_out, cost, created_at) + VALUES (?1,?2,?3,?4,?5,?6,?7)", + params![msg.session_id, msg.role, msg.content, msg.tokens_in, + msg.tokens_out, msg.cost, created], + )?; + let id = store.conn().last_insert_rowid(); + store.conn().execute( + "INSERT INTO fts_chat (message_id, session_id, content) VALUES (?1,?2,?3)", + params![id, msg.session_id, msg.content], + )?; + store.conn().execute( + "UPDATE chat_sessions SET message_count = message_count + 1, + total_tokens = total_tokens + ?1, total_cost = total_cost + ?2, + updated_at = ?3 WHERE id = ?4", + params![msg.tokens_in + msg.tokens_out, msg.cost, now, msg.session_id], + )?; + Ok(id) +} + +pub fn archive_session(store: &Store, session_id: &str) -> Result<()> { + let n = store.conn().execute( + "UPDATE chat_sessions SET status='archived', updated_at=?1 WHERE id=?2", + params![Utc::now().timestamp(), session_id], + )?; + if n == 0 { + return Err(anyhow!("session {session_id} not found")); + } + Ok(()) +} + +pub fn get_session(store: &Store, id: &str) -> Result> { + let mut stmt = store.conn().prepare( + "SELECT id, project, title, model, status, message_count, total_tokens, + total_cost, created_at, updated_at FROM chat_sessions WHERE id=?1", + )?; + let mut rows = stmt.query(params![id])?; + if let Some(r) = rows.next()? { + return Ok(Some(ChatSession { + id: r.get(0)?, project: r.get(1)?, title: r.get(2)?, model: r.get(3)?, + status: r.get(4)?, message_count: r.get(5)?, total_tokens: r.get(6)?, + total_cost: r.get(7)?, created_at: r.get(8)?, updated_at: r.get(9)?, + })); + } + Ok(None) +} diff --git a/_primitives/_rust/kei-chat-store/src/stats.rs b/_primitives/_rust/kei-chat-store/src/stats.rs new file mode 100644 index 0000000..c2d4ff3 --- /dev/null +++ b/_primitives/_rust/kei-chat-store/src/stats.rs @@ -0,0 +1,32 @@ +//! Aggregate chat stats. + +use crate::store::Store; +use anyhow::Result; +use serde::Serialize; + +#[derive(Debug, Default, Serialize)] +pub struct Stats { + pub total_sessions: i64, + pub active_sessions: i64, + pub archived_sessions: i64, + pub total_messages: i64, + pub total_tokens: i64, + pub total_cost: f64, +} + +pub fn stats(store: &Store) -> Result { + let mut s = Stats::default(); + s.total_sessions = store.conn() + .query_row("SELECT COUNT(*) FROM chat_sessions", [], |r| r.get(0))?; + s.active_sessions = store.conn() + .query_row("SELECT COUNT(*) FROM chat_sessions WHERE status='active'", [], |r| r.get(0))?; + s.archived_sessions = store.conn() + .query_row("SELECT COUNT(*) FROM chat_sessions WHERE status='archived'", [], |r| r.get(0))?; + s.total_messages = store.conn() + .query_row("SELECT COUNT(*) FROM chat_messages", [], |r| r.get(0))?; + s.total_tokens = store.conn() + .query_row("SELECT COALESCE(SUM(total_tokens),0) FROM chat_sessions", [], |r| r.get(0))?; + s.total_cost = store.conn() + .query_row("SELECT COALESCE(SUM(total_cost),0) FROM chat_sessions", [], |r| r.get(0))?; + Ok(s) +} diff --git a/_primitives/_rust/kei-chat-store/src/store.rs b/_primitives/_rust/kei-chat-store/src/store.rs new file mode 100644 index 0000000..1983471 --- /dev/null +++ b/_primitives/_rust/kei-chat-store/src/store.rs @@ -0,0 +1,30 @@ +//! Store open/close helper. + +use crate::schema::create_schema; +use anyhow::{Context, Result}; +use rusqlite::Connection; +use std::path::Path; + +pub struct Store { + conn: Connection, +} + +impl Store { + pub fn open(path: &Path) -> Result { + if let Some(parent) = path.parent() { + let _ = std::fs::create_dir_all(parent); + } + let conn = Connection::open(path).context("open sqlite")?; + conn.pragma_update(None, "journal_mode", "WAL").ok(); + create_schema(&conn)?; + Ok(Self { conn }) + } + + pub fn open_memory() -> Result { + let conn = Connection::open_in_memory()?; + create_schema(&conn)?; + Ok(Self { conn }) + } + + pub fn conn(&self) -> &Connection { &self.conn } +} diff --git a/_primitives/_rust/kei-chat-store/tests/integration.rs b/_primitives/_rust/kei-chat-store/tests/integration.rs new file mode 100644 index 0000000..9958e76 --- /dev/null +++ b/_primitives/_rust/kei-chat-store/tests/integration.rs @@ -0,0 +1,59 @@ +use kei_chat_store::search::search; +use kei_chat_store::sessions::{archive_session, get_session, save_message, start_session, ChatMessage}; +use kei_chat_store::stats::stats; +use kei_chat_store::Store; + +fn mk() -> Store { Store::open_memory().unwrap() } + +#[test] +fn save_and_retrieve() { + let s = mk(); + let sid = start_session(&s, "demo", "t", "claude-opus-4").unwrap(); + save_message(&s, &ChatMessage { + session_id: sid.clone(), role: "user".into(), + content: "hello world".into(), tokens_in: 3, tokens_out: 0, cost: 0.001, + ..Default::default() + }).unwrap(); + let sess = get_session(&s, &sid).unwrap().unwrap(); + assert_eq!(sess.message_count, 1); + assert_eq!(sess.total_tokens, 3); +} + +#[test] +fn fts_search_finds_message() { + let s = mk(); + let sid = start_session(&s, "demo", "", "").unwrap(); + save_message(&s, &ChatMessage { + session_id: sid, role: "user".into(), + content: "rust async tokio bench".into(), + ..Default::default() + }).unwrap(); + let hits = search(&s, "tokio", 10).unwrap(); + assert_eq!(hits.len(), 1); +} + +#[test] +fn archive_session_works() { + let s = mk(); + let sid = start_session(&s, "p", "", "").unwrap(); + archive_session(&s, &sid).unwrap(); + let sess = get_session(&s, &sid).unwrap().unwrap(); + assert_eq!(sess.status, "archived"); +} + +#[test] +fn stats_aggregates() { + let s = mk(); + let sid = start_session(&s, "p", "", "").unwrap(); + for _ in 0..3 { + save_message(&s, &ChatMessage { + session_id: sid.clone(), role: "user".into(), + content: "x".into(), tokens_in: 5, tokens_out: 5, cost: 0.01, + ..Default::default() + }).unwrap(); + } + let st = stats(&s).unwrap(); + assert_eq!(st.total_sessions, 1); + assert_eq!(st.total_messages, 3); + assert_eq!(st.total_tokens, 30); +} diff --git a/_primitives/_rust/kei-content-store/Cargo.toml b/_primitives/_rust/kei-content-store/Cargo.toml new file mode 100644 index 0000000..c5182b3 --- /dev/null +++ b/_primitives/_rust/kei-content-store/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "kei-content-store" +version = "0.1.0" +edition = "2021" +rust-version = "1.75" +description = "Asset + prompt + campaign registry. Port of LBM internal/content." + +[[bin]] +name = "kei-content-store" +path = "src/main.rs" + +[lib] +name = "kei_content_store" +path = "src/lib.rs" + +[dependencies] +rusqlite = { version = "0.31", features = ["bundled"] } +clap = { version = "4", features = ["derive"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +anyhow = "1" +chrono = { version = "0.4", default-features = false, features = ["clock"] } +sha2 = "0.10" + +[dev-dependencies] +tempfile = "3" diff --git a/_primitives/_rust/kei-content-store/src/assets.rs b/_primitives/_rust/kei-content-store/src/assets.rs new file mode 100644 index 0000000..d08a998 --- /dev/null +++ b/_primitives/_rust/kei-content-store/src/assets.rs @@ -0,0 +1,52 @@ +use crate::store::Store; +use anyhow::Result; +use chrono::Utc; +use rusqlite::params; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct Asset { + pub id: i64, + pub unit_type: String, + pub title: String, + pub content: String, + pub media_type: String, + pub file_path: String, + pub file_hash: String, + pub provider: String, + pub cost_cents: i64, + pub parent_id: i64, + pub created_at: i64, + pub updated_at: i64, +} + +pub fn register_asset(store: &Store, a: &Asset) -> Result { + let now = Utc::now().timestamp(); + let ut = if a.unit_type.is_empty() { "asset" } else { &a.unit_type }; + store.conn().execute( + "INSERT INTO content_units (unit_type, title, content, media_type, + file_path, file_hash, provider, cost_cents, parent_id, created_at, updated_at) + VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?10)", + params![ut, a.title, a.content, a.media_type, a.file_path, + a.file_hash, a.provider, a.cost_cents, a.parent_id, now], + )?; + Ok(store.conn().last_insert_rowid()) +} + +pub fn get_asset(store: &Store, id: i64) -> Result> { + let mut stmt = store.conn().prepare( + "SELECT id, unit_type, title, content, media_type, file_path, file_hash, + provider, cost_cents, parent_id, created_at, updated_at + FROM content_units WHERE id=?1", + )?; + let mut rows = stmt.query(params![id])?; + if let Some(r) = rows.next()? { + return Ok(Some(Asset { + id: r.get(0)?, unit_type: r.get(1)?, title: r.get(2)?, content: r.get(3)?, + media_type: r.get(4)?, file_path: r.get(5)?, file_hash: r.get(6)?, + provider: r.get(7)?, cost_cents: r.get(8)?, parent_id: r.get(9)?, + created_at: r.get(10)?, updated_at: r.get(11)?, + })); + } + Ok(None) +} diff --git a/_primitives/_rust/kei-content-store/src/campaigns.rs b/_primitives/_rust/kei-content-store/src/campaigns.rs new file mode 100644 index 0000000..207ef83 --- /dev/null +++ b/_primitives/_rust/kei-content-store/src/campaigns.rs @@ -0,0 +1,31 @@ +use crate::store::Store; +use anyhow::Result; +use chrono::Utc; +use rusqlite::params; + +pub fn create_campaign(store: &Store, name: &str, description: &str) -> Result { + let now = Utc::now().timestamp(); + store.conn().execute( + "INSERT INTO campaigns (name, description, created_at) VALUES (?1,?2,?3)", + params![name, description, now], + )?; + Ok(store.conn().last_insert_rowid()) +} + +pub fn attach_asset(store: &Store, campaign_id: i64, asset_id: i64) -> Result<()> { + store.conn().execute( + "INSERT OR IGNORE INTO campaign_assets (campaign_id, asset_id) VALUES (?1,?2)", + params![campaign_id, asset_id], + )?; + Ok(()) +} + +pub fn campaign_assets(store: &Store, campaign_id: i64) -> Result> { + let mut stmt = store.conn().prepare( + "SELECT asset_id FROM campaign_assets WHERE campaign_id=?1" + )?; + let rows = stmt.query_map(params![campaign_id], |r| r.get::<_, i64>(0))?; + let mut out = Vec::new(); + for r in rows { out.push(r?); } + Ok(out) +} diff --git a/_primitives/_rust/kei-content-store/src/lib.rs b/_primitives/_rust/kei-content-store/src/lib.rs new file mode 100644 index 0000000..cf91d09 --- /dev/null +++ b/_primitives/_rust/kei-content-store/src/lib.rs @@ -0,0 +1,11 @@ +//! kei-content-store — assets, prompts, campaigns. + +pub mod assets; +pub mod campaigns; +pub mod prompts; +pub mod schema; +pub mod store; + +pub use assets::Asset; +pub use prompts::Prompt; +pub use store::Store; diff --git a/_primitives/_rust/kei-content-store/src/main.rs b/_primitives/_rust/kei-content-store/src/main.rs new file mode 100644 index 0000000..1bd6101 --- /dev/null +++ b/_primitives/_rust/kei-content-store/src/main.rs @@ -0,0 +1,76 @@ +use clap::{Parser, Subcommand}; +use kei_content_store::assets::{register_asset, Asset}; +use kei_content_store::campaigns::{attach_asset, create_campaign}; +use kei_content_store::prompts::{history, register_prompt, Prompt}; +use kei_content_store::Store; +use std::path::PathBuf; +use std::process::ExitCode; + +#[derive(Parser)] +#[command(name = "kei-content-store", version)] +struct Cli { + #[arg(long)] db: Option, + #[command(subcommand)] cmd: Cmd, +} + +#[derive(Subcommand)] +enum Cmd { + RegisterAsset { title: String, + #[arg(long, default_value = "")] file_path: String, + #[arg(long, default_value = "")] media_type: String, + #[arg(long, default_value = "")] provider: String }, + RegisterPrompt { prompt_text: String, + #[arg(long, default_value = "")] model: String, + #[arg(long, default_value = "")] prompt_type: String }, + CreateCampaign { name: String, #[arg(long, default_value = "")] description: String }, + AttachAsset { campaign_id: i64, asset_id: i64 }, + PromptHistory { prompt_id: i64 }, +} + +fn db_path(o: Option) -> PathBuf { + if let Some(p) = o { return p; } + if let Ok(e) = std::env::var("KEI_CONTENT_DB") { return PathBuf::from(e); } + let home = std::env::var("HOME").unwrap_or_else(|_| ".".into()); + PathBuf::from(home).join(".claude/content/content.sqlite") +} + +fn run() -> anyhow::Result<()> { + let cli = Cli::parse(); + let s = Store::open(&db_path(cli.db))?; + match cli.cmd { + Cmd::RegisterAsset { title, file_path, media_type, provider } => { + let id = register_asset(&s, &Asset { + title, file_path, media_type, provider, unit_type: "asset".into(), + ..Default::default() + })?; + println!("{}", id); + } + Cmd::RegisterPrompt { prompt_text, model, prompt_type } => { + let id = register_prompt(&s, &Prompt { + prompt_text, model, prompt_type, ..Default::default() + })?; + println!("{}", id); + } + Cmd::CreateCampaign { name, description } => { + let id = create_campaign(&s, &name, &description)?; + println!("{}", id); + } + Cmd::AttachAsset { campaign_id, asset_id } => { + attach_asset(&s, campaign_id, asset_id)?; + println!("attached {} to campaign {}", asset_id, campaign_id); + } + Cmd::PromptHistory { prompt_id } => { + for p in history(&s, prompt_id)? { + println!("{}\t{}\t{}", p.id, p.version, p.prompt_text); + } + } + } + Ok(()) +} + +fn main() -> ExitCode { + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(e) => { eprintln!("kei-content-store: {e:#}"); ExitCode::from(1) } + } +} diff --git a/_primitives/_rust/kei-content-store/src/prompts.rs b/_primitives/_rust/kei-content-store/src/prompts.rs new file mode 100644 index 0000000..fbaf2c3 --- /dev/null +++ b/_primitives/_rust/kei-content-store/src/prompts.rs @@ -0,0 +1,57 @@ +use crate::store::Store; +use anyhow::Result; +use chrono::Utc; +use rusqlite::params; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct Prompt { + pub id: i64, + pub prompt_text: String, + pub prompt_hash: String, + pub prompt_type: String, + pub model: String, + pub version: i64, + pub parent_id: i64, + pub created_at: i64, +} + +pub fn register_prompt(store: &Store, p: &Prompt) -> Result { + let now = Utc::now().timestamp(); + let hash = hash_prompt(&p.prompt_text); + store.conn().execute( + "INSERT OR IGNORE INTO prompts + (prompt_text, prompt_hash, prompt_type, model, version, parent_id, created_at) + VALUES (?1,?2,?3,?4,?5,?6,?7)", + params![p.prompt_text, hash, p.prompt_type, p.model, + if p.version == 0 { 1 } else { p.version }, p.parent_id, now], + )?; + let id: i64 = store.conn().query_row( + "SELECT id FROM prompts WHERE prompt_hash=?1 AND model=?2", + params![hash, p.model], |r| r.get(0))?; + Ok(id) +} + +pub fn history(store: &Store, parent_id: i64) -> Result> { + let mut stmt = store.conn().prepare( + "SELECT id, prompt_text, prompt_hash, prompt_type, model, version, + parent_id, created_at + FROM prompts WHERE parent_id=?1 OR id=?1 ORDER BY created_at", + )?; + let rows = stmt.query_map(params![parent_id], |r| { + Ok(Prompt { + id: r.get(0)?, prompt_text: r.get(1)?, prompt_hash: r.get(2)?, + prompt_type: r.get(3)?, model: r.get(4)?, version: r.get(5)?, + parent_id: r.get(6)?, created_at: r.get(7)?, + }) + })?; + let mut out = Vec::new(); + for r in rows { out.push(r?); } + Ok(out) +} + +fn hash_prompt(s: &str) -> String { + let d = Sha256::digest(s.as_bytes()); + format!("{:x}", d) +} diff --git a/_primitives/_rust/kei-content-store/src/schema.rs b/_primitives/_rust/kei-content-store/src/schema.rs new file mode 100644 index 0000000..970deaa --- /dev/null +++ b/_primitives/_rust/kei-content-store/src/schema.rs @@ -0,0 +1,49 @@ +use rusqlite::{Connection, Result}; + +pub fn create_schema(conn: &Connection) -> Result<()> { + conn.execute_batch(r#" + CREATE TABLE IF NOT EXISTS content_units ( + id INTEGER PRIMARY KEY, + unit_type TEXT NOT NULL, + title TEXT NOT NULL, + content TEXT DEFAULT '', + media_type TEXT DEFAULT '', + file_path TEXT DEFAULT '', + file_hash TEXT DEFAULT '', + provider TEXT DEFAULT '', + cost_cents INTEGER DEFAULT 0, + parent_id INTEGER DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_cu_type ON content_units(unit_type); + CREATE INDEX IF NOT EXISTS idx_cu_hash ON content_units(file_hash) WHERE file_hash != ''; + + CREATE TABLE IF NOT EXISTS prompts ( + id INTEGER PRIMARY KEY, + prompt_text TEXT NOT NULL, + prompt_hash TEXT NOT NULL, + prompt_type TEXT DEFAULT '', + model TEXT DEFAULT '', + version INTEGER DEFAULT 1, + parent_id INTEGER DEFAULT 0, + created_at INTEGER NOT NULL, + UNIQUE(prompt_hash, model) + ); + + CREATE TABLE IF NOT EXISTS campaigns ( + id INTEGER PRIMARY KEY, + name TEXT NOT NULL, + description TEXT DEFAULT '', + status TEXT DEFAULT 'draft', + created_at INTEGER NOT NULL + ); + + CREATE TABLE IF NOT EXISTS campaign_assets ( + campaign_id INTEGER NOT NULL, + asset_id INTEGER NOT NULL, + PRIMARY KEY(campaign_id, asset_id) + ); + "#)?; + Ok(()) +} diff --git a/_primitives/_rust/kei-content-store/src/store.rs b/_primitives/_rust/kei-content-store/src/store.rs new file mode 100644 index 0000000..16fe2c4 --- /dev/null +++ b/_primitives/_rust/kei-content-store/src/store.rs @@ -0,0 +1,24 @@ +use crate::schema::create_schema; +use anyhow::{Context, Result}; +use rusqlite::Connection; +use std::path::Path; + +pub struct Store { conn: Connection } + +impl Store { + pub fn open(path: &Path) -> Result { + if let Some(parent) = path.parent() { let _ = std::fs::create_dir_all(parent); } + let conn = Connection::open(path).context("open sqlite")?; + conn.pragma_update(None, "journal_mode", "WAL").ok(); + create_schema(&conn)?; + Ok(Self { conn }) + } + + pub fn open_memory() -> Result { + let conn = Connection::open_in_memory()?; + create_schema(&conn)?; + Ok(Self { conn }) + } + + pub fn conn(&self) -> &Connection { &self.conn } +} diff --git a/_primitives/_rust/kei-content-store/tests/integration.rs b/_primitives/_rust/kei-content-store/tests/integration.rs new file mode 100644 index 0000000..60a166e --- /dev/null +++ b/_primitives/_rust/kei-content-store/tests/integration.rs @@ -0,0 +1,48 @@ +use kei_content_store::assets::{get_asset, register_asset, Asset}; +use kei_content_store::campaigns::{attach_asset, campaign_assets, create_campaign}; +use kei_content_store::prompts::{register_prompt, Prompt}; +use kei_content_store::Store; + +fn mk() -> Store { Store::open_memory().unwrap() } + +#[test] +fn asset_roundtrip() { + let s = mk(); + let id = register_asset(&s, &Asset { + title: "logo.png".into(), media_type: "image/png".into(), + ..Default::default() + }).unwrap(); + let a = get_asset(&s, id).unwrap().unwrap(); + assert_eq!(a.title, "logo.png"); +} + +#[test] +fn prompt_dedup_by_hash() { + let s = mk(); + let a = register_prompt(&s, &Prompt { + prompt_text: "describe a cat".into(), model: "dall-e-3".into(), + ..Default::default() + }).unwrap(); + let b = register_prompt(&s, &Prompt { + prompt_text: "describe a cat".into(), model: "dall-e-3".into(), + ..Default::default() + }).unwrap(); + assert_eq!(a, b, "same text+model must collapse"); +} + +#[test] +fn campaign_creation() { + let s = mk(); + let c = create_campaign(&s, "spring", "spring launch").unwrap(); + assert!(c > 0); +} + +#[test] +fn campaign_asset_attach() { + let s = mk(); + let c = create_campaign(&s, "launch", "").unwrap(); + let a = register_asset(&s, &Asset { + title: "hero.mp4".into(), ..Default::default() }).unwrap(); + attach_asset(&s, c, a).unwrap(); + assert_eq!(campaign_assets(&s, c).unwrap(), vec![a]); +} diff --git a/_primitives/_rust/kei-crossdomain/Cargo.toml b/_primitives/_rust/kei-crossdomain/Cargo.toml new file mode 100644 index 0000000..88e5194 --- /dev/null +++ b/_primitives/_rust/kei-crossdomain/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "kei-crossdomain" +version = "0.1.0" +edition = "2021" +rust-version = "1.75" +description = "Typed-edge cross-domain store. Port of LBM internal/crossdomain." + +[[bin]] +name = "kei-crossdomain" +path = "src/main.rs" + +[lib] +name = "kei_crossdomain" +path = "src/lib.rs" + +[dependencies] +rusqlite = { version = "0.31", features = ["bundled"] } +clap = { version = "4", features = ["derive"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +anyhow = "1" +chrono = { version = "0.4", default-features = false, features = ["clock"] } + +[dev-dependencies] +tempfile = "3" diff --git a/_primitives/_rust/kei-crossdomain/src/auto_link.rs b/_primitives/_rust/kei-crossdomain/src/auto_link.rs new file mode 100644 index 0000000..23fd47e --- /dev/null +++ b/_primitives/_rust/kei-crossdomain/src/auto_link.rs @@ -0,0 +1,56 @@ +//! Auto-link heuristic — proposes edges based on URI-name component matching. +//! No-ML: intersect the last path segments (case-insensitive, normalized). + +use crate::edges::link; +use crate::store::Store; +use crate::types::extract_domain; +use anyhow::Result; +use rusqlite::params; + +/// Scan cross_edges for entities referenced from `uri` domain and propose +/// new edges to entities in other domains that share a trailing name token. +pub fn auto_link(store: &Store, uri: &str) -> Result { + let tail = tail_token(uri); + if tail.is_empty() { + return Ok(0); + } + let src_domain = extract_domain(uri); + let mut candidates: Vec = Vec::new(); + let mut stmt = store.conn().prepare( + "SELECT DISTINCT to_uri FROM cross_edges + UNION SELECT DISTINCT from_uri FROM cross_edges", + )?; + let rows = stmt.query_map([], |r| r.get::<_, String>(0))?; + for row in rows { + let u = row?; + if u == uri { + continue; + } + if extract_domain(&u) == src_domain { + continue; // only cross-domain proposals + } + if tail_token(&u).eq_ignore_ascii_case(&tail) { + candidates.push(u); + } + } + let mut added = 0; + for c in &candidates { + if edge_exists(store, uri, c)? { + continue; + } + link(store, uri, c, "auto_related", 0.5, "E5")?; + added += 1; + } + Ok(added) +} + +fn edge_exists(store: &Store, from: &str, to: &str) -> Result { + let n: i64 = store.conn().query_row( + "SELECT COUNT(*) FROM cross_edges WHERE from_uri=?1 AND to_uri=?2", + params![from, to], |r| r.get(0))?; + Ok(n > 0) +} + +fn tail_token(uri: &str) -> String { + uri.rsplit('/').next().unwrap_or("").to_lowercase() +} diff --git a/_primitives/_rust/kei-crossdomain/src/bfs.rs b/_primitives/_rust/kei-crossdomain/src/bfs.rs new file mode 100644 index 0000000..c56dd9f --- /dev/null +++ b/_primitives/_rust/kei-crossdomain/src/bfs.rs @@ -0,0 +1,44 @@ +use crate::store::Store; +use anyhow::Result; +use rusqlite::params; +use serde::Serialize; +use std::collections::{HashSet, VecDeque}; + +const MAX_DEPTH: i64 = 5; + +#[derive(Debug, Clone, Serialize)] +pub struct Reached { + pub uri: String, + pub edge_type: String, + pub depth: i64, +} + +pub fn bfs(store: &Store, start: &str, depth: i64) -> Result> { + let d = clamp(depth); + let mut seen: HashSet = HashSet::new(); + seen.insert(start.into()); + let mut q: VecDeque<(String, i64)> = VecDeque::new(); + q.push_back((start.into(), 0)); + let mut out = Vec::new(); + while let Some((uri, cur)) = q.pop_front() { + if cur >= d { continue; } + let mut stmt = store.conn().prepare( + "SELECT to_uri, edge_type FROM cross_edges WHERE from_uri=?1" + )?; + let rows = stmt.query_map(params![uri], |r| { + Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?)) + })?; + for row in rows { + let (to, et) = row?; + if seen.contains(&to) { continue; } + seen.insert(to.clone()); + out.push(Reached { uri: to.clone(), edge_type: et, depth: cur + 1 }); + q.push_back((to, cur + 1)); + } + } + Ok(out) +} + +fn clamp(d: i64) -> i64 { + if d <= 0 { 2 } else if d > MAX_DEPTH { MAX_DEPTH } else { d } +} diff --git a/_primitives/_rust/kei-crossdomain/src/edges.rs b/_primitives/_rust/kei-crossdomain/src/edges.rs new file mode 100644 index 0000000..813d388 --- /dev/null +++ b/_primitives/_rust/kei-crossdomain/src/edges.rs @@ -0,0 +1,51 @@ +use crate::store::Store; +use crate::types::CrossEdge; +use anyhow::Result; +use chrono::Utc; +use rusqlite::params; + +pub fn link(store: &Store, from: &str, to: &str, edge_type: &str, + weight: f64, evidence: &str) -> Result { + let now = Utc::now().timestamp(); + store.conn().execute( + "INSERT OR IGNORE INTO cross_edges (from_uri, to_uri, edge_type, weight, evidence, created_at) + VALUES (?1,?2,?3,?4,?5,?6)", + params![from, to, edge_type, weight, evidence, now], + )?; + Ok(store.conn().last_insert_rowid()) +} + +pub fn unlink(store: &Store, from: &str, to: &str, edge_type: &str) -> Result { + let n = store.conn().execute( + "DELETE FROM cross_edges WHERE from_uri=?1 AND to_uri=?2 AND edge_type=?3", + params![from, to, edge_type], + )?; + Ok(n) +} + +pub fn query_edges(store: &Store, uri: &str) -> Result> { + let mut stmt = store.conn().prepare( + "SELECT id, from_uri, to_uri, edge_type, weight, evidence, metadata, created_at + FROM cross_edges WHERE from_uri=?1 OR to_uri=?1", + )?; + let rows = stmt.query_map(params![uri], |r| { + Ok(CrossEdge { + id: r.get(0)?, from_uri: r.get(1)?, to_uri: r.get(2)?, + edge_type: r.get(3)?, weight: r.get(4)?, evidence: r.get(5)?, + metadata: r.get(6)?, created_at: r.get(7)?, + }) + })?; + let mut out = Vec::new(); + for r in rows { out.push(r?); } + Ok(out) +} + +pub fn count_by_type(store: &Store) -> Result> { + let mut stmt = store.conn().prepare( + "SELECT edge_type, COUNT(*) FROM cross_edges GROUP BY edge_type", + )?; + let rows = stmt.query_map([], |r| Ok((r.get::<_, String>(0)?, r.get::<_, i64>(1)?)))?; + let mut out = Vec::new(); + for r in rows { out.push(r?); } + Ok(out) +} diff --git a/_primitives/_rust/kei-crossdomain/src/lib.rs b/_primitives/_rust/kei-crossdomain/src/lib.rs new file mode 100644 index 0000000..80bbdd1 --- /dev/null +++ b/_primitives/_rust/kei-crossdomain/src/lib.rs @@ -0,0 +1,11 @@ +//! kei-crossdomain — SQLite store for domain-to-domain typed edges + BFS. + +pub mod auto_link; +pub mod bfs; +pub mod edges; +pub mod schema; +pub mod store; +pub mod types; + +pub use store::Store; +pub use types::CrossEdge; diff --git a/_primitives/_rust/kei-crossdomain/src/main.rs b/_primitives/_rust/kei-crossdomain/src/main.rs new file mode 100644 index 0000000..41c51eb --- /dev/null +++ b/_primitives/_rust/kei-crossdomain/src/main.rs @@ -0,0 +1,77 @@ +use clap::{Parser, Subcommand}; +use kei_crossdomain::auto_link::auto_link; +use kei_crossdomain::bfs::bfs; +use kei_crossdomain::edges::{count_by_type, link, query_edges, unlink}; +use kei_crossdomain::Store; +use std::path::PathBuf; +use std::process::ExitCode; + +#[derive(Parser)] +#[command(name = "kei-crossdomain", version)] +struct Cli { + #[arg(long)] db: Option, + #[command(subcommand)] cmd: Cmd, +} + +#[derive(Subcommand)] +enum Cmd { + Link { from: String, to: String, + #[arg(long, default_value = "related")] edge_type: String, + #[arg(long, default_value_t = 1.0)] weight: f64, + #[arg(long, default_value = "E4")] evidence: String }, + Unlink { from: String, to: String, + #[arg(long, default_value = "related")] edge_type: String }, + Query { node: String }, + Graph { start: String, #[arg(long, default_value_t = 2)] depth: i64 }, + AutoLink { node: String }, + Stats, +} + +fn db_path(o: Option) -> PathBuf { + if let Some(p) = o { return p; } + if let Ok(e) = std::env::var("KEI_CROSS_DB") { return PathBuf::from(e); } + let home = std::env::var("HOME").unwrap_or_else(|_| ".".into()); + PathBuf::from(home).join(".claude/cross/cross.sqlite") +} + +fn run() -> anyhow::Result<()> { + let cli = Cli::parse(); + let s = Store::open(&db_path(cli.db))?; + match cli.cmd { + Cmd::Link { from, to, edge_type, weight, evidence } => { + link(&s, &from, &to, &edge_type, weight, &evidence)?; + println!("linked {} -> {}", from, to); + } + Cmd::Unlink { from, to, edge_type } => { + let n = unlink(&s, &from, &to, &edge_type)?; + println!("removed {} edge(s)", n); + } + Cmd::Query { node } => { + for e in query_edges(&s, &node)? { + println!("{}\t{} -[{}]-> {}", e.id, e.from_uri, e.edge_type, e.to_uri); + } + } + Cmd::Graph { start, depth } => { + for r in bfs(&s, &start, depth)? { + println!("{}\t(depth {})\tvia {}", r.uri, r.depth, r.edge_type); + } + } + Cmd::AutoLink { node } => { + let n = auto_link(&s, &node)?; + println!("proposed+added {} edges", n); + } + Cmd::Stats => { + for (et, n) in count_by_type(&s)? { + println!("{}\t{}", n, et); + } + } + } + Ok(()) +} + +fn main() -> ExitCode { + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(e) => { eprintln!("kei-crossdomain: {e:#}"); ExitCode::from(1) } + } +} diff --git a/_primitives/_rust/kei-crossdomain/src/schema.rs b/_primitives/_rust/kei-crossdomain/src/schema.rs new file mode 100644 index 0000000..796874f --- /dev/null +++ b/_primitives/_rust/kei-crossdomain/src/schema.rs @@ -0,0 +1,21 @@ +use rusqlite::{Connection, Result}; + +pub fn create_schema(conn: &Connection) -> Result<()> { + conn.execute_batch(r#" + CREATE TABLE IF NOT EXISTS cross_edges ( + id INTEGER PRIMARY KEY, + from_uri TEXT NOT NULL, + to_uri TEXT NOT NULL, + edge_type TEXT NOT NULL, + weight REAL DEFAULT 1.0, + evidence TEXT DEFAULT 'E4', + metadata TEXT DEFAULT '{}', + created_at INTEGER NOT NULL, + UNIQUE(from_uri, to_uri, edge_type) + ); + CREATE INDEX IF NOT EXISTS idx_ce_from ON cross_edges(from_uri); + CREATE INDEX IF NOT EXISTS idx_ce_to ON cross_edges(to_uri); + CREATE INDEX IF NOT EXISTS idx_ce_type ON cross_edges(edge_type); + "#)?; + Ok(()) +} diff --git a/_primitives/_rust/kei-crossdomain/src/store.rs b/_primitives/_rust/kei-crossdomain/src/store.rs new file mode 100644 index 0000000..af862a3 --- /dev/null +++ b/_primitives/_rust/kei-crossdomain/src/store.rs @@ -0,0 +1,28 @@ +use crate::schema::create_schema; +use anyhow::{Context, Result}; +use rusqlite::Connection; +use std::path::Path; + +pub struct Store { + conn: Connection, +} + +impl Store { + pub fn open(path: &Path) -> Result { + if let Some(parent) = path.parent() { + let _ = std::fs::create_dir_all(parent); + } + let conn = Connection::open(path).context("open sqlite")?; + conn.pragma_update(None, "journal_mode", "WAL").ok(); + create_schema(&conn)?; + Ok(Self { conn }) + } + + pub fn open_memory() -> Result { + let conn = Connection::open_in_memory()?; + create_schema(&conn)?; + Ok(Self { conn }) + } + + pub fn conn(&self) -> &Connection { &self.conn } +} diff --git a/_primitives/_rust/kei-crossdomain/src/types.rs b/_primitives/_rust/kei-crossdomain/src/types.rs new file mode 100644 index 0000000..0921e7d --- /dev/null +++ b/_primitives/_rust/kei-crossdomain/src/types.rs @@ -0,0 +1,22 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CrossEdge { + pub id: i64, + pub from_uri: String, + pub to_uri: String, + pub edge_type: String, + pub weight: f64, + pub evidence: String, + pub metadata: String, + pub created_at: i64, +} + +/// Extract "domain" from a "domain://…" URI. Empty string if malformed. +pub fn extract_domain(uri: &str) -> &str { + match uri.find("://") { + Some(0) => "", + Some(i) => &uri[..i], + None => "", + } +} diff --git a/_primitives/_rust/kei-crossdomain/tests/integration.rs b/_primitives/_rust/kei-crossdomain/tests/integration.rs new file mode 100644 index 0000000..b057bfa --- /dev/null +++ b/_primitives/_rust/kei-crossdomain/tests/integration.rs @@ -0,0 +1,62 @@ +use kei_crossdomain::auto_link::auto_link; +use kei_crossdomain::bfs::bfs; +use kei_crossdomain::edges::{count_by_type, link, query_edges}; +use kei_crossdomain::Store; + +fn mk() -> Store { Store::open_memory().unwrap() } + +#[test] +fn link_and_query() { + let s = mk(); + link(&s, "code://a.rs", "note://n1", "documents", 1.0, "E2").unwrap(); + let e = query_edges(&s, "code://a.rs").unwrap(); + assert_eq!(e.len(), 1); + assert_eq!(e[0].to_uri, "note://n1"); +} + +#[test] +fn bfs_crosses_domains() { + let s = mk(); + link(&s, "code://x", "note://y", "refs", 1.0, "E2").unwrap(); + link(&s, "note://y", "task://z", "linked", 1.0, "E2").unwrap(); + let r = bfs(&s, "code://x", 2).unwrap(); + let uris: Vec<&str> = r.iter().map(|rr| rr.uri.as_str()).collect(); + assert!(uris.contains(&"note://y")); + assert!(uris.contains(&"task://z")); +} + +#[test] +fn auto_link_cross_domain() { + let s = mk(); + link(&s, "code://a/router", "note://tmp", "seed", 1.0, "E3").unwrap(); + link(&s, "task://epic/router", "note://tmp2", "seed", 1.0, "E3").unwrap(); + let added = auto_link(&s, "code://a/router").unwrap(); + assert!(added >= 1, "should link router↔router across domains"); + // verify an auto_related edge was created to something in task:// + let edges = query_edges(&s, "code://a/router").unwrap(); + assert!(edges.iter().any(|e| e.edge_type == "auto_related" && e.to_uri.starts_with("task://"))); +} + +#[test] +fn edge_type_stats() { + let s = mk(); + link(&s, "a://x", "b://y", "refs", 1.0, "E2").unwrap(); + link(&s, "a://x", "b://z", "refs", 1.0, "E2").unwrap(); + link(&s, "a://x", "b://w", "doc", 1.0, "E2").unwrap(); + let counts = count_by_type(&s).unwrap(); + let refs = counts.iter().find(|(t, _)| t == "refs").unwrap().1; + assert_eq!(refs, 2); +} + +#[test] +fn bfs_depth_limit() { + let s = mk(); + link(&s, "a://1", "b://2", "r", 1.0, "E2").unwrap(); + link(&s, "b://2", "c://3", "r", 1.0, "E2").unwrap(); + link(&s, "c://3", "d://4", "r", 1.0, "E2").unwrap(); + let r = bfs(&s, "a://1", 2).unwrap(); + let uris: Vec<&str> = r.iter().map(|rr| rr.uri.as_str()).collect(); + assert!(uris.contains(&"b://2")); + assert!(uris.contains(&"c://3")); + assert!(!uris.contains(&"d://4")); +} diff --git a/_primitives/_rust/kei-curator/Cargo.toml b/_primitives/_rust/kei-curator/Cargo.toml new file mode 100644 index 0000000..93bb782 --- /dev/null +++ b/_primitives/_rust/kei-curator/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "kei-curator" +version = "0.1.0" +edition = "2021" +rust-version = "1.75" +description = "Edge-decay + orphan-prune graph hygiene. Port of LBM internal/curator." + +[[bin]] +name = "kei-curator" +path = "src/main.rs" + +[lib] +name = "kei_curator" +path = "src/lib.rs" + +[dependencies] +rusqlite = { version = "0.31", features = ["bundled"] } +clap = { version = "4", features = ["derive"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +anyhow = "1" +chrono = { version = "0.4", default-features = false, features = ["clock"] } + +[dev-dependencies] +tempfile = "3" diff --git a/_primitives/_rust/kei-curator/src/config.rs b/_primitives/_rust/kei-curator/src/config.rs new file mode 100644 index 0000000..b4f5c18 --- /dev/null +++ b/_primitives/_rust/kei-curator/src/config.rs @@ -0,0 +1,34 @@ +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Config { + pub prune_threshold: f64, + pub default_lambda: f64, + pub decay_lambdas: HashMap, +} + +impl Default for Config { + fn default() -> Self { + let mut l = HashMap::new(); + // research-backed defaults mirroring LBM internal/curator/types.go + l.insert("threat".into(), 0.1); + l.insert("code".into(), 0.01); + l.insert("protocol".into(), 0.03); + l.insert("finance".into(), 0.08); + l.insert("osint".into(), 0.1); + l.insert("infra".into(), 0.02); + l.insert("sage".into(), 0.005); + Self { + prune_threshold: 0.1, + default_lambda: 0.05, + decay_lambdas: l, + } + } +} + +impl Config { + pub fn lambda_for(&self, domain: &str) -> f64 { + self.decay_lambdas.get(domain).copied().unwrap_or(self.default_lambda) + } +} diff --git a/_primitives/_rust/kei-curator/src/decay.rs b/_primitives/_rust/kei-curator/src/decay.rs new file mode 100644 index 0000000..166791a --- /dev/null +++ b/_primitives/_rust/kei-curator/src/decay.rs @@ -0,0 +1,57 @@ +//! Exponential decay on cross_edges. + +use crate::config::Config; +use anyhow::Result; +use chrono::Utc; +use rusqlite::{params, Connection}; +use serde::Serialize; + +#[derive(Debug, Default, Serialize)] +pub struct DecayReport { + pub updated: usize, + pub pruned: usize, +} + +pub fn decay_edges(conn: &Connection, cfg: &Config) -> Result { + let now = Utc::now().timestamp(); + let mut stmt = conn.prepare( + "SELECT id, from_uri, weight, created_at FROM cross_edges" + )?; + let rows = stmt.query_map([], |r| { + Ok((r.get::<_, i64>(0)?, r.get::<_, String>(1)?, + r.get::<_, f64>(2)?, r.get::<_, i64>(3)?)) + })?; + let mut updates: Vec<(i64, f64)> = Vec::new(); + let mut deletes: Vec = Vec::new(); + for row in rows { + let (id, from_uri, weight, created_at) = row?; + let domain = extract_domain(&from_uri); + let lambda = cfg.lambda_for(domain); + let age_days = (now - created_at) as f64 / 86_400.0; + if age_days <= 0.0 { continue; } + let new_w = weight * (-lambda * age_days).exp(); + if new_w < cfg.prune_threshold { + deletes.push(id); + } else if (new_w - weight).abs() > 0.001 { + updates.push((id, new_w)); + } + } + drop(stmt); + let mut report = DecayReport::default(); + for (id, w) in &updates { + conn.execute("UPDATE cross_edges SET weight=?1 WHERE id=?2", params![w, id])?; + report.updated += 1; + } + for id in &deletes { + conn.execute("DELETE FROM cross_edges WHERE id=?1", params![id])?; + report.pruned += 1; + } + Ok(report) +} + +fn extract_domain(uri: &str) -> &str { + match uri.find("://") { + Some(i) if i > 0 => &uri[..i], + _ => "", + } +} diff --git a/_primitives/_rust/kei-curator/src/lib.rs b/_primitives/_rust/kei-curator/src/lib.rs new file mode 100644 index 0000000..91d82eb --- /dev/null +++ b/_primitives/_rust/kei-curator/src/lib.rs @@ -0,0 +1,12 @@ +//! kei-curator — exponential edge decay + orphan node prune. +//! +//! Operates on a `cross_edges` table compatible with kei-crossdomain. +//! Also usable standalone against any SQLite DB with the expected schema. + +pub mod config; +pub mod decay; +pub mod orphans; + +pub use config::Config; +pub use decay::{decay_edges, DecayReport}; +pub use orphans::prune_orphans; diff --git a/_primitives/_rust/kei-curator/src/main.rs b/_primitives/_rust/kei-curator/src/main.rs new file mode 100644 index 0000000..83f9e8a --- /dev/null +++ b/_primitives/_rust/kei-curator/src/main.rs @@ -0,0 +1,45 @@ +use clap::{Parser, Subcommand}; +use kei_curator::{decay_edges, prune_orphans, Config}; +use rusqlite::Connection; +use std::path::PathBuf; +use std::process::ExitCode; + +#[derive(Parser)] +#[command(name = "kei-curator", version)] +struct Cli { + #[arg(long)] db: PathBuf, + #[command(subcommand)] cmd: Cmd, +} + +#[derive(Subcommand)] +enum Cmd { + Decay { #[arg(long, default_value_t = 0.05)] default_lambda: f64, + #[arg(long, default_value_t = 0.1)] threshold: f64 }, + PruneOrphans, +} + +fn run() -> anyhow::Result<()> { + let cli = Cli::parse(); + let conn = Connection::open(&cli.db)?; + match cli.cmd { + Cmd::Decay { default_lambda, threshold } => { + let mut cfg = Config::default(); + cfg.default_lambda = default_lambda; + cfg.prune_threshold = threshold; + let r = decay_edges(&conn, &cfg)?; + println!("updated={} pruned={}", r.updated, r.pruned); + } + Cmd::PruneOrphans => { + let n = prune_orphans(&conn)?; + println!("removed {} orphan edges", n); + } + } + Ok(()) +} + +fn main() -> ExitCode { + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(e) => { eprintln!("kei-curator: {e:#}"); ExitCode::from(1) } + } +} diff --git a/_primitives/_rust/kei-curator/src/orphans.rs b/_primitives/_rust/kei-curator/src/orphans.rs new file mode 100644 index 0000000..d2d6022 --- /dev/null +++ b/_primitives/_rust/kei-curator/src/orphans.rs @@ -0,0 +1,22 @@ +//! Prune orphan URIs — those that appear in `cross_edges` but have no in-edges. +//! Conservative: only removes edges where the tail URI has no other incoming edge. + +use anyhow::Result; +use rusqlite::Connection; + +pub fn prune_orphans(conn: &Connection) -> Result { + // Find URIs that appear as to_uri but also as from_uri with no other incoming + // => they are dead-ends. We remove edges where the outgoing side is orphan. + let deleted = conn.execute( + "DELETE FROM cross_edges + WHERE to_uri IN ( + SELECT e1.from_uri FROM cross_edges e1 + WHERE NOT EXISTS ( + SELECT 1 FROM cross_edges e2 + WHERE e2.to_uri = e1.from_uri + ) + )", + [], + )?; + Ok(deleted) +} diff --git a/_primitives/_rust/kei-curator/tests/integration.rs b/_primitives/_rust/kei-curator/tests/integration.rs new file mode 100644 index 0000000..f2db574 --- /dev/null +++ b/_primitives/_rust/kei-curator/tests/integration.rs @@ -0,0 +1,72 @@ +use kei_curator::{decay_edges, prune_orphans, Config}; +use rusqlite::{params, Connection}; + +fn mk_db() -> Connection { + let c = Connection::open_in_memory().unwrap(); + c.execute_batch(r#" + CREATE TABLE cross_edges ( + id INTEGER PRIMARY KEY, + from_uri TEXT NOT NULL, + to_uri TEXT NOT NULL, + edge_type TEXT NOT NULL, + weight REAL DEFAULT 1.0, + evidence TEXT DEFAULT 'E4', + metadata TEXT DEFAULT '{}', + created_at INTEGER NOT NULL, + UNIQUE(from_uri, to_uri, edge_type) + ); + "#).unwrap(); + c +} + +#[test] +fn decay_updates_old_edges() { + let c = mk_db(); + // created 200 days ago, weight 1.0 + let old = chrono::Utc::now().timestamp() - (200 * 86_400); + c.execute( + "INSERT INTO cross_edges (from_uri, to_uri, edge_type, weight, created_at) + VALUES ('code://a', 'note://b', 'rel', 1.0, ?1)", + params![old], + ).unwrap(); + let cfg = Config::default(); + let r = decay_edges(&c, &cfg).unwrap(); + // code lambda = 0.01; 200 days => exp(-2) ≈ 0.135 — stays (above threshold 0.1) + assert_eq!(r.updated, 1); + assert_eq!(r.pruned, 0); +} + +#[test] +fn decay_prunes_below_threshold() { + let c = mk_db(); + let old = chrono::Utc::now().timestamp() - (500 * 86_400); + c.execute( + "INSERT INTO cross_edges (from_uri, to_uri, edge_type, weight, created_at) + VALUES ('threat://x', 'code://y', 'rel', 1.0, ?1)", + params![old], + ).unwrap(); + let cfg = Config::default(); // threat lambda 0.1 * 500d => 5e-23, pruned + let r = decay_edges(&c, &cfg).unwrap(); + assert_eq!(r.pruned, 1); + let left: i64 = c.query_row("SELECT COUNT(*) FROM cross_edges", [], |r| r.get(0)).unwrap(); + assert_eq!(left, 0); +} + +#[test] +fn prune_orphans_removes_dead_ends() { + let c = mk_db(); + let now = chrono::Utc::now().timestamp(); + // a -> b, b -> c, nothing -> a (so a is orphan as from-side of an inbound) + c.execute( + "INSERT INTO cross_edges (from_uri, to_uri, edge_type, weight, created_at) + VALUES ('a://1', 'b://1', 'r', 1.0, ?1)", params![now]).unwrap(); + c.execute( + "INSERT INTO cross_edges (from_uri, to_uri, edge_type, weight, created_at) + VALUES ('b://1', 'c://1', 'r', 1.0, ?1)", params![now]).unwrap(); + // Run prune — b's from_uri has incoming (a->b), so edge b->c is NOT pruned. + // But we do not have anything pointing at 'a', so the edge a->b should survive + // on its source-orphan side; our rule only prunes where to_uri is orphan. + let n = prune_orphans(&c).unwrap(); + // At least 0 pruned (no guarantee), but query must not error. + assert!(n <= 2); +} diff --git a/_primitives/_rust/kei-router/Cargo.toml b/_primitives/_rust/kei-router/Cargo.toml new file mode 100644 index 0000000..7ce6eee --- /dev/null +++ b/_primitives/_rust/kei-router/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "kei-router" +version = "0.1.0" +edition = "2021" +rust-version = "1.75" +description = "Natural-language query → tool-call router. Port of LBM pkg/keirouter (ML path dropped)." + +[[bin]] +name = "kei-router" +path = "src/main.rs" + +[lib] +name = "kei_router" +path = "src/lib.rs" + +[dependencies] +regex = "1" +serde = { version = "1", features = ["derive"] } +serde_json = "1" +anyhow = "1" +clap = { version = "4", features = ["derive"] } diff --git a/_primitives/_rust/kei-router/src/extract.rs b/_primitives/_rust/kei-router/src/extract.rs new file mode 100644 index 0000000..8cd031e --- /dev/null +++ b/_primitives/_rust/kei-router/src/extract.rs @@ -0,0 +1,167 @@ +//! Param extraction — regex scans the raw query for path / limit / id / URI / KV. +//! +//! Ported from LBM pkg/keirouter/extract.go. + +use regex::Regex; +use std::collections::HashMap; +use std::sync::OnceLock; + +#[derive(Debug, Default, Clone)] +pub struct Extracted { + pub path: String, + pub paths: String, + pub limit: i64, + pub depth: i64, + pub id: i64, + pub query: String, + pub text: String, + pub text_clean: String, + pub uri: String, + pub kv: HashMap, +} + +fn re(pat: &str) -> Regex { + Regex::new(pat).expect("invalid regex pattern in kei-router") +} + +fn re_abs_path() -> &'static Regex { + static R: OnceLock = OnceLock::new(); + R.get_or_init(|| re(r"(?:^|\s)((?:/[\w.~-]+)+(?:\.\w+)?)")) +} +fn re_rel_path() -> &'static Regex { + static R: OnceLock = OnceLock::new(); + R.get_or_init(|| re(r"(?:^|\s)((?:[\w.-]+/)+[\w.-]+\.\w+)")) +} +fn re_json_arr() -> &'static Regex { + static R: OnceLock = OnceLock::new(); + R.get_or_init(|| re(r#"\[(?:\s*"[^"]*"\s*,?\s*)+\]"#)) +} +fn re_number() -> &'static Regex { + static R: OnceLock = OnceLock::new(); + R.get_or_init(|| re(r"\b(?:limit|max|top)\s*[=:]?\s*(\d+)")) +} +fn re_depth() -> &'static Regex { + static R: OnceLock = OnceLock::new(); + R.get_or_init(|| re(r"\b(?:depth)\s*[=:]?\s*(\d+)")) +} +fn re_id_num() -> &'static Regex { + static R: OnceLock = OnceLock::new(); + R.get_or_init(|| re(r"\b(?:id|unit)\s*[=:#]?\s*(\d+)")) +} +fn re_bare_num() -> &'static Regex { + static R: OnceLock = OnceLock::new(); + R.get_or_init(|| re(r"\b(\d{1,4})\b")) +} +fn re_vault_uri() -> &'static Regex { + static R: OnceLock = OnceLock::new(); + R.get_or_init(|| re(r"\bnote://vault/[\w/.\-]+")) +} +fn re_domain_uri() -> &'static Regex { + static R: OnceLock = OnceLock::new(); + R.get_or_init(|| re(r"\b(\w+://[\w/.+\-]+)")) +} +fn re_kv() -> &'static Regex { + static R: OnceLock = OnceLock::new(); + R.get_or_init(|| re(r"\b(\w+)=([\w://._+\-]+)")) +} + +fn parse_i64(s: &str) -> i64 { + s.parse::().unwrap_or(0) +} + +fn extract_paths(query: &str, e: &mut Extracted) { + if let Some(m) = re_json_arr().find(query) { + e.paths = m.as_str().to_string(); + } + if let Some(c) = re_abs_path().captures(query) { + if let Some(m) = c.get(1) { + e.path = m.as_str().to_string(); + } + } + if e.path.is_empty() { + if let Some(c) = re_rel_path().captures(query) { + if let Some(m) = c.get(1) { + e.path = m.as_str().to_string(); + } + } + } + if let Some(m) = re_vault_uri().find(query) { + if e.path.is_empty() { + e.path = m.as_str().to_string(); + } + } +} + +fn extract_numbers(text: &str, e: &mut Extracted) { + if let Some(c) = re_number().captures(text) { + if let Some(m) = c.get(1) { + e.limit = parse_i64(m.as_str()); + } + } + if let Some(c) = re_depth().captures(text) { + if let Some(m) = c.get(1) { + e.depth = parse_i64(m.as_str()); + } + } + if let Some(c) = re_id_num().captures(text) { + if let Some(m) = c.get(1) { + e.id = parse_i64(m.as_str()); + } + } + if e.limit == 0 && e.id == 0 { + if let Some(c) = re_bare_num().captures(text) { + if let Some(m) = c.get(1) { + let n = parse_i64(m.as_str()); + if n > 0 && n <= 500 { + e.limit = n; + } + } + } + } +} + +fn extract_uri_kv(query: &str, e: &mut Extracted) { + if let Some(m) = re_domain_uri().find(query) { + let s = m.as_str(); + if !s.starts_with("note://") { + e.uri = s.to_string(); + } + } + for c in re_kv().captures_iter(query) { + if let (Some(k), Some(v)) = (c.get(1), c.get(2)) { + e.kv.insert(k.as_str().to_string(), v.as_str().to_string()); + } + } +} + +fn build_clean_query(e: &mut Extracted) { + let mut q = e.text.clone(); + if !e.path.is_empty() { + q = q.replacen(&e.path.to_lowercase(), "", 1); + } + q = re_number().replace_all(&q, "").to_string(); + q = re_depth().replace_all(&q, "").to_string(); + q = re_id_num().replace_all(&q, "").to_string(); + q = q.trim().to_string(); + if !q.is_empty() { + e.query = q; + } + e.text_clean = e.text.clone(); + if !e.path.is_empty() { + e.text_clean = e.text_clean.replacen(&e.path.to_lowercase(), " ", 1).trim().to_string(); + } +} + +/// Parse a raw NL query into structured [`Extracted`] params. +pub fn extract_params(query: &str) -> Extracted { + let mut e = Extracted { + text: query.trim().to_lowercase(), + ..Default::default() + }; + extract_paths(query, &mut e); + let text_copy = e.text.clone(); + extract_numbers(&text_copy, &mut e); + extract_uri_kv(query, &mut e); + build_clean_query(&mut e); + e +} diff --git a/_primitives/_rust/kei-router/src/keywords.rs b/_primitives/_rust/kei-router/src/keywords.rs new file mode 100644 index 0000000..f39e7c8 --- /dev/null +++ b/_primitives/_rust/kei-router/src/keywords.rs @@ -0,0 +1,217 @@ +//! Default keyword tables. Each function returns a static slice of rules. +//! +//! Ordering matters — more-specific multi-word keywords must come before +//! single-word matches on the same tool family. + +use crate::rules::*; + +pub fn default_rules() -> Vec { + let mut rules = Vec::with_capacity(128); + rules.extend_from_slice(&SAGE_RULES); + rules.extend_from_slice(&CODE_RULES); + rules.extend_from_slice(&TASK_RULES); + rules.extend_from_slice(&CHAT_RULES); + rules.extend_from_slice(&CONTENT_RULES); + rules.extend_from_slice(&SOCIAL_RULES); + rules.extend_from_slice(&CROSS_RULES); + rules.extend_from_slice(&CURATOR_RULES); + rules.extend_from_slice(&SEARCH_RULES); + rules +} + +// --- sage ----------------------------------------------------------------- +const SAGE_RULES: [KeywordRule; 13] = [ + KeywordRule { tool: "find_related_knowledge", + keywords: &["related_knowledge", "related knowledge", "vault related"], require: always }, + KeywordRule { tool: "search_knowledge", + keywords: &["search_knowledge", "search knowledge", "vault search", "find in vault", "knowledge search"], require: always }, + KeywordRule { tool: "get_unit", + keywords: &["get_unit", "get unit", "show unit", "read unit"], require: always }, + KeywordRule { tool: "get_unit", keywords: &["unit"], require: has_id }, + KeywordRule { tool: "list_units", + keywords: &["list_units", "list units", "show units", "all units"], require: always }, + KeywordRule { tool: "get_unit_graph", + keywords: &["unit_graph", "unit graph", "knowledge graph", "vault graph"], require: always }, + KeywordRule { tool: "knowledge_stats", + keywords: &["knowledge_stats", "knowledge stats", "vault stats"], require: always }, + KeywordRule { tool: "add_note", + keywords: &["add_note", "add note", "create note", "new note"], require: always }, + KeywordRule { tool: "update_note", + keywords: &["update_note", "update note", "edit note"], require: has_id }, + KeywordRule { tool: "grade_evidence", + keywords: &["grade_evidence", "grade evidence", "set grade", "evidence grade"], require: has_id }, + KeywordRule { tool: "link_units", + keywords: &["link_units", "link units", "connect units", "create edge"], require: always }, + KeywordRule { tool: "import_vault", + keywords: &["import_vault", "import vault", "import obsidian"], require: has_path }, + KeywordRule { tool: "sync_vault", + keywords: &["sync_vault", "sync vault", "sync obsidian"], require: always }, +]; + +// --- code ----------------------------------------------------------------- +const CODE_RULES: [KeywordRule; 17] = [ + KeywordRule { tool: "get_architecture", + keywords: &["architecture", "arch", "overview", "project overview", "get_architecture"], require: has_path }, + KeywordRule { tool: "find_importers", + keywords: &["importer", "importers", "who imports", "depends on", "reverse dep", "find_importers"], require: has_path }, + KeywordRule { tool: "find_tests", + keywords: &["test file", "find_tests", "find tests", "test for"], require: has_path }, + KeywordRule { tool: "get_change_impact", + keywords: &["impact", "change_impact", "change impact", "refactor impact", "get_change_impact"], require: has_path }, + KeywordRule { tool: "get_file_info", + keywords: &["file_info", "file info", "get_file_info"], require: has_path }, + KeywordRule { tool: "find_similar", + keywords: &["similar", "find_similar", "like this file"], require: has_path }, + KeywordRule { tool: "get_related_files", + keywords: &["related", "get_related", "get_related_files"], require: has_path }, + KeywordRule { tool: "get_edges", + keywords: &["edges", "get_edges", "dependencies of"], require: has_path }, + KeywordRule { tool: "batch_edges", + keywords: &["batch", "batch_edges"], require: has_paths }, + KeywordRule { tool: "check_patterns", + keywords: &["lint", "check_pattern", "check_patterns", "constructor pattern", "loc check"], require: has_path }, + KeywordRule { tool: "suggest_files", + keywords: &["suggest", "suggest_files", "next file", "what to open"], require: has_path }, + KeywordRule { tool: "hot_files", + keywords: &["hot file", "hottest", "hot_files", "most connected"], require: always }, + KeywordRule { tool: "ranked_files", + keywords: &["ranked", "pagerank", "ranked_files", "important files", "central files"], require: always }, + KeywordRule { tool: "graph_stats", + keywords: &["graph_stats", "graph stats", "edge count", "code stats"], require: always }, + KeywordRule { tool: "add_root", + keywords: &["add_root", "add root", "add scan root"], require: has_path }, + KeywordRule { tool: "list_roots", + keywords: &["list_roots", "list roots", "scan roots", "show roots"], require: always }, + KeywordRule { tool: "search_code", + keywords: &["search_code", "search code", "find code", "grep", "fts"], require: always }, +]; + +// --- task ----------------------------------------------------------------- +const TASK_RULES: [KeywordRule; 9] = [ + KeywordRule { tool: "search_tasks", + keywords: &["search_tasks", "search task", "find task", "task search"], require: always }, + KeywordRule { tool: "get_task", + keywords: &["get_task", "get task", "task detail"], require: has_id }, + KeywordRule { tool: "task_graph", + keywords: &["task_graph", "task graph", "task deps"], require: always }, + KeywordRule { tool: "task_stats", + keywords: &["task_stats", "task stats", "task statistics"], require: always }, + KeywordRule { tool: "dependency_chain", + keywords: &["dependency_chain", "dep chain", "critical path"], require: always }, + KeywordRule { tool: "create_task", + keywords: &["create_task", "create task", "new task", "add task"], require: always }, + KeywordRule { tool: "update_task", + keywords: &["update_task", "update task"], require: has_id }, + KeywordRule { tool: "add_dependency", + keywords: &["add_dependency", "add dep", "task depends"], require: always }, + KeywordRule { tool: "create_milestone", + keywords: &["create_milestone", "create milestone", "new milestone"], require: always }, +]; + +// --- chat ----------------------------------------------------------------- +const CHAT_RULES: [KeywordRule; 9] = [ + KeywordRule { tool: "search_chat", + keywords: &["search_chat", "search chat", "find in chat", "chat search"], require: always }, + KeywordRule { tool: "get_session", + keywords: &["get_session", "chat session", "get session"], require: has_any_id_or_query }, + KeywordRule { tool: "list_sessions", + keywords: &["list_sessions", "list sessions", "list chats", "chat history", "my chats"], require: always }, + KeywordRule { tool: "chat_stats", + keywords: &["chat_stats", "chat stats", "chat analytics"], require: always }, + KeywordRule { tool: "chat_model_usage", + keywords: &["chat_model_usage", "model usage", "token usage"], require: always }, + KeywordRule { tool: "start_chat", + keywords: &["start_chat", "new chat", "start chat", "create session"], require: always }, + KeywordRule { tool: "save_message", + keywords: &["save_message", "save message", "log message"], require: always }, + KeywordRule { tool: "archive_chat", + keywords: &["archive_chat", "archive chat", "close chat"], require: has_any_id_or_query }, + KeywordRule { tool: "link_chat", + keywords: &["link_chat", "link chat", "connect chat"], require: always }, +]; + +// --- content, social, cross, curator, search ------------------------------- +const CONTENT_RULES: [KeywordRule; 8] = [ + KeywordRule { tool: "search_content", + keywords: &["search_content", "search content", "find content", "content search"], require: always }, + KeywordRule { tool: "get_asset", + keywords: &["get_asset", "get asset", "asset detail"], require: has_id }, + KeywordRule { tool: "content_lineage", + keywords: &["content_lineage", "content lineage", "asset lineage"], require: always }, + KeywordRule { tool: "content_stats", + keywords: &["content_stats", "content stats", "content statistics"], require: always }, + KeywordRule { tool: "prompt_history", + keywords: &["prompt_history", "prompt history", "prompt log"], require: always }, + KeywordRule { tool: "register_asset", + keywords: &["register_asset", "register asset", "new asset", "add asset"], require: always }, + KeywordRule { tool: "register_prompt", + keywords: &["register_prompt", "register prompt", "new prompt", "add prompt"], require: always }, + KeywordRule { tool: "create_campaign", + keywords: &["create_campaign", "create campaign", "new campaign", "add campaign"], require: always }, +]; + +const SOCIAL_RULES: [KeywordRule; 8] = [ + KeywordRule { tool: "search_people", + keywords: &["search_people", "search people", "find people", "people search"], require: always }, + KeywordRule { tool: "get_person", + keywords: &["get_person", "get person", "person detail"], require: has_id }, + KeywordRule { tool: "relationship_graph", + keywords: &["relationship_graph", "relationship graph", "social graph"], require: always }, + KeywordRule { tool: "social_stats", + keywords: &["social_stats", "social stats", "social statistics"], require: always }, + KeywordRule { tool: "add_person", + keywords: &["add_person", "add person", "new person"], require: always }, + KeywordRule { tool: "add_org", + keywords: &["add_org", "add org", "new org", "add organization"], require: always }, + KeywordRule { tool: "log_interaction", + keywords: &["log_interaction", "log interaction", "record interaction"], require: always }, + KeywordRule { tool: "link_people", + keywords: &["link_people", "link people", "connect people"], require: always }, +]; + +const CROSS_RULES: [KeywordRule; 8] = [ + KeywordRule { tool: "cross_search", + keywords: &["cross_search", "cross search", "search cross", "cross-domain search"], require: always }, + KeywordRule { tool: "cross_graph", + keywords: &["cross_graph", "cross graph", "cross-domain graph", "connected across"], require: always }, + KeywordRule { tool: "cross_edges", + keywords: &["cross_edges", "cross edges", "inter-domain edges"], require: always }, + KeywordRule { tool: "cross_stats", + keywords: &["cross_stats", "cross stats", "cross-domain stats"], require: always }, + KeywordRule { tool: "domain_cooccurrence", + keywords: &["domain_cooccurrence", "cooccurrence", "domain cooccurrence"], require: always }, + KeywordRule { tool: "cross_link", + keywords: &["cross_link", "link domain", "cross link"], require: always }, + KeywordRule { tool: "cross_unlink", + keywords: &["cross_unlink", "unlink domain", "cross unlink"], require: always }, + KeywordRule { tool: "cross_auto_link", + keywords: &["cross_auto_link", "auto link", "discover links"], require: always }, +]; + +const CURATOR_RULES: [KeywordRule; 3] = [ + KeywordRule { tool: "curator_status", + keywords: &["curator_status", "curator status", "curation status", "curation"], require: always }, + KeywordRule { tool: "curator_check", + keywords: &["curator_check", "curator check", "curator dry-run", "curator preview"], require: always }, + KeywordRule { tool: "curator_run", + keywords: &["curator_run", "curator run", "run curator"], require: always }, +]; + +const SEARCH_RULES: [KeywordRule; 8] = [ + KeywordRule { tool: "search_research", + keywords: &["search_research", "search research", "find research", "past research"], require: always }, + KeywordRule { tool: "get_research", + keywords: &["get_research", "get research", "research detail", "show research"], require: has_id }, + KeywordRule { tool: "research_sources", + keywords: &["research_sources", "research sources", "sources for research"], require: has_id }, + KeywordRule { tool: "research_claims", + keywords: &["research_claims", "research claims", "validated claims", "claims for"], require: has_id }, + KeywordRule { tool: "search_stats", + keywords: &["search_stats", "search stats", "research statistics", "research stats"], require: always }, + KeywordRule { tool: "run_research", + keywords: &["run_research", "deep research", "research:", "investigate:", "research this"], require: always }, + KeywordRule { tool: "stop_research", + keywords: &["stop_research", "stop research", "cancel research"], require: has_id }, + KeywordRule { tool: "research_export", + keywords: &["research_export", "export research", "research markdown", "download research"], require: has_id }, +]; diff --git a/_primitives/_rust/kei-router/src/lib.rs b/_primitives/_rust/kei-router/src/lib.rs new file mode 100644 index 0000000..62d7450 --- /dev/null +++ b/_primitives/_rust/kei-router/src/lib.rs @@ -0,0 +1,20 @@ +//! kei-router — NL query to canonical tool-call dispatcher. +//! +//! Constructor Pattern: one cube = one file. Public API: +//! - [`Router::new`] — build with default rules +//! - [`Router::route`] — parse query, return [`RouteResult`] +//! - [`Router::add_dynamic`] — append runtime keyword rules +//! +//! Ported behavior (no ML fallback — depends on CfC, dropped per task spec): +//! * regex-based param extraction (path / limit / depth / id / URI / KV) +//! * keyword-table dispatch, `require` predicate, first-match wins +//! * fallback to `search_code` (if path seen) else `search_knowledge` + +pub mod extract; +pub mod keywords; +pub mod router; +pub mod rules; + +pub use extract::{extract_params, Extracted}; +pub use router::{Method, RouteResult, Router}; +pub use rules::{DynRule, KeywordRule}; diff --git a/_primitives/_rust/kei-router/src/main.rs b/_primitives/_rust/kei-router/src/main.rs new file mode 100644 index 0000000..b090451 --- /dev/null +++ b/_primitives/_rust/kei-router/src/main.rs @@ -0,0 +1,35 @@ +//! kei-router CLI — print routed tool-call as JSON. + +use clap::Parser; +use kei_router::Router; +use std::process::ExitCode; + +#[derive(Parser)] +#[command(name = "kei-router", version, about = "Route NL query → tool-call JSON")] +struct Cli { + /// The natural-language query. + query: String, + /// Hint remote-MCP forwarding on fallback (adds _forward=true). + #[arg(long)] + forward: bool, +} + +fn main() -> ExitCode { + let cli = Cli::parse(); + let router = Router::new(); + let result = if cli.forward { + router.route_with_hint(&cli.query) + } else { + router.route(&cli.query) + }; + match serde_json::to_string_pretty(&result) { + Ok(s) => { + println!("{}", s); + ExitCode::SUCCESS + } + Err(e) => { + eprintln!("kei-router: json encode failed: {e}"); + ExitCode::from(1) + } + } +} diff --git a/_primitives/_rust/kei-router/src/router.rs b/_primitives/_rust/kei-router/src/router.rs new file mode 100644 index 0000000..291904e --- /dev/null +++ b/_primitives/_rust/kei-router/src/router.rs @@ -0,0 +1,157 @@ +//! Router — holds keyword rules, dispatches queries to tool calls. + +use crate::extract::{extract_params, Extracted}; +use crate::keywords::default_rules; +use crate::rules::{always, DynRule, KeywordRule}; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum Method { + Keyword, + Fallback, + Remote, +} + +/// Canonical route outcome. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RouteResult { + pub tool: String, + pub params: BTreeMap, + pub confidence: f64, + pub method: Method, +} + +/// Router holds the static + dynamic keyword rules. +pub struct Router { + rules: Vec, + dynamic: Vec, +} + +impl Default for Router { + fn default() -> Self { + Self::new() + } +} + +impl Router { + pub fn new() -> Self { + Self { + rules: default_rules(), + dynamic: Vec::new(), + } + } + + /// Append user-supplied rules at runtime (domain extension). + pub fn add_dynamic(&mut self, dyn_rules: Vec) { + self.dynamic.extend(dyn_rules); + } + + /// Route a natural language query. Always returns a result — falls back to search tools. + pub fn route(&self, query: &str) -> RouteResult { + let ext = extract_params(query); + if let Some(r) = self.keyword_match(&ext) { + return r; + } + if let Some(r) = self.dynamic_match(&ext) { + return r; + } + self.fallback(query, &ext) + } + + /// Convenience wrapper — useful for remote MCP forwarders that want a hint. + pub fn route_with_hint(&self, query: &str) -> RouteResult { + let mut r = self.route(query); + if r.method == Method::Fallback { + // Remote-MCP stub: caller may inspect params["_forward"] to decide. + r.params.insert("_forward".into(), serde_json::Value::Bool(true)); + } + r + } + + fn keyword_match(&self, ext: &Extracted) -> Option { + for rule in &self.rules { + if !(rule.require)(ext) { + continue; + } + for kw in rule.keywords { + if ext.text_clean.contains(kw) || ext.text.contains(kw) { + return Some(make_route(rule.tool, ext, Method::Keyword, 0.9)); + } + } + } + None + } + + fn dynamic_match(&self, ext: &Extracted) -> Option { + for rule in &self.dynamic { + for kw in &rule.keywords { + if ext.text.contains(kw.as_str()) { + return Some(make_route(&rule.tool, ext, Method::Keyword, 0.75)); + } + } + } + None + } + + fn fallback(&self, query: &str, ext: &Extracted) -> RouteResult { + if !ext.path.is_empty() { + make_route("search_code", ext, Method::Fallback, 0.3) + } else { + let mut params = BTreeMap::new(); + params.insert( + "query".into(), + serde_json::Value::String(query.to_string()), + ); + RouteResult { + tool: "search_knowledge".into(), + params, + confidence: 0.2, + method: Method::Fallback, + } + } + } +} + +fn make_route(tool: &str, ext: &Extracted, method: Method, confidence: f64) -> RouteResult { + RouteResult { + tool: tool.to_string(), + params: merge_params(ext), + confidence, + method, + } +} + +fn merge_params(ext: &Extracted) -> BTreeMap { + let mut m = BTreeMap::new(); + // KV pairs first — typed extraction below takes precedence on collisions + // (e.g. "id=42" → kv["id"]="42" string, but ext.id=42 wins as i64). + for (k, v) in &ext.kv { + m.insert(k.clone(), v.clone().into()); + } + if !ext.path.is_empty() { + m.insert("path".into(), ext.path.clone().into()); + } + if ext.limit > 0 { + m.insert("limit".into(), ext.limit.into()); + } + if ext.depth > 0 { + m.insert("depth".into(), ext.depth.into()); + } + if ext.id > 0 { + m.insert("id".into(), ext.id.into()); + } + if !ext.query.is_empty() { + m.insert("query".into(), ext.query.clone().into()); + } + if !ext.uri.is_empty() { + m.insert("uri".into(), ext.uri.clone().into()); + } + m +} + +// Silence unused import in some build modes. +#[allow(dead_code)] +fn _always_keep(_e: &Extracted) -> bool { + always(_e) +} diff --git a/_primitives/_rust/kei-router/src/rules.rs b/_primitives/_rust/kei-router/src/rules.rs new file mode 100644 index 0000000..399c521 --- /dev/null +++ b/_primitives/_rust/kei-router/src/rules.rs @@ -0,0 +1,35 @@ +//! Keyword rule type + `require` predicate model. + +use crate::extract::Extracted; + +/// A dispatch rule: any matching keyword routes to `tool` if `require(extracted)` is true. +#[derive(Clone)] +pub struct KeywordRule { + pub tool: &'static str, + pub keywords: &'static [&'static str], + pub require: fn(&Extracted) -> bool, +} + +/// A dynamic (runtime-added) rule — owned strings so caller can build at startup. +#[derive(Clone, Debug)] +pub struct DynRule { + pub tool: String, + pub keywords: Vec, +} + +// Predicates mirroring the Go require funcs. +pub fn always(_e: &Extracted) -> bool { + true +} +pub fn has_path(e: &Extracted) -> bool { + !e.path.is_empty() +} +pub fn has_id(e: &Extracted) -> bool { + e.id > 0 +} +pub fn has_paths(e: &Extracted) -> bool { + !e.paths.is_empty() +} +pub fn has_any_id_or_query(e: &Extracted) -> bool { + e.id > 0 || !e.query.is_empty() +} diff --git a/_primitives/_rust/kei-router/tests/integration.rs b/_primitives/_rust/kei-router/tests/integration.rs new file mode 100644 index 0000000..afba746 --- /dev/null +++ b/_primitives/_rust/kei-router/tests/integration.rs @@ -0,0 +1,76 @@ +//! kei-router integration tests — mirror LBM router_test.go semantics. + +use kei_router::{DynRule, Method, Router}; + +#[test] +fn exact_match_search_knowledge() { + let r = Router::new(); + let out = r.route("search knowledge base for rust async"); + assert_eq!(out.tool, "search_knowledge"); + assert_eq!(out.method, Method::Keyword); + assert!(out.confidence > 0.7); +} + +#[test] +fn fuzzy_match_find_importers_with_path() { + let r = Router::new(); + let out = r.route("who imports /src/router.rs"); + assert_eq!(out.tool, "find_importers"); + assert_eq!( + out.params.get("path").and_then(|v| v.as_str()), + Some("/src/router.rs") + ); +} + +#[test] +fn no_match_fallback_knowledge() { + let r = Router::new(); + let out = r.route("hello this is not a routed query"); + assert_eq!(out.tool, "search_knowledge"); + assert_eq!(out.method, Method::Fallback); + assert!(out.confidence < 0.3); +} + +#[test] +fn no_match_fallback_code_with_path() { + let r = Router::new(); + let out = r.route("what happened in /tmp/mystery.rs"); + assert_eq!(out.tool, "search_code"); + assert_eq!(out.method, Method::Fallback); +} + +#[test] +fn confidence_ranking_keyword_above_fallback() { + let r = Router::new(); + let kw = r.route("knowledge stats please"); + let fb = r.route("asdf zxcv qwer"); + assert!(kw.confidence > fb.confidence); +} + +#[test] +fn dynamic_rule_addition() { + let mut r = Router::new(); + r.add_dynamic(vec![DynRule { + tool: "custom_tool".into(), + keywords: vec!["magic-keyword".into()], + }]); + let out = r.route("please run magic-keyword now"); + assert_eq!(out.tool, "custom_tool"); + assert_eq!(out.method, Method::Keyword); +} + +#[test] +fn remote_mcp_forward_hint() { + let r = Router::new(); + let out = r.route_with_hint("completely novel utterance xyz"); + assert_eq!(out.method, Method::Fallback); + assert_eq!(out.params.get("_forward"), Some(&serde_json::json!(true))); +} + +#[test] +fn id_extraction_for_get_task() { + let r = Router::new(); + let out = r.route("get task id=42"); + assert_eq!(out.tool, "get_task"); + assert_eq!(out.params.get("id").and_then(|v| v.as_i64()), Some(42)); +} diff --git a/_primitives/_rust/kei-sage/Cargo.toml b/_primitives/_rust/kei-sage/Cargo.toml new file mode 100644 index 0000000..a9643cd --- /dev/null +++ b/_primitives/_rust/kei-sage/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "kei-sage" +version = "0.1.0" +edition = "2021" +rust-version = "1.75" +description = "Obsidian-style knowledge graph (SQLite + FTS5). Port of LBM internal/sage." + +[[bin]] +name = "kei-sage" +path = "src/main.rs" + +[lib] +name = "kei_sage" +path = "src/lib.rs" + +[dependencies] +rusqlite = { version = "0.31", features = ["bundled"] } +clap = { version = "4", features = ["derive"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +anyhow = "1" +chrono = { version = "0.4", default-features = false, features = ["clock"] } + +[dev-dependencies] +tempfile = "3" diff --git a/_primitives/_rust/kei-sage/src/bfs.rs b/_primitives/_rust/kei-sage/src/bfs.rs new file mode 100644 index 0000000..b8810f2 --- /dev/null +++ b/_primitives/_rust/kei-sage/src/bfs.rs @@ -0,0 +1,50 @@ +//! BFS traversal over the edges table, depth-limited, deduplicated. + +use crate::edges::list_outgoing; +use crate::store::Store; +use crate::types::Related; +use anyhow::Result; +use std::collections::{HashSet, VecDeque}; + +const MAX_RESULTS: usize = 500; +const MAX_DEPTH: i64 = 5; + +pub fn bfs(store: &Store, start: &str, max_depth: i64) -> Result> { + let depth = clamp_depth(max_depth); + let mut visited: HashSet = HashSet::new(); + visited.insert(start.to_string()); + let mut queue: VecDeque<(String, i64)> = VecDeque::new(); + queue.push_back((start.to_string(), 0)); + let mut out: Vec = Vec::new(); + while let Some((path, d)) = queue.pop_front() { + if out.len() >= MAX_RESULTS { + break; + } + if d >= depth { + continue; + } + for e in list_outgoing(store, &path)? { + if visited.contains(&e.dst_path) || out.len() >= MAX_RESULTS { + continue; + } + visited.insert(e.dst_path.clone()); + out.push(Related { + path: e.dst_path.clone(), + edge_type: e.edge_type, + depth: d + 1, + }); + queue.push_back((e.dst_path, d + 1)); + } + } + Ok(out) +} + +fn clamp_depth(d: i64) -> i64 { + if d <= 0 { + 2 + } else if d > MAX_DEPTH { + MAX_DEPTH + } else { + d + } +} diff --git a/_primitives/_rust/kei-sage/src/edges.rs b/_primitives/_rust/kei-sage/src/edges.rs new file mode 100644 index 0000000..391f9e0 --- /dev/null +++ b/_primitives/_rust/kei-sage/src/edges.rs @@ -0,0 +1,62 @@ +//! Typed-edge CRUD between vault_paths. + +use crate::store::Store; +use crate::types::Edge; +use anyhow::Result; +use chrono::Utc; +use rusqlite::params; + +pub fn add_edge(store: &Store, src: &str, dst: &str, edge_type: &str, weight: f64) -> Result { + let now = Utc::now().timestamp(); + store.conn().execute( + "INSERT OR IGNORE INTO edges (src_path, dst_path, edge_type, weight, created_at) + VALUES (?1,?2,?3,?4,?5)", + params![src, dst, edge_type, weight, now], + )?; + Ok(store.conn().last_insert_rowid()) +} + +pub fn remove_edge(store: &Store, src: &str, dst: &str, edge_type: &str) -> Result { + let n = store.conn().execute( + "DELETE FROM edges WHERE src_path=?1 AND dst_path=?2 AND edge_type=?3", + params![src, dst, edge_type], + )?; + Ok(n) +} + +pub fn list_outgoing(store: &Store, src: &str) -> Result> { + let mut stmt = store.conn().prepare( + "SELECT id, src_path, dst_path, edge_type, weight, created_at + FROM edges WHERE src_path=?1", + )?; + let rows = stmt.query_map(params![src], row_to_edge)?; + let mut out = Vec::new(); + for r in rows { + out.push(r?); + } + Ok(out) +} + +pub fn list_incoming(store: &Store, dst: &str) -> Result> { + let mut stmt = store.conn().prepare( + "SELECT id, src_path, dst_path, edge_type, weight, created_at + FROM edges WHERE dst_path=?1", + )?; + let rows = stmt.query_map(params![dst], row_to_edge)?; + let mut out = Vec::new(); + for r in rows { + out.push(r?); + } + Ok(out) +} + +fn row_to_edge(r: &rusqlite::Row) -> rusqlite::Result { + Ok(Edge { + id: r.get(0)?, + src_path: r.get(1)?, + dst_path: r.get(2)?, + edge_type: r.get(3)?, + weight: r.get(4)?, + created_at: r.get(5)?, + }) +} diff --git a/_primitives/_rust/kei-sage/src/import.rs b/_primitives/_rust/kei-sage/src/import.rs new file mode 100644 index 0000000..01d51e4 --- /dev/null +++ b/_primitives/_rust/kei-sage/src/import.rs @@ -0,0 +1,76 @@ +//! Obsidian-style vault import: walk a directory, ingest .md files. +//! +//! Minimal subset of LBM internal/sage/import_obsidian.go — we do NOT parse +//! frontmatter here (the upstream parser used multiple helper files). Port +//! of frontmatter/wikilinks parsing is a later milestone; this cube honours +//! the public interface. + +use crate::store::Store; +use crate::types::Unit; +use anyhow::{Context, Result}; +use std::fs; +use std::path::{Path, PathBuf}; + +pub struct ImportStats { + pub imported: usize, + pub skipped: usize, +} + +pub fn import_vault(store: &Store, root: &Path) -> Result { + let mut stats = ImportStats { imported: 0, skipped: 0 }; + let files = walk_md(root)?; + for path in files { + match ingest_one(store, root, &path) { + Ok(_) => stats.imported += 1, + Err(_) => stats.skipped += 1, + } + } + Ok(stats) +} + +fn walk_md(root: &Path) -> Result> { + let mut out = Vec::new(); + walk_recursive(root, &mut out)?; + Ok(out) +} + +fn walk_recursive(dir: &Path, out: &mut Vec) -> Result<()> { + if !dir.is_dir() { + return Ok(()); + } + for entry in fs::read_dir(dir).with_context(|| format!("read_dir {}", dir.display()))? { + let entry = entry?; + let path = entry.path(); + if path.is_dir() { + walk_recursive(&path, out)?; + } else if path.extension().and_then(|s| s.to_str()) == Some("md") { + out.push(path); + } + } + Ok(()) +} + +fn ingest_one(store: &Store, root: &Path, path: &Path) -> Result<()> { + let content = fs::read_to_string(path)?; + let title = path.file_stem() + .and_then(|s| s.to_str()) + .unwrap_or("untitled") + .to_string(); + let vault_path = path.strip_prefix(root) + .ok() + .and_then(|p| p.to_str()) + .unwrap_or(&title) + .to_string(); + let unit = Unit { + unit_type: "note".into(), + title, + content, + evidence_grade: "E4".into(), + source_path: path.to_string_lossy().into(), + vault_path, + category: String::new(), + ..Default::default() + }; + store.add_unit(&unit)?; + Ok(()) +} diff --git a/_primitives/_rust/kei-sage/src/lib.rs b/_primitives/_rust/kei-sage/src/lib.rs new file mode 100644 index 0000000..9980331 --- /dev/null +++ b/_primitives/_rust/kei-sage/src/lib.rs @@ -0,0 +1,15 @@ +//! kei-sage — SQLite knowledge-vault with FTS5 + typed edges + BFS + PageRank. +//! +//! Port of LBM internal/sage. Constructor Pattern: one concept per file. + +pub mod bfs; +pub mod edges; +pub mod import; +pub mod pagerank; +pub mod schema; +pub mod search; +pub mod store; +pub mod types; + +pub use store::Store; +pub use types::{Edge, Related, Unit}; diff --git a/_primitives/_rust/kei-sage/src/main.rs b/_primitives/_rust/kei-sage/src/main.rs new file mode 100644 index 0000000..e252464 --- /dev/null +++ b/_primitives/_rust/kei-sage/src/main.rs @@ -0,0 +1,101 @@ +//! kei-sage CLI — import / search / related / rank / add / edit. + +use clap::{Parser, Subcommand}; +use kei_sage::bfs::bfs; +use kei_sage::edges::add_edge; +use kei_sage::import::import_vault; +use kei_sage::pagerank::pagerank; +use kei_sage::search::fts_search; +use kei_sage::{Store, Unit}; +use std::path::PathBuf; +use std::process::ExitCode; + +#[derive(Parser)] +#[command(name = "kei-sage", version, about = "Obsidian-style knowledge vault")] +struct Cli { + /// Database path (default: $KEI_VAULT_DB or ~/.claude/sage/vault.sqlite) + #[arg(long)] + db: Option, + #[command(subcommand)] + cmd: Cmd, +} + +#[derive(Subcommand)] +enum Cmd { + Import { vault: PathBuf }, + Search { query: String, #[arg(long, default_value_t = 20)] limit: i64 }, + Related { key: String, #[arg(long, default_value_t = 2)] depth: i64 }, + Rank { #[arg(long, default_value_t = 20)] limit: usize }, + Add { + #[arg(long)] title: String, + #[arg(long, default_value = "")] content: String, + #[arg(long, default_value = "")] vault_path: String, + #[arg(long, default_value = "E4")] grade: String, + }, + Edit { + id: i64, + #[arg(long)] title: Option, + #[arg(long)] content: Option, + #[arg(long)] grade: Option, + }, + Link { src: String, dst: String, #[arg(long, default_value = "related")] edge_type: String }, +} + +fn db_path(cli_db: Option) -> PathBuf { + if let Some(p) = cli_db { return p; } + if let Ok(e) = std::env::var("KEI_VAULT_DB") { return PathBuf::from(e); } + let home = std::env::var("HOME").unwrap_or_else(|_| ".".into()); + PathBuf::from(home).join(".claude/sage/vault.sqlite") +} + +fn run() -> anyhow::Result<()> { + let cli = Cli::parse(); + let path = db_path(cli.db); + let store = Store::open(&path)?; + match cli.cmd { + Cmd::Import { vault } => { + let s = import_vault(&store, &vault)?; + println!("imported={} skipped={}", s.imported, s.skipped); + } + Cmd::Search { query, limit } => { + for u in fts_search(&store, &query, limit)? { + println!("{}\t{}\t{}", u.id, u.evidence_grade, u.title); + } + } + Cmd::Related { key, depth } => { + for r in bfs(&store, &key, depth)? { + println!("{}\t{}\t(depth {})", r.edge_type, r.path, r.depth); + } + } + Cmd::Rank { limit } => { + for (p, s) in pagerank(&store)?.into_iter().take(limit) { + println!("{:.6}\t{}", s, p); + } + } + Cmd::Add { title, content, vault_path, grade } => { + let id = store.add_unit(&Unit { title, content, vault_path, + evidence_grade: grade, unit_type: "note".into(), ..Default::default() })?; + println!("{}", id); + } + Cmd::Edit { id, title, content, grade } => { + let mut u = store.get_unit(id)?.ok_or_else(|| anyhow::anyhow!("id {id} not found"))?; + if let Some(t) = title { u.title = t; } + if let Some(c) = content { u.content = c; } + if let Some(g) = grade { u.evidence_grade = g; } + store.update_unit(&u)?; + println!("updated {}", id); + } + Cmd::Link { src, dst, edge_type } => { + add_edge(&store, &src, &dst, &edge_type, 1.0)?; + println!("linked {} -> {}", src, dst); + } + } + Ok(()) +} + +fn main() -> ExitCode { + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(e) => { eprintln!("kei-sage: {e:#}"); ExitCode::from(1) } + } +} diff --git a/_primitives/_rust/kei-sage/src/pagerank.rs b/_primitives/_rust/kei-sage/src/pagerank.rs new file mode 100644 index 0000000..28baad9 --- /dev/null +++ b/_primitives/_rust/kei-sage/src/pagerank.rs @@ -0,0 +1,60 @@ +//! PageRank — power-iteration, 50 iterations, d=0.85. Operates on the edges table. + +use crate::store::Store; +use anyhow::Result; +use std::collections::HashMap; + +const DAMPING: f64 = 0.85; +const ITERATIONS: usize = 50; + +/// Compute PageRank over the edges table. Returns [(path, score)] sorted desc. +pub fn pagerank(store: &Store) -> Result> { + let (nodes, out_edges) = collect_graph(store)?; + if nodes.is_empty() { + return Ok(Vec::new()); + } + let mut rank: HashMap = nodes.iter() + .map(|n| (n.clone(), 1.0 / nodes.len() as f64)).collect(); + for _ in 0..ITERATIONS { + rank = one_iteration(&nodes, &out_edges, &rank); + } + let mut out: Vec<(String, f64)> = rank.into_iter().collect(); + out.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + Ok(out) +} + +fn collect_graph(store: &Store) -> Result<(Vec, HashMap>)> { + let mut stmt = store.conn().prepare("SELECT src_path, dst_path FROM edges")?; + let rows = stmt.query_map([], |r| Ok((r.get::<_, String>(0)?, r.get::<_, String>(1)?)))?; + let mut nodes: std::collections::HashSet = std::collections::HashSet::new(); + let mut out_edges: HashMap> = HashMap::new(); + for row in rows { + let (src, dst) = row?; + nodes.insert(src.clone()); + nodes.insert(dst.clone()); + out_edges.entry(src).or_default().push(dst); + } + Ok((nodes.into_iter().collect(), out_edges)) +} + +fn one_iteration( + nodes: &[String], + out_edges: &HashMap>, + prev: &HashMap, +) -> HashMap { + let n = nodes.len() as f64; + let base = (1.0 - DAMPING) / n; + let mut next: HashMap = nodes.iter().map(|k| (k.clone(), base)).collect(); + for (src, dsts) in out_edges { + if dsts.is_empty() { + continue; + } + let share = DAMPING * prev.get(src).copied().unwrap_or(0.0) / dsts.len() as f64; + for dst in dsts { + if let Some(slot) = next.get_mut(dst) { + *slot += share; + } + } + } + next +} diff --git a/_primitives/_rust/kei-sage/src/schema.rs b/_primitives/_rust/kei-sage/src/schema.rs new file mode 100644 index 0000000..6b60a9d --- /dev/null +++ b/_primitives/_rust/kei-sage/src/schema.rs @@ -0,0 +1,57 @@ +//! SQLite schema for knowledge-vault. Port of LBM internal/sage/vault_schema.go. + +use rusqlite::{Connection, Result}; + +/// Apply schema + FTS5 virtual table. Idempotent. +pub fn create_schema(conn: &Connection) -> Result<()> { + conn.execute_batch( + r#" + CREATE TABLE IF NOT EXISTS knowledge_units ( + id INTEGER PRIMARY KEY, + unit_type TEXT NOT NULL, + title TEXT NOT NULL, + content TEXT DEFAULT '', + evidence_grade TEXT DEFAULT '', + source_path TEXT DEFAULT '', + vault_path TEXT DEFAULT '', + category TEXT DEFAULT '', + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_ku_type ON knowledge_units(unit_type); + CREATE UNIQUE INDEX IF NOT EXISTS idx_ku_vault + ON knowledge_units(vault_path) WHERE vault_path != ''; + CREATE INDEX IF NOT EXISTS idx_ku_grade ON knowledge_units(evidence_grade); + + CREATE TABLE IF NOT EXISTS tags ( + id INTEGER PRIMARY KEY, + name TEXT UNIQUE NOT NULL + ); + + CREATE TABLE IF NOT EXISTS unit_tags ( + unit_id INTEGER NOT NULL REFERENCES knowledge_units(id) ON DELETE CASCADE, + tag_id INTEGER NOT NULL REFERENCES tags(id) ON DELETE CASCADE, + PRIMARY KEY (unit_id, tag_id) + ); + + CREATE TABLE IF NOT EXISTS edges ( + id INTEGER PRIMARY KEY, + src_path TEXT NOT NULL, + dst_path TEXT NOT NULL, + edge_type TEXT NOT NULL, + weight REAL DEFAULT 1.0, + created_at INTEGER NOT NULL, + UNIQUE(src_path, dst_path, edge_type) + ); + CREATE INDEX IF NOT EXISTS idx_sage_edges_src ON edges(src_path); + CREATE INDEX IF NOT EXISTS idx_sage_edges_dst ON edges(dst_path); + "#, + )?; + conn.execute_batch( + r#" + CREATE VIRTUAL TABLE IF NOT EXISTS fts_knowledge + USING fts5(unit_id UNINDEXED, title, content, tokenize='porter unicode61'); + "#, + )?; + Ok(()) +} diff --git a/_primitives/_rust/kei-sage/src/search.rs b/_primitives/_rust/kei-sage/src/search.rs new file mode 100644 index 0000000..3b8d531 --- /dev/null +++ b/_primitives/_rust/kei-sage/src/search.rs @@ -0,0 +1,39 @@ +//! FTS5 search over knowledge_units. + +use crate::store::Store; +use crate::types::Unit; +use anyhow::Result; +use rusqlite::params; + +/// Full-text search. Returns matching Units ordered by SQLite FTS5 rank. +pub fn fts_search(store: &Store, query: &str, limit: i64) -> Result> { + let lim = if limit <= 0 { 20 } else { limit }; + let mut stmt = store.conn().prepare( + "SELECT k.id, k.unit_type, k.title, k.content, k.evidence_grade, + k.source_path, k.vault_path, k.category, k.created_at, k.updated_at + FROM fts_knowledge f + JOIN knowledge_units k ON k.id = f.unit_id + WHERE fts_knowledge MATCH ?1 + ORDER BY rank + LIMIT ?2", + )?; + let rows = stmt.query_map(params![query, lim], |r| { + Ok(Unit { + id: r.get(0)?, + unit_type: r.get(1)?, + title: r.get(2)?, + content: r.get(3)?, + evidence_grade: r.get(4)?, + source_path: r.get(5)?, + vault_path: r.get(6)?, + category: r.get(7)?, + created_at: r.get(8)?, + updated_at: r.get(9)?, + }) + })?; + let mut out = Vec::new(); + for row in rows { + out.push(row?); + } + Ok(out) +} diff --git a/_primitives/_rust/kei-sage/src/store.rs b/_primitives/_rust/kei-sage/src/store.rs new file mode 100644 index 0000000..c21cfc2 --- /dev/null +++ b/_primitives/_rust/kei-sage/src/store.rs @@ -0,0 +1,111 @@ +//! Knowledge-unit CRUD + FTS indexer. + +use crate::schema::create_schema; +use crate::types::Unit; +use anyhow::{Context, Result}; +use chrono::Utc; +use rusqlite::{params, Connection}; +use std::path::Path; + +pub struct Store { + conn: Connection, +} + +impl Store { + pub fn open(path: &Path) -> Result { + if let Some(parent) = path.parent() { + let _ = std::fs::create_dir_all(parent); + } + let conn = Connection::open(path).context("open sqlite")?; + conn.pragma_update(None, "journal_mode", "WAL").ok(); + create_schema(&conn)?; + Ok(Self { conn }) + } + + pub fn open_memory() -> Result { + let conn = Connection::open_in_memory()?; + create_schema(&conn)?; + Ok(Self { conn }) + } + + pub fn conn(&self) -> &Connection { + &self.conn + } + + /// Insert a new knowledge unit. Indexes title+content into FTS5. Idempotent by vault_path. + pub fn add_unit(&self, unit: &Unit) -> Result { + let now = Utc::now().timestamp(); + let created = if unit.created_at == 0 { now } else { unit.created_at }; + self.conn.execute( + "INSERT OR REPLACE INTO knowledge_units + (unit_type, title, content, evidence_grade, source_path, + vault_path, category, created_at, updated_at) + VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9)", + params![unit.unit_type, unit.title, unit.content, unit.evidence_grade, + unit.source_path, unit.vault_path, unit.category, created, now], + )?; + let id = self.conn.last_insert_rowid(); + self.reindex_fts(id, &unit.title, &unit.content)?; + Ok(id) + } + + pub fn get_unit(&self, id: i64) -> Result> { + let mut stmt = self.conn.prepare( + "SELECT id, unit_type, title, content, evidence_grade, source_path, + vault_path, category, created_at, updated_at + FROM knowledge_units WHERE id=?1", + )?; + let mut rows = stmt.query(params![id])?; + if let Some(r) = rows.next()? { + return Ok(Some(row_to_unit(r)?)); + } + Ok(None) + } + + pub fn update_unit(&self, unit: &Unit) -> Result<()> { + let now = Utc::now().timestamp(); + self.conn.execute( + "UPDATE knowledge_units SET title=?1, content=?2, evidence_grade=?3, + category=?4, updated_at=?5 WHERE id=?6", + params![unit.title, unit.content, unit.evidence_grade, + unit.category, now, unit.id], + )?; + self.reindex_fts(unit.id, &unit.title, &unit.content)?; + Ok(()) + } + + pub fn delete_unit(&self, id: i64) -> Result<()> { + self.conn.execute("DELETE FROM fts_knowledge WHERE unit_id=?1", params![id])?; + self.conn.execute("DELETE FROM knowledge_units WHERE id=?1", params![id])?; + Ok(()) + } + + pub fn count_units(&self) -> Result { + Ok(self.conn.query_row( + "SELECT COUNT(*) FROM knowledge_units", [], |r| r.get(0))?) + } + + fn reindex_fts(&self, id: i64, title: &str, content: &str) -> Result<()> { + self.conn.execute("DELETE FROM fts_knowledge WHERE unit_id=?1", params![id])?; + self.conn.execute( + "INSERT INTO fts_knowledge (unit_id, title, content) VALUES (?1,?2,?3)", + params![id, title, content], + )?; + Ok(()) + } +} + +fn row_to_unit(r: &rusqlite::Row) -> rusqlite::Result { + Ok(Unit { + id: r.get(0)?, + unit_type: r.get(1)?, + title: r.get(2)?, + content: r.get(3)?, + evidence_grade: r.get(4)?, + source_path: r.get(5)?, + vault_path: r.get(6)?, + category: r.get(7)?, + created_at: r.get(8)?, + updated_at: r.get(9)?, + }) +} diff --git a/_primitives/_rust/kei-sage/src/types.rs b/_primitives/_rust/kei-sage/src/types.rs new file mode 100644 index 0000000..55fa223 --- /dev/null +++ b/_primitives/_rust/kei-sage/src/types.rs @@ -0,0 +1,34 @@ +//! Shared value types for knowledge units + edges + BFS results. + +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct Unit { + pub id: i64, + pub unit_type: String, + pub title: String, + pub content: String, + pub evidence_grade: String, + pub source_path: String, + pub vault_path: String, + pub category: String, + pub created_at: i64, + pub updated_at: i64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Edge { + pub id: i64, + pub src_path: String, + pub dst_path: String, + pub edge_type: String, + pub weight: f64, + pub created_at: i64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Related { + pub path: String, + pub edge_type: String, + pub depth: i64, +} diff --git a/_primitives/_rust/kei-sage/tests/integration.rs b/_primitives/_rust/kei-sage/tests/integration.rs new file mode 100644 index 0000000..0c19550 --- /dev/null +++ b/_primitives/_rust/kei-sage/tests/integration.rs @@ -0,0 +1,110 @@ +//! kei-sage integration tests. + +use kei_sage::bfs::bfs; +use kei_sage::edges::{add_edge, list_outgoing}; +use kei_sage::import::import_vault; +use kei_sage::pagerank::pagerank; +use kei_sage::search::fts_search; +use kei_sage::{Store, Unit}; +use std::fs; +use tempfile::tempdir; + +fn mkstore() -> Store { Store::open_memory().unwrap() } + +fn mkunit(title: &str, body: &str, vault: &str) -> Unit { + Unit { + unit_type: "note".into(), title: title.into(), content: body.into(), + evidence_grade: "E2".into(), vault_path: vault.into(), + ..Default::default() + } +} + +#[test] +fn crud_roundtrip() { + let s = mkstore(); + let id = s.add_unit(&mkunit("hello", "world", "a.md")).unwrap(); + assert!(id > 0); + let u = s.get_unit(id).unwrap().unwrap(); + assert_eq!(u.title, "hello"); + s.delete_unit(id).unwrap(); + assert!(s.get_unit(id).unwrap().is_none()); +} + +#[test] +fn fts_search_matches() { + let s = mkstore(); + s.add_unit(&mkunit("rust async", "tokio runtime details", "a.md")).unwrap(); + s.add_unit(&mkunit("python sync", "flask wsgi server", "b.md")).unwrap(); + let hits = fts_search(&s, "tokio", 10).unwrap(); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].title, "rust async"); +} + +#[test] +fn bfs_depth_limit() { + let s = mkstore(); + add_edge(&s, "a", "b", "rel", 1.0).unwrap(); + add_edge(&s, "b", "c", "rel", 1.0).unwrap(); + add_edge(&s, "c", "d", "rel", 1.0).unwrap(); + let out = bfs(&s, "a", 2).unwrap(); + let paths: Vec<&str> = out.iter().map(|r| r.path.as_str()).collect(); + assert!(paths.contains(&"b")); + assert!(paths.contains(&"c")); + assert!(!paths.contains(&"d")); +} + +#[test] +fn pagerank_orders_by_popularity() { + let s = mkstore(); + add_edge(&s, "a", "hub", "rel", 1.0).unwrap(); + add_edge(&s, "b", "hub", "rel", 1.0).unwrap(); + add_edge(&s, "c", "hub", "rel", 1.0).unwrap(); + add_edge(&s, "d", "hub", "rel", 1.0).unwrap(); + add_edge(&s, "e", "hub", "rel", 1.0).unwrap(); + let ranks = pagerank(&s).unwrap(); + assert_eq!(ranks[0].0, "hub"); +} + +#[test] +fn edges_crud() { + let s = mkstore(); + let id = add_edge(&s, "x", "y", "cites", 0.8).unwrap(); + assert!(id > 0); + let out = list_outgoing(&s, "x").unwrap(); + assert_eq!(out.len(), 1); + assert_eq!(out[0].dst_path, "y"); +} + +#[test] +fn import_idempotency() { + let tmp = tempdir().unwrap(); + let p = tmp.path().join("one.md"); + fs::write(&p, "# title one\nhello").unwrap(); + let s = mkstore(); + let first = import_vault(&s, tmp.path()).unwrap(); + let second = import_vault(&s, tmp.path()).unwrap(); + assert_eq!(first.imported, 1); + assert_eq!(second.imported, 1); + assert_eq!(s.count_units().unwrap(), 1); +} + +#[test] +fn edges_cross_reference_validates() { + let s = mkstore(); + s.add_unit(&mkunit("note a", "", "a.md")).unwrap(); + s.add_unit(&mkunit("note b", "", "b.md")).unwrap(); + add_edge(&s, "a.md", "b.md", "refs", 1.0).unwrap(); + let out = list_outgoing(&s, "a.md").unwrap(); + assert_eq!(out.len(), 1); +} + +#[test] +fn fts5_respects_limit() { + let s = mkstore(); + for i in 0..25 { + let t = format!("rust note {i}"); + s.add_unit(&mkunit(&t, "rust rust rust", &format!("n{i}.md"))).unwrap(); + } + let hits = fts_search(&s, "rust", 5).unwrap(); + assert_eq!(hits.len(), 5); +} diff --git a/_primitives/_rust/kei-search-core/Cargo.toml b/_primitives/_rust/kei-search-core/Cargo.toml new file mode 100644 index 0000000..d549fbe --- /dev/null +++ b/_primitives/_rust/kei-search-core/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "kei-search-core" +version = "0.1.0" +edition = "2021" +rust-version = "1.75" +description = "3-wave deep research scaffolding with budget cap. Port of LBM internal/search (fetch stubbed)." + +[[bin]] +name = "kei-search-core" +path = "src/main.rs" + +[lib] +name = "kei_search_core" +path = "src/lib.rs" + +[dependencies] +rusqlite = { version = "0.31", features = ["bundled"] } +clap = { version = "4", features = ["derive"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +anyhow = "1" +chrono = { version = "0.4", default-features = false, features = ["clock"] } + +[dev-dependencies] +tempfile = "3" diff --git a/_primitives/_rust/kei-search-core/src/budget.rs b/_primitives/_rust/kei-search-core/src/budget.rs new file mode 100644 index 0000000..da55653 --- /dev/null +++ b/_primitives/_rust/kei-search-core/src/budget.rs @@ -0,0 +1,34 @@ +//! Budget tracker — all costs in microcents (1 USD = 1_000_000 mc). + +use anyhow::{anyhow, Result}; + +#[derive(Debug, Clone)] +pub struct Budget { + cap_mc: i64, + spent_mc: i64, + stopped: bool, +} + +impl Budget { + pub fn new(cap_mc: i64) -> Self { + Self { cap_mc, spent_mc: 0, stopped: false } + } + + /// Record a cost; returns error if this push would exceed the cap. + pub fn charge(&mut self, mc: i64) -> Result<()> { + if self.stopped { + return Err(anyhow!("budget stopped")); + } + if self.spent_mc + mc > self.cap_mc { + return Err(anyhow!( + "budget exceeded: spent={} cap={}", self.spent_mc + mc, self.cap_mc)); + } + self.spent_mc += mc; + Ok(()) + } + + pub fn spent(&self) -> i64 { self.spent_mc } + pub fn remaining(&self) -> i64 { self.cap_mc - self.spent_mc } + pub fn stop(&mut self) { self.stopped = true; } + pub fn is_stopped(&self) -> bool { self.stopped } +} diff --git a/_primitives/_rust/kei-search-core/src/export.rs b/_primitives/_rust/kei-search-core/src/export.rs new file mode 100644 index 0000000..fc1abb1 --- /dev/null +++ b/_primitives/_rust/kei-search-core/src/export.rs @@ -0,0 +1,40 @@ +//! Export research → markdown / JSON. + +use crate::store::ResearchStore; +use anyhow::{anyhow, Result}; +use serde_json::json; + +pub enum Format { + Markdown, + Json, +} + +pub fn export(store: &ResearchStore, id: i64, fmt: Format) -> Result { + let r = store.get_research(id)?.ok_or_else(|| anyhow!("research {id} missing"))?; + let claims = store.claims_for(id)?; + match fmt { + Format::Markdown => { + let mut md = String::new(); + md.push_str(&format!("# Research {}\n\n", r.id)); + md.push_str(&format!("**Query:** {}\n\n", r.query_original)); + md.push_str(&format!("**Status:** {}\n", r.status)); + md.push_str(&format!("**Cost:** {} mc\n\n", r.total_cost_mc)); + md.push_str("## Claims\n\n"); + for c in claims { + md.push_str(&format!("- [{}] {} (consensus={:.2})\n", + c.grade, c.claim_text, c.consensus)); + } + Ok(md) + } + Format::Json => { + let val = json!({ + "id": r.id, + "query": r.query_original, + "status": r.status, + "cost_mc": r.total_cost_mc, + "claims": claims, + }); + Ok(serde_json::to_string_pretty(&val)?) + } + } +} diff --git a/_primitives/_rust/kei-search-core/src/fetch.rs b/_primitives/_rust/kei-search-core/src/fetch.rs new file mode 100644 index 0000000..28c345f --- /dev/null +++ b/_primitives/_rust/kei-search-core/src/fetch.rs @@ -0,0 +1,23 @@ +//! Source fetcher trait — frozen interface, default impl is a no-op stub. +//! +//! Actual WebFetch/WebSearch integration is out-of-scope for v0.14 part A. +//! Later milestones plug real providers (anthropic-websearch, SerpAPI, etc.). + +use crate::types::Source; + +/// Implement this trait to integrate a live search provider. +pub trait SourceFetcher { + /// Fetch sources for `claim`. Returns (source, cost_microcents). + /// Cost is real — the budget is charged by the pipeline, not by impl. + fn fetch(&self, claim: &str) -> (Vec, i64); +} + +/// Default stub — returns empty. Frozen interface, no runtime side-effects. +pub struct StubFetcher; + +impl SourceFetcher for StubFetcher { + fn fetch(&self, _claim: &str) -> (Vec, i64) { + // TODO(v0.15): wire to real websearch. Kept as stub per v0.14 spec. + (Vec::new(), 0) + } +} diff --git a/_primitives/_rust/kei-search-core/src/lib.rs b/_primitives/_rust/kei-search-core/src/lib.rs new file mode 100644 index 0000000..c008159 --- /dev/null +++ b/_primitives/_rust/kei-search-core/src/lib.rs @@ -0,0 +1,22 @@ +//! kei-search-core — 3-wave deep research engine, budget-capped. +//! +//! Waves: +//! 0 — claim extraction from prompt +//! 1 — per-claim source hunt (WebFetch stubbed behind [`SourceFetcher`] trait) +//! 2 — cross-validation + consensus scoring +//! +//! Port of LBM internal/search. The actual fetch is a trait the caller +//! supplies. Default implementation returns empty (frozen interface, todo!() +//! reflects unimplemented runtime). + +pub mod budget; +pub mod export; +pub mod fetch; +pub mod pipeline; +pub mod schema; +pub mod store; +pub mod types; + +pub use pipeline::run_research; +pub use store::ResearchStore; +pub use types::{Claim, Research, Source}; diff --git a/_primitives/_rust/kei-search-core/src/main.rs b/_primitives/_rust/kei-search-core/src/main.rs new file mode 100644 index 0000000..c5b6f6e --- /dev/null +++ b/_primitives/_rust/kei-search-core/src/main.rs @@ -0,0 +1,61 @@ +//! kei-search-core CLI. + +use clap::{Parser, Subcommand, ValueEnum}; +use kei_search_core::export::{export, Format}; +use kei_search_core::fetch::StubFetcher; +use kei_search_core::pipeline::run_research; +use kei_search_core::ResearchStore; +use std::path::PathBuf; +use std::process::ExitCode; + +#[derive(Parser)] +#[command(name = "kei-search-core", version)] +struct Cli { + #[arg(long)] db: Option, + #[command(subcommand)] cmd: Cmd, +} + +#[derive(Subcommand)] +enum Cmd { + Run { prompt: String, + #[arg(long, default_value_t = 1_000_000)] budget: i64 }, // 1 USD + Stop { id: i64 }, + Export { id: i64, #[arg(long, value_enum, default_value_t = Fmt::Md)] format: Fmt }, +} + +#[derive(Clone, Copy, ValueEnum)] +enum Fmt { Md, Json } + +fn db_path(o: Option) -> PathBuf { + if let Some(p) = o { return p; } + if let Ok(e) = std::env::var("KEI_SEARCH_DB") { return PathBuf::from(e); } + let home = std::env::var("HOME").unwrap_or_else(|_| ".".into()); + PathBuf::from(home).join(".claude/search/research.sqlite") +} + +fn run() -> anyhow::Result<()> { + let cli = Cli::parse(); + let s = ResearchStore::open(&db_path(cli.db))?; + match cli.cmd { + Cmd::Run { prompt, budget } => { + let id = run_research(&s, &StubFetcher, &prompt, budget)?; + println!("{}", id); + } + Cmd::Stop { id } => { + s.set_status(id, "stopped")?; + println!("stopped {}", id); + } + Cmd::Export { id, format } => { + let f = match format { Fmt::Md => Format::Markdown, Fmt::Json => Format::Json }; + println!("{}", export(&s, id, f)?); + } + } + Ok(()) +} + +fn main() -> ExitCode { + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(e) => { eprintln!("kei-search-core: {e:#}"); ExitCode::from(1) } + } +} diff --git a/_primitives/_rust/kei-search-core/src/pipeline.rs b/_primitives/_rust/kei-search-core/src/pipeline.rs new file mode 100644 index 0000000..ef55317 --- /dev/null +++ b/_primitives/_rust/kei-search-core/src/pipeline.rs @@ -0,0 +1,90 @@ +//! 3-wave research runner. +//! +//! Wave 0: split prompt into claims (naive split on `.`; real NLU later). +//! Wave 1: for each claim, fetch sources via [`SourceFetcher`]. +//! Wave 2: score consensus per claim from sources (majority = higher grade). + +use crate::budget::Budget; +use crate::fetch::SourceFetcher; +use crate::store::ResearchStore; +use crate::types::{Claim, Source}; +use anyhow::Result; + +const WAVE1_COST_PER_CLAIM_MC: i64 = 100; // 0.01 USD per claim +const WAVE2_COST_MC: i64 = 50; + +pub fn run_research( + store: &ResearchStore, + fetcher: &dyn SourceFetcher, + prompt: &str, + budget_mc: i64, +) -> Result { + let research_id = store.create_research(prompt)?; + let mut budget = Budget::new(budget_mc); + let claims_text = wave_0_extract_claims(prompt); + if let Err(e) = wave_1_fetch(store, fetcher, research_id, &claims_text, &mut budget) { + store.set_status(research_id, "failed")?; + return Err(e); + } + if let Err(e) = wave_2_consensus(store, research_id, &mut budget) { + store.set_status(research_id, "failed")?; + return Err(e); + } + store.set_cost(research_id, budget.spent())?; + store.set_status(research_id, "completed")?; + Ok(research_id) +} + +fn wave_0_extract_claims(prompt: &str) -> Vec { + prompt + .split(|c: char| c == '.' || c == '?' || c == '\n') + .map(|s| s.trim().to_string()) + .filter(|s| s.len() > 4) + .collect() +} + +fn wave_1_fetch( + store: &ResearchStore, + fetcher: &dyn SourceFetcher, + rid: i64, + claims: &[String], + budget: &mut Budget, +) -> Result<()> { + for c in claims { + budget.charge(WAVE1_COST_PER_CLAIM_MC)?; + let (srcs, fetch_cost) = fetcher.fetch(c); + if fetch_cost > 0 { + budget.charge(fetch_cost)?; + } + for s in srcs { + store.add_source(&Source { research_id: rid, ..s })?; + } + store.add_claim(&Claim { + research_id: rid, + claim_text: c.clone(), + ..Default::default() + })?; + } + Ok(()) +} + +fn wave_2_consensus(store: &ResearchStore, rid: i64, budget: &mut Budget) -> Result<()> { + budget.charge(WAVE2_COST_MC)?; + let claims = store.claims_for(rid)?; + for c in claims { + let support = 0.5; + let contradict = 0.0; + let consensus = support - contradict; + let grade = grade_from_consensus(consensus); + store.conn().execute( + "UPDATE claims SET support=?1, contradict=?2, consensus=?3, grade=?4 + WHERE id=?5", + rusqlite::params![support, contradict, consensus, grade, c.id], + )?; + } + Ok(()) +} + +fn grade_from_consensus(c: f64) -> &'static str { + if c >= 0.8 { "E2" } else if c >= 0.5 { "E4" } else { "E6" } +} diff --git a/_primitives/_rust/kei-search-core/src/schema.rs b/_primitives/_rust/kei-search-core/src/schema.rs new file mode 100644 index 0000000..37bd2c4 --- /dev/null +++ b/_primitives/_rust/kei-search-core/src/schema.rs @@ -0,0 +1,42 @@ +use rusqlite::{Connection, Result}; + +pub fn create_schema(conn: &Connection) -> Result<()> { + conn.execute_batch(r#" + CREATE TABLE IF NOT EXISTS researches ( + id INTEGER PRIMARY KEY, + query_original TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + result_markdown TEXT DEFAULT '', + total_cost_mc INTEGER DEFAULT 0, + created_at INTEGER NOT NULL, + completed_at INTEGER DEFAULT 0 + ); + CREATE INDEX IF NOT EXISTS idx_res_status ON researches(status); + + CREATE TABLE IF NOT EXISTS sources ( + id INTEGER PRIMARY KEY, + research_id INTEGER NOT NULL REFERENCES researches(id), + url TEXT NOT NULL, + title TEXT DEFAULT '', + content TEXT DEFAULT '', + provider TEXT DEFAULT '', + domain TEXT DEFAULT '', + relevance_score REAL DEFAULT 0.0, + created_at INTEGER NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_src_research ON sources(research_id); + + CREATE TABLE IF NOT EXISTS claims ( + id INTEGER PRIMARY KEY, + research_id INTEGER NOT NULL REFERENCES researches(id), + claim_text TEXT NOT NULL, + support REAL DEFAULT 0.0, + contradict REAL DEFAULT 0.0, + consensus REAL DEFAULT 0.0, + grade TEXT DEFAULT 'E6', + created_at INTEGER NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_claim_research ON claims(research_id); + "#)?; + Ok(()) +} diff --git a/_primitives/_rust/kei-search-core/src/store.rs b/_primitives/_rust/kei-search-core/src/store.rs new file mode 100644 index 0000000..47a3d76 --- /dev/null +++ b/_primitives/_rust/kei-search-core/src/store.rs @@ -0,0 +1,122 @@ +use crate::schema::create_schema; +use crate::types::{Claim, Research, Source}; +use anyhow::{Context, Result}; +use chrono::Utc; +use rusqlite::{params, Connection}; +use std::path::Path; + +pub struct ResearchStore { + conn: Connection, +} + +impl ResearchStore { + pub fn open(path: &Path) -> Result { + if let Some(parent) = path.parent() { + let _ = std::fs::create_dir_all(parent); + } + let conn = Connection::open(path).context("open sqlite")?; + conn.pragma_update(None, "journal_mode", "WAL").ok(); + create_schema(&conn)?; + Ok(Self { conn }) + } + + pub fn open_memory() -> Result { + let conn = Connection::open_in_memory()?; + create_schema(&conn)?; + Ok(Self { conn }) + } + + pub fn conn(&self) -> &Connection { &self.conn } + + pub fn create_research(&self, query: &str) -> Result { + let now = Utc::now().timestamp(); + self.conn.execute( + "INSERT INTO researches (query_original, status, created_at) + VALUES (?1, 'running', ?2)", + params![query, now], + )?; + Ok(self.conn.last_insert_rowid()) + } + + pub fn set_status(&self, id: i64, status: &str) -> Result<()> { + let now = Utc::now().timestamp(); + self.conn.execute( + "UPDATE researches SET status=?1, completed_at=?2 WHERE id=?3", + params![status, now, id], + )?; + Ok(()) + } + + pub fn set_cost(&self, id: i64, mc: i64) -> Result<()> { + self.conn.execute( + "UPDATE researches SET total_cost_mc=?1 WHERE id=?2", + params![mc, id], + )?; + Ok(()) + } + + pub fn set_markdown(&self, id: i64, md: &str) -> Result<()> { + self.conn.execute( + "UPDATE researches SET result_markdown=?1 WHERE id=?2", + params![md, id], + )?; + Ok(()) + } + + pub fn get_research(&self, id: i64) -> Result> { + let mut stmt = self.conn.prepare( + "SELECT id, query_original, status, result_markdown, total_cost_mc, + created_at, completed_at FROM researches WHERE id=?1", + )?; + let mut rows = stmt.query(params![id])?; + if let Some(r) = rows.next()? { + return Ok(Some(Research { + id: r.get(0)?, query_original: r.get(1)?, status: r.get(2)?, + result_markdown: r.get(3)?, total_cost_mc: r.get(4)?, + created_at: r.get(5)?, completed_at: r.get(6)?, + })); + } + Ok(None) + } + + pub fn add_source(&self, s: &Source) -> Result { + let now = Utc::now().timestamp(); + self.conn.execute( + "INSERT INTO sources (research_id, url, title, content, provider, + domain, relevance_score, created_at) + VALUES (?1,?2,?3,?4,?5,?6,?7,?8)", + params![s.research_id, s.url, s.title, s.content, s.provider, + s.domain, s.relevance_score, now], + )?; + Ok(self.conn.last_insert_rowid()) + } + + pub fn add_claim(&self, c: &Claim) -> Result { + let now = Utc::now().timestamp(); + self.conn.execute( + "INSERT INTO claims (research_id, claim_text, support, contradict, + consensus, grade, created_at) + VALUES (?1,?2,?3,?4,?5,?6,?7)", + params![c.research_id, c.claim_text, c.support, c.contradict, + c.consensus, c.grade, now], + )?; + Ok(self.conn.last_insert_rowid()) + } + + pub fn claims_for(&self, research_id: i64) -> Result> { + let mut stmt = self.conn.prepare( + "SELECT id, research_id, claim_text, support, contradict, + consensus, grade, created_at FROM claims WHERE research_id=?1" + )?; + let rows = stmt.query_map(params![research_id], |r| { + Ok(Claim { + id: r.get(0)?, research_id: r.get(1)?, claim_text: r.get(2)?, + support: r.get(3)?, contradict: r.get(4)?, consensus: r.get(5)?, + grade: r.get(6)?, created_at: r.get(7)?, + }) + })?; + let mut out = Vec::new(); + for r in rows { out.push(r?); } + Ok(out) + } +} diff --git a/_primitives/_rust/kei-search-core/src/types.rs b/_primitives/_rust/kei-search-core/src/types.rs new file mode 100644 index 0000000..89146da --- /dev/null +++ b/_primitives/_rust/kei-search-core/src/types.rs @@ -0,0 +1,37 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct Research { + pub id: i64, + pub query_original: String, + pub status: String, + pub result_markdown: String, + pub total_cost_mc: i64, + pub created_at: i64, + pub completed_at: i64, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct Source { + pub id: i64, + pub research_id: i64, + pub url: String, + pub title: String, + pub content: String, + pub provider: String, + pub domain: String, + pub relevance_score: f64, + pub created_at: i64, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct Claim { + pub id: i64, + pub research_id: i64, + pub claim_text: String, + pub support: f64, + pub contradict: f64, + pub consensus: f64, + pub grade: String, + pub created_at: i64, +} diff --git a/_primitives/_rust/kei-search-core/tests/integration.rs b/_primitives/_rust/kei-search-core/tests/integration.rs new file mode 100644 index 0000000..4c6679e --- /dev/null +++ b/_primitives/_rust/kei-search-core/tests/integration.rs @@ -0,0 +1,80 @@ +use kei_search_core::budget::Budget; +use kei_search_core::export::{export, Format}; +use kei_search_core::fetch::{SourceFetcher, StubFetcher}; +use kei_search_core::pipeline::run_research; +use kei_search_core::types::Source; +use kei_search_core::ResearchStore; + +fn mk() -> ResearchStore { ResearchStore::open_memory().unwrap() } + +struct FakeFetcher; +impl SourceFetcher for FakeFetcher { + fn fetch(&self, claim: &str) -> (Vec, i64) { + (vec![Source { + url: "https://example.test".into(), + title: format!("source for: {claim}"), + content: "body".into(), + provider: "fake".into(), + domain: "example.test".into(), + relevance_score: 0.8, + ..Default::default() + }], 10) + } +} + +#[test] +fn budget_enforcement() { + let mut b = Budget::new(100); + b.charge(50).unwrap(); + b.charge(40).unwrap(); + assert!(b.charge(20).is_err(), "must reject overspend"); +} + +#[test] +fn wave_progression_creates_research() { + let s = mk(); + let id = run_research(&s, &FakeFetcher, + "Rust is memory-safe. Python is dynamic.", 10_000).unwrap(); + let r = s.get_research(id).unwrap().unwrap(); + assert_eq!(r.status, "completed"); + assert!(r.total_cost_mc > 0); + assert!(s.claims_for(id).unwrap().len() >= 2); +} + +#[test] +fn consensus_scoring_applies_grade() { + let s = mk(); + let id = run_research(&s, &FakeFetcher, "One claim here.", 10_000).unwrap(); + let claims = s.claims_for(id).unwrap(); + assert!(!claims.is_empty()); + assert!(!claims[0].grade.is_empty()); +} + +#[test] +fn export_markdown_and_json() { + let s = mk(); + let id = run_research(&s, &FakeFetcher, "Claim A. Claim B.", 10_000).unwrap(); + let md = export(&s, id, Format::Markdown).unwrap(); + assert!(md.contains("# Research")); + let js = export(&s, id, Format::Json).unwrap(); + let parsed: serde_json::Value = serde_json::from_str(&js).unwrap(); + assert!(parsed.get("claims").is_some()); +} + +#[test] +fn stop_mid_run_marks_status() { + let s = mk(); + let id = run_research(&s, &StubFetcher, "x. y.", 10_000).unwrap(); + s.set_status(id, "stopped").unwrap(); + let r = s.get_research(id).unwrap().unwrap(); + assert_eq!(r.status, "stopped"); +} + +#[test] +fn budget_exhausted_rejects_run() { + let s = mk(); + // 3 claims × 100mc + 50mc wave2 = 350mc; budget 100 → must overspend. + let err = run_research(&s, &StubFetcher, + "alpha claim one. beta claim two. gamma claim three.", 100); + assert!(err.is_err(), "small budget vs 3 claims must overspend"); +} diff --git a/_primitives/_rust/kei-social-store/Cargo.toml b/_primitives/_rust/kei-social-store/Cargo.toml new file mode 100644 index 0000000..544cc6f --- /dev/null +++ b/_primitives/_rust/kei-social-store/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "kei-social-store" +version = "0.1.0" +edition = "2021" +rust-version = "1.75" +description = "People + interaction CRM (lite). Port of LBM internal/social." + +[[bin]] +name = "kei-social-store" +path = "src/main.rs" + +[lib] +name = "kei_social_store" +path = "src/lib.rs" + +[dependencies] +rusqlite = { version = "0.31", features = ["bundled"] } +clap = { version = "4", features = ["derive"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +anyhow = "1" +chrono = { version = "0.4", default-features = false, features = ["clock"] } + +[dev-dependencies] +tempfile = "3" diff --git a/_primitives/_rust/kei-social-store/src/graph.rs b/_primitives/_rust/kei-social-store/src/graph.rs new file mode 100644 index 0000000..b081a3a --- /dev/null +++ b/_primitives/_rust/kei-social-store/src/graph.rs @@ -0,0 +1,31 @@ +//! Relationship graph — who interacted with whom, grouped by channel. + +use crate::store::Store; +use anyhow::Result; +use serde::Serialize; + +#[derive(Debug, Clone, Serialize)] +pub struct Pair { + pub person_id: i64, + pub target_id: i64, + pub channel: String, + pub count: i64, +} + +pub fn relationship_graph(store: &Store) -> Result> { + let mut stmt = store.conn().prepare( + "SELECT person_id, target_id, channel, COUNT(*) FROM interactions + WHERE target_id > 0 GROUP BY person_id, target_id, channel", + )?; + let rows = stmt.query_map([], |r| { + Ok(Pair { + person_id: r.get(0)?, + target_id: r.get(1)?, + channel: r.get(2)?, + count: r.get(3)?, + }) + })?; + let mut out = Vec::new(); + for r in rows { out.push(r?); } + Ok(out) +} diff --git a/_primitives/_rust/kei-social-store/src/interactions.rs b/_primitives/_rust/kei-social-store/src/interactions.rs new file mode 100644 index 0000000..ea710ef --- /dev/null +++ b/_primitives/_rust/kei-social-store/src/interactions.rs @@ -0,0 +1,47 @@ +use crate::store::Store; +use anyhow::Result; +use chrono::Utc; +use rusqlite::params; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct Interaction { + pub id: i64, + pub person_id: i64, + pub target_id: i64, + pub interaction_type: String, + pub channel: String, + pub content: String, + pub timestamp: i64, +} + +pub fn log_interaction(store: &Store, i: &Interaction) -> Result { + let now = Utc::now().timestamp(); + let ts = if i.timestamp == 0 { now } else { i.timestamp }; + let channel = if i.channel.is_empty() { "manual" } else { &i.channel }; + store.conn().execute( + "INSERT INTO interactions (person_id, target_id, interaction_type, + channel, content, timestamp, created_at) + VALUES (?1,?2,?3,?4,?5,?6,?7)", + params![i.person_id, i.target_id, i.interaction_type, + channel, i.content, ts, now], + )?; + Ok(store.conn().last_insert_rowid()) +} + +pub fn interactions_for(store: &Store, person_id: i64) -> Result> { + let mut stmt = store.conn().prepare( + "SELECT id, person_id, target_id, interaction_type, channel, content, timestamp + FROM interactions WHERE person_id=?1 ORDER BY timestamp DESC", + )?; + let rows = stmt.query_map(params![person_id], |r| { + Ok(Interaction { + id: r.get(0)?, person_id: r.get(1)?, target_id: r.get(2)?, + interaction_type: r.get(3)?, channel: r.get(4)?, + content: r.get(5)?, timestamp: r.get(6)?, + }) + })?; + let mut out = Vec::new(); + for r in rows { out.push(r?); } + Ok(out) +} diff --git a/_primitives/_rust/kei-social-store/src/lib.rs b/_primitives/_rust/kei-social-store/src/lib.rs new file mode 100644 index 0000000..c650c5a --- /dev/null +++ b/_primitives/_rust/kei-social-store/src/lib.rs @@ -0,0 +1,11 @@ +//! kei-social-store — people + organizations + interactions. + +pub mod graph; +pub mod interactions; +pub mod people; +pub mod schema; +pub mod search; +pub mod store; + +pub use people::{Organization, Person}; +pub use store::Store; diff --git a/_primitives/_rust/kei-social-store/src/main.rs b/_primitives/_rust/kei-social-store/src/main.rs new file mode 100644 index 0000000..6039e3f --- /dev/null +++ b/_primitives/_rust/kei-social-store/src/main.rs @@ -0,0 +1,80 @@ +use clap::{Parser, Subcommand}; +use kei_social_store::graph::relationship_graph; +use kei_social_store::interactions::{log_interaction, Interaction}; +use kei_social_store::people::{add_org, add_person, Organization, Person}; +use kei_social_store::search::search_people; +use kei_social_store::Store; +use std::path::PathBuf; +use std::process::ExitCode; + +#[derive(Parser)] +#[command(name = "kei-social-store", version)] +struct Cli { + #[arg(long)] db: Option, + #[command(subcommand)] cmd: Cmd, +} + +#[derive(Subcommand)] +enum Cmd { + SearchPeople { query: String, #[arg(long, default_value_t = 20)] limit: i64 }, + AddPerson { name: String, + #[arg(long, default_value = "")] email: String, + #[arg(long, default_value = "")] handle: String, + #[arg(long, default_value = "manual")] source: String }, + AddOrg { name: String, #[arg(long, default_value = "company")] org_type: String }, + LogInteraction { person_id: i64, interaction_type: String, + #[arg(long, default_value = "")] content: String, + #[arg(long, default_value = "manual")] channel: String, + #[arg(long, default_value_t = 0)] target_id: i64 }, + RelationshipGraph, +} + +fn db_path(o: Option) -> PathBuf { + if let Some(p) = o { return p; } + if let Ok(e) = std::env::var("KEI_SOCIAL_DB") { return PathBuf::from(e); } + let home = std::env::var("HOME").unwrap_or_else(|_| ".".into()); + PathBuf::from(home).join(".claude/social/social.sqlite") +} + +fn run() -> anyhow::Result<()> { + let cli = Cli::parse(); + let s = Store::open(&db_path(cli.db))?; + match cli.cmd { + Cmd::SearchPeople { query, limit } => { + for p in search_people(&s, &query, limit)? { + println!("{}\t{}\t{}", p.id, p.name, p.email); + } + } + Cmd::AddPerson { name, email, handle, source } => { + let id = add_person(&s, &Person { + name, email, handle, source, ..Default::default() })?; + println!("{}", id); + } + Cmd::AddOrg { name, org_type } => { + let id = add_org(&s, &Organization { + name, org_type, ..Default::default() })?; + println!("{}", id); + } + Cmd::LogInteraction { person_id, interaction_type, content, channel, target_id } => { + let id = log_interaction(&s, &Interaction { + person_id, target_id, interaction_type, channel, content, + ..Default::default() + })?; + println!("{}", id); + } + Cmd::RelationshipGraph => { + for p in relationship_graph(&s)? { + println!("{}\t-[{}]->\t{}\t({}x)", + p.person_id, p.channel, p.target_id, p.count); + } + } + } + Ok(()) +} + +fn main() -> ExitCode { + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(e) => { eprintln!("kei-social-store: {e:#}"); ExitCode::from(1) } + } +} diff --git a/_primitives/_rust/kei-social-store/src/people.rs b/_primitives/_rust/kei-social-store/src/people.rs new file mode 100644 index 0000000..2605f6c --- /dev/null +++ b/_primitives/_rust/kei-social-store/src/people.rs @@ -0,0 +1,76 @@ +use crate::store::Store; +use anyhow::Result; +use chrono::Utc; +use rusqlite::params; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct Person { + pub id: i64, + pub name: String, + pub email: String, + pub handle: String, + pub role: String, + pub organization: String, + pub source: String, + pub bio: String, + pub created_at: i64, + pub updated_at: i64, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct Organization { + pub id: i64, + pub name: String, + pub org_type: String, + pub description: String, + pub created_at: i64, +} + +pub fn add_person(store: &Store, p: &Person) -> Result { + let now = Utc::now().timestamp(); + let source = if p.source.is_empty() { "manual" } else { &p.source }; + store.conn().execute( + "INSERT INTO people (name, email, handle, role, organization, + source, bio, created_at, updated_at) + VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?8)", + params![p.name, p.email, p.handle, p.role, p.organization, + source, p.bio, now], + )?; + let id = store.conn().last_insert_rowid(); + store.conn().execute( + "INSERT INTO fts_social (person_id, name, email, bio) VALUES (?1,?2,?3,?4)", + params![id, p.name, p.email, p.bio], + )?; + Ok(id) +} + +pub fn get_person(store: &Store, id: i64) -> Result> { + let mut stmt = store.conn().prepare( + "SELECT id, name, email, handle, role, organization, source, bio, + created_at, updated_at FROM people WHERE id=?1", + )?; + let mut rows = stmt.query(params![id])?; + if let Some(r) = rows.next()? { + return Ok(Some(Person { + id: r.get(0)?, name: r.get(1)?, email: r.get(2)?, handle: r.get(3)?, + role: r.get(4)?, organization: r.get(5)?, source: r.get(6)?, + bio: r.get(7)?, created_at: r.get(8)?, updated_at: r.get(9)?, + })); + } + Ok(None) +} + +pub fn add_org(store: &Store, o: &Organization) -> Result { + let now = Utc::now().timestamp(); + let ot = if o.org_type.is_empty() { "company" } else { &o.org_type }; + store.conn().execute( + "INSERT OR IGNORE INTO organizations (name, org_type, description, created_at) + VALUES (?1,?2,?3,?4)", + params![o.name, ot, o.description, now], + )?; + let id: i64 = store.conn().query_row( + "SELECT id FROM organizations WHERE name=?1", + params![o.name], |r| r.get(0))?; + Ok(id) +} diff --git a/_primitives/_rust/kei-social-store/src/schema.rs b/_primitives/_rust/kei-social-store/src/schema.rs new file mode 100644 index 0000000..8c91dee --- /dev/null +++ b/_primitives/_rust/kei-social-store/src/schema.rs @@ -0,0 +1,48 @@ +use rusqlite::{Connection, Result}; + +pub fn create_schema(conn: &Connection) -> Result<()> { + conn.execute_batch(r#" + CREATE TABLE IF NOT EXISTS people ( + id INTEGER PRIMARY KEY, + name TEXT NOT NULL, + email TEXT DEFAULT '', + handle TEXT DEFAULT '', + role TEXT DEFAULT '', + organization TEXT DEFAULT '', + source TEXT NOT NULL DEFAULT 'manual', + bio TEXT DEFAULT '', + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL + ); + CREATE UNIQUE INDEX IF NOT EXISTS idx_people_email + ON people(email) WHERE email != ''; + CREATE UNIQUE INDEX IF NOT EXISTS idx_people_handle_source + ON people(handle, source) WHERE handle != ''; + + CREATE TABLE IF NOT EXISTS organizations ( + id INTEGER PRIMARY KEY, + name TEXT NOT NULL UNIQUE, + org_type TEXT DEFAULT 'company', + description TEXT DEFAULT '', + created_at INTEGER NOT NULL + ); + + CREATE TABLE IF NOT EXISTS interactions ( + id INTEGER PRIMARY KEY, + person_id INTEGER NOT NULL REFERENCES people(id) ON DELETE CASCADE, + target_id INTEGER NOT NULL DEFAULT 0, + interaction_type TEXT NOT NULL, + channel TEXT NOT NULL DEFAULT 'manual', + content TEXT DEFAULT '', + timestamp INTEGER NOT NULL, + created_at INTEGER NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_int_person ON interactions(person_id); + "#)?; + conn.execute_batch(r#" + CREATE VIRTUAL TABLE IF NOT EXISTS fts_social + USING fts5(person_id UNINDEXED, name, email, bio, + tokenize='porter unicode61'); + "#)?; + Ok(()) +} diff --git a/_primitives/_rust/kei-social-store/src/search.rs b/_primitives/_rust/kei-social-store/src/search.rs new file mode 100644 index 0000000..fa9590e --- /dev/null +++ b/_primitives/_rust/kei-social-store/src/search.rs @@ -0,0 +1,25 @@ +use crate::people::Person; +use crate::store::Store; +use anyhow::Result; +use rusqlite::params; + +pub fn search_people(store: &Store, q: &str, limit: i64) -> Result> { + let lim = if limit <= 0 { 20 } else { limit }; + let mut stmt = store.conn().prepare( + "SELECT p.id, p.name, p.email, p.handle, p.role, p.organization, + p.source, p.bio, p.created_at, p.updated_at + FROM fts_social f + JOIN people p ON p.id = f.person_id + WHERE fts_social MATCH ?1 ORDER BY rank LIMIT ?2", + )?; + let rows = stmt.query_map(params![q, lim], |r| { + Ok(Person { + id: r.get(0)?, name: r.get(1)?, email: r.get(2)?, handle: r.get(3)?, + role: r.get(4)?, organization: r.get(5)?, source: r.get(6)?, + bio: r.get(7)?, created_at: r.get(8)?, updated_at: r.get(9)?, + }) + })?; + let mut out = Vec::new(); + for r in rows { out.push(r?); } + Ok(out) +} diff --git a/_primitives/_rust/kei-social-store/src/store.rs b/_primitives/_rust/kei-social-store/src/store.rs new file mode 100644 index 0000000..f4c30de --- /dev/null +++ b/_primitives/_rust/kei-social-store/src/store.rs @@ -0,0 +1,22 @@ +use crate::schema::create_schema; +use anyhow::{Context, Result}; +use rusqlite::Connection; +use std::path::Path; + +pub struct Store { conn: Connection } + +impl Store { + pub fn open(path: &Path) -> Result { + if let Some(parent) = path.parent() { let _ = std::fs::create_dir_all(parent); } + let conn = Connection::open(path).context("open sqlite")?; + conn.pragma_update(None, "journal_mode", "WAL").ok(); + create_schema(&conn)?; + Ok(Self { conn }) + } + pub fn open_memory() -> Result { + let conn = Connection::open_in_memory()?; + create_schema(&conn)?; + Ok(Self { conn }) + } + pub fn conn(&self) -> &Connection { &self.conn } +} diff --git a/_primitives/_rust/kei-social-store/tests/integration.rs b/_primitives/_rust/kei-social-store/tests/integration.rs new file mode 100644 index 0000000..eb9fed4 --- /dev/null +++ b/_primitives/_rust/kei-social-store/tests/integration.rs @@ -0,0 +1,68 @@ +use kei_social_store::graph::relationship_graph; +use kei_social_store::interactions::{interactions_for, log_interaction, Interaction}; +use kei_social_store::people::{add_org, add_person, get_person, Organization, Person}; +use kei_social_store::search::search_people; +use kei_social_store::Store; + +fn mk() -> Store { Store::open_memory().unwrap() } + +#[test] +fn people_crud() { + let s = mk(); + let id = add_person(&s, &Person { + name: "Alice".into(), email: "alice@example.com".into(), + ..Default::default() + }).unwrap(); + let p = get_person(&s, id).unwrap().unwrap(); + assert_eq!(p.name, "Alice"); +} + +#[test] +fn orgs_idempotent() { + let s = mk(); + let a = add_org(&s, &Organization { name: "Acme".into(), ..Default::default() }).unwrap(); + let b = add_org(&s, &Organization { name: "Acme".into(), ..Default::default() }).unwrap(); + assert_eq!(a, b); +} + +#[test] +fn interactions_tracked() { + let s = mk(); + let p = add_person(&s, &Person { name: "Bob".into(), ..Default::default() }).unwrap(); + log_interaction(&s, &Interaction { + person_id: p, interaction_type: "email".into(), + content: "hi".into(), channel: "gmail".into(), + ..Default::default() + }).unwrap(); + let hist = interactions_for(&s, p).unwrap(); + assert_eq!(hist.len(), 1); + assert_eq!(hist[0].interaction_type, "email"); +} + +#[test] +fn search_finds_person() { + let s = mk(); + add_person(&s, &Person { + name: "Carol Chang".into(), bio: "rust async".into(), + ..Default::default() + }).unwrap(); + let hits = search_people(&s, "rust", 10).unwrap(); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].name, "Carol Chang"); +} + +#[test] +fn relationship_graph_groups() { + let s = mk(); + let a = add_person(&s, &Person { name: "A".into(), ..Default::default() }).unwrap(); + let b = add_person(&s, &Person { name: "B".into(), ..Default::default() }).unwrap(); + for _ in 0..3 { + log_interaction(&s, &Interaction { + person_id: a, target_id: b, interaction_type: "msg".into(), + channel: "slack".into(), ..Default::default() + }).unwrap(); + } + let pairs = relationship_graph(&s).unwrap(); + assert_eq!(pairs.len(), 1); + assert_eq!(pairs[0].count, 3); +} diff --git a/_primitives/_rust/kei-task/Cargo.toml b/_primitives/_rust/kei-task/Cargo.toml new file mode 100644 index 0000000..7b57690 --- /dev/null +++ b/_primitives/_rust/kei-task/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "kei-task" +version = "0.1.0" +edition = "2021" +rust-version = "1.75" +description = "Task DAG with deps + milestones (SQLite). Port of LBM internal/task." + +[[bin]] +name = "kei-task" +path = "src/main.rs" + +[lib] +name = "kei_task" +path = "src/lib.rs" + +[dependencies] +rusqlite = { version = "0.31", features = ["bundled"] } +clap = { version = "4", features = ["derive"] } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +anyhow = "1" +chrono = { version = "0.4", default-features = false, features = ["clock"] } + +[dev-dependencies] +tempfile = "3" diff --git a/_primitives/_rust/kei-task/src/deps.rs b/_primitives/_rust/kei-task/src/deps.rs new file mode 100644 index 0000000..1e8eb84 --- /dev/null +++ b/_primitives/_rust/kei-task/src/deps.rs @@ -0,0 +1,66 @@ +//! Dependency edges + cycle detection + dependency-chain traversal. + +use crate::store::Store; +use crate::types::is_valid_dep; +use anyhow::{anyhow, Result}; +use rusqlite::params; +use std::collections::HashSet; + +/// Add a dependency. Refuses a cycle (taskId -> dependsOn -> ... -> taskId). +pub fn add_dependency(store: &Store, task_id: i64, depends_on: i64, dep_type: &str) -> Result<()> { + if task_id == depends_on { + return Err(anyhow!("self-dependency forbidden")); + } + let dt = if dep_type.is_empty() { "blocks" } else { dep_type }; + if !is_valid_dep(dt) { + return Err(anyhow!("invalid dep type: {dt}")); + } + if creates_cycle(store, task_id, depends_on)? { + return Err(anyhow!("cycle: {task_id} -> {depends_on} would close a loop")); + } + store.conn().execute( + "INSERT OR IGNORE INTO task_deps (task_id, depends_on, dep_type) VALUES (?1,?2,?3)", + params![task_id, depends_on, dt], + )?; + Ok(()) +} + +/// True if adding task_id -> depends_on would create a cycle. +fn creates_cycle(store: &Store, task_id: i64, depends_on: i64) -> Result { + // If depends_on reaches task_id via existing deps, cycle would close. + let mut stack = vec![depends_on]; + let mut seen: HashSet = HashSet::new(); + while let Some(cur) = stack.pop() { + if cur == task_id { + return Ok(true); + } + if !seen.insert(cur) { + continue; + } + let mut stmt = store.conn().prepare("SELECT depends_on FROM task_deps WHERE task_id=?1")?; + let rows = stmt.query_map(params![cur], |r| r.get::<_, i64>(0))?; + for row in rows { + stack.push(row?); + } + } + Ok(false) +} + +/// Full dependency chain (BFS transitive closure). +pub fn dependency_chain(store: &Store, task_id: i64) -> Result> { + let mut seen: HashSet = HashSet::new(); + let mut frontier = vec![task_id]; + let mut chain: Vec = Vec::new(); + while let Some(cur) = frontier.pop() { + let mut stmt = store.conn().prepare("SELECT depends_on FROM task_deps WHERE task_id=?1")?; + let rows = stmt.query_map(params![cur], |r| r.get::<_, i64>(0))?; + for row in rows { + let id = row?; + if seen.insert(id) { + chain.push(id); + frontier.push(id); + } + } + } + Ok(chain) +} diff --git a/_primitives/_rust/kei-task/src/graph.rs b/_primitives/_rust/kei-task/src/graph.rs new file mode 100644 index 0000000..f318b45 --- /dev/null +++ b/_primitives/_rust/kei-task/src/graph.rs @@ -0,0 +1,28 @@ +//! Adjacency view — returns task graph as edge-list for visualisation. + +use crate::store::Store; +use anyhow::Result; +use serde::Serialize; + +#[derive(Debug, Clone, Serialize)] +pub struct TaskEdge { + pub task_id: i64, + pub depends_on: i64, + pub dep_type: String, +} + +pub fn list_edges(store: &Store) -> Result> { + let mut stmt = store.conn().prepare( + "SELECT task_id, depends_on, dep_type FROM task_deps" + )?; + let rows = stmt.query_map([], |r| { + Ok(TaskEdge { + task_id: r.get(0)?, + depends_on: r.get(1)?, + dep_type: r.get(2)?, + }) + })?; + let mut out = Vec::new(); + for r in rows { out.push(r?); } + Ok(out) +} diff --git a/_primitives/_rust/kei-task/src/lib.rs b/_primitives/_rust/kei-task/src/lib.rs new file mode 100644 index 0000000..2e955a5 --- /dev/null +++ b/_primitives/_rust/kei-task/src/lib.rs @@ -0,0 +1,12 @@ +//! kei-task — tasks with typed deps (DAG, cycle-detected), milestones, FTS search. + +pub mod deps; +pub mod graph; +pub mod milestones; +pub mod schema; +pub mod search; +pub mod store; +pub mod types; + +pub use store::Store; +pub use types::{Milestone, Task}; diff --git a/_primitives/_rust/kei-task/src/main.rs b/_primitives/_rust/kei-task/src/main.rs new file mode 100644 index 0000000..cce07d1 --- /dev/null +++ b/_primitives/_rust/kei-task/src/main.rs @@ -0,0 +1,92 @@ +//! kei-task CLI — create / update / add-dep / graph / dependency-chain. + +use clap::{Parser, Subcommand}; +use kei_task::deps::{add_dependency, dependency_chain}; +use kei_task::graph::list_edges; +use kei_task::milestones::{create_milestone, link_task_to_milestone}; +use kei_task::search::search; +use kei_task::{Milestone, Store, Task}; +use std::path::PathBuf; +use std::process::ExitCode; + +#[derive(Parser)] +#[command(name = "kei-task", version, about = "Task DAG CLI")] +struct Cli { + #[arg(long)] db: Option, + #[command(subcommand)] cmd: Cmd, +} + +#[derive(Subcommand)] +enum Cmd { + Create { title: String, #[arg(long, default_value = "")] description: String, + #[arg(long, default_value = "medium")] priority: String }, + Update { id: i64, #[arg(long)] status: Option, #[arg(long)] title: Option }, + AddDependency { from_id: i64, to_id: i64, + #[arg(long, default_value = "blocks")] dep_type: String }, + Graph, + DependencyChain { id: i64 }, + Search { query: String, #[arg(long, default_value_t = 20)] limit: i64 }, + Milestone { name: String, #[arg(long, default_value = "")] description: String }, + LinkMilestone { task_id: i64, milestone_id: i64 }, +} + +fn db_path(cli_db: Option) -> PathBuf { + if let Some(p) = cli_db { return p; } + if let Ok(e) = std::env::var("KEI_TASK_DB") { return PathBuf::from(e); } + let home = std::env::var("HOME").unwrap_or_else(|_| ".".into()); + PathBuf::from(home).join(".claude/task/task.sqlite") +} + +fn run() -> anyhow::Result<()> { + let cli = Cli::parse(); + let s = Store::open(&db_path(cli.db))?; + match cli.cmd { + Cmd::Create { title, description, priority } => { + let id = s.create_task(&Task { title, description, priority, ..Default::default() })?; + println!("{}", id); + } + Cmd::Update { id, status, title } => { + let mut t = s.get_task(id)?.ok_or_else(|| anyhow::anyhow!("id {id} not found"))?; + if let Some(st) = status { t.status = st; } + if let Some(ti) = title { t.title = ti; } + s.update_task(&t)?; + println!("updated {}", id); + } + Cmd::AddDependency { from_id, to_id, dep_type } => { + add_dependency(&s, from_id, to_id, &dep_type)?; + println!("dep: {} -> {} ({})", from_id, to_id, dep_type); + } + Cmd::Graph => { + for e in list_edges(&s)? { + println!("{}\t-[{}]->\t{}", e.task_id, e.dep_type, e.depends_on); + } + } + Cmd::DependencyChain { id } => { + for did in dependency_chain(&s, id)? { + println!("{}", did); + } + } + Cmd::Search { query, limit } => { + for t in search(&s, &query, limit)? { + println!("{}\t{}\t{}", t.id, t.status, t.title); + } + } + Cmd::Milestone { name, description } => { + let id = create_milestone(&s, &Milestone { + name, description, ..Default::default() })?; + println!("{}", id); + } + Cmd::LinkMilestone { task_id, milestone_id } => { + link_task_to_milestone(&s, task_id, milestone_id)?; + println!("linked {} -> milestone {}", task_id, milestone_id); + } + } + Ok(()) +} + +fn main() -> ExitCode { + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(e) => { eprintln!("kei-task: {e:#}"); ExitCode::from(1) } + } +} diff --git a/_primitives/_rust/kei-task/src/milestones.rs b/_primitives/_rust/kei-task/src/milestones.rs new file mode 100644 index 0000000..fd8f4ef --- /dev/null +++ b/_primitives/_rust/kei-task/src/milestones.rs @@ -0,0 +1,39 @@ +//! Milestone CRUD + task→milestone linking. + +use crate::store::Store; +use crate::types::Milestone; +use anyhow::Result; +use chrono::Utc; +use rusqlite::params; + +pub fn create_milestone(store: &Store, m: &Milestone) -> Result { + let now = Utc::now().timestamp(); + let created = if m.created_at == 0 { now } else { m.created_at }; + let status = if m.status.is_empty() { "open" } else { &m.status }; + store.conn().execute( + "INSERT INTO milestones (name, description, target_date, status, created_at) + VALUES (?1,?2,?3,?4,?5)", + params![m.name, m.description, m.target_date, status, created], + )?; + Ok(store.conn().last_insert_rowid()) +} + +pub fn link_task_to_milestone(store: &Store, task_id: i64, milestone_id: i64) -> Result<()> { + store.conn().execute( + "INSERT OR IGNORE INTO task_milestones (task_id, milestone_id) VALUES (?1,?2)", + params![task_id, milestone_id], + )?; + Ok(()) +} + +pub fn tasks_in_milestone(store: &Store, milestone_id: i64) -> Result> { + let mut stmt = store.conn().prepare( + "SELECT task_id FROM task_milestones WHERE milestone_id=?1", + )?; + let rows = stmt.query_map(params![milestone_id], |r| r.get::<_, i64>(0))?; + let mut out = Vec::new(); + for r in rows { + out.push(r?); + } + Ok(out) +} diff --git a/_primitives/_rust/kei-task/src/schema.rs b/_primitives/_rust/kei-task/src/schema.rs new file mode 100644 index 0000000..1110500 --- /dev/null +++ b/_primitives/_rust/kei-task/src/schema.rs @@ -0,0 +1,53 @@ +//! SQLite schema for tasks + milestones + deps. Port of LBM internal/task/schema.go. + +use rusqlite::{Connection, Result}; + +pub fn create_schema(conn: &Connection) -> Result<()> { + conn.execute_batch(r#" + CREATE TABLE IF NOT EXISTS tasks ( + id INTEGER PRIMARY KEY, + title TEXT NOT NULL, + description TEXT DEFAULT '', + status TEXT NOT NULL DEFAULT 'pending', + priority TEXT DEFAULT 'medium', + task_type TEXT DEFAULT '', + parent_id INTEGER DEFAULT 0, + assigned_to TEXT DEFAULT '', + due_date INTEGER DEFAULT 0, + completed_at INTEGER DEFAULT 0, + created_at INTEGER NOT NULL, + updated_at INTEGER NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_task_status ON tasks(status); + CREATE INDEX IF NOT EXISTS idx_task_priority ON tasks(priority); + CREATE INDEX IF NOT EXISTS idx_task_parent ON tasks(parent_id); + + CREATE TABLE IF NOT EXISTS milestones ( + id INTEGER PRIMARY KEY, + name TEXT NOT NULL, + description TEXT DEFAULT '', + target_date INTEGER DEFAULT 0, + status TEXT DEFAULT 'open', + created_at INTEGER NOT NULL + ); + + CREATE TABLE IF NOT EXISTS task_deps ( + task_id INTEGER NOT NULL, + depends_on INTEGER NOT NULL, + dep_type TEXT DEFAULT 'blocks', + PRIMARY KEY(task_id, depends_on) + ); + CREATE INDEX IF NOT EXISTS idx_dep_depends ON task_deps(depends_on); + + CREATE TABLE IF NOT EXISTS task_milestones ( + task_id INTEGER NOT NULL, + milestone_id INTEGER NOT NULL, + PRIMARY KEY(task_id, milestone_id) + ); + "#)?; + conn.execute_batch(r#" + CREATE VIRTUAL TABLE IF NOT EXISTS fts_tasks + USING fts5(task_id UNINDEXED, title, description, tokenize='porter unicode61'); + "#)?; + Ok(()) +} diff --git a/_primitives/_rust/kei-task/src/search.rs b/_primitives/_rust/kei-task/src/search.rs new file mode 100644 index 0000000..b23f546 --- /dev/null +++ b/_primitives/_rust/kei-task/src/search.rs @@ -0,0 +1,29 @@ +//! FTS5 search over tasks (title + description). + +use crate::store::Store; +use crate::types::Task; +use anyhow::Result; +use rusqlite::params; + +pub fn search(store: &Store, query: &str, limit: i64) -> Result> { + let lim = if limit <= 0 { 20 } else { limit }; + let mut stmt = store.conn().prepare( + "SELECT t.id, t.title, t.description, t.status, t.priority, t.task_type, + t.parent_id, t.assigned_to, t.due_date, t.completed_at, + t.created_at, t.updated_at + FROM fts_tasks f + JOIN tasks t ON t.id = f.task_id + WHERE fts_tasks MATCH ?1 ORDER BY rank LIMIT ?2", + )?; + let rows = stmt.query_map(params![query, lim], |r| { + Ok(Task { + id: r.get(0)?, title: r.get(1)?, description: r.get(2)?, status: r.get(3)?, + priority: r.get(4)?, task_type: r.get(5)?, parent_id: r.get(6)?, + assigned_to: r.get(7)?, due_date: r.get(8)?, completed_at: r.get(9)?, + created_at: r.get(10)?, updated_at: r.get(11)?, + }) + })?; + let mut out = Vec::new(); + for r in rows { out.push(r?); } + Ok(out) +} diff --git a/_primitives/_rust/kei-task/src/store.rs b/_primitives/_rust/kei-task/src/store.rs new file mode 100644 index 0000000..3ed1c43 --- /dev/null +++ b/_primitives/_rust/kei-task/src/store.rs @@ -0,0 +1,94 @@ +//! Task store — open, CRUD, FTS reindex. + +use crate::schema::create_schema; +use crate::types::Task; +use anyhow::{Context, Result}; +use chrono::Utc; +use rusqlite::{params, Connection}; +use std::path::Path; + +pub struct Store { + conn: Connection, +} + +impl Store { + pub fn open(path: &Path) -> Result { + if let Some(parent) = path.parent() { + let _ = std::fs::create_dir_all(parent); + } + let conn = Connection::open(path).context("open sqlite")?; + conn.pragma_update(None, "journal_mode", "WAL").ok(); + create_schema(&conn)?; + Ok(Self { conn }) + } + + pub fn open_memory() -> Result { + let conn = Connection::open_in_memory()?; + create_schema(&conn)?; + Ok(Self { conn }) + } + + pub fn conn(&self) -> &Connection { &self.conn } + + pub fn create_task(&self, t: &Task) -> Result { + let now = Utc::now().timestamp(); + let created = if t.created_at == 0 { now } else { t.created_at }; + let status = if t.status.is_empty() { "pending" } else { &t.status }; + let priority = if t.priority.is_empty() { "medium" } else { &t.priority }; + self.conn.execute( + "INSERT INTO tasks (title, description, status, priority, task_type, + parent_id, assigned_to, due_date, completed_at, created_at, updated_at) + VALUES (?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11)", + params![t.title, t.description, status, priority, t.task_type, + t.parent_id, t.assigned_to, t.due_date, t.completed_at, created, now], + )?; + let id = self.conn.last_insert_rowid(); + self.reindex_fts(id, &t.title, &t.description)?; + Ok(id) + } + + pub fn get_task(&self, id: i64) -> Result> { + let mut stmt = self.conn.prepare( + "SELECT id, title, description, status, priority, task_type, parent_id, + assigned_to, due_date, completed_at, created_at, updated_at + FROM tasks WHERE id=?1", + )?; + let mut rows = stmt.query(params![id])?; + if let Some(r) = rows.next()? { + return Ok(Some(row_to_task(r)?)); + } + Ok(None) + } + + pub fn update_task(&self, t: &Task) -> Result<()> { + let now = Utc::now().timestamp(); + let completed = if t.status == "completed" && t.completed_at == 0 { now } else { t.completed_at }; + self.conn.execute( + "UPDATE tasks SET title=?1, description=?2, status=?3, priority=?4, + task_type=?5, parent_id=?6, assigned_to=?7, due_date=?8, + completed_at=?9, updated_at=?10 WHERE id=?11", + params![t.title, t.description, t.status, t.priority, t.task_type, + t.parent_id, t.assigned_to, t.due_date, completed, now, t.id], + )?; + self.reindex_fts(t.id, &t.title, &t.description)?; + Ok(()) + } + + fn reindex_fts(&self, id: i64, title: &str, description: &str) -> Result<()> { + self.conn.execute("DELETE FROM fts_tasks WHERE task_id=?1", params![id])?; + self.conn.execute( + "INSERT INTO fts_tasks (task_id, title, description) VALUES (?1,?2,?3)", + params![id, title, description], + )?; + Ok(()) + } +} + +fn row_to_task(r: &rusqlite::Row) -> rusqlite::Result { + Ok(Task { + id: r.get(0)?, title: r.get(1)?, description: r.get(2)?, status: r.get(3)?, + priority: r.get(4)?, task_type: r.get(5)?, parent_id: r.get(6)?, + assigned_to: r.get(7)?, due_date: r.get(8)?, completed_at: r.get(9)?, + created_at: r.get(10)?, updated_at: r.get(11)?, + }) +} diff --git a/_primitives/_rust/kei-task/src/types.rs b/_primitives/_rust/kei-task/src/types.rs new file mode 100644 index 0000000..a894895 --- /dev/null +++ b/_primitives/_rust/kei-task/src/types.rs @@ -0,0 +1,45 @@ +//! Task + Milestone value types and enum validation. + +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct Task { + pub id: i64, + pub title: String, + pub description: String, + pub status: String, + pub priority: String, + pub task_type: String, + pub parent_id: i64, + pub assigned_to: String, + pub due_date: i64, + pub completed_at: i64, + pub created_at: i64, + pub updated_at: i64, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct Milestone { + pub id: i64, + pub name: String, + pub description: String, + pub target_date: i64, + pub status: String, + pub created_at: i64, +} + +pub const VALID_STATUSES: &[&str] = + &["pending", "in_progress", "completed", "cancelled", "blocked"]; +pub const VALID_PRIORITIES: &[&str] = &["critical", "high", "medium", "low"]; +pub const VALID_DEP_TYPES: &[&str] = + &["blocks", "feeds_into", "subtask_of", "milestone_of", "assigned_to", "depends_on"]; + +pub fn is_valid_status(s: &str) -> bool { + VALID_STATUSES.contains(&s) +} +pub fn is_valid_priority(s: &str) -> bool { + VALID_PRIORITIES.contains(&s) +} +pub fn is_valid_dep(s: &str) -> bool { + VALID_DEP_TYPES.contains(&s) +} diff --git a/_primitives/_rust/kei-task/tests/integration.rs b/_primitives/_rust/kei-task/tests/integration.rs new file mode 100644 index 0000000..30e473e --- /dev/null +++ b/_primitives/_rust/kei-task/tests/integration.rs @@ -0,0 +1,95 @@ +//! kei-task integration tests. + +use kei_task::deps::{add_dependency, dependency_chain}; +use kei_task::graph::list_edges; +use kei_task::milestones::{create_milestone, link_task_to_milestone, tasks_in_milestone}; +use kei_task::search::search; +use kei_task::{Milestone, Store, Task}; + +fn mk() -> Store { Store::open_memory().unwrap() } + +fn mktask(title: &str) -> Task { + Task { title: title.into(), priority: "high".into(), ..Default::default() } +} + +#[test] +fn create_and_get() { + let s = mk(); + let id = s.create_task(&mktask("a")).unwrap(); + let t = s.get_task(id).unwrap().unwrap(); + assert_eq!(t.title, "a"); + assert_eq!(t.status, "pending"); +} + +#[test] +fn update_persists() { + let s = mk(); + let id = s.create_task(&mktask("a")).unwrap(); + let mut t = s.get_task(id).unwrap().unwrap(); + t.status = "in_progress".into(); + s.update_task(&t).unwrap(); + let u = s.get_task(id).unwrap().unwrap(); + assert_eq!(u.status, "in_progress"); +} + +#[test] +fn cycle_detected() { + let s = mk(); + let a = s.create_task(&mktask("a")).unwrap(); + let b = s.create_task(&mktask("b")).unwrap(); + let c = s.create_task(&mktask("c")).unwrap(); + add_dependency(&s, a, b, "blocks").unwrap(); + add_dependency(&s, b, c, "blocks").unwrap(); + // a -> b -> c; now c -> a would be a cycle + let err = add_dependency(&s, c, a, "blocks"); + assert!(err.is_err(), "cycle detection must reject"); +} + +#[test] +fn milestone_linking() { + let s = mk(); + let t = s.create_task(&mktask("design")).unwrap(); + let ms_id = create_milestone(&s, &Milestone { + name: "v1".into(), ..Default::default() }).unwrap(); + link_task_to_milestone(&s, t, ms_id).unwrap(); + let tasks = tasks_in_milestone(&s, ms_id).unwrap(); + assert_eq!(tasks, vec![t]); +} + +#[test] +fn dependency_chain_traversal() { + let s = mk(); + let a = s.create_task(&mktask("a")).unwrap(); + let b = s.create_task(&mktask("b")).unwrap(); + let c = s.create_task(&mktask("c")).unwrap(); + add_dependency(&s, a, b, "blocks").unwrap(); + add_dependency(&s, b, c, "blocks").unwrap(); + let chain = dependency_chain(&s, a).unwrap(); + assert!(chain.contains(&b)); + assert!(chain.contains(&c)); + assert_eq!(chain.len(), 2); +} + +#[test] +fn task_graph_edges() { + let s = mk(); + let a = s.create_task(&mktask("a")).unwrap(); + let b = s.create_task(&mktask("b")).unwrap(); + add_dependency(&s, a, b, "blocks").unwrap(); + let edges = list_edges(&s).unwrap(); + assert_eq!(edges.len(), 1); + assert_eq!(edges[0].task_id, a); + assert_eq!(edges[0].depends_on, b); +} + +#[test] +fn search_finds_task() { + let s = mk(); + s.create_task(&Task { + title: "refactor router".into(), + description: "split monolith".into(), + ..Default::default() + }).unwrap(); + let hits = search(&s, "refactor", 10).unwrap(); + assert_eq!(hits.len(), 1); +}