Merge pull request #20 from smartscanapp/update/store

Update store and removed retriever (breaking changes)

Author: dev-diaries41

core/src/main/java/com/fpf/smartscansdk/core/embeddings/FileEmbeddingRetriever.kt

@@ -12,6 +12,7 @@ import java.io.RandomAccessFile
 import java.nio.ByteBuffer
 import java.nio.ByteOrder
 import java.nio.channels.FileChannel
+import kotlin.collections.map
 
 class FileEmbeddingStore(
     private val file: File,
@@ -25,13 +26,10 @@ class FileEmbeddingStore(
     }
 
     private var cache: LinkedHashMap<Long, Embedding>? = null
+    private var cachedIds: List<Long>? = null
 
     override val exists: Boolean get() = file.exists()
 
-    override val isCached: Boolean
-        get() = cache != null
-
-
     // prevent OOM in FileEmbeddingStore.save() by batching writes
     private suspend fun save(embeddingsList: List<Embedding>): Unit = withContext(Dispatchers.IO) {
         if (embeddingsList.isEmpty()) return@withContext
@@ -99,6 +97,19 @@ class FileEmbeddingStore(
         }
     }
 
+    suspend fun get(ids: List<Long>): List<Embedding> = withContext(Dispatchers.IO) {
+        val map = cache ?: run {
+            val all = get()
+            LinkedHashMap(all.associateBy { it.id })
+        }
+        val embeddings = mutableListOf<Embedding>()
+
+        for (id in ids) {
+            map.get(id)?.let { embeddings.add(it) }
+        }
+        embeddings
+    }
+
     override suspend fun add(newEmbeddings: List<Embedding>): Unit = withContext(Dispatchers.IO) {
         if (newEmbeddings.isEmpty()) return@withContext
 
@@ -185,29 +196,42 @@ class FileEmbeddingStore(
         }
     }
 
-    suspend fun get(ids: List<Long>): List<Embedding> = withContext(Dispatchers.IO) {
-        val map = cache ?: run {
-            val all = get()
-            LinkedHashMap(all.associateBy { it.id })
-        }
-        val embeddings = mutableListOf<Embedding>()
 
-        for (id in ids) {
-            map.get(id)?.let { embeddings.add(it) }
-        }
-        embeddings
+    override fun clear(){
+        cache = null
     }
 
-    suspend fun get(id: Long): Embedding? = withContext(Dispatchers.IO) {
-        val map = cache ?: run {
-            val all = get()
-            LinkedHashMap(all.associateBy { it.id })
+
+    override suspend fun query(embedding: FloatArray, topK: Int, threshold: Float): List<Embedding> {
+
+        cachedIds = null // clear on new search
+
+        val storedEmbeddings = get()
+
+        if (storedEmbeddings.isEmpty()) return emptyList()
+
+        val similarities = getSimilarities(embedding, storedEmbeddings.map { it.embeddings })
+        val resultIndices = getTopN(similarities, topK, threshold)
+
+        if (resultIndices.isEmpty()) return emptyList()
+
+        val idsToCache = mutableListOf<Long>()
+        val results = resultIndices.map{idx ->
+            idsToCache.add( storedEmbeddings[idx].id)
+            storedEmbeddings[idx]
         }
-        map.get(id)
+        cachedIds = idsToCache
+        return results
     }
 
-    override fun clear(){
-        cache = null
+    suspend fun query(start: Int, end: Int): List<Embedding> {
+        val ids = cachedIds ?: return emptyList()
+        val s = start.coerceAtLeast(0)
+        val e = end.coerceAtMost(ids.size)
+        if (s >= e) return emptyList()
+
+        val batch = get(ids.subList(s, e))
+        return batch
     }
 
 }

core/src/main/java/com/fpf/smartscansdk/core/embeddings/FileEmbeddingStore.kt

@@ -3,10 +3,15 @@ package com.fpf.smartscansdk.core.embeddings
 import com.fpf.smartscansdk.core.data.Embedding
 
 interface IEmbeddingStore {
-    val isCached: Boolean
     val exists: Boolean
     suspend fun add(newEmbeddings: List<Embedding>)
     suspend fun remove(ids: List<Long>)
     suspend fun get(): List<Embedding>
     fun clear()
+
+    suspend fun query(
+        embedding: FloatArray,
+        topK: Int,
+        threshold: Float
+    ): List<Embedding>
 }