From 31a2d894ac4316f8c8b03ae2c0aa3a88c798b22d Mon Sep 17 00:00:00 2001 From: raulraja Date: Thu, 20 Jun 2024 10:45:43 +0200 Subject: [PATCH 1/2] Add Document structure with source to VectorStore --- .../xef/conversation/Conversation.kt | 12 --------- .../xef/store/CombinedVectorStore.kt | 7 +++-- .../functional/xef/store/LocalVectorStore.kt | 24 ++++++++++------- .../xebia/functional/xef/store/VectorStore.kt | 27 ++++++++++++++----- .../xef/vectorstore/VectorStoreExample.kt | 18 +++++++++++++ .../xef/store/PostgreSQLVectorStore.kt | 15 +++++++---- .../src/test/kotlin/xef/PGVectorStoreSpec.kt | 14 +++++----- 7 files changed, 75 insertions(+), 42 deletions(-) create mode 100644 examples/src/main/kotlin/com/xebia/functional/xef/vectorstore/VectorStoreExample.kt diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/conversation/Conversation.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/conversation/Conversation.kt index 509354857..7841a31c9 100644 --- a/core/src/commonMain/kotlin/com/xebia/functional/xef/conversation/Conversation.kt +++ b/core/src/commonMain/kotlin/com/xebia/functional/xef/conversation/Conversation.kt @@ -16,18 +16,6 @@ constructor( val conversationId: ConversationId? = ConversationId(UUID.generateUUID().toString()) ) { - @AiDsl - @JvmSynthetic - suspend fun addContext(vararg docs: String) { - store.addTexts(docs.toList()) - } - - @AiDsl - @JvmSynthetic - suspend fun addContext(docs: Iterable): Unit { - store.addTexts(docs.toList()) - } - companion object { @JvmSynthetic diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/store/CombinedVectorStore.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/store/CombinedVectorStore.kt index 4e10fb1de..4499a20d6 100644 --- a/core/src/commonMain/kotlin/com/xebia/functional/xef/store/CombinedVectorStore.kt +++ b/core/src/commonMain/kotlin/com/xebia/functional/xef/store/CombinedVectorStore.kt @@ -27,7 +27,7 @@ class CombinedVectorStore(private val top: VectorStore, private val bottom: Vect .reversed() } - override suspend fun similaritySearch(query: String, limit: Int): List { + override suspend fun similaritySearch(query: String, limit: Int): List { val topResults = top.similaritySearch(query, limit) return when { topResults.size >= limit -> topResults @@ -35,7 +35,10 @@ class CombinedVectorStore(private val top: VectorStore, private val bottom: Vect } } - override suspend fun similaritySearchByVector(embedding: Embedding, limit: Int): List { + override suspend fun similaritySearchByVector( + embedding: Embedding, + limit: Int + ): List { val topResults = top.similaritySearchByVector(embedding, limit) return when { topResults.size >= limit -> topResults diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/store/LocalVectorStore.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/store/LocalVectorStore.kt index 67908be24..d3799d40b 100644 --- a/core/src/commonMain/kotlin/com/xebia/functional/xef/store/LocalVectorStore.kt +++ b/core/src/commonMain/kotlin/com/xebia/functional/xef/store/LocalVectorStore.kt @@ -15,7 +15,7 @@ import kotlin.math.sqrt private data class State( val orderedMemories: Map>, - val documents: List, + val documents: List, val precomputedEmbeddings: Map ) { companion object { @@ -75,26 +75,30 @@ private constructor( .reversed() } - override suspend fun addTexts(texts: List) { + override suspend fun addDocuments(texts: List) { + val docsAsJson = texts.map { it.content } val embeddingsList = - embeddings.embedDocuments(texts, embeddingRequestModel = embeddingRequestModel) + embeddings.embedDocuments(docsAsJson, embeddingRequestModel = embeddingRequestModel) state.getAndUpdate { prevState -> - val newEmbeddings = prevState.precomputedEmbeddings + texts.zip(embeddingsList) + val newEmbeddings = prevState.precomputedEmbeddings + docsAsJson.zip(embeddingsList) State(prevState.orderedMemories, prevState.documents + texts, newEmbeddings) } } - override suspend fun similaritySearch(query: String, limit: Int): List { + override suspend fun similaritySearch(query: String, limit: Int): List { val queryEmbedding = embeddings.embedQuery(query, embeddingRequestModel = embeddingRequestModel).firstOrNull() return queryEmbedding?.let { similaritySearchByVector(it, limit) }.orEmpty() } - override suspend fun similaritySearchByVector(embedding: Embedding, limit: Int): List { + override suspend fun similaritySearchByVector( + embedding: Embedding, + limit: Int + ): List { val state0 = state.get() return state0.documents .asSequence() - .mapNotNull { doc -> state0.precomputedEmbeddings[doc]?.let { doc to it } } + .mapNotNull { doc -> state0.precomputedEmbeddings[doc.content]?.let { doc to it } } .map { (doc, e) -> doc to embedding.cosineSimilarity(e) } .sortedByDescending { (_, similarity) -> similarity } .take(limit) @@ -103,9 +107,9 @@ private constructor( } private fun Embedding.cosineSimilarity(other: Embedding): Double { - val dotProduct = this.embedding.zip(other.embedding).sumOf { (a, b) -> (a * b).toDouble() } - val magnitudeA = sqrt(this.embedding.sumOf { (it * it).toDouble() }) - val magnitudeB = sqrt(other.embedding.sumOf { (it * it).toDouble() }) + val dotProduct = this.embedding.zip(other.embedding).sumOf { (a, b) -> (a * b) } + val magnitudeA = sqrt(this.embedding.sumOf { (it * it) }) + val magnitudeB = sqrt(other.embedding.sumOf { (it * it) }) return dotProduct / (magnitudeA * magnitudeB) } } diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/store/VectorStore.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/store/VectorStore.kt index 65d13dce8..5585ec549 100644 --- a/core/src/commonMain/kotlin/com/xebia/functional/xef/store/VectorStore.kt +++ b/core/src/commonMain/kotlin/com/xebia/functional/xef/store/VectorStore.kt @@ -3,10 +3,22 @@ package com.xebia.functional.xef.store import arrow.atomic.AtomicInt import com.xebia.functional.openai.generated.model.CreateChatCompletionRequestModel import com.xebia.functional.openai.generated.model.Embedding +import com.xebia.functional.xef.Config import kotlin.jvm.JvmStatic +import kotlinx.serialization.Serializable interface VectorStore { + @Serializable + data class Document(val content: String, val source: String) { + fun toJson(): String = Config.DEFAULT.json.encodeToString(serializer(), this) + + companion object { + fun fromJson(json: String): Document = + Config.DEFAULT.json.decodeFromString(serializer(), json) + } + } + val indexValue: AtomicInt fun incrementIndexAndGet(): Int = indexValue.addAndGet(1) @@ -27,9 +39,9 @@ interface VectorStore { * @param texts list of text to add to the vector store * @return a list of IDs from adding the texts to the vector store */ - suspend fun addTexts(texts: List) + suspend fun addDocuments(texts: List) - suspend fun addText(texts: String) = addTexts(listOf(texts)) + suspend fun addDocument(texts: Document) = addDocuments(listOf(texts)) /** * Return the docs most similar to the query @@ -38,7 +50,7 @@ interface VectorStore { * @param limit number of documents to return * @return a list of Documents most similar to query */ - suspend fun similaritySearch(query: String, limit: Int): List + suspend fun similaritySearch(query: String, limit: Int): List /** * Return the docs most similar to the embedding @@ -47,7 +59,7 @@ interface VectorStore { * @param limit number of documents to return * @return list of Documents most similar to the embedding */ - suspend fun similaritySearchByVector(embedding: Embedding, limit: Int): List + suspend fun similaritySearchByVector(embedding: Embedding, limit: Int): List companion object { @JvmStatic @@ -65,14 +77,15 @@ interface VectorStore { limitTokens: Int ): List = emptyList() - override suspend fun addTexts(texts: List) {} + override suspend fun addDocuments(texts: List) {} - override suspend fun similaritySearch(query: String, limit: Int): List = emptyList() + override suspend fun similaritySearch(query: String, limit: Int): List = + emptyList() override suspend fun similaritySearchByVector( embedding: Embedding, limit: Int - ): List = emptyList() + ): List = emptyList() } } } diff --git a/examples/src/main/kotlin/com/xebia/functional/xef/vectorstore/VectorStoreExample.kt b/examples/src/main/kotlin/com/xebia/functional/xef/vectorstore/VectorStoreExample.kt new file mode 100644 index 000000000..10cfb9552 --- /dev/null +++ b/examples/src/main/kotlin/com/xebia/functional/xef/vectorstore/VectorStoreExample.kt @@ -0,0 +1,18 @@ +package com.xebia.functional.xef.vectorstore + +import com.xebia.functional.xef.OpenAI +import com.xebia.functional.xef.store.LocalVectorStore +import com.xebia.functional.xef.store.VectorStore.Document + +suspend fun main() { + val embeddings = OpenAI().embeddings + val vectorStore = LocalVectorStore(embeddings) + val helloDoc = Document("Hello, how are you?", "source1") + val unrelatedDoc = Document("Unrelated text", "source2") + vectorStore.addDocuments(listOf(helloDoc, unrelatedDoc)) + val maybeHelloDoc = vectorStore.similaritySearch("Hello", 1).first() + assert(maybeHelloDoc == helloDoc) { "Expected $helloDoc but got $maybeHelloDoc" } + val maybeUnrelatedDoc = vectorStore.similaritySearch("Unrelated", 1).first() + assert(maybeUnrelatedDoc == unrelatedDoc) { "Expected $unrelatedDoc but got $maybeUnrelatedDoc" } + println("All expected documents found!") +} diff --git a/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/PostgreSQLVectorStore.kt b/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/PostgreSQLVectorStore.kt index cd32df067..6ae46dc4e 100644 --- a/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/PostgreSQLVectorStore.kt +++ b/integrations/postgresql/src/main/kotlin/com/xebia/functional/xef/store/PostgreSQLVectorStore.kt @@ -84,9 +84,10 @@ class PGVectorStore( } } - override suspend fun addTexts(texts: List): Unit = + override suspend fun addDocuments(texts: List): Unit = dataSource.connection { - val embeddings = embeddings.embedDocuments(texts, chunkSize, embeddingRequestModel) + val docsContent = texts.map { it.content } + val embeddings = embeddings.embedDocuments(docsContent, chunkSize, embeddingRequestModel) val collection = getCollection(collectionName) texts.zip(embeddings) { text, embedding -> val uuid = UUID.generateUUID() @@ -94,12 +95,12 @@ class PGVectorStore( bind(uuid.toString()) bind(collection.uuid.toString()) bind(embedding.embedding.toString()) - bind(text) + bind(text.toJson()) } } } - override suspend fun similaritySearch(query: String, limit: Int): List = + override suspend fun similaritySearch(query: String, limit: Int): List = dataSource.connection { val collection = getCollection(collectionName) @@ -123,10 +124,12 @@ class PGVectorStore( } ) { string() + }.map { json -> + VectorStore.Document.fromJson(json) } } - override suspend fun similaritySearchByVector(embedding: Embedding, limit: Int): List = + override suspend fun similaritySearchByVector(embedding: Embedding, limit: Int): List = dataSource.connection { val collection = getCollection(collectionName) queryAsList( @@ -138,6 +141,8 @@ class PGVectorStore( } ) { string() + }.map { json -> + VectorStore.Document.fromJson(json) } } diff --git a/integrations/postgresql/src/test/kotlin/xef/PGVectorStoreSpec.kt b/integrations/postgresql/src/test/kotlin/xef/PGVectorStoreSpec.kt index c7e5ab83b..2ee1dbfd0 100644 --- a/integrations/postgresql/src/test/kotlin/xef/PGVectorStoreSpec.kt +++ b/integrations/postgresql/src/test/kotlin/xef/PGVectorStoreSpec.kt @@ -4,6 +4,7 @@ import com.xebia.functional.openai.generated.model.CreateChatCompletionRequestMo import com.xebia.functional.openai.generated.model.CreateEmbeddingRequestModel import com.xebia.functional.openai.generated.model.Embedding import com.xebia.functional.xef.store.PGVectorStore +import com.xebia.functional.xef.store.VectorStore import com.xebia.functional.xef.store.migrations.runDatabaseMigrations import com.xebia.functional.xef.store.postgresql.PGDistanceStrategy import com.zaxxer.hikari.HikariConfig @@ -17,7 +18,6 @@ import io.kotest.matchers.shouldBe import org.junit.jupiter.api.assertThrows import org.testcontainers.containers.PostgreSQLContainer import org.testcontainers.utility.DockerImageName -import kotlin.coroutines.coroutineContext val postgres: PostgreSQLContainer = PostgreSQLContainer( @@ -67,10 +67,12 @@ class PGVectorStoreSpec : postgresVector.createCollection() } + val docs = listOf(VectorStore.Document(content = "foo", source = "tests"), VectorStore.Document(content = "bar", source = "tests")) + "initialDbSetup should configure the DB properly" { pg().initialDbSetup() } "addTexts should fail with a CollectionNotFoundError if collection isn't present in the DB" { - assertThrows { pg().addTexts(listOf("foo", "bar")) }.message shouldBe + assertThrows { pg().addDocuments(docs) }.message shouldBe "Collection 'test_collection' not found" } @@ -82,13 +84,13 @@ class PGVectorStoreSpec : "createCollection should create collection" { pg().createCollection() } "addTexts should not fail now that we created the collection" { - pg().addTexts(listOf("foo", "bar")) + pg().addDocuments(docs) } "similaritySearchByVector should return both documents" { - pg().addTexts(listOf("bar", "foo")) + pg().addDocuments(docs.reversed()) pg().similaritySearchByVector(Embedding(0, listOf(4.0, 5.0, 6.0), Embedding.Object.embedding), 2) shouldBe - listOf("bar", "foo") + docs.reversed() } "similaritySearch should return 2 documents" { @@ -104,7 +106,7 @@ class PGVectorStoreSpec : pg().similaritySearchByVector( Embedding(0, listOf(1.0, 2.0, 3.0), Embedding.Object.embedding), 1 - ) shouldBe listOf("foo") + ) shouldBe listOf(docs[0]) } "the added memories sorted by index should be obtained in the same order" { From cc0f725063fb2afe642b0b50f9dc59c5ec75b0ac Mon Sep 17 00:00:00 2001 From: raulraja Date: Thu, 20 Jun 2024 10:56:25 +0200 Subject: [PATCH 2/2] clean up --- .../com/xebia/functional/xef/store/LocalVectorStore.kt | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/store/LocalVectorStore.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/store/LocalVectorStore.kt index d3799d40b..fb4c98b33 100644 --- a/core/src/commonMain/kotlin/com/xebia/functional/xef/store/LocalVectorStore.kt +++ b/core/src/commonMain/kotlin/com/xebia/functional/xef/store/LocalVectorStore.kt @@ -76,11 +76,11 @@ private constructor( } override suspend fun addDocuments(texts: List) { - val docsAsJson = texts.map { it.content } + val docsContent = texts.map { it.content } val embeddingsList = - embeddings.embedDocuments(docsAsJson, embeddingRequestModel = embeddingRequestModel) + embeddings.embedDocuments(docsContent, embeddingRequestModel = embeddingRequestModel) state.getAndUpdate { prevState -> - val newEmbeddings = prevState.precomputedEmbeddings + docsAsJson.zip(embeddingsList) + val newEmbeddings = prevState.precomputedEmbeddings + docsContent.zip(embeddingsList) State(prevState.orderedMemories, prevState.documents + texts, newEmbeddings) } }