Lucene as vector store

core/build.gradle.kts

@@ -81,6 +81,7 @@ kotlin {
         implementation(libs.jtokk.it)
         implementation(libs.scrape.it)
         implementation(libs.rss.reader)
+        implementation(libs.apache.lucene)
       }
     }
 

core/src/jvmMain/kotlin/com/xebia/functional/vectorstores/Lucene.kt

@@ -0,0 +1,79 @@
+package com.xebia.functional.vectorstores
+
+import com.xebia.functional.embeddings.Embedding
+import com.xebia.functional.embeddings.Embeddings
+import com.xebia.functional.llm.openai.EmbeddingModel
+import com.xebia.functional.llm.openai.RequestConfig
+import org.apache.lucene.document.Document
+import org.apache.lucene.document.Field
+import org.apache.lucene.document.KnnFloatVectorField
+import org.apache.lucene.document.TextField
+import org.apache.lucene.index.*
+import org.apache.lucene.search.FuzzyQuery
+import org.apache.lucene.search.IndexSearcher
+import org.apache.lucene.search.KnnFloatVectorQuery
+import org.apache.lucene.search.Query
+import org.apache.lucene.store.Directory
+import org.apache.lucene.store.MMapDirectory
+import java.nio.file.Path
+
+open class Lucene(
+    private val writer: IndexWriter,
+    private val searcher: IndexSearcher,
+    private val embeddings: Embeddings,
+    private val similarity: VectorSimilarityFunction = VectorSimilarityFunction.EUCLIDEAN
+): VectorStore, AutoCloseable {
+
+    constructor(
+        writer: IndexWriter,
+        embeddings: Embeddings,
+        similarity: VectorSimilarityFunction = VectorSimilarityFunction.EUCLIDEAN
+    ): this(writer, IndexSearcher(DirectoryReader.open(writer)), embeddings, similarity)
+
+    private val requestConfig =
+        RequestConfig(EmbeddingModel.TextEmbeddingAda002, RequestConfig.Companion.User("user"))
+
+    override suspend fun addTexts(texts: List<String>) = texts.forEach {
+        val embedding = embeddings.embedQuery(it, requestConfig)
+        val doc = Document().apply {
+            add(TextField("contents",  it, Field.Store.YES))
+            add(KnnFloatVectorField("embedding", embedding.toFloatArray(), similarity))
+        }
+        writer.addDocument(doc)
+    }
+
+    override suspend fun similaritySearch(query: String, limit: Int): List<String> =
+        search(FuzzyQuery(Term("contents", query)), limit)
+
+    override suspend fun similaritySearchByVector(embedding: Embedding, limit: Int): List<String> =
+        search(KnnFloatVectorQuery("embedding", embedding.data.toFloatArray(), limit), limit)
+
+    private fun search(q: Query, limit: Int): List<String> = searcher.search(q, limit).scoreDocs.map {
+        searcher.storedFields().document(it.doc).get("contents")
+    }
+
+    override fun close() {
+        writer.close()
+    }
+}
+
+class DirectoryLucene(
+    private val directory: Directory,
+    writerConfig: IndexWriterConfig = IndexWriterConfig(),
+    embeddings: Embeddings,
+    similarity: VectorSimilarityFunction = VectorSimilarityFunction.EUCLIDEAN
+): Lucene(IndexWriter(directory, writerConfig), embeddings, similarity) {
+    override fun close() {
+        super.close()
+        directory.close()
+    }
+}
+
+fun InMemoryLucene(
+    path: Path,
+    writerConfig: IndexWriterConfig = IndexWriterConfig(),
+    embeddings: Embeddings,
+    similarity: VectorSimilarityFunction = VectorSimilarityFunction.EUCLIDEAN
+): DirectoryLucene = DirectoryLucene(MMapDirectory(path), writerConfig, embeddings, similarity)
+
+fun List<Embedding>.toFloatArray(): FloatArray = flatMap { it.data }.toFloatArray()

gradle/libs.versions.toml

@@ -33,6 +33,7 @@ munitCatsEffect = "1.0.5"
 scrapeit = "1.1.5"
 jtokkit = "0.4.0"
 rssreader = "3.4.4"
+lucene = "9.6.0"
 
 [libraries]
 arrow-core = { module = "io.arrow-kt:arrow-core", version.ref = "arrow" }
@@ -89,6 +90,7 @@ scrape-it-browser-fetcher = { module = "it.skrape:skrapeit-browser-fetcher", ver
 scrape-it-async-fetcher = { module = "it.skrape:skrapeit-asyn-fetcher", version.ref = "scrapeit" }
 jtokk-it = { module = "com.knuddels:jtokkit", version.ref = "jtokkit" }
 rss-reader = { module = "com.apptasticsoftware:rssreader", version.ref = "rssreader" }
+apache-lucene = { module = "org.apache.lucene:lucene-core", version.ref = "lucene" }
 
 [bundles]
 arrow = [