remove :storage-ksqlite (conversation/message) and related tests; decouple schema from journal
ci / JVM build + tests (push) Successful in 6m6s
ci / JVM build + tests (push) Successful in 6m6s
This commit is contained in:
@@ -0,0 +1,32 @@
|
||||
plugins {
|
||||
alias(libs.plugins.kotlin.multiplatform)
|
||||
}
|
||||
|
||||
// KMP-реализация :vector-index-api (MutableVectorIndexStore) поверх ksqlite.
|
||||
// Brute-force ANN: SELECT всех записей + cosine similarity в Kotlin. Persistence —
|
||||
// обычная SQLite БД. Подходит для small-to-medium масштабов (≤10K записей на
|
||||
// embedding ~512d); для больших датасетов — sqlite-vec или JVector.
|
||||
//
|
||||
// Цели сборки — только те, для которых ksqlite 0.1.2 опубликован в Maven Central:
|
||||
// jvm + linuxX64/Arm64 + mingwX64. Apple/iOS/tvOS/watchOS — НЕ публикуются;
|
||||
// для них использовать JVM-only `:vector-index-jvector`.
|
||||
|
||||
kotlin {
|
||||
jvmToolchain(21)
|
||||
jvm()
|
||||
linuxX64()
|
||||
linuxArm64()
|
||||
mingwX64()
|
||||
|
||||
sourceSets {
|
||||
commonMain.dependencies {
|
||||
// ksqlite 0.1.2 опубликован в Maven Central.
|
||||
implementation("pw.binom.db:ksqlite:0.1.2")
|
||||
api(project(":vector-index-api"))
|
||||
}
|
||||
commonTest.dependencies {
|
||||
implementation(kotlin("test"))
|
||||
implementation(libs.kotlinx.coroutines.test)
|
||||
}
|
||||
}
|
||||
}
|
||||
+245
@@ -0,0 +1,245 @@
|
||||
package pw.binom.agentik.vectorindex.ksqlite
|
||||
|
||||
import pw.binom.agentik.vectorindex.MutableVectorIndexStore
|
||||
import pw.binom.agentik.vectorindex.VectorIndex
|
||||
import pw.binom.agentik.vectorindex.VectorIndexAlreadyExistsException
|
||||
import pw.binom.agentik.vectorindex.VectorSearchResult
|
||||
import pw.binom.db.ksqlite.SQLiteConnection
|
||||
import pw.binom.db.ksqlite.SQLitePreparedStatement
|
||||
import pw.binom.db.ksqlite.SQLiteResultSet
|
||||
import kotlin.uuid.Uuid
|
||||
import kotlinx.coroutines.Dispatchers
|
||||
import kotlinx.coroutines.sync.Mutex
|
||||
import kotlinx.coroutines.sync.withLock
|
||||
import kotlinx.coroutines.withContext
|
||||
import pw.binom.agentik.vectorindex.ksqlite.utils.toFloatArray
|
||||
import pw.binom.agentik.vectorindex.ksqlite.utils.toLittleEndianBytes
|
||||
import kotlin.math.sqrt
|
||||
|
||||
/**
|
||||
* ksqlite-реализация [MutableVectorIndexStore].
|
||||
*
|
||||
* **Стратегия поиска: brute-force cosine similarity.**
|
||||
* Один SELECT всех записей → декодирование BLOB → cosine sim в Kotlin →
|
||||
* sort desc → top-[limit]. Просто и KMP-совместимо. Для масштабов >10K
|
||||
* записей становится узким местом — переезжать на `:vector-index-jvector`
|
||||
* или sqlite-vec (см. KDoc [Schema]).
|
||||
*
|
||||
* ## Lifecycle соединения
|
||||
*
|
||||
* Семантика владения connection'ом идентична другим ksqlite-store'ам
|
||||
* (`pw.binom.agentik.journal.ksqlite.KsqliteJournalStore` и т.п.):
|
||||
* - `KsqliteVectorIndexStore(dimension, path)` — открывает файловое
|
||||
* соединение, прогоняет [Schema.migrate], закрывает в [close].
|
||||
* - `KsqliteVectorIndexStore(dimension, connection)` — внешнее соединение,
|
||||
* store НЕ закрывает его в [close].
|
||||
* - `KsqliteVectorIndexStore.memory(dimension, name)` — in-memory, мигрирует,
|
||||
* закрывает в [close].
|
||||
*/
|
||||
class KsqliteVectorIndexStore private constructor(
|
||||
override val dimension: Int,
|
||||
private val connection: SQLiteConnection,
|
||||
private val ownsConnection: Boolean,
|
||||
) : MutableVectorIndexStore {
|
||||
|
||||
init {
|
||||
require(dimension > 0) { "dimension must be positive, got $dimension" }
|
||||
Schema.migrate(connection)
|
||||
}
|
||||
|
||||
/**
|
||||
* Открывает файловое соединение через [SQLiteConnection.open] и берёт
|
||||
* на себя его закрытие в [close].
|
||||
*/
|
||||
constructor(dimension: Int, path: String) : this(
|
||||
dimension = dimension,
|
||||
connection = SQLiteConnection.open(path = path),
|
||||
ownsConnection = true,
|
||||
)
|
||||
|
||||
/**
|
||||
* Внешнее соединение — store НЕ закрывает его в [close].
|
||||
*/
|
||||
constructor(dimension: Int, connection: SQLiteConnection) : this(
|
||||
dimension = dimension,
|
||||
connection = connection,
|
||||
ownsConnection = false,
|
||||
)
|
||||
|
||||
private val mutex = Mutex()
|
||||
|
||||
// pre-prepare (см. KDoc KsqliteJournalStore — почему это критично против
|
||||
// SIGSEGV в StmtHolder.finalize на закрытой connection).
|
||||
private val existsStmt: SQLitePreparedStatement = connection.prepare(
|
||||
"SELECT 1 FROM ${Schema.TABLE_VECTOR_INDEX} WHERE ${Schema.COL_ID} = ?"
|
||||
)
|
||||
private val insertStmt: SQLitePreparedStatement = connection.prepare(
|
||||
"""
|
||||
INSERT INTO ${Schema.TABLE_VECTOR_INDEX}
|
||||
(${Schema.COL_ID}, ${Schema.COL_DIMENSION},
|
||||
${Schema.COL_EMBEDDING}, ${Schema.COL_PAYLOAD})
|
||||
VALUES (?, ?, ?, ?)
|
||||
""".trimIndent()
|
||||
)
|
||||
private val deleteStmt: SQLitePreparedStatement = connection.prepare(
|
||||
"DELETE FROM ${Schema.TABLE_VECTOR_INDEX} WHERE ${Schema.COL_ID} = ?"
|
||||
)
|
||||
private val clearStmt: SQLitePreparedStatement = connection.prepare(
|
||||
"DELETE FROM ${Schema.TABLE_VECTOR_INDEX}"
|
||||
)
|
||||
private val countStmt: SQLitePreparedStatement = connection.prepare(
|
||||
"SELECT COUNT(*) FROM ${Schema.TABLE_VECTOR_INDEX}"
|
||||
)
|
||||
private val allStmt: SQLitePreparedStatement = connection.prepare(
|
||||
"""
|
||||
SELECT ${Schema.COL_ID}, ${Schema.COL_EMBEDDING}, ${Schema.COL_PAYLOAD}
|
||||
FROM ${Schema.TABLE_VECTOR_INDEX}
|
||||
""".trimIndent()
|
||||
)
|
||||
|
||||
override suspend fun getSize(): Long = withContext(Dispatchers.Default) {
|
||||
mutex.withLock {
|
||||
countStmt.reset()
|
||||
countStmt.clearBindings()
|
||||
countStmt.executeQuery().use { rs ->
|
||||
if (rs.next()) rs.getLong(0) ?: 0L else 0L
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
override suspend fun add(
|
||||
id: String?,
|
||||
embedding: FloatArray,
|
||||
payload: String?,
|
||||
): VectorIndex = withContext(Dispatchers.Default) {
|
||||
require(embedding.size == dimension) {
|
||||
"embedding size ${embedding.size} != dimension $dimension"
|
||||
}
|
||||
val resolvedId = id ?: Uuid.random().toString()
|
||||
mutex.withLock {
|
||||
if (execExists(resolvedId)) {
|
||||
throw VectorIndexAlreadyExistsException()
|
||||
}
|
||||
insertStmt.reset()
|
||||
insertStmt.clearBindings()
|
||||
insertStmt.bindText(1, resolvedId)
|
||||
insertStmt.bindInt(2, dimension)
|
||||
insertStmt.bindBlob(3, embedding.toLittleEndianBytes())
|
||||
if (payload != null) insertStmt.bindText(4, payload) else insertStmt.bindNull(4)
|
||||
insertStmt.executeUpdate()
|
||||
VectorIndex(id = resolvedId, embedding = embedding, payload = payload)
|
||||
}
|
||||
}
|
||||
|
||||
override suspend fun delete(id: String): Boolean = withContext(Dispatchers.Default) {
|
||||
mutex.withLock {
|
||||
deleteStmt.reset()
|
||||
deleteStmt.clearBindings()
|
||||
deleteStmt.bindText(1, id)
|
||||
deleteStmt.executeUpdate() > 0
|
||||
}
|
||||
}
|
||||
|
||||
override suspend fun clear(): Long = withContext(Dispatchers.Default) {
|
||||
mutex.withLock {
|
||||
val before = getSizeNoLock()
|
||||
clearStmt.reset()
|
||||
clearStmt.clearBindings()
|
||||
clearStmt.executeUpdate()
|
||||
before
|
||||
}
|
||||
}
|
||||
|
||||
override suspend fun search(
|
||||
embedding: FloatArray,
|
||||
limit: Int,
|
||||
): List<VectorSearchResult> = withContext(Dispatchers.Default) {
|
||||
require(embedding.size == dimension) {
|
||||
"query size ${embedding.size} != dimension $dimension"
|
||||
}
|
||||
if (limit <= 0) return@withContext emptyList()
|
||||
mutex.withLock {
|
||||
allStmt.reset()
|
||||
allStmt.clearBindings()
|
||||
val scored = mutableListOf<ScoredRow>()
|
||||
allStmt.executeQuery().use { rs ->
|
||||
while (rs.next()) {
|
||||
val row = rs.toRow() ?: continue
|
||||
val score = cosineSimilarity(embedding, row.embedding)
|
||||
scored.add(ScoredRow(row, score))
|
||||
}
|
||||
}
|
||||
scored.sortByDescending { it.score }
|
||||
scored.take(limit).map { it.toResult() }
|
||||
}
|
||||
}
|
||||
|
||||
override fun close() {
|
||||
existsStmt.close()
|
||||
insertStmt.close()
|
||||
deleteStmt.close()
|
||||
clearStmt.close()
|
||||
countStmt.close()
|
||||
allStmt.close()
|
||||
if (ownsConnection) connection.close()
|
||||
}
|
||||
|
||||
companion object {
|
||||
fun memory(dimension: Int, name: String? = null): KsqliteVectorIndexStore =
|
||||
KsqliteVectorIndexStore(
|
||||
dimension = dimension,
|
||||
connection = SQLiteConnection.memory(name),
|
||||
ownsConnection = true,
|
||||
)
|
||||
}
|
||||
|
||||
private fun execExists(id: String): Boolean {
|
||||
existsStmt.reset()
|
||||
existsStmt.clearBindings()
|
||||
existsStmt.bindText(1, id)
|
||||
existsStmt.executeQuery().use { rs -> return rs.next() }
|
||||
}
|
||||
|
||||
private fun getSizeNoLock(): Long {
|
||||
countStmt.reset()
|
||||
countStmt.clearBindings()
|
||||
countStmt.executeQuery().use { rs ->
|
||||
if (rs.next()) return rs.getLong(0) ?: 0L
|
||||
}
|
||||
return 0L
|
||||
}
|
||||
|
||||
private fun SQLiteResultSet.toRow(): Row? {
|
||||
val id = getText(0) ?: return null
|
||||
val blob = getBlob(1) ?: return null
|
||||
val payload = getText(2)
|
||||
return Row(id = id, embedding = blob.toFloatArray(), payload = payload)
|
||||
}
|
||||
|
||||
private data class Row(val id: String, val embedding: FloatArray, val payload: String?)
|
||||
|
||||
private data class ScoredRow(val row: Row, val score: Float) {
|
||||
fun toResult(): VectorSearchResult = VectorSearchResult(
|
||||
index = VectorIndex(id = row.id, embedding = row.embedding, payload = row.payload),
|
||||
score = score,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Cosine similarity двух одинаковой длины векторов.
|
||||
* Возвращает 0, если один из векторов — нулевой (безопасно для пустых запросов).
|
||||
*/
|
||||
private fun cosineSimilarity(a: FloatArray, b: FloatArray): Float {
|
||||
require(a.size == b.size) { "dimension mismatch: ${a.size} vs ${b.size}" }
|
||||
var dot = 0f
|
||||
var normA = 0f
|
||||
var normB = 0f
|
||||
for (i in a.indices) {
|
||||
dot += a[i] * b[i]
|
||||
normA += a[i] * a[i]
|
||||
normB += b[i] * b[i]
|
||||
}
|
||||
val denom = sqrt(normA) * sqrt(normB)
|
||||
return if (denom == 0f) 0f else dot / denom
|
||||
}
|
||||
+78
@@ -0,0 +1,78 @@
|
||||
package pw.binom.agentik.vectorindex.ksqlite
|
||||
|
||||
import pw.binom.db.ksqlite.SQLiteConnection
|
||||
|
||||
/**
|
||||
* Имена таблиц/колонок для ksqlite-бэкенда `:vector-index-api`.
|
||||
*
|
||||
* Хранит embeddings как BLOB (raw float32 LE, см. [pw.binom.agentik.vectorindex.ksqlite.utils.toLittleEndianBytes])
|
||||
* + dimension для sanity-check на insert + опциональный payload как TEXT.
|
||||
*
|
||||
* Brute-force search не использует индексов на стороне SQLite — все embeddings
|
||||
* загружаются одним SELECT и cosine similarity считается в Kotlin. Для
|
||||
* масштабов >10K записей это становится узким местом; тогда мигрировать на
|
||||
* `:vector-index-jvector` или sqlite-vec.
|
||||
*/
|
||||
internal object Schema {
|
||||
|
||||
/** Версия схемы. Увеличивать при ЛЮБОМ изменении DDL. */
|
||||
const val CURRENT_VERSION: Int = 1
|
||||
|
||||
const val TABLE_VECTOR_INDEX = "vector_index"
|
||||
|
||||
const val COL_ID = "id"
|
||||
const val COL_DIMENSION = "dimension"
|
||||
const val COL_EMBEDDING = "embedding"
|
||||
const val COL_PAYLOAD = "payload"
|
||||
|
||||
private val v1Ddl = """
|
||||
CREATE TABLE IF NOT EXISTS $TABLE_VECTOR_INDEX (
|
||||
$COL_ID TEXT NOT NULL PRIMARY KEY,
|
||||
$COL_DIMENSION INTEGER NOT NULL,
|
||||
$COL_EMBEDDING BLOB NOT NULL,
|
||||
$COL_PAYLOAD TEXT
|
||||
);
|
||||
""".trimIndent()
|
||||
|
||||
/**
|
||||
* Прогоняет миграцию схемы до [CURRENT_VERSION] на пустой или существующей БД.
|
||||
*
|
||||
* Версия хранится в `PRAGMA user_version` (стандартный SQLite-механизм,
|
||||
* 32-bit int в заголовке БД — без своей таблицы). Каждая миграция —
|
||||
* блок DDL под номером `fromV+1`, выполняется в транзакции. Если миграция
|
||||
* упадёт посередине — `ROLLBACK` оставит БД на предыдущей версии.
|
||||
*
|
||||
* Идемпотентен: повторный вызов на уже мигрированной БД — no-op.
|
||||
*/
|
||||
fun migrate(conn: SQLiteConnection) {
|
||||
val current = readUserVersion(conn)
|
||||
if (current >= CURRENT_VERSION) return
|
||||
|
||||
conn.exec("BEGIN")
|
||||
try {
|
||||
if (current < 1) {
|
||||
conn.exec(v1Ddl)
|
||||
}
|
||||
writeUserVersion(conn, CURRENT_VERSION)
|
||||
conn.exec("COMMIT")
|
||||
} catch (t: Throwable) {
|
||||
runCatching { conn.exec("ROLLBACK") }
|
||||
throw t
|
||||
}
|
||||
}
|
||||
|
||||
private fun readUserVersion(conn: SQLiteConnection): Int {
|
||||
conn.prepare("PRAGMA user_version").use { stmt ->
|
||||
stmt.executeQuery().use { rs ->
|
||||
if (rs.next()) return rs.getLong(0)?.toInt() ?: 0
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
private fun writeUserVersion(conn: SQLiteConnection, version: Int) {
|
||||
// SQLite PRAGMA с literal-аргументом нельзя параметризовать через `?`,
|
||||
// поэтому собираем SQL строкой (значение контролируемое, не user input).
|
||||
conn.exec("PRAGMA user_version = $version")
|
||||
}
|
||||
}
|
||||
+36
@@ -0,0 +1,36 @@
|
||||
package pw.binom.agentik.vectorindex.ksqlite.utils
|
||||
|
||||
/**
|
||||
* Multiplatform-safe сериализация [FloatArray] ↔ [ByteArray] в little-endian.
|
||||
*
|
||||
* Используется для хранения embeddings в SQLite BLOB-колонке. Float занимает
|
||||
* ровно 4 байта (IEEE-754 binary32). На JVM little-endian совпадает с native
|
||||
* byte order; на linux/macos/ios/mingw — то же самое. Выбран little-endian
|
||||
* как наиболее распространённый вариант и для совместимости с sqlite-vec
|
||||
* (см. `pw.binom.db.ksqlite.SQLitePreparedStatement.bindVector`).
|
||||
*/
|
||||
internal fun FloatArray.toLittleEndianBytes(): ByteArray {
|
||||
val out = ByteArray(size * 4)
|
||||
for (i in indices) {
|
||||
val bits = this[i].toRawBits()
|
||||
out[i * 4 + 0] = bits.toByte()
|
||||
out[i * 4 + 1] = (bits ushr 8).toByte()
|
||||
out[i * 4 + 2] = (bits ushr 16).toByte()
|
||||
out[i * 4 + 3] = (bits ushr 24).toByte()
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
internal fun ByteArray.toFloatArray(): FloatArray {
|
||||
require(size % 4 == 0) { "byte array size $size not a multiple of 4" }
|
||||
val out = FloatArray(size / 4)
|
||||
for (i in out.indices) {
|
||||
val off = i * 4
|
||||
val bits = (this[off].toInt() and 0xFF).toLong() or
|
||||
((this[off + 1].toInt() and 0xFF).toLong() shl 8) or
|
||||
((this[off + 2].toInt() and 0xFF).toLong() shl 16) or
|
||||
((this[off + 3].toInt() and 0xFF).toLong() shl 24)
|
||||
out[i] = Float.fromBits(bits.toInt())
|
||||
}
|
||||
return out
|
||||
}
|
||||
+179
@@ -0,0 +1,179 @@
|
||||
package pw.binom.agentik.vectorindex.ksqlite
|
||||
|
||||
import kotlinx.coroutines.test.runTest
|
||||
import pw.binom.agentik.vectorindex.VectorIndexAlreadyExistsException
|
||||
import pw.binom.agentik.vectorindex.ksqlite.utils.toFloatArray
|
||||
import pw.binom.db.ksqlite.SQLiteConnection
|
||||
import kotlin.test.Test
|
||||
import kotlin.test.assertEquals
|
||||
import kotlin.test.assertFailsWith
|
||||
import kotlin.test.assertNotNull
|
||||
import kotlin.test.assertTrue
|
||||
|
||||
class KsqliteVectorIndexStoreTest {
|
||||
|
||||
private fun vec(vararg values: Float) = values
|
||||
|
||||
private fun newStore(dim: Int = 4): KsqliteVectorIndexStore =
|
||||
KsqliteVectorIndexStore.memory(dimension = dim, name = "test-${kotlin.random.Random.nextLong()}")
|
||||
|
||||
@Test
|
||||
fun `add with explicit id and search returns it as top-1`() = runTest {
|
||||
val store = newStore(4)
|
||||
try {
|
||||
val a = store.add(id = "a", embedding = vec(1f, 0f, 0f, 0f), payload = "first")
|
||||
assertEquals("a", a.id)
|
||||
assertEquals("first", a.payload)
|
||||
|
||||
val results = store.search(embedding = vec(1f, 0f, 0f, 0f), limit = 1)
|
||||
assertEquals(1, results.size)
|
||||
assertEquals("a", results[0].index.id)
|
||||
assertTrue(results[0].score > 0.99f, "score=${results[0].score} should be ~1.0")
|
||||
} finally {
|
||||
store.close()
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `add with null id auto-generates`() = runTest {
|
||||
val store = newStore(2)
|
||||
try {
|
||||
val r1 = store.add(id = null, embedding = vec(1f, 0f), payload = null)
|
||||
val r2 = store.add(id = null, embedding = vec(0f, 1f), payload = null)
|
||||
assertNotNull(r1.id)
|
||||
assertNotNull(r2.id)
|
||||
assertTrue(r1.id != r2.id, "auto-generated ids must differ")
|
||||
assertEquals(2L, store.getSize())
|
||||
} finally {
|
||||
store.close()
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `add with duplicate explicit id throws`() = runTest {
|
||||
val store = newStore(2)
|
||||
try {
|
||||
store.add(id = "dup", embedding = vec(1f, 0f), payload = null)
|
||||
assertFailsWith<VectorIndexAlreadyExistsException> {
|
||||
store.add(id = "dup", embedding = vec(0f, 1f), payload = null)
|
||||
}
|
||||
} finally {
|
||||
store.close()
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `delete removes and clears count`() = runTest {
|
||||
val store = newStore(2)
|
||||
try {
|
||||
store.add(id = "x", embedding = vec(1f, 0f), payload = null)
|
||||
store.add(id = "y", embedding = vec(0f, 1f), payload = null)
|
||||
assertEquals(2L, store.getSize())
|
||||
|
||||
assertTrue(store.delete("x"))
|
||||
assertEquals(1L, store.getSize())
|
||||
assertEquals(false, store.delete("x"), "second delete returns false")
|
||||
|
||||
// После delete можно заново add с тем же id.
|
||||
store.add(id = "x", embedding = vec(-1f, 0f), payload = null)
|
||||
assertEquals(2L, store.getSize())
|
||||
} finally {
|
||||
store.close()
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `clear empties and returns count`() = runTest {
|
||||
val store = newStore(2)
|
||||
try {
|
||||
store.add(id = "a", embedding = vec(1f, 0f), payload = null)
|
||||
store.add(id = "b", embedding = vec(0f, 1f), payload = null)
|
||||
store.add(id = "c", embedding = vec(-1f, 0f), payload = null)
|
||||
|
||||
val cleared = store.clear()
|
||||
assertEquals(3L, cleared)
|
||||
assertEquals(0L, store.getSize())
|
||||
assertEquals(emptyList(), store.search(vec(1f, 0f), limit = 5))
|
||||
} finally {
|
||||
store.close()
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `search returns top-K sorted by descending score`() = runTest {
|
||||
val store = newStore(3)
|
||||
try {
|
||||
store.add(id = "exact", embedding = vec(1f, 0f, 0f), payload = null)
|
||||
store.add(id = "noise1", embedding = vec(-1f, 0f, 0f), payload = null)
|
||||
store.add(id = "noise2", embedding = vec(0f, 1f, 0f), payload = null)
|
||||
|
||||
val results = store.search(embedding = vec(1f, 0f, 0f), limit = 3)
|
||||
assertEquals(3, results.size)
|
||||
assertTrue(results[0].score >= results[1].score)
|
||||
assertTrue(results[1].score >= results[2].score)
|
||||
assertEquals("exact", results[0].index.id)
|
||||
} finally {
|
||||
store.close()
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `search with limit greater than index returns all entries`() = runTest {
|
||||
val store = newStore(2)
|
||||
try {
|
||||
store.add(id = "a", embedding = vec(1f, 0f), payload = null)
|
||||
store.add(id = "b", embedding = vec(0f, 1f), payload = null)
|
||||
val results = store.search(vec(1f, 0f), limit = 100)
|
||||
assertEquals(2, results.size)
|
||||
} finally {
|
||||
store.close()
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `payload round-trips through search`() = runTest {
|
||||
val store = newStore(2)
|
||||
try {
|
||||
store.add(id = "p", embedding = vec(1f, 0f), payload = """{"k":"v"}""")
|
||||
val results = store.search(vec(1f, 0f), limit = 1)
|
||||
assertEquals(1, results.size)
|
||||
assertEquals("p", results[0].index.id)
|
||||
assertEquals("""{"k":"v"}""", results[0].index.payload)
|
||||
} finally {
|
||||
store.close()
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `embedding BLOB round-trips bit-exact through raw SQL`() = runTest {
|
||||
// Открываем свой connection, создаём store, пишем, потом читаем через
|
||||
// raw SQL на том же connection — проверяем что FloatArray<->BLOB encoding
|
||||
// round-trip'ит без потерь.
|
||||
val conn = SQLiteConnection.memory("roundtrip-${kotlin.random.Random.nextLong()}")
|
||||
try {
|
||||
val store = KsqliteVectorIndexStore(dimension = 3, connection = conn)
|
||||
store.add(id = "p", embedding = vec(1.5f, -2.25f, 3.875f), payload = null)
|
||||
store.close()
|
||||
|
||||
val stmt = conn.prepare(
|
||||
"SELECT id, embedding, payload FROM vector_index WHERE id = ?"
|
||||
)
|
||||
stmt.bindText(1, "p")
|
||||
stmt.executeQuery().use { rs ->
|
||||
assertTrue(rs.next())
|
||||
assertEquals("p", rs.getText(0))
|
||||
val blob = rs.getBlob(1)!!
|
||||
val floats = blob.toFloatArray()
|
||||
assertEquals(3, floats.size)
|
||||
// Float сравниваем по toRawBits из-за float-округления.
|
||||
assertEquals(1.5f.toRawBits(), floats[0].toRawBits())
|
||||
assertEquals((-2.25f).toRawBits(), floats[1].toRawBits())
|
||||
assertEquals(3.875f.toRawBits(), floats[2].toRawBits())
|
||||
assertEquals(null, rs.getText(2))
|
||||
}
|
||||
stmt.close()
|
||||
} finally {
|
||||
conn.close()
|
||||
}
|
||||
}
|
||||
}
|
||||
+92
@@ -0,0 +1,92 @@
|
||||
package pw.binom.agentik.vectorindex.ksqlite
|
||||
|
||||
import kotlinx.coroutines.test.runTest
|
||||
import pw.binom.db.ksqlite.SQLiteConnection
|
||||
import kotlin.test.Test
|
||||
import kotlin.test.assertEquals
|
||||
import kotlin.test.assertTrue
|
||||
|
||||
/**
|
||||
* Тесты Schema.migrate(): idempotency, current_user_version, table existence.
|
||||
*
|
||||
* Принцип: каждый тест использует свежий in-memory connection, чтобы не
|
||||
* зависеть от порядка выполнения.
|
||||
*/
|
||||
class SchemaMigrationTest {
|
||||
|
||||
private fun freshConn(name: String = "mig-${kotlin.random.Random.nextLong()}"): SQLiteConnection =
|
||||
SQLiteConnection.memory(name)
|
||||
|
||||
@Test
|
||||
fun `fresh DB gets the table and CURRENT_VERSION`() = runTest {
|
||||
val conn = freshConn()
|
||||
try {
|
||||
Schema.migrate(conn)
|
||||
|
||||
assertTrue(tableExists(conn, Schema.TABLE_VECTOR_INDEX), "vector_index table should exist")
|
||||
assertEquals(Schema.CURRENT_VERSION, readUserVersion(conn))
|
||||
} finally {
|
||||
conn.close()
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `migrate is idempotent on already-migrated DB`() = runTest {
|
||||
val conn = freshConn()
|
||||
try {
|
||||
Schema.migrate(conn)
|
||||
// Добавим запись, чтобы убедиться что вторая migrate ничего не сломала.
|
||||
val store = KsqliteVectorIndexStore(dimension = 2, connection = conn)
|
||||
store.add(id = "x", embedding = floatArrayOf(1f, 0f), payload = null)
|
||||
store.close()
|
||||
|
||||
// Повторный migrate должен быть no-op.
|
||||
Schema.migrate(conn)
|
||||
assertEquals(Schema.CURRENT_VERSION, readUserVersion(conn))
|
||||
|
||||
// Запись должна быть на месте.
|
||||
val checkStmt = conn.prepare("SELECT COUNT(*) FROM ${Schema.TABLE_VECTOR_INDEX}")
|
||||
checkStmt.executeQuery().use { rs ->
|
||||
assertTrue(rs.next())
|
||||
assertEquals(1L, rs.getLong(0))
|
||||
}
|
||||
checkStmt.close()
|
||||
} finally {
|
||||
conn.close()
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `migrate leaves user_version alone if already at CURRENT_VERSION`() = runTest {
|
||||
val conn = freshConn()
|
||||
try {
|
||||
// Имитируем БД, которая уже мигрирована (выставляем user_version вручную).
|
||||
conn.exec("PRAGMA user_version = ${Schema.CURRENT_VERSION}")
|
||||
assertEquals(Schema.CURRENT_VERSION, readUserVersion(conn))
|
||||
|
||||
// migrate должен быть no-op — таблица НЕ создаётся.
|
||||
Schema.migrate(conn)
|
||||
assertEquals(false, tableExists(conn, Schema.TABLE_VECTOR_INDEX),
|
||||
"migrate on already-current DB must not create tables")
|
||||
} finally {
|
||||
conn.close()
|
||||
}
|
||||
}
|
||||
|
||||
private fun tableExists(conn: SQLiteConnection, name: String): Boolean {
|
||||
val stmt = conn.prepare(
|
||||
"SELECT 1 FROM sqlite_master WHERE type='table' AND name=?"
|
||||
)
|
||||
stmt.bindText(1, name)
|
||||
stmt.executeQuery().use { rs -> return rs.next() }
|
||||
}
|
||||
|
||||
private fun readUserVersion(conn: SQLiteConnection): Int {
|
||||
conn.prepare("PRAGMA user_version").use { stmt ->
|
||||
stmt.executeQuery().use { rs ->
|
||||
if (rs.next()) return rs.getLong(0)?.toInt() ?: 0
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user