core: каркас сборки, embedder, схема БД, 4 теста

2026-10-02 01:06:06 +03:00
parent cefc17d782
commit 700508dc8e
13 changed files with 585 additions and 0 deletions
@@ -0,0 +1,11 @@
+plugins {
+    kotlin("jvm")
+}
+
+dependencies {
+    implementation("pw.binom.db:ksqlite:0.1.4")
+    implementation("pw.binom.ai.embeddingtext:api:5")
+    runtimeOnly("pw.binom.ai.embeddingtext:siglip-jvm:5")
+
+    testImplementation(kotlin("test"))
+}
@@ -0,0 +1,34 @@
+package memo.core
+
+import pw.binom.db.ksqlite.SQLiteConnection
+
+class Db(val path: String) : AutoCloseable {
+
+    val conn: SQLiteConnection = SQLiteConnection.open(path)
+
+    init {
+        conn.exec("PRAGMA journal_mode=WAL")
+        conn.exec("PRAGMA synchronous=NORMAL")
+    }
+
+    fun init() {
+        conn.exec(
+            "CREATE TABLE IF NOT EXISTS files(" +
+                "path TEXT PRIMARY KEY, mtime REAL, size INTEGER, hash TEXT, indexed_at REAL" +
+                ")"
+        )
+        conn.exec(
+            "CREATE TABLE IF NOT EXISTS chunks(" +
+                "id INTEGER PRIMARY KEY, path TEXT NOT NULL, heading TEXT, " +
+                "line INTEGER NOT NULL, ord INTEGER NOT NULL, text TEXT NOT NULL, hash TEXT NOT NULL" +
+                ")"
+        )
+        conn.exec("CREATE INDEX IF NOT EXISTS idx_chunks_path ON chunks(path)")
+        conn.exec("CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(text, heading, tokenize='unicode61')")
+        conn.exec("CREATE VIRTUAL TABLE IF NOT EXISTS chunks_vec USING vec0(embedding float[768])")
+    }
+
+    override fun close() {
+        conn.close()
+    }
+}
@@ -0,0 +1,46 @@
+package memo.core
+
+import pw.binom.voice.embeddingtext.TextEmbeddingExtractor
+
+class Embedder(
+    private val modelPath: String,
+    private val tokenizerPath: String,
+) : AutoCloseable {
+
+    private var extractor: TextEmbeddingExtractor? = null
+    private val factory: Class<*>? = runCatching {
+        Class.forName("pw.binom.voice.embeddingtext.Siglip2TextExtractorFactoryKt")
+    }.getOrNull()
+    private val createMethod = factory?.methods?.firstOrNull {
+        it.name == "createSiglip2TextExtractor" &&
+            it.parameterTypes.size == 2 &&
+            it.parameterTypes[0] == String::class.java &&
+            it.parameterTypes[1] == String::class.java
+    }
+
+    private fun obtain(): TextEmbeddingExtractor {
+        extractor?.let { return it }
+        val method = createMethod ?: error(
+                "createSiglip2TextExtractor is not available on classpath; " +
+                    "ensure pw.binom.ai.embeddingtext:siglip-jvm is on the runtime classpath"
+            )
+        val created = method.invoke(null, modelPath, tokenizerPath) as TextEmbeddingExtractor
+        extractor = created
+        return created
+    }
+
+    fun embed(text: String): FloatArray {
+        val ex = obtain()
+        val values = ex.embed(text).values
+        if (values.size != 768) {
+            throw IllegalStateException("expected 768 dims, got ${values.size}")
+        }
+        return values
+    }
+
+    override fun close() {
+        val current = extractor ?: return
+        extractor = null
+        current.close()
+    }
+}
@@ -0,0 +1,143 @@
+package memo.core
+
+import java.io.File
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertTrue
+import kotlin.test.fail
+import pw.binom.db.ksqlite.SQLiteConnection
+
+class CoreSmokeTest {
+
+    @Test
+    fun ksqliteSmoke() {
+        val file = File.createTempFile("memo-ksqlite-", ".db")
+        file.deleteOnExit()
+        val conn = SQLiteConnection.open(file.absolutePath)
+        try {
+            conn.exec("CREATE TABLE t(a INTEGER)")
+            conn.exec("INSERT INTO t(a) VALUES (42)")
+
+            val stmt = conn.prepare("SELECT a FROM t")
+            try {
+                val rs = stmt.executeQuery()
+                try {
+                    assertTrue(rs.next(), "expected one row")
+                    assertEquals(42L, rs.getLong(0)!!)
+                } finally {
+                    rs.close()
+                }
+            } finally {
+                stmt.close()
+            }
+        } finally {
+            conn.close()
+            file.delete()
+        }
+    }
+
+    @Test
+    fun vec0KnnRoundTrip() {
+        val file = File.createTempFile("memo-vec0-", ".db")
+        file.deleteOnExit()
+        val conn = SQLiteConnection.open(file.absolutePath)
+        try {
+            conn.exec("CREATE VIRTUAL TABLE v USING vec0(embedding float[4])")
+
+            val insert = conn.prepare("INSERT INTO v(rowid, embedding) VALUES (?, ?)")
+            try {
+                insert.bindLong(1, 1L)
+                insert.bindVector(2, floatArrayOf(1.0f, 0.0f, 0.0f, 0.0f))
+                insert.executeUpdate()
+
+                insert.reset()
+                insert.bindLong(1, 2L)
+                insert.bindVector(2, floatArrayOf(0.0f, 1.0f, 0.0f, 0.0f))
+                insert.executeUpdate()
+
+                insert.reset()
+                insert.bindLong(1, 3L)
+                insert.bindVector(2, floatArrayOf(0.0f, 0.0f, 1.0f, 0.0f))
+                insert.executeUpdate()
+            } finally {
+                insert.close()
+            }
+
+            val query = conn.prepare("SELECT rowid FROM v WHERE embedding MATCH ? ORDER BY distance LIMIT 1")
+            try {
+                query.bindVector(1, floatArrayOf(0.1f, 0.9f, 0.0f, 0.0f))
+                val rs = query.executeQuery()
+                try {
+                    assertTrue(rs.next(), "expected one match")
+                    assertEquals(2L, rs.getLong(0)!!)
+                } finally {
+                    rs.close()
+                }
+            } finally {
+                query.close()
+            }
+        } finally {
+            conn.close()
+            file.delete()
+        }
+    }
+
+    @Test
+    fun fts5FindsCyrillic() {
+        val file = File.createTempFile("memo-fts-", ".db")
+        file.deleteOnExit()
+        val conn = SQLiteConnection.open(file.absolutePath)
+        try {
+            conn.exec("CREATE VIRTUAL TABLE docs USING fts5(text, tokenize='unicode61')")
+
+            val insert = conn.prepare("INSERT INTO docs(rowid, text) VALUES (?, ?)")
+            try {
+                insert.bindLong(1, 1L)
+                insert.bindText(2, "внутренний домен траефик")
+                insert.executeUpdate()
+            } finally {
+                insert.close()
+            }
+
+            val query = conn.prepare("SELECT rowid FROM docs WHERE docs MATCH ?")
+            try {
+                query.bindText(1, "траефик")
+                val rs = query.executeQuery()
+                try {
+                    var hits = 0
+                    var lastRowid = -1L
+                    while (rs.next()) {
+                        hits++
+                        lastRowid = rs.getLong(0)!!
+                    }
+                    assertEquals(1, hits, "expected exactly one FTS5 hit")
+                    assertEquals(1L, lastRowid)
+                } finally {
+                    rs.close()
+                }
+            } finally {
+                query.close()
+            }
+        } finally {
+            conn.close()
+            file.delete()
+        }
+    }
+
+    @Test
+    fun embedderProduces768() {
+        val modelDir = System.getenv("MEMO_MODEL_DIR") ?: "/root/WORK/memo/models/siglip2"
+        val modelPath = "$modelDir/text_model_int8.onnx"
+        val tokenizerPath = "$modelDir/tokenizer.model"
+
+        if (!File(modelPath).exists() || !File(tokenizerPath).exists()) {
+            fail("модель не найдена: $modelDir")
+        }
+
+        Embedder(modelPath, tokenizerPath).use { embedder ->
+            val v = embedder.embed("привет мир")
+            assertEquals(768, v.size)
+            assertTrue(v.none { it.isNaN() }, "embedding contains NaN")
+        }
+    }
+}