cli: index/search/status/mcp-probe

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"))
 }

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<String>): 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<String>): 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<String>): 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<String>) {
+    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 <path>
+          memo search <path> <query> [--k N] [--mode hybrid|lex|vec] [--json]
+          memo status <path>
+          memo mcp-probe --tool <name> [--args <json>]
+        """.trimIndent()
+    )
+}
+
+private fun modelPaths(): Pair<String, String> {
+    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<File> {
+    if (!base.isDirectory) return emptyList()
+    val candidates = LinkedHashSet<File>()
+    candidates.add(base)
+    val q = ArrayDeque<Pair<File, Int>>()
+    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<Hit>()
+    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<File> {
+    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<Hit> {
+    val targets = resolveSearchTargets(File(path))
+    if (targets.isEmpty()) return emptyList()
+    val (modelPath, tokenizerPath) = modelPaths()
+    val allHits = ArrayList<Hit>()
+    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<Hit>) {
+    println(formatHitsJson(hits))
+}
+
+private fun formatHitsJson(hits: List<Hit>): 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<Hit>): 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<String, Any?> {
+    val map = HashMap<String, Any?>()
+    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<Any?, Int> {
+    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()
+}

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")
+    }
+}

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")