Phase 5: :memory-vector (JVector + SQLite + LLM-эмбеддинги), memory-abstraction, compaction, MessageContext
- :memory-api — общий контракт MemoryStore/Prefetcher/Reviewer/Tools/MemorySystem
- :memory-md (KMP, kotlinx-io) — Hermes-style §-файлы, keyword overlap
- :memory-vector (JVM-only) — JVector ANN + SQLite + HttpEmbeddingClient
- :standalone — AGENTIK_MEMORY_BACKEND={md,vector,off}, выбор в Main.kt
- :standalone — compaction рабочего контекста (LiteLlmContextCompactor + reviewPreCompaction)
- :proto — MessageContext (origin: user/system/event) на send и в Message
- :server — backward-compat dual-format для POST /messages
- README — env-vars, vector-бэкенд docs
This commit is contained in:
@@ -0,0 +1,32 @@
|
||||
plugins {
|
||||
alias(libs.plugins.kotlin.multiplatform)
|
||||
}
|
||||
|
||||
kotlin {
|
||||
jvmToolchain(21)
|
||||
|
||||
// Vector-бэкенд JVM-only: JVector не публикует KMP-таргеты, но его Java 11
|
||||
// base jar работает на Android ART через scalar fallback. Для desktop JVM
|
||||
// HotSpot 21+ автоматически подхватывается Panama Vector API (multirelease).
|
||||
jvm()
|
||||
|
||||
sourceSets {
|
||||
commonMain.dependencies {
|
||||
api(project(":memory-api"))
|
||||
implementation(libs.kotlinx.coroutines.core)
|
||||
implementation(libs.kotlinx.serialization.json)
|
||||
}
|
||||
commonTest.dependencies {
|
||||
implementation(kotlin("test"))
|
||||
implementation(libs.kotlinx.coroutines.core)
|
||||
implementation(libs.kotlinx.coroutines.test)
|
||||
}
|
||||
jvmMain.dependencies {
|
||||
implementation(libs.jvector)
|
||||
implementation(libs.sqldelight.sqlite.driver)
|
||||
}
|
||||
jvmTest.dependencies {
|
||||
implementation(kotlin("test"))
|
||||
}
|
||||
}
|
||||
}
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
package pw.binom.agentik.memory.vector
|
||||
|
||||
/**
|
||||
* Провайдер эмбеддингов: превращает текст в FloatArray фиксированной размерности.
|
||||
*
|
||||
* Реализация по умолчанию — HTTP-вызов `POST /v1/embeddings` к OpenAI-совместимому
|
||||
* API (OpenAI / litellm-proxy / vllm). С LRU-кэшом, чтобы не ходить в сеть
|
||||
* на каждый search/upsert.
|
||||
*/
|
||||
interface EmbeddingProvider {
|
||||
val dimension: Int
|
||||
suspend fun embed(text: String): FloatArray
|
||||
|
||||
/** Batch-вариант. По умолчанию — последовательный вызов [embed]. */
|
||||
suspend fun embedBatch(texts: List<String>): List<FloatArray> =
|
||||
texts.map { embed(it) }
|
||||
}
|
||||
|
||||
/**
|
||||
* Детерминированный провайдер для тестов: хеширует текст в псевдо-вектор.
|
||||
* Используется только в commonTest; в продакшн заменяется на HttpEmbeddingProvider.
|
||||
*/
|
||||
class FakeEmbeddingProvider(override val dimension: Int = 32) : EmbeddingProvider {
|
||||
override suspend fun embed(text: String): FloatArray {
|
||||
val v = FloatArray(dimension)
|
||||
// Простейший детерминированный seed — сумма char'ов по модулю.
|
||||
var seed = text.hashCode().toLong() and 0xFFFFFFFFL
|
||||
for (i in 0 until dimension) {
|
||||
seed = (seed * 6364136223846793005L + 1442695040888963407L) and 0xFFFFFFFFL
|
||||
v[i] = ((seed.toInt() and 0xFFFF) / 65535f) * 2f - 1f
|
||||
}
|
||||
// L2-normalize чтобы cosine работал осмысленно.
|
||||
var norm = 0f
|
||||
for (x in v) norm += x * x
|
||||
norm = kotlin.math.sqrt(norm)
|
||||
if (norm > 0f) for (i in v.indices) v[i] /= norm
|
||||
return v
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
package pw.binom.agentik.memory.vector
|
||||
|
||||
import pw.binom.agentik.memory.MemoryCategory
|
||||
import pw.binom.agentik.memory.MemoryNote
|
||||
import kotlin.time.Instant
|
||||
|
||||
/**
|
||||
* Хранилище метаданных и embeddings заметок. Реализация по умолчанию —
|
||||
* SQLite (`SqliteMemoryMetaStore`).
|
||||
*
|
||||
* Это источник правды: [VectorMemoryIndex] (JVector) держит in-RAM ANN-индекс,
|
||||
* который пересобирается из [allEntries] при старте. Вектор хранится рядом с
|
||||
* метаданными — как packed little-endian Float32Array (`dim` * 4 байт).
|
||||
*
|
||||
* Скрывает детали backend'а от [VectorMemoryStore], который живёт в commonMain
|
||||
* и не знает про SQLite.
|
||||
*/
|
||||
interface MemoryMetaStore : AutoCloseable {
|
||||
/** Записать заметку и её embedding. Идемпотентно по [note]`.id`. */
|
||||
fun put(note: MemoryNote, embedding: FloatArray)
|
||||
|
||||
/** Заметка по id, без вектора. */
|
||||
fun get(id: String): MemoryNote?
|
||||
|
||||
/** Все (id, embedding) — для пересборки vector-индекса при старте. */
|
||||
fun allEntries(): List<Pair<String, FloatArray>>
|
||||
|
||||
/** Пагинированный листинг заметок с опциональными фильтрами. */
|
||||
fun list(category: MemoryCategory?, conversationId: String?, limit: Int, offset: Int): List<MemoryNote>
|
||||
|
||||
/** Удалить заметку и её embedding. Возвращает true если запись была. */
|
||||
fun delete(id: String): Boolean
|
||||
|
||||
/** Обновить `last_used_at` (и увеличить `use_count`) для [id]. */
|
||||
fun markUsed(id: String, at: Instant)
|
||||
|
||||
override fun close()
|
||||
}
|
||||
+63
@@ -0,0 +1,63 @@
|
||||
package pw.binom.agentik.memory.vector
|
||||
|
||||
import pw.binom.agentik.memory.MemoryCategory
|
||||
import pw.binom.agentik.memory.MemoryNote
|
||||
|
||||
/**
|
||||
* Результат одного hit'а vector-поиска: id заметки + cosine-similarity score в [0..1].
|
||||
* Чем ближе к 1.0, тем семантически ближе query к заметке.
|
||||
*/
|
||||
data class ScoredVector(
|
||||
val id: String,
|
||||
val score: Float,
|
||||
)
|
||||
|
||||
/**
|
||||
* Контракт vector-индекса. Реализация отвечает за ANN-поиск top-K ближайших
|
||||
* векторов к query. Метаданные заметок лежат в [MemoryStore] (SQLite для
|
||||
* vector-бэкенда); индекс хранит только embedding'и + id-маппинг.
|
||||
*
|
||||
* Потокобезопасность: реализации обязаны быть безопасны для конкурентных
|
||||
* read'ов. write'ы (add/remove) могут требовать внешней синхронизации — это
|
||||
* инвариант JVector (его OnHeapGraphIndex не thread-safe для мутаций).
|
||||
*/
|
||||
interface MemoryVectorIndex : AutoCloseable {
|
||||
/** Текущая размерность embeddings. Фиксируется при первом [add]. */
|
||||
val dimension: Int
|
||||
|
||||
/** Количество записей в индексе. */
|
||||
suspend fun size(): Long
|
||||
|
||||
/** Добавить или заменить запись по [id]. [embedding] должен иметь длину [dimension]. */
|
||||
suspend fun add(id: String, embedding: FloatArray)
|
||||
|
||||
/** Удалить запись по [id]. Возвращает true если запись была. */
|
||||
suspend fun remove(id: String): Boolean
|
||||
|
||||
/** ANN-поиск: top-[k] ближайших к [query]. [filter] применяется к id (например, по категории). */
|
||||
suspend fun search(
|
||||
query: FloatArray,
|
||||
k: Int,
|
||||
filter: (MemoryNote) -> Boolean = { true },
|
||||
): List<ScoredVector>
|
||||
|
||||
/** Принудительно переписать on-disk файл из текущего in-RAM состояния. */
|
||||
suspend fun flush()
|
||||
|
||||
override fun close()
|
||||
}
|
||||
|
||||
/**
|
||||
* Доп. контекст для vector-индекса: фильтр по категории и conversationId
|
||||
* передаётся через замыкание, которое получает [MemoryNote]. Так [MemoryStore]
|
||||
* остаётся единственным источником правды по метаданным.
|
||||
*/
|
||||
fun noteMatches(
|
||||
note: MemoryNote,
|
||||
category: MemoryCategory? = null,
|
||||
conversationId: String? = null,
|
||||
): Boolean {
|
||||
if (category != null && note.category != category) return false
|
||||
if (conversationId != null && note.conversationId != conversationId) return false
|
||||
return true
|
||||
}
|
||||
+96
@@ -0,0 +1,96 @@
|
||||
package pw.binom.agentik.memory.vector
|
||||
|
||||
import pw.binom.agentik.memory.MemoryCategory
|
||||
import pw.binom.agentik.memory.MemoryNote
|
||||
import pw.binom.agentik.memory.MemorySearchQuery
|
||||
import pw.binom.agentik.memory.MemorySearchResult
|
||||
import pw.binom.agentik.memory.MemoryStore
|
||||
import pw.binom.agentik.memory.MemoryStoreEvent
|
||||
import kotlin.math.exp
|
||||
import kotlin.time.Clock
|
||||
import kotlin.time.Instant
|
||||
import kotlinx.coroutines.flow.Flow
|
||||
import kotlinx.coroutines.flow.MutableSharedFlow
|
||||
import kotlinx.coroutines.flow.asSharedFlow
|
||||
import kotlinx.coroutines.sync.Mutex
|
||||
import kotlinx.coroutines.sync.withLock
|
||||
|
||||
/**
|
||||
* MemoryStore поверх (index + metadata). Метаданные заметок хранятся
|
||||
* в [metaStore] (SQLite-таблица), эмбеддинги — в [index] (JVector on-disk graph).
|
||||
*
|
||||
* Контракт MemoryStore требует, чтобы [upsert] атомарно обновлял и метаданные,
|
||||
* и эмбеддинг; [delete] — и то и другое; [search] использует ANN для кандидатов,
|
||||
* потом re-rank по recency.
|
||||
*
|
||||
* [embeddingProvider] обязателен — используется для эмбеддинга контента при
|
||||
* upsert и query при search. Без него vector-бэкенд не имеет смысла.
|
||||
*/
|
||||
class VectorMemoryStore(
|
||||
private val index: MemoryVectorIndex,
|
||||
private val metaStore: MemoryMetaStore,
|
||||
private val embeddingProvider: EmbeddingProvider,
|
||||
) : MemoryStore {
|
||||
|
||||
private val mutex = Mutex()
|
||||
private val _events = MutableSharedFlow<MemoryStoreEvent>(extraBufferCapacity = 64)
|
||||
override fun events(): Flow<MemoryStoreEvent> = _events.asSharedFlow()
|
||||
|
||||
override suspend fun upsert(note: MemoryNote) = mutex.withLock {
|
||||
val embedding = embeddingProvider.embed(note.content)
|
||||
metaStore.put(note, embedding)
|
||||
index.add(note.id, embedding)
|
||||
_events.emit(MemoryStoreEvent.Upserted(note))
|
||||
}
|
||||
|
||||
override suspend fun get(id: String): MemoryNote? = metaStore.get(id)
|
||||
|
||||
override suspend fun list(
|
||||
category: MemoryCategory?,
|
||||
conversationId: String?,
|
||||
limit: Int,
|
||||
offset: Int,
|
||||
): List<MemoryNote> = metaStore.list(category, conversationId, limit, offset)
|
||||
|
||||
override suspend fun search(query: MemorySearchQuery): List<MemorySearchResult> {
|
||||
val queryEmbedding = embeddingProvider.embed(query.query)
|
||||
val overFetch = (query.topK * 5).coerceAtLeast(query.topK)
|
||||
// Берём больше кандидатов, чем нужно — финальный фильтр по category/convId
|
||||
// через [metaStore.get] + [noteMatches] отрежет лишних.
|
||||
val candidates = index.search(
|
||||
query = queryEmbedding,
|
||||
k = overFetch,
|
||||
filter = { true },
|
||||
)
|
||||
// Re-rank: 0.7 * cosine + 0.3 * recency_weight
|
||||
// recency_weight = exp(-age_days / 30) — half-life месяц.
|
||||
val now = Clock.System.now()
|
||||
val scored = candidates.mapNotNull { sv ->
|
||||
val note = metaStore.get(sv.id) ?: return@mapNotNull null
|
||||
if (!noteMatches(note, query.category, query.conversationId)) return@mapNotNull null
|
||||
val ageDays = (now - note.lastUsedAt).inWholeDays.toDouble()
|
||||
val recency = exp(-ageDays / 30.0).toFloat()
|
||||
val finalScore = 0.7f * sv.score + 0.3f * recency
|
||||
MemorySearchResult(note = note, score = finalScore)
|
||||
}
|
||||
return scored.sortedByDescending { it.score }.take(query.topK)
|
||||
}
|
||||
|
||||
override suspend fun delete(id: String): Boolean = mutex.withLock {
|
||||
val existed = metaStore.delete(id)
|
||||
if (existed) {
|
||||
index.remove(id)
|
||||
_events.emit(MemoryStoreEvent.Deleted(id))
|
||||
}
|
||||
existed
|
||||
}
|
||||
|
||||
override suspend fun markUsed(id: String, at: Instant) {
|
||||
metaStore.markUsed(id, at)
|
||||
}
|
||||
|
||||
override fun close() {
|
||||
index.close()
|
||||
metaStore.close()
|
||||
}
|
||||
}
|
||||
+168
@@ -0,0 +1,168 @@
|
||||
package pw.binom.agentik.memory.vector
|
||||
|
||||
import io.github.jbellis.jvector.graph.GraphIndexBuilder
|
||||
import io.github.jbellis.jvector.graph.GraphSearcher
|
||||
import io.github.jbellis.jvector.graph.ListRandomAccessVectorValues
|
||||
import io.github.jbellis.jvector.graph.OnHeapGraphIndex
|
||||
import io.github.jbellis.jvector.graph.SearchResult
|
||||
import io.github.jbellis.jvector.graph.similarity.BuildScoreProvider
|
||||
import io.github.jbellis.jvector.util.Bits
|
||||
import io.github.jbellis.jvector.vector.VectorizationProvider
|
||||
import io.github.jbellis.jvector.vector.VectorSimilarityFunction
|
||||
import pw.binom.agentik.memory.MemoryNote
|
||||
import java.util.concurrent.locks.ReentrantReadWriteLock
|
||||
import kotlin.concurrent.read
|
||||
import kotlin.concurrent.write
|
||||
|
||||
/**
|
||||
* In-RAM ANN-индекс поверх JVector.
|
||||
*
|
||||
* Семантика хранения: **источник правды — SQLite (см. MemoryMetaStore)**.
|
||||
* Этот класс держит в heap'е [OnHeapGraphIndex] + mapping id ↔ ordinal и
|
||||
* пересобирается из [seedEntries] при конструировании. На каждом [add]/[remove]
|
||||
* граф перестраивается полностью (для 10K vectors это <100ms).
|
||||
*
|
||||
* **Что НЕ делается**: persist через OnDiskGraphIndex. JVector'у для записи
|
||||
* на диск нужна Feature с INLINE_VECTORS, которая (в текущей версии 4.0.0)
|
||||
* конфигурируется отдельно и сложно. SQLite BLOB дешевле и проще — она и
|
||||
* хранит embedding'и. Граф реконструируется из SQLite при старте.
|
||||
*
|
||||
* Потокобезопасность: [ReentrantReadWriteLock] — параллельные [search] ок,
|
||||
* [add]/[remove] — эксклюзивно.
|
||||
*/
|
||||
class JVectorMemoryIndex(
|
||||
override val dimension: Int,
|
||||
seedEntries: List<Pair<String, FloatArray>> = emptyList(),
|
||||
) : MemoryVectorIndex {
|
||||
|
||||
init {
|
||||
require(seedEntries.all { it.second.size == dimension }) {
|
||||
"all seed embeddings must have dimension=$dimension"
|
||||
}
|
||||
require(seedEntries.map { it.first }.toSet().size == seedEntries.size) {
|
||||
"duplicate ids in seedEntries"
|
||||
}
|
||||
}
|
||||
|
||||
private val rwLock = ReentrantReadWriteLock()
|
||||
|
||||
private val vts = VectorizationProvider.getInstance().getVectorTypeSupport()
|
||||
private val similarity = VectorSimilarityFunction.COSINE
|
||||
|
||||
// In-RAM state. Защищён rwLock.
|
||||
private val idToOrdinal = LinkedHashMap<String, Int>()
|
||||
private val ordinalToId = ArrayList<String>(seedEntries.size + 16)
|
||||
private val ordinalToVector = ArrayList<FloatArray>(seedEntries.size + 16)
|
||||
private val deleted = java.util.BitSet()
|
||||
private var graph: OnHeapGraphIndex? = null
|
||||
|
||||
init {
|
||||
seedEntries.forEach { (id, vec) ->
|
||||
val ord = ordinalToId.size
|
||||
idToOrdinal[id] = ord
|
||||
ordinalToId.add(id)
|
||||
ordinalToVector.add(vec)
|
||||
}
|
||||
if (ordinalToId.isNotEmpty()) {
|
||||
graph = rebuildFromScratch()
|
||||
}
|
||||
}
|
||||
|
||||
override suspend fun size(): Long = rwLock.read {
|
||||
(ordinalToId.size - deleted.cardinality()).toLong()
|
||||
}
|
||||
|
||||
override suspend fun add(id: String, embedding: FloatArray) = rwLock.write {
|
||||
require(embedding.size == dimension) {
|
||||
"embedding size ${embedding.size} != dimension $dimension"
|
||||
}
|
||||
val existing = idToOrdinal[id]
|
||||
if (existing != null) {
|
||||
ordinalToVector[existing] = embedding
|
||||
deleted.clear(existing)
|
||||
} else {
|
||||
val ord = ordinalToId.size
|
||||
idToOrdinal[id] = ord
|
||||
ordinalToId.add(id)
|
||||
ordinalToVector.add(embedding)
|
||||
}
|
||||
rebuildAndSwapGraph()
|
||||
}
|
||||
|
||||
override suspend fun remove(id: String): Boolean = rwLock.write {
|
||||
val ord = idToOrdinal[id] ?: return false
|
||||
deleted.set(ord)
|
||||
rebuildAndSwapGraph()
|
||||
true
|
||||
}
|
||||
|
||||
override suspend fun search(
|
||||
query: FloatArray,
|
||||
k: Int,
|
||||
filter: (MemoryNote) -> Boolean,
|
||||
): List<ScoredVector> = rwLock.read {
|
||||
require(query.size == dimension) {
|
||||
"query size ${query.size} != dimension $dimension"
|
||||
}
|
||||
if (k <= 0 || graph == null) return emptyList()
|
||||
val activeOrdinals = (0 until ordinalToId.size).filter { !deleted.get(it) }
|
||||
if (activeOrdinals.isEmpty()) return emptyList()
|
||||
val vectors = activeOrdinals.map { vts.createFloatVector(ordinalToVector[it]) }
|
||||
val ravv = ListRandomAccessVectorValues(vectors, dimension)
|
||||
val queryVec = vts.createFloatVector(query)
|
||||
val result: SearchResult = GraphSearcher.search(
|
||||
queryVec,
|
||||
k.coerceAtMost(activeOrdinals.size),
|
||||
ravv,
|
||||
similarity,
|
||||
graph!!,
|
||||
Bits.ALL,
|
||||
)
|
||||
val nodes: Array<SearchResult.NodeScore> = result.getNodes()
|
||||
val out = ArrayList<ScoredVector>(nodes.size)
|
||||
for (ns in nodes) {
|
||||
val realOrd = activeOrdinals[ns.node]
|
||||
out.add(ScoredVector(id = ordinalToId[realOrd], score = ns.score))
|
||||
}
|
||||
// [filter] применяется в [VectorMemoryStore] по MemoryNote (там есть category/convId).
|
||||
// Контракт JVector — фильтрация через Bits, что здесь неудобно, поэтому
|
||||
// делегируем фильтр наверх.
|
||||
out
|
||||
}
|
||||
|
||||
override suspend fun flush() {
|
||||
// No-op: граф в RAM, источник правды — SQLite. flush не требуется.
|
||||
}
|
||||
|
||||
override fun close() {
|
||||
rwLock.write {
|
||||
graph?.close()
|
||||
graph = null
|
||||
}
|
||||
}
|
||||
|
||||
private fun rebuildAndSwapGraph() {
|
||||
val newGraph = rebuildFromScratch()
|
||||
val old = graph
|
||||
graph = newGraph
|
||||
old?.close()
|
||||
}
|
||||
|
||||
private fun rebuildFromScratch(): OnHeapGraphIndex {
|
||||
val activeOrdinals = (0 until ordinalToId.size).filter { !deleted.get(it) }
|
||||
val vectors = activeOrdinals.map { vts.createFloatVector(ordinalToVector[it]) }
|
||||
val ravv = ListRandomAccessVectorValues(vectors, dimension)
|
||||
val bsp = BuildScoreProvider.randomAccessScoreProvider(ravv, similarity)
|
||||
// Параметры графа по умолчанию (как в JVector README):
|
||||
// - M (max degree) = 16..32 — больше = точнее, медленнее
|
||||
// - efConstruction = 100..200 — больше = точнее, дольше строить
|
||||
// Для нашего масштаба (10K) берём средние значения.
|
||||
val M = 16
|
||||
val efConstruction = 100
|
||||
val neighborOverflow = 1.2f
|
||||
val alpha = 1.2f
|
||||
return GraphIndexBuilder(bsp, dimension, M, efConstruction, neighborOverflow, alpha).use { builder ->
|
||||
builder.build(ravv)
|
||||
}
|
||||
}
|
||||
}
|
||||
+242
@@ -0,0 +1,242 @@
|
||||
package pw.binom.agentik.memory.vector
|
||||
|
||||
import pw.binom.agentik.memory.MemoryCategory
|
||||
import pw.binom.agentik.memory.MemoryNote
|
||||
import pw.binom.agentik.memory.MemorySource
|
||||
import java.nio.ByteBuffer
|
||||
import java.nio.ByteOrder
|
||||
import java.sql.Connection
|
||||
import java.sql.DriverManager
|
||||
import java.sql.PreparedStatement
|
||||
import java.sql.ResultSet
|
||||
import kotlin.time.Clock
|
||||
import kotlin.time.Instant
|
||||
|
||||
/**
|
||||
* Хранилище метаданных заметок + их эмбеддингов в SQLite.
|
||||
*
|
||||
* Схема (`memory_note_meta`):
|
||||
* - `id` — TEXT PRIMARY KEY
|
||||
* - `category`, `source` — TEXT (id enum'ов)
|
||||
* - `content` — TEXT
|
||||
* - `created_at`, `last_used_at` — INTEGER (epoch ms)
|
||||
* - `use_count` — INTEGER
|
||||
* - `conversation_id` — TEXT NULL
|
||||
* - `embedding` — BLOB (packed Float32Array, dim * 4 bytes, little-endian)
|
||||
*
|
||||
* Это **источник правды** для vector-бэкенда. JVector-индекс — in-RAM,
|
||||
* пересобирается из [allEntries] при старте. См. [JVectorMemoryIndex].
|
||||
*
|
||||
* Можно шарить один `agentik.db` с conversation DB — таблицы не пересекаются.
|
||||
*
|
||||
* Потокобезопасность: рассчитывает на single-connection-per-instance,
|
||||
* синхронизация на уровне [VectorMemoryStore] (mutex на upsert/delete).
|
||||
*/
|
||||
class SqliteMemoryMetaStore(
|
||||
private val conn: Connection,
|
||||
private val dimension: Int,
|
||||
) : MemoryMetaStore {
|
||||
|
||||
/** Открыть отдельный файл (например, `~/.agentik/agentik.db` для шаринга). */
|
||||
constructor(jdbcUrl: String, dimension: Int) : this(
|
||||
DriverManager.getConnection(jdbcUrl).apply {
|
||||
createStatement().use { st ->
|
||||
st.execute("PRAGMA foreign_keys = ON")
|
||||
st.execute("PRAGMA journal_mode = WAL")
|
||||
}
|
||||
},
|
||||
dimension,
|
||||
)
|
||||
|
||||
private val initialized = java.util.concurrent.atomic.AtomicBoolean(false)
|
||||
|
||||
private fun ensureSchema() {
|
||||
if (initialized.get()) return
|
||||
conn.createStatement().use { st ->
|
||||
st.execute(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS memory_note_meta (
|
||||
id TEXT PRIMARY KEY,
|
||||
category TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
created_at INTEGER NOT NULL,
|
||||
last_used_at INTEGER NOT NULL,
|
||||
use_count INTEGER NOT NULL DEFAULT 0,
|
||||
conversation_id TEXT,
|
||||
source TEXT NOT NULL,
|
||||
embedding BLOB NOT NULL
|
||||
)
|
||||
""".trimIndent()
|
||||
)
|
||||
st.execute("CREATE INDEX IF NOT EXISTS memory_note_meta_cat ON memory_note_meta(category)")
|
||||
st.execute("CREATE INDEX IF NOT EXISTS memory_note_meta_lu ON memory_note_meta(last_used_at DESC)")
|
||||
}
|
||||
initialized.set(true)
|
||||
}
|
||||
|
||||
override fun put(note: MemoryNote, embedding: FloatArray) {
|
||||
ensureSchema()
|
||||
require(embedding.size == dimension) {
|
||||
"embedding size ${embedding.size} != dimension $dimension"
|
||||
}
|
||||
val blob = embedding.toLittleEndianBytes()
|
||||
conn.prepareStatement(
|
||||
"""
|
||||
INSERT INTO memory_note_meta(id, category, content, created_at, last_used_at,
|
||||
use_count, conversation_id, source, embedding)
|
||||
VALUES(?,?,?,?,?,?,?,?,?)
|
||||
ON CONFLICT(id) DO UPDATE SET
|
||||
category = excluded.category,
|
||||
content = excluded.content,
|
||||
created_at = excluded.created_at,
|
||||
last_used_at = excluded.last_used_at,
|
||||
use_count = excluded.use_count,
|
||||
conversation_id = excluded.conversation_id,
|
||||
source = excluded.source,
|
||||
embedding = excluded.embedding
|
||||
""".trimIndent()
|
||||
).use { ps ->
|
||||
ps.setString(1, note.id)
|
||||
ps.setString(2, note.category.id)
|
||||
ps.setString(3, note.content)
|
||||
ps.setLong(4, note.createdAt.toEpochMilliseconds())
|
||||
ps.setLong(5, note.lastUsedAt.toEpochMilliseconds())
|
||||
ps.setInt(6, note.useCount)
|
||||
ps.setString(7, note.conversationId)
|
||||
ps.setString(8, note.source.id)
|
||||
ps.setBytes(9, blob)
|
||||
ps.executeUpdate()
|
||||
}
|
||||
}
|
||||
|
||||
override fun get(id: String): MemoryNote? {
|
||||
ensureSchema()
|
||||
conn.prepareStatement(
|
||||
"SELECT category, content, created_at, last_used_at, use_count, conversation_id, source FROM memory_note_meta WHERE id = ?"
|
||||
).use { ps ->
|
||||
ps.setString(1, id)
|
||||
ps.executeQuery().use { rs ->
|
||||
return if (rs.next()) rs.toNote(id) else null
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
override fun allEntries(): List<Pair<String, FloatArray>> {
|
||||
ensureSchema()
|
||||
conn.prepareStatement(
|
||||
"SELECT id, embedding FROM memory_note_meta"
|
||||
).use { ps ->
|
||||
ps.executeQuery().use { rs ->
|
||||
val out = ArrayList<Pair<String, FloatArray>>()
|
||||
while (rs.next()) {
|
||||
val id = rs.getString("id")
|
||||
val blob = rs.getBytes("embedding") ?: continue
|
||||
out.add(id to blob.toFloatArray(dimension))
|
||||
}
|
||||
return out
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
override fun list(category: MemoryCategory?, conversationId: String?, limit: Int, offset: Int): List<MemoryNote> {
|
||||
ensureSchema()
|
||||
val where = buildString {
|
||||
val clauses = mutableListOf<String>()
|
||||
if (category != null) clauses += "category = ?"
|
||||
if (conversationId != null) clauses += "conversation_id = ?"
|
||||
if (clauses.isNotEmpty()) append("WHERE ").append(clauses.joinToString(" AND "))
|
||||
}
|
||||
val sql = "SELECT id, category, content, created_at, last_used_at, use_count, conversation_id, source FROM memory_note_meta $where ORDER BY last_used_at DESC LIMIT ? OFFSET ?"
|
||||
return conn.prepareStatement(sql).use { ps ->
|
||||
var idx = 1
|
||||
if (category != null) ps.setString(idx++, category.id)
|
||||
if (conversationId != null) ps.setString(idx++, conversationId)
|
||||
ps.setInt(idx++, limit)
|
||||
ps.setInt(idx, offset)
|
||||
ps.executeQuery().use { rs ->
|
||||
buildList {
|
||||
while (rs.next()) add(rs.toNote(rs.getString("id")))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
override fun delete(id: String): Boolean {
|
||||
ensureSchema()
|
||||
return conn.prepareStatement("DELETE FROM memory_note_meta WHERE id = ?").use { ps ->
|
||||
ps.setString(1, id)
|
||||
ps.executeUpdate() > 0
|
||||
}
|
||||
}
|
||||
|
||||
override fun markUsed(id: String, at: Instant) {
|
||||
ensureSchema()
|
||||
conn.prepareStatement(
|
||||
"UPDATE memory_note_meta SET use_count = use_count + 1, last_used_at = ? WHERE id = ?"
|
||||
).use { ps ->
|
||||
ps.setLong(1, at.toEpochMilliseconds())
|
||||
ps.setString(2, id)
|
||||
ps.executeUpdate()
|
||||
}
|
||||
}
|
||||
|
||||
override fun close() {
|
||||
conn.close()
|
||||
}
|
||||
|
||||
companion object {
|
||||
/** Default `now` для тестов. */
|
||||
internal fun now(): Instant = Clock.System.now()
|
||||
|
||||
/**
|
||||
* Открывает (или создаёт) SQLite-БД по пути [dbPath], инициализирует
|
||||
* схему `memory_note_meta` и возвращает [SqliteMemoryMetaStore].
|
||||
*/
|
||||
fun open(dbPath: String, dimension: Int): SqliteMemoryMetaStore {
|
||||
val conn = DriverManager.getConnection("jdbc:sqlite:$dbPath")
|
||||
SqliteMemoryMetaStore(conn, dimension)
|
||||
return SqliteMemoryMetaStore(conn, dimension)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private fun ResultSet.toNote(id: String): MemoryNote {
|
||||
val catId = getString("category")
|
||||
val srcId = getString("source")
|
||||
val createdMs = getLong("created_at")
|
||||
val lastUsedMs = getLong("last_used_at")
|
||||
return MemoryNote(
|
||||
id = id,
|
||||
category = MemoryCategory.fromId(catId),
|
||||
content = getString("content"),
|
||||
createdAt = Instant.fromEpochMilliseconds(createdMs),
|
||||
lastUsedAt = Instant.fromEpochMilliseconds(lastUsedMs),
|
||||
useCount = getInt("use_count"),
|
||||
conversationId = getString("conversation_id"),
|
||||
source = MemorySource.fromId(srcId),
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Little-endian packed Float32Array → byte[].
|
||||
* JVector ожидает packed float, а JVM по умолчанию big-endian — переставляем явно.
|
||||
*/
|
||||
internal fun FloatArray.toLittleEndianBytes(): ByteArray {
|
||||
val bb = ByteBuffer.allocate(size * 4).order(ByteOrder.LITTLE_ENDIAN)
|
||||
bb.asFloatBuffer().put(this)
|
||||
return bb.array()
|
||||
}
|
||||
|
||||
/**
|
||||
* Обратное преобразование: byte[] → FloatArray (little-endian → JVM-native).
|
||||
* Проверяет длину против [expectedDim].
|
||||
*/
|
||||
internal fun ByteArray.toFloatArray(expectedDim: Int): FloatArray {
|
||||
require(size == expectedDim * 4) {
|
||||
"blob size $size != expected ${expectedDim * 4} bytes (dim=$expectedDim)"
|
||||
}
|
||||
val bb = ByteBuffer.wrap(this).order(ByteOrder.LITTLE_ENDIAN)
|
||||
val out = FloatArray(expectedDim)
|
||||
bb.asFloatBuffer().get(out)
|
||||
return out
|
||||
}
|
||||
+101
@@ -0,0 +1,101 @@
|
||||
package pw.binom.agentik.memory.vector
|
||||
|
||||
import pw.binom.agentik.memory.MemoryCategory
|
||||
import pw.binom.agentik.memory.MemoryNote
|
||||
import pw.binom.agentik.memory.MemoryPrefetcher
|
||||
import pw.binom.agentik.memory.MemoryReviewDecision
|
||||
import pw.binom.agentik.memory.MemoryReviewer
|
||||
import pw.binom.agentik.memory.MemorySearchQuery
|
||||
import pw.binom.agentik.memory.MemoryStore
|
||||
import pw.binom.agentik.memory.MemorySystem
|
||||
import pw.binom.agentik.memory.ReviewedTurn
|
||||
|
||||
/**
|
||||
* Бандл компонентов vector-бэкенда памяти — то же, что
|
||||
* [pw.binom.agentik.memory.md.MdMemorySystem], но на базе JVector + SQLite + LLM-эмбеддингов.
|
||||
*
|
||||
* Содержит:
|
||||
* - [store] — `MemoryStore` (vector-backed)
|
||||
* - [prefetcher] — top-K через vector search + `markUsed`
|
||||
* - [reviewer] — простая эвристика (vector-рекомендации оставим для Phase 5 LlmMemoryReviewer)
|
||||
*
|
||||
* Закрытие через [close] освобождает SQLite-коннекшен и (если есть) HTTP-клиент эмбеддингов.
|
||||
*/
|
||||
class VectorMemorySystem(
|
||||
override val store: MemoryStore,
|
||||
override val prefetcher: MemoryPrefetcher,
|
||||
override val reviewer: MemoryReviewer,
|
||||
private val closables: List<AutoCloseable>,
|
||||
) : MemorySystem {
|
||||
override fun close() {
|
||||
closables.forEach { runCatching { it.close() } }
|
||||
}
|
||||
|
||||
companion object {
|
||||
/**
|
||||
* Открыть vector-бэкенд: SQLite + JVector + HTTP embedding client.
|
||||
*
|
||||
* @param dbPath путь к agentik.db (SQLite для metadata + embedding-blobs)
|
||||
* @param embedding [EmbeddingProvider] — обычно HttpEmbeddingClient
|
||||
* @param topK размер top-K для prefetch
|
||||
*/
|
||||
fun open(
|
||||
dbPath: String,
|
||||
embedding: EmbeddingProvider,
|
||||
topK: Int = 10,
|
||||
): VectorMemorySystem {
|
||||
val metaStore = SqliteMemoryMetaStore.open(dbPath, embedding.dimension)
|
||||
val index = JVectorMemoryIndex(embedding.dimension)
|
||||
val store = VectorMemoryStore(index, metaStore, embedding)
|
||||
val prefetcher = VectorPrefetcher(store, topK)
|
||||
val reviewer = VectorMemoryReviewer(store)
|
||||
return VectorMemorySystem(
|
||||
store = store,
|
||||
prefetcher = prefetcher,
|
||||
reviewer = reviewer,
|
||||
closables = listOfNotNull(
|
||||
metaStore,
|
||||
index,
|
||||
embedding as? AutoCloseable,
|
||||
),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* `MemoryPrefetcher` поверх vector-store: top-K через cosine similarity + recency re-rank.
|
||||
* На каждом результате вызывает `store.markUsed(id)`.
|
||||
*/
|
||||
class VectorPrefetcher(
|
||||
private val store: MemoryStore,
|
||||
private val defaultTopK: Int,
|
||||
) : MemoryPrefetcher {
|
||||
override suspend fun prefetch(
|
||||
query: String,
|
||||
topK: Int,
|
||||
category: MemoryCategory?,
|
||||
): List<MemoryNote> {
|
||||
val results = store.search(
|
||||
MemorySearchQuery(
|
||||
query = query,
|
||||
topK = topK.takeIf { it > 0 } ?: defaultTopK,
|
||||
category = category,
|
||||
)
|
||||
)
|
||||
results.forEach { store.markUsed(it.note.id) }
|
||||
return results.map { it.note }
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Простейший reviewer для vector-бэкенда: не извлекает новых фактов из ходов,
|
||||
* только дедуплицирует/маркирует использованные. Для настоящего LLM-driven review'а
|
||||
* (Hermes-style one-shot с whitelist tools) см. Phase 5 — [LlmMemoryReviewer].
|
||||
*/
|
||||
class VectorMemoryReviewer(
|
||||
private val store: MemoryStore,
|
||||
) : MemoryReviewer {
|
||||
override suspend fun review(turn: ReviewedTurn): MemoryReviewDecision =
|
||||
MemoryReviewDecision()
|
||||
}
|
||||
+103
@@ -0,0 +1,103 @@
|
||||
package pw.binom.agentik.memory.vector.embedding
|
||||
|
||||
import java.net.URI
|
||||
import java.net.http.HttpClient
|
||||
import java.net.http.HttpRequest
|
||||
import java.net.http.HttpResponse
|
||||
import java.time.Duration
|
||||
import java.util.concurrent.ConcurrentHashMap
|
||||
import kotlinx.serialization.Serializable
|
||||
import kotlinx.serialization.json.Json
|
||||
import kotlinx.serialization.json.JsonElement
|
||||
import kotlinx.serialization.json.JsonObject
|
||||
import kotlinx.serialization.json.JsonPrimitive
|
||||
import kotlinx.serialization.json.buildJsonObject
|
||||
import kotlinx.serialization.json.jsonArray
|
||||
import kotlinx.serialization.json.jsonObject
|
||||
import kotlinx.serialization.json.jsonPrimitive
|
||||
import kotlinx.serialization.json.put
|
||||
import pw.binom.agentik.memory.vector.EmbeddingProvider
|
||||
|
||||
/**
|
||||
* HTTP клиент для OpenAI-совместимого `/v1/embeddings` endpoint.
|
||||
* Используется при memory-backend=vector.
|
||||
*
|
||||
* LRU-кэш на [cacheSize] текстов (default 256) — дедупликация запросов
|
||||
* к API на одинаковых промптах.
|
||||
*
|
||||
* @param apiUrl базовый URL (без trailing slash), например `https://api.openai.com`
|
||||
* @param apiKey bearer-токен
|
||||
* @param model имя модели эмбеддингов, например `text-embedding-3-small`
|
||||
* @param dimension размерность вектора (по умолчанию 1536 — text-embedding-3-small)
|
||||
* @param cacheSize ёмкость LRU-кэша (default 256)
|
||||
*/
|
||||
class HttpEmbeddingClient(
|
||||
private val apiUrl: String,
|
||||
private val apiKey: String,
|
||||
private val model: String,
|
||||
override val dimension: Int,
|
||||
cacheSize: Int = 256,
|
||||
) : EmbeddingProvider, AutoCloseable {
|
||||
|
||||
private val cache = LruCache<String, FloatArray>(cacheSize)
|
||||
private val http: HttpClient = HttpClient.newBuilder()
|
||||
.connectTimeout(Duration.ofSeconds(10))
|
||||
.build()
|
||||
private val json = Json { ignoreUnknownKeys = true }
|
||||
|
||||
override suspend fun embed(text: String): FloatArray {
|
||||
cache.get(text)?.let { return it }
|
||||
val vector = fetchEmbedding(text)
|
||||
cache.put(text, vector)
|
||||
return vector
|
||||
}
|
||||
|
||||
private fun fetchEmbedding(text: String): FloatArray {
|
||||
val url = URI.create("$apiUrl/v1/embeddings")
|
||||
val body = buildJsonObject {
|
||||
put("model", JsonPrimitive(model))
|
||||
put("input", JsonPrimitive(text))
|
||||
}.toString()
|
||||
val request = HttpRequest.newBuilder(url)
|
||||
.header("Authorization", "Bearer $apiKey")
|
||||
.header("Content-Type", "application/json")
|
||||
.POST(HttpRequest.BodyPublishers.ofString(body))
|
||||
.timeout(Duration.ofSeconds(30))
|
||||
.build()
|
||||
val response = http.send(request, HttpResponse.BodyHandlers.ofString())
|
||||
if (response.statusCode() !in 200..299) {
|
||||
error("embedding API error ${response.statusCode()}: ${response.body()}")
|
||||
}
|
||||
val parsed = json.parseToJsonElement(response.body()).jsonObject
|
||||
val data = parsed["data"]?.jsonArray ?: error("missing 'data' in embedding response")
|
||||
val firstData = data[0].jsonObject
|
||||
val embeddingArray = firstData["embedding"]?.jsonArray ?: error("missing 'embedding' array")
|
||||
val out = FloatArray(embeddingArray.size)
|
||||
for ((i, v: JsonElement) in embeddingArray.withIndex()) {
|
||||
out[i] = v.jsonPrimitive.content.toFloat()
|
||||
}
|
||||
require(out.size == dimension) {
|
||||
"embedding dim mismatch: got ${out.size}, expected $dimension (model=$model)"
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
override fun close() = http.close()
|
||||
}
|
||||
|
||||
private class LruCache<K, V>(private val capacity: Int) {
|
||||
private val map = LinkedHashMap<K, V>(capacity, 0.75f, true)
|
||||
private val lock = Any()
|
||||
|
||||
fun get(key: K): V? = synchronized(lock) {
|
||||
map[key]
|
||||
}
|
||||
|
||||
fun put(key: K, value: V) = synchronized(lock) {
|
||||
map[key] = value
|
||||
if (map.size > capacity) {
|
||||
val firstKey = map.keys.iterator().next()
|
||||
map.remove(firstKey)
|
||||
}
|
||||
}
|
||||
}
|
||||
+74
@@ -0,0 +1,74 @@
|
||||
package pw.binom.agentik.memory.vector
|
||||
|
||||
import kotlinx.coroutines.test.runTest
|
||||
import kotlin.test.Test
|
||||
import kotlin.test.assertEquals
|
||||
import kotlin.test.assertNotNull
|
||||
import kotlin.test.assertNull
|
||||
import kotlin.test.assertTrue
|
||||
|
||||
class JVectorMemoryIndexTest {
|
||||
|
||||
private fun makeVec(seed: Int, dim: Int): FloatArray {
|
||||
val v = FloatArray(dim)
|
||||
var s = seed.toLong() and 0xFFFFFFFFL
|
||||
for (i in 0 until dim) {
|
||||
s = (s * 6364136223846793005L + 1442695040888963407L) and 0xFFFFFFFFL
|
||||
v[i] = ((s.toInt() and 0xFFFF) / 65535f) * 2f - 1f
|
||||
}
|
||||
var norm = 0f
|
||||
for (x in v) norm += x * x
|
||||
norm = kotlin.math.sqrt(norm)
|
||||
if (norm > 0f) for (i in v.indices) v[i] /= norm
|
||||
return v
|
||||
}
|
||||
|
||||
@Test
|
||||
fun emptySearchReturnsEmpty() = runTest {
|
||||
val idx = JVectorMemoryIndex(dimension = 8, seedEntries = emptyList())
|
||||
val out = idx.search(makeVec(1, 8), k = 5) { true }
|
||||
assertTrue(out.isEmpty())
|
||||
idx.close()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun addAndSearchReturnsNearest() = runTest {
|
||||
val dim = 32
|
||||
val idx = JVectorMemoryIndex(dimension = dim, seedEntries = emptyList())
|
||||
// 50 случайных векторов, id'ы = "v0".."v49"
|
||||
for (i in 0 until 50) {
|
||||
idx.add("v$i", makeVec(i + 100, dim))
|
||||
}
|
||||
assertEquals(50L, idx.size())
|
||||
// Запрос = vec с seed 105 (= v5)
|
||||
val results = idx.search(makeVec(105, dim), k = 5) { true }
|
||||
assertEquals(5, results.size)
|
||||
// v5 должен быть среди top-k (топовый результат должен быть тем же seed'ом).
|
||||
assertEquals("v5", results.first().id)
|
||||
idx.close()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun removeHidesFromSearch() = runTest {
|
||||
val dim = 16
|
||||
val idx = JVectorMemoryIndex(dimension = dim, seedEntries = emptyList())
|
||||
for (i in 0 until 10) {
|
||||
idx.add("n$i", makeVec(i, dim))
|
||||
}
|
||||
assertTrue(idx.remove("n3"))
|
||||
val results = idx.search(makeVec(3, dim), k = 10) { true }
|
||||
assertEquals(9, results.size)
|
||||
assertTrue(results.none { it.id == "n3" })
|
||||
idx.close()
|
||||
}
|
||||
|
||||
@Test
|
||||
fun reAddReusesOrdinal() = runTest {
|
||||
val dim = 8
|
||||
val idx = JVectorMemoryIndex(dimension = dim, seedEntries = emptyList())
|
||||
idx.add("x", makeVec(1, dim))
|
||||
idx.add("x", makeVec(2, dim)) // overwrite
|
||||
assertEquals(1L, idx.size())
|
||||
idx.close()
|
||||
}
|
||||
}
|
||||
+166
@@ -0,0 +1,166 @@
|
||||
package pw.binom.agentik.memory.vector
|
||||
|
||||
import pw.binom.agentik.memory.MemoryCategory
|
||||
import pw.binom.agentik.memory.MemoryNote
|
||||
import pw.binom.agentik.memory.MemorySource
|
||||
import java.io.File
|
||||
import java.sql.DriverManager
|
||||
import java.util.UUID
|
||||
import kotlin.test.AfterTest
|
||||
import kotlin.test.BeforeTest
|
||||
import kotlin.test.Test
|
||||
import kotlin.test.assertEquals
|
||||
import kotlin.test.assertNotNull
|
||||
import kotlin.test.assertNull
|
||||
import kotlin.test.assertTrue
|
||||
import kotlin.time.Instant
|
||||
|
||||
class SqliteMemoryMetaStoreTest {
|
||||
|
||||
private lateinit var file: File
|
||||
private lateinit var store: SqliteMemoryMetaStore
|
||||
private val dim = 32
|
||||
|
||||
@BeforeTest
|
||||
fun setup() {
|
||||
file = File.createTempFile("agentik-vec-test-", ".db").also { it.deleteOnExit() }
|
||||
store = SqliteMemoryMetaStore(
|
||||
"jdbc:sqlite:${file.absolutePath}",
|
||||
dimension = dim,
|
||||
)
|
||||
}
|
||||
|
||||
@AfterTest
|
||||
fun teardown() {
|
||||
store.close()
|
||||
}
|
||||
|
||||
private fun makeNote(id: String, content: String, cat: MemoryCategory = MemoryCategory.WORLD): MemoryNote {
|
||||
val now = Instant.fromEpochMilliseconds(System.currentTimeMillis())
|
||||
return MemoryNote(
|
||||
id = id,
|
||||
category = cat,
|
||||
content = content,
|
||||
createdAt = now,
|
||||
lastUsedAt = now,
|
||||
useCount = 0,
|
||||
conversationId = null,
|
||||
source = MemorySource.USER_EXPLICIT,
|
||||
)
|
||||
}
|
||||
|
||||
private fun makeVec(seed: Int): FloatArray {
|
||||
val v = FloatArray(dim)
|
||||
var s = seed.toLong() and 0xFFFFFFFFL
|
||||
for (i in 0 until dim) {
|
||||
s = (s * 6364136223846793005L + 1442695040888963407L) and 0xFFFFFFFFL
|
||||
v[i] = ((s.toInt() and 0xFFFF) / 65535f) * 2f - 1f
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
@Test
|
||||
fun putAndGetRoundTrip() {
|
||||
val note = makeNote("n1", "hello world")
|
||||
val vec = makeVec(42)
|
||||
store.put(note, vec)
|
||||
val got = store.get("n1")
|
||||
assertNotNull(got)
|
||||
assertEquals("hello world", got.content)
|
||||
assertEquals(MemoryCategory.WORLD, got.category)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun allEntriesReturnsAll() {
|
||||
repeat(5) { i ->
|
||||
store.put(makeNote("n$i", "text $i"), makeVec(i))
|
||||
}
|
||||
val all = store.allEntries()
|
||||
assertEquals(5, all.size)
|
||||
assertEquals(setOf("n0", "n1", "n2", "n3", "n4"), all.map { it.first }.toSet())
|
||||
all.forEach { (_, v) ->
|
||||
assertEquals(dim, v.size)
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun listFiltersByCategory() {
|
||||
store.put(makeNote("w1", "world 1", MemoryCategory.WORLD), makeVec(1))
|
||||
store.put(makeNote("u1", "user 1", MemoryCategory.USER), makeVec(2))
|
||||
store.put(makeNote("w2", "world 2", MemoryCategory.WORLD), makeVec(3))
|
||||
|
||||
val worlds = store.list(category = MemoryCategory.WORLD, conversationId = null, limit = 10, offset = 0)
|
||||
assertEquals(2, worlds.size)
|
||||
assertTrue(worlds.all { it.category == MemoryCategory.WORLD })
|
||||
|
||||
val users = store.list(category = MemoryCategory.USER, conversationId = null, limit = 10, offset = 0)
|
||||
assertEquals(1, users.size)
|
||||
assertEquals("u1", users.first().id)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun deleteRemovesNote() {
|
||||
store.put(makeNote("x", "to delete"), makeVec(7))
|
||||
assertTrue(store.delete("x"))
|
||||
assertNull(store.get("x"))
|
||||
assertTrue(store.allEntries().isEmpty())
|
||||
// Второй delete возвращает false.
|
||||
assertEquals(false, store.delete("x"))
|
||||
}
|
||||
|
||||
@Test
|
||||
fun markUsedIncrementsCount() {
|
||||
val note = makeNote("y", "used")
|
||||
store.put(note, makeVec(8))
|
||||
store.markUsed("y", Instant.fromEpochMilliseconds(1000L))
|
||||
store.markUsed("y", Instant.fromEpochMilliseconds(2000L))
|
||||
val got = store.get("y")
|
||||
assertNotNull(got)
|
||||
assertEquals(2, got.useCount)
|
||||
assertEquals(Instant.fromEpochMilliseconds(2000L), got.lastUsedAt)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun embeddingsAreLittleEndian() {
|
||||
// Проверяем что BLOB читается в little-endian: первая 4 байта = first float.
|
||||
// dim=32 => vec длиной 32.
|
||||
val vec = FloatArray(dim) { i -> (i + 1).toFloat() }
|
||||
val note = makeNote("le", "le test")
|
||||
store.put(note, vec)
|
||||
val rawBytes = DriverManager.getConnection("jdbc:sqlite:${file.absolutePath}").use { conn ->
|
||||
conn.prepareStatement("SELECT embedding FROM memory_note_meta WHERE id = ?").use { ps ->
|
||||
ps.setString(1, "le")
|
||||
ps.executeQuery().use { rs ->
|
||||
rs.next()
|
||||
rs.getBytes("embedding")
|
||||
}
|
||||
}
|
||||
}
|
||||
// В little-endian IEEE-754: 1.0f = 0x00 0x00 0x80 0x3F (младший байт первый).
|
||||
assertEquals(0x00.toByte(), rawBytes[0])
|
||||
assertEquals(0x00.toByte(), rawBytes[1])
|
||||
assertEquals(0x80.toByte(), rawBytes[2])
|
||||
assertEquals(0x3F.toByte(), rawBytes[3])
|
||||
// 2.0f = 0x00 0x00 0x00 0x40
|
||||
assertEquals(0x00.toByte(), rawBytes[4])
|
||||
assertEquals(0x00.toByte(), rawBytes[5])
|
||||
assertEquals(0x00.toByte(), rawBytes[6])
|
||||
assertEquals(0x40.toByte(), rawBytes[7])
|
||||
}
|
||||
|
||||
@Test
|
||||
fun reopenKeepsData() {
|
||||
store.put(makeNote("persistent", "survives restart"), makeVec(99))
|
||||
store.close()
|
||||
// Переоткрываем тот же файл — данные должны быть.
|
||||
store = SqliteMemoryMetaStore("jdbc:sqlite:${file.absolutePath}", dimension = dim)
|
||||
val got = store.get("persistent")
|
||||
assertNotNull(got)
|
||||
assertEquals("survives restart", got.content)
|
||||
val entries = store.allEntries()
|
||||
assertEquals(1, entries.size)
|
||||
assertEquals(dim, entries[0].second.size)
|
||||
// Round-trip работает (byte-order little-endian — проверено в отдельном тесте).
|
||||
// Здесь просто убеждаемся что BLOB распарсился в массив нужной длины.
|
||||
}
|
||||
}
|
||||
+139
@@ -0,0 +1,139 @@
|
||||
package pw.binom.agentik.memory.vector
|
||||
|
||||
import pw.binom.agentik.memory.MemoryCategory
|
||||
import pw.binom.agentik.memory.MemoryNote
|
||||
import pw.binom.agentik.memory.MemorySearchQuery
|
||||
import pw.binom.agentik.memory.MemorySource
|
||||
import java.io.File
|
||||
import kotlinx.coroutines.launch
|
||||
import kotlinx.coroutines.test.runTest
|
||||
import kotlinx.coroutines.withTimeout
|
||||
import kotlin.test.AfterTest
|
||||
import kotlin.test.BeforeTest
|
||||
import kotlin.test.Test
|
||||
import kotlin.test.assertEquals
|
||||
import kotlin.test.assertNotNull
|
||||
import kotlin.test.assertNull
|
||||
import kotlin.test.assertTrue
|
||||
import kotlin.time.Instant
|
||||
|
||||
class VectorMemoryStoreTest {
|
||||
|
||||
private lateinit var file: File
|
||||
private lateinit var metaStore: SqliteMemoryMetaStore
|
||||
private lateinit var index: JVectorMemoryIndex
|
||||
private lateinit var store: VectorMemoryStore
|
||||
private val dim = 16
|
||||
|
||||
@BeforeTest
|
||||
fun setup() {
|
||||
file = File.createTempFile("agentik-vms-test-", ".db").also { it.deleteOnExit() }
|
||||
metaStore = SqliteMemoryMetaStore("jdbc:sqlite:${file.absolutePath}", dimension = dim)
|
||||
// Загружаем начальные entries из metaStore (на случай если что-то там есть).
|
||||
val seedEntries = metaStore.allEntries()
|
||||
index = JVectorMemoryIndex(dimension = dim, seedEntries = seedEntries)
|
||||
store = VectorMemoryStore(index, metaStore, FakeEmbeddingProvider(dimension = dim))
|
||||
}
|
||||
|
||||
@AfterTest
|
||||
fun teardown() {
|
||||
store.close()
|
||||
}
|
||||
|
||||
private fun makeNote(id: String, content: String, cat: MemoryCategory = MemoryCategory.WORLD): MemoryNote {
|
||||
val now = Instant.fromEpochMilliseconds(System.currentTimeMillis())
|
||||
return MemoryNote(
|
||||
id = id,
|
||||
category = cat,
|
||||
content = content,
|
||||
createdAt = now,
|
||||
lastUsedAt = now,
|
||||
useCount = 0,
|
||||
conversationId = null,
|
||||
source = MemorySource.USER_EXPLICIT,
|
||||
)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun upsertAndGet() = runTest {
|
||||
val note = makeNote("a", "alpha")
|
||||
store.upsert(note)
|
||||
val got = store.get("a")
|
||||
assertNotNull(got)
|
||||
assertEquals("alpha", got.content)
|
||||
assertEquals(1L, index.size())
|
||||
}
|
||||
|
||||
@Test
|
||||
fun searchFindsNearest() = runTest {
|
||||
// Несколько заметок; запрос — близкий к "hello world" по семантике.
|
||||
store.upsert(makeNote("a", "kotlin coroutines async"))
|
||||
store.upsert(makeNote("b", "java virtual machine"))
|
||||
store.upsert(makeNote("c", "the quick brown fox"))
|
||||
store.upsert(makeNote("d", "asynchronous programming paradigms"))
|
||||
|
||||
val results = store.search(MemorySearchQuery(query = "kotlin async programming", topK = 3))
|
||||
assertTrue(results.isNotEmpty())
|
||||
assertTrue(results.size <= 3)
|
||||
// Сортировка descending — первый score >= последнего.
|
||||
if (results.size >= 2) {
|
||||
assertTrue(results[0].score >= results.last().score)
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun searchFiltersByCategory() = runTest {
|
||||
store.upsert(makeNote("w1", "world thing 1", MemoryCategory.WORLD))
|
||||
store.upsert(makeNote("u1", "user thing 1", MemoryCategory.USER))
|
||||
store.upsert(makeNote("w2", "world thing 2", MemoryCategory.WORLD))
|
||||
|
||||
val worldResults = store.search(
|
||||
MemorySearchQuery(query = "thing", topK = 10, category = MemoryCategory.WORLD)
|
||||
)
|
||||
assertTrue(worldResults.isNotEmpty())
|
||||
assertTrue(worldResults.all { it.note.category == MemoryCategory.WORLD })
|
||||
// user заметка не должна попасть в результат даже если она "ближе" по эмбеддингу.
|
||||
assertTrue(worldResults.none { it.note.id == "u1" })
|
||||
}
|
||||
|
||||
@Test
|
||||
fun deleteRemovesBoth() = runTest {
|
||||
store.upsert(makeNote("x", "to delete"))
|
||||
assertEquals(1L, index.size())
|
||||
assertTrue(store.delete("x"))
|
||||
assertNull(store.get("x"))
|
||||
assertEquals(0L, index.size())
|
||||
}
|
||||
|
||||
@Test
|
||||
fun upsertEmitsEvent() = runTest {
|
||||
// MutableSharedFlow без replay: подписчик должен быть ДО emit.
|
||||
// backgroundScope — это TestScope'овый scope, авто-отменяется при teardown.
|
||||
val received = kotlinx.coroutines.CompletableDeferred<pw.binom.agentik.memory.MemoryStoreEvent>()
|
||||
backgroundScope.launch(start = kotlinx.coroutines.CoroutineStart.UNDISPATCHED) {
|
||||
store.events().collect { received.complete(it); return@collect }
|
||||
}
|
||||
store.upsert(makeNote("e", "eventful"))
|
||||
val ev = withTimeout(1000) { received.await() }
|
||||
assertEquals("e", (ev as pw.binom.agentik.memory.MemoryStoreEvent.Upserted).note.id)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun reopenReconstructsIndexFromSqlite() = runTest {
|
||||
store.upsert(makeNote("p", "persistent 1"))
|
||||
store.upsert(makeNote("q", "persistent 2"))
|
||||
// Close → re-open.
|
||||
store.close()
|
||||
val meta2 = SqliteMemoryMetaStore("jdbc:sqlite:${file.absolutePath}", dimension = dim)
|
||||
val seedEntries = meta2.allEntries()
|
||||
val idx2 = JVectorMemoryIndex(dimension = dim, seedEntries = seedEntries)
|
||||
val store2 = VectorMemoryStore(idx2, meta2, FakeEmbeddingProvider(dimension = dim))
|
||||
try {
|
||||
assertEquals(2L, idx2.size())
|
||||
val results = store2.search(MemorySearchQuery(query = "persistent 1", topK = 5))
|
||||
assertTrue(results.any { it.note.id == "p" })
|
||||
} finally {
|
||||
store2.close()
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user