diff --git a/memo-cli/build.gradle.kts b/memo-cli/build.gradle.kts index 9bc4506..7aa9889 100644 --- a/memo-cli/build.gradle.kts +++ b/memo-cli/build.gradle.kts @@ -1,3 +1,14 @@ plugins { kotlin("jvm") + application +} + +application { + mainClass.set("memo.cli.MainKt") + applicationName = "memo" +} + +dependencies { + implementation(project(":memo-core")) + testImplementation(kotlin("test")) } \ No newline at end of file diff --git a/memo-cli/src/main/kotlin/memo/cli/Main.kt b/memo-cli/src/main/kotlin/memo/cli/Main.kt new file mode 100644 index 0000000..3dc28bc --- /dev/null +++ b/memo-cli/src/main/kotlin/memo/cli/Main.kt @@ -0,0 +1,575 @@ +package memo.cli + +import memo.core.Db +import memo.core.Embedder +import memo.core.Hit +import memo.core.Indexer +import memo.core.RefreshHook +import memo.core.SearchMode +import memo.core.Searcher +import java.io.File + +sealed interface Cmd +data class IndexCmd(val path: String) : Cmd +data class SearchCmd( + val path: String, + val query: String, + val k: Int, + val mode: SearchMode, + val json: Boolean, +) : Cmd +data class StatusCmd(val path: String) : Cmd +data class McpProbeCmd(val tool: String, val argsJson: String) : Cmd +data object HelpCmd : Cmd + +fun parseArgs(args: Array): Cmd { + if (args.isEmpty()) return HelpCmd + return when (val cmd = args[0]) { + "index" -> { + val path = args.getOrNull(1) ?: return HelpCmd + IndexCmd(path) + } + "search" -> parseSearch(args.drop(1)) + "status" -> { + val path = args.getOrNull(1) ?: return HelpCmd + StatusCmd(path) + } + "mcp-probe" -> parseMcpProbe(args.drop(1)) + "--help", "-h" -> HelpCmd + else -> HelpCmd + } +} + +private fun parseSearch(rest: List): Cmd { + if (rest.size < 2) return HelpCmd + val path = rest[0] + val query = rest[1] + var k = 8 + var mode = SearchMode.HYBRID + var json = false + var i = 2 + while (i < rest.size) { + when (val a = rest[i]) { + "--k" -> { + val v = rest.getOrNull(i + 1) ?: return HelpCmd + val n = v.toIntOrNull() ?: return HelpCmd + k = n + i += 2 + } + "--mode" -> { + val v = rest.getOrNull(i + 1) ?: return HelpCmd + mode = when (v) { + "hybrid" -> SearchMode.HYBRID + "lex" -> SearchMode.LEX + "vec" -> SearchMode.VEC + else -> return HelpCmd + } + i += 2 + } + "--json" -> { + json = true + i += 1 + } + else -> return HelpCmd + } + } + return SearchCmd(path, query, k, mode, json) +} + +private fun parseMcpProbe(rest: List): Cmd { + var tool: String? = null + var argsJson = "{}" + var i = 0 + while (i < rest.size) { + when (rest[i]) { + "--tool" -> { + tool = rest.getOrNull(i + 1) ?: return HelpCmd + i += 2 + } + "--args" -> { + argsJson = rest.getOrNull(i + 1) ?: return HelpCmd + i += 2 + } + else -> return HelpCmd + } + } + val t = tool ?: return HelpCmd + return McpProbeCmd(t, argsJson) +} + +fun main(args: Array) { + when (val cmd = parseArgs(args)) { + is IndexCmd -> runIndex(cmd) + is SearchCmd -> runSearch(cmd) + is StatusCmd -> runStatus(cmd) + is McpProbeCmd -> runMcpProbe(cmd) + HelpCmd -> printHelp() + } +} + +private fun printHelp() { + println( + """ + usage: + memo index + memo search [--k N] [--mode hybrid|lex|vec] [--json] + memo status + memo mcp-probe --tool [--args ] + """.trimIndent() + ) +} + +private fun modelPaths(): Pair { + val modelDir = System.getenv("MEMO_MODEL_DIR") ?: "/root/WORK/memo/models/siglip2" + return "$modelDir/text_model_int8.onnx" to "$modelDir/tokenizer.model" +} + +private fun runIndex(cmd: IndexCmd) { + val raw = File(cmd.path) + val base = if (raw.name == ".memo") raw.parentFile ?: raw else raw + val collections = discoverCollections(base) + if (collections.isEmpty()) { + System.err.println("коллекции не найдены в ${base.absolutePath}") + return + } + val (modelPath, tokenizerPath) = modelPaths() + var totalUpdated = 0 + for (coll in collections) { + val memoDir = File(coll, ".memo") + memoDir.mkdirs() + val dbPath = File(memoDir, "index.db").absolutePath + val totalFiles = coll.walkTopDown() + .maxDepth(8) + .count { it.isFile && it.extension == "md" } + val db = Db(dbPath) + try { + db.init() + val embedder = Embedder(modelPath, tokenizerPath) + try { + val updated = Indexer(db, embedder).indexTree(coll) + println("индексировано: $updated обновлено, $totalFiles файлов всего") + totalUpdated += updated + } finally { + embedder.close() + } + } finally { + db.close() + } + } + println("итого: $totalUpdated обновлено в ${collections.size} коллекциях") +} + +private fun discoverCollections(base: File): List { + if (!base.isDirectory) return emptyList() + val candidates = LinkedHashSet() + candidates.add(base) + val q = ArrayDeque>() + q.addLast(base to 0) + while (q.isNotEmpty()) { + val (d, depth) = q.removeFirst() + if (depth >= 2) continue + val children = d.listFiles() ?: continue + for (c in children) { + if (c.isDirectory && !c.name.startsWith(".")) { + candidates.add(c) + q.addLast(c to depth + 1) + } + } + } + return candidates.filter { d -> + d.walkTopDown() + .maxDepth(8) + .any { it.isFile && it.extension == "md" } + } +} + +private fun runSearch(cmd: SearchCmd) { + val targets = resolveSearchTargets(File(cmd.path)) + if (targets.isEmpty()) { + if (cmd.json) println("[]") + else println("коллекции не найдены") + return + } + val (modelPath, tokenizerPath) = modelPaths() + val allHits = ArrayList() + for (coll in targets) { + val dbPath = File(coll, ".memo/index.db") + if (!dbPath.exists()) continue + runCatching { + val db = Db(dbPath.absolutePath) + try { + val embedder = Embedder(modelPath, tokenizerPath) + try { + val refresh = RefreshHook { root -> + val innerDbPath = File(root, ".memo/index.db") + if (innerDbPath.exists()) { + val innerDb = Db(innerDbPath.absolutePath) + try { + Indexer(innerDb, embedder).indexTree(root) + } finally { + innerDb.close() + } + } + } + val hits = Searcher(db, embedder, refresh) + .search(coll, cmd.query, cmd.k, cmd.mode) + allHits.addAll(hits) + } finally { + embedder.close() + } + } finally { + db.close() + } + } + } + val merged = allHits.sortedByDescending { it.score }.take(cmd.k) + if (cmd.json) { + printJsonHits(merged) + } else { + for (h in merged) { + val text = if (h.text.length > 300) h.text.substring(0, 300) else h.text + println(String.format("%.3f %s:%d %s", h.score, h.path, h.line, h.heading)) + for (line in text.lines()) { + println(" $line") + } + } + } +} + +private fun resolveSearchTargets(path: File): List { + if (path.isFile && path.extension == "md") { + return listOfNotNull(path.parentFile) + } + if (File(path, ".memo/index.db").exists()) { + return listOf(path) + } + val subs = path.listFiles() + ?.filter { it.isDirectory && File(it, ".memo/index.db").exists() } + ?: emptyList() + return subs +} + +private fun runStatus(cmd: StatusCmd) { + val base = File(cmd.path) + val collections = if (File(base, ".memo/index.db").exists()) { + listOf(base) + } else { + discoverCollections(base) + } + if (collections.isEmpty()) { + println("коллекции не найдены") + return + } + for (coll in collections) { + val dbPath = File(coll, ".memo/index.db") + if (!dbPath.exists()) { + println("${coll.name}: нет индекса") + continue + } + val db = Db(dbPath.absolutePath) + try { + val fileCount = countInt(db, "SELECT COUNT(*) FROM files") + val chunkCount = countInt(db, "SELECT COUNT(*) FROM chunks") + val indexedAt = maxDouble(db, "SELECT indexed_at FROM files WHERE indexed_at IS NOT NULL") + val date = if (indexedAt != null) { + java.time.Instant.ofEpochSecond(indexedAt.toLong()).toString() + } else { + "—" + } + println("$coll: файлов $fileCount, чанков $chunkCount, индекс $date") + } finally { + db.close() + } + } +} + +private fun countInt(db: Db, sql: String): Int { + val stmt = db.conn.prepare(sql) + return try { + val rs = stmt.executeQuery() + try { + if (rs.next()) rs.getInt(0) ?: 0 else 0 + } finally { + rs.close() + } + } finally { + stmt.close() + } +} + +private fun maxDouble(db: Db, sql: String): Double? { + val stmt = db.conn.prepare(sql) + return try { + val rs = stmt.executeQuery() + try { + var best: Double? = null + while (rs.next()) { + val v = rs.getDouble(0) ?: continue + if (best == null || v > best) best = v + } + best + } finally { + rs.close() + } + } finally { + stmt.close() + } +} + +private fun runMcpProbe(cmd: McpProbeCmd) { + when (cmd.tool) { + "memo_reindex" -> printMcpError("not implemented yet") + "memo_search" -> { + val args = parseJsonArgs(cmd.argsJson) + val path = (args["path"] as? String) ?: "" + val query = (args["query"] as? String) ?: "" + val k = (args["k"] as? Number)?.toInt() ?: 8 + val mode = when (args["mode"] as? String) { + "lex" -> SearchMode.LEX + "vec" -> SearchMode.VEC + else -> SearchMode.HYBRID + } + val jsonOutput = (args["json"] as? Boolean) ?: true + val hits = performSearch(path, query, k, mode) + val text = if (jsonOutput) formatHitsJson(hits) else formatHitsHuman(hits) + printMcpOk(text) + } + "memo_status" -> { + val args = parseJsonArgs(cmd.argsJson) + val path = (args["path"] as? String) ?: "" + val text = formatStatusInner(path) + printMcpOk(text) + } + else -> printMcpError("unknown tool: ${cmd.tool}") + } +} + +private fun performSearch(path: String, query: String, k: Int, mode: SearchMode): List { + val targets = resolveSearchTargets(File(path)) + if (targets.isEmpty()) return emptyList() + val (modelPath, tokenizerPath) = modelPaths() + val allHits = ArrayList() + for (coll in targets) { + val dbPath = File(coll, ".memo/index.db") + if (!dbPath.exists()) continue + runCatching { + val db = Db(dbPath.absolutePath) + try { + val embedder = Embedder(modelPath, tokenizerPath) + try { + val refresh = RefreshHook { root -> + val innerDbPath = File(root, ".memo/index.db") + if (innerDbPath.exists()) { + val innerDb = Db(innerDbPath.absolutePath) + try { + Indexer(innerDb, embedder).indexTree(root) + } finally { + innerDb.close() + } + } + } + val hits = Searcher(db, embedder, refresh) + .search(coll, query, k, mode) + allHits.addAll(hits) + } finally { + embedder.close() + } + } finally { + db.close() + } + } + } + return allHits.sortedByDescending { it.score }.take(k) +} + +private fun formatStatusInner(path: String): String { + val base = File(path) + val collections = if (File(base, ".memo/index.db").exists()) { + listOf(base) + } else { + discoverCollections(base) + } + if (collections.isEmpty()) return "коллекции не найдены" + val sb = StringBuilder() + for (coll in collections) { + val dbPath = File(coll, ".memo/index.db") + if (!dbPath.exists()) { + sb.append("${coll.name}: нет индекса\n") + continue + } + val db = Db(dbPath.absolutePath) + try { + val fileCount = countInt(db, "SELECT COUNT(*) FROM files") + val chunkCount = countInt(db, "SELECT COUNT(*) FROM chunks") + val indexedAt = maxDouble(db, "SELECT indexed_at FROM files WHERE indexed_at IS NOT NULL") + val date = if (indexedAt != null) { + java.time.Instant.ofEpochSecond(indexedAt.toLong()).toString() + } else { + "—" + } + sb.append("$coll: файлов $fileCount, чанков $chunkCount, индекс $date\n") + } finally { + db.close() + } + } + return sb.toString().trimEnd('\n') +} + +private fun printJsonHits(hits: List) { + println(formatHitsJson(hits)) +} + +private fun formatHitsJson(hits: List): String { + val sb = StringBuilder() + sb.append('[') + for ((i, h) in hits.withIndex()) { + if (i > 0) sb.append(',') + sb.append('{') + sb.append("\"path\":").append(jsonStr(h.path)).append(',') + sb.append("\"line\":").append(h.line).append(',') + sb.append("\"heading\":").append(jsonStr(h.heading)).append(',') + sb.append("\"score\":").append("%.6f".format(h.score)).append(',') + sb.append("\"text\":").append(jsonStr(h.text)) + sb.append('}') + } + sb.append(']') + return sb.toString() +} + +private fun formatHitsHuman(hits: List): String { + val sb = StringBuilder() + for (h in hits) { + val text = if (h.text.length > 300) h.text.substring(0, 300) else h.text + sb.append(String.format("%.3f %s:%d %s", h.score, h.path, h.line, h.heading)) + sb.append('\n') + for (line in text.lines()) { + sb.append(" ").append(line).append('\n') + } + } + return sb.toString().trimEnd('\n') +} + +private fun jsonStr(s: String): String { + val sb = StringBuilder(s.length + 2) + sb.append('"') + for (c in s) { + when (c) { + '"' -> sb.append("\\\"") + '\\' -> sb.append("\\\\") + '\n' -> sb.append("\\n") + '\r' -> sb.append("\\r") + '\t' -> sb.append("\\t") + '\b' -> sb.append("\\b") + '\u000C' -> sb.append("\\f") + else -> if (c.code < 0x20) { + sb.append("\\u%04x".format(c.code)) + } else { + sb.append(c) + } + } + } + sb.append('"') + return sb.toString() +} + +private fun printMcpOk(text: String) { + println("""{"content":[{"type":"text","text":${jsonStr(text)}}],"isError":false}""") +} + +private fun printMcpError(text: String) { + println("""{"content":[{"type":"text","text":${jsonStr(text)}}],"isError":true}""") +} + +private fun parseJsonArgs(json: String): Map { + val map = HashMap() + val trimmed = json.trim() + if (trimmed.isEmpty() || trimmed == "{}") return map + val inner = trimmed.trim().removePrefix("{").removeSuffix("}") + if (inner.isBlank()) return map + var i = 0 + while (i < inner.length) { + while (i < inner.length && inner[i].isWhitespace()) i++ + if (i >= inner.length) break + if (inner[i] != '"') return map + val keyEnd = readJsonString(inner, i) + val key = unescapeJson(inner.substring(i + 1, keyEnd)) + i = keyEnd + 1 + while (i < inner.length && inner[i].isWhitespace()) i++ + if (i >= inner.length || inner[i] != ':') return map + i++ + while (i < inner.length && inner[i].isWhitespace()) i++ + if (i >= inner.length) return map + val (value, next) = readJsonValue(inner, i) + map[key] = value + i = next + while (i < inner.length && inner[i].isWhitespace()) i++ + if (i < inner.length && inner[i] == ',') i++ + } + return map +} + +private fun readJsonString(s: String, start: Int): Int { + var i = start + 1 + while (i < s.length) { + if (s[i] == '\\') { + i += 2 + continue + } + if (s[i] == '"') return i + i++ + } + return s.length +} + +private fun readJsonValue(s: String, start: Int): Pair { + val c = s[start] + return when { + c == '"' -> { + val end = readJsonString(s, start) + unescapeJson(s.substring(start + 1, end)) to (end + 1) + } + c == 't' && s.regionMatches(start, "true", 0, 4) -> true to (start + 4) + c == 'f' && s.regionMatches(start, "false", 0, 5) -> false to (start + 5) + c == 'n' && s.regionMatches(start, "null", 0, 4) -> null to (start + 4) + c == '-' || c.isDigit() -> { + var i = start + while (i < s.length && (s[i].isDigit() || s[i] == '-' || s[i] == '+' || s[i] == '.' || s[i] == 'e' || s[i] == 'E')) i++ + val raw = s.substring(start, i) + val num = raw.toDoubleOrNull() ?: raw + num to i + } + else -> null to (start + 1) + } +} + +private fun unescapeJson(s: String): String { + val sb = StringBuilder(s.length) + var i = 0 + while (i < s.length) { + val c = s[i] + if (c == '\\' && i + 1 < s.length) { + when (s[i + 1]) { + '"' -> sb.append('"') + '\\' -> sb.append('\\') + '/' -> sb.append('/') + 'n' -> sb.append('\n') + 'r' -> sb.append('\r') + 't' -> sb.append('\t') + 'b' -> sb.append('\b') + 'f' -> sb.append('\u000C') + 'u' -> { + if (i + 5 < s.length) { + val hex = s.substring(i + 2, i + 6) + sb.append(hex.toInt(16).toChar()) + i += 4 + } + } + } + i += 2 + } else { + sb.append(c) + i++ + } + } + return sb.toString() +} \ No newline at end of file diff --git a/memo-cli/src/test/kotlin/memo/cli/CliArgsTest.kt b/memo-cli/src/test/kotlin/memo/cli/CliArgsTest.kt new file mode 100644 index 0000000..b4a4499 --- /dev/null +++ b/memo-cli/src/test/kotlin/memo/cli/CliArgsTest.kt @@ -0,0 +1,48 @@ +package memo.cli + +import memo.core.SearchMode +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertTrue + +class CliArgsTest { + + @Test + fun indexCommandParsed() { + val cmd = parseArgs(arrayOf("index", "/tmp/x")) + assertTrue(cmd is IndexCmd, "ожидался IndexCmd, получено $cmd") + assertEquals("/tmp/x", cmd.path) + } + + @Test + fun searchDefaultsAreK8HybridNoJson() { + val cmd = parseArgs(arrayOf("search", "/tmp/x", "вопрос")) + assertTrue(cmd is SearchCmd, "ожидался SearchCmd, получено $cmd") + assertEquals(8, cmd.k) + assertEquals(SearchMode.HYBRID, cmd.mode) + assertEquals(false, cmd.json) + } + + @Test + fun searchFlagsParsed() { + val cmd = parseArgs( + arrayOf("search", "/tmp/x", "вопрос", "--k", "3", "--mode", "lex", "--json") + ) + assertTrue(cmd is SearchCmd, "ожидался SearchCmd, получено $cmd") + assertEquals(3, cmd.k) + assertEquals(SearchMode.LEX, cmd.mode) + assertEquals(true, cmd.json) + } + + @Test + fun badKValueFallsBackToHelp() { + val cmd = parseArgs(arrayOf("search", "/tmp/x", "вопрос", "--k", "abc")) + assertTrue(cmd is HelpCmd, "ожидался HelpCmd, получено $cmd") + } + + @Test + fun unknownCommandFallsBackToHelp() { + val cmd = parseArgs(arrayOf("сломать")) + assertTrue(cmd is HelpCmd, "ожидался HelpCmd, получено $cmd") + } +} \ No newline at end of file diff --git a/memo-core/build.gradle.kts b/memo-core/build.gradle.kts index 7e69942..92d7c86 100644 --- a/memo-core/build.gradle.kts +++ b/memo-core/build.gradle.kts @@ -3,7 +3,7 @@ plugins { } dependencies { - implementation("pw.binom.db:ksqlite:0.1.4") + api("pw.binom.db:ksqlite:0.1.4") implementation("pw.binom.ai.embeddingtext:api:5") runtimeOnly("pw.binom.ai.embeddingtext:siglip-jvm:5")