Improve Lucene similarity search (#108)

Author: serras

gradle/libs.versions.toml

@@ -96,7 +96,8 @@ skrape = { module = "it.skrape:skrapeit", version.ref = "scrapeit" }
 skrape-browser-fetcher = { module = "it.skrape:skrapeit-browser-fetcher", version.ref = "scrapeit" }
 skrape-async-fetcher = { module = "it.skrape:skrapeit-asyn-fetcher", version.ref = "scrapeit" }
 rss-reader = { module = "com.apptasticsoftware:rssreader", version.ref = "rssreader" }
-apache-lucene = { module = "org.apache.lucene:lucene-core", version.ref = "lucene" }
+lucene-core = { module = "org.apache.lucene:lucene-core", version.ref = "lucene" }
+lucene-queries = { module = "org.apache.lucene:lucene-queries", version.ref = "lucene" }
 assertj = { module = "org.assertj:assertj-core", version.ref = "assertj" }
 apache-pdf-box = { module = "org.apache.pdfbox:pdfbox", version.ref = "pdfbox" }
 jdbc-mysql-connector = { module = "mysql:mysql-connector-java", version.ref = "mysql" }

integrations/lucene/build.gradle.kts

@@ -19,7 +19,8 @@ java {
 
 dependencies {
     implementation(projects.xefCore)
-    api(libs.apache.lucene)
+    api(libs.lucene.core)
+    api(libs.lucene.queries)
 }
 
 tasks.withType<AbstractPublishToMaven> {

integrations/lucene/src/main/kotlin/com/xebia/functional/xef/vectorstores/Lucene.kt

@@ -6,32 +6,22 @@ import com.xebia.functional.xef.embeddings.Embedding
 import com.xebia.functional.xef.embeddings.Embeddings
 import com.xebia.functional.xef.llm.openai.EmbeddingModel
 import com.xebia.functional.xef.llm.openai.RequestConfig
+import org.apache.lucene.analysis.standard.StandardAnalyzer
+import org.apache.lucene.document.*
 import java.nio.file.Path
-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.queries.mlt.MoreLikeThis
+import org.apache.lucene.search.*
 import org.apache.lucene.store.Directory
 import org.apache.lucene.store.MMapDirectory
+import java.io.StringReader
 
 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"))
 
@@ -48,16 +38,27 @@ open class Lucene(
     writer.commit()
   }
 
-  override suspend fun similaritySearch(query: String, limit: Int): List<String> =
-    search(FuzzyQuery(Term("contents", query)), limit)
+  override suspend fun similaritySearch(query: String, limit: Int): List<String> {
+    val reader = DirectoryReader.open(writer)
+    val mlt = MoreLikeThis(reader)
+    mlt.analyzer = StandardAnalyzer()
+    mlt.minTermFreq = 1
+    mlt.minDocFreq = 1
+    mlt.minWordLen = 3
+    val luceneQuery = mlt.like("contents", StringReader(query))
+    val searcher = IndexSearcher(reader)
+    return IndexSearcher(reader).search(luceneQuery, limit).extract(searcher)
+  }
 
   override suspend fun similaritySearchByVector(embedding: Embedding, limit: Int): List<String> {
     requireNotNull(embeddings) { "no embeddings were computed for this model" }
-    return search(KnnFloatVectorQuery("embedding", embedding.data.toFloatArray(), limit), limit)
+    val luceneQuery = KnnFloatVectorQuery("embedding", embedding.data.toFloatArray(), limit)
+    val searcher = IndexSearcher(DirectoryReader.open(writer))
+    return searcher.search(luceneQuery, limit).extract(searcher)
   }
 
-  private fun search(q: Query, limit: Int): List<String> =
-    searcher.search(q, limit).scoreDocs.map {
+  fun TopDocs.extract(searcher: IndexSearcher): List<String> =
+    scoreDocs.map {
       searcher.storedFields().document(it.doc).get("contents")
     }