From ac686d8849af845937e297d79f005007ce020f60 Mon Sep 17 00:00:00 2001 From: Simon Vergauwen Date: Mon, 22 May 2023 14:17:02 +0200 Subject: [PATCH 1/2] Fix scoping contextScope --- .../com/xebia/functional/xef/auto/AI.kt | 42 +++++++++++-------- .../xef/vectorstores/LocalVectorStore.kt | 12 +++--- 2 files changed, 29 insertions(+), 25 deletions(-) diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/auto/AI.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/auto/AI.kt index 5d584a19f..aa02beefc 100644 --- a/core/src/commonMain/kotlin/com/xebia/functional/xef/auto/AI.kt +++ b/core/src/commonMain/kotlin/com/xebia/functional/xef/auto/AI.kt @@ -5,7 +5,6 @@ import arrow.core.left import arrow.core.raise.Raise import arrow.core.raise.recover import arrow.core.right -import arrow.fx.coroutines.Resource import arrow.fx.coroutines.ResourceScope import arrow.fx.coroutines.resourceScope import com.xebia.functional.xef.AIError @@ -26,7 +25,8 @@ import kotlinx.serialization.DeserializationStrategy import kotlinx.serialization.descriptors.SerialDescriptor import kotlinx.serialization.json.JsonObject -@DslMarker annotation class AiDsl +@DslMarker +annotation class AiDsl data class SerializationConfig( val jsonSchema: JsonObject, @@ -147,38 +147,44 @@ class AIScope( * } * ``` */ - @AiDsl @JvmName("invokeAI") suspend operator fun AI.invoke(): A = invoke(this@AIScope) + @AiDsl + @JvmName("invokeAI") + suspend operator fun AI.invoke(): A = invoke(this@AIScope) @AiDsl suspend fun extendContext(vararg docs: String) { context.addTexts(docs.toList()) } + /** + * Creates a new scoped [VectorStore] using [store], which is scoped to the [block] lambda. + * The [block] also runs on a _nested_ [resourceScope], + * meaning that all additional resources created within [block] will be finalized after [block] finishes. + */ @AiDsl suspend fun contextScope( - store: suspend (Embeddings) -> Resource, + store: suspend ResourceScope.(Embeddings) -> VectorStore, block: AI - ): A { - val newStore = store(embeddings).bind() - return AIScope( - openAIClient, - CombinedVectorStore(newStore, context), - embeddings, - logger, - this, - this - ) - .block() + ): A = resourceScope { + val newStore = store(this@AIScope.embeddings) + AIScope( + this@AIScope.openAIClient, + CombinedVectorStore(newStore, this@AIScope.context), + this@AIScope.embeddings, + this@AIScope.logger, + this, + this@AIScope + ).block() } @AiDsl suspend fun contextScope(block: AI): A = contextScope(LocalVectorStoreBuilder, block) - /** Add new [docs] to the [context], and then executes the [scope]. */ + /** Add new [docs] to the [context], and then executes the [block]. */ @AiDsl - suspend fun contextScope(docs: List, scope: suspend AIScope.() -> A): A = + suspend fun contextScope(docs: List, block: AI): A = contextScope { extendContext(*docs.toTypedArray()) - scope(this) + block(this) } } diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/vectorstores/LocalVectorStore.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/vectorstores/LocalVectorStore.kt index 70f35338c..6004a8b41 100644 --- a/core/src/commonMain/kotlin/com/xebia/functional/xef/vectorstores/LocalVectorStore.kt +++ b/core/src/commonMain/kotlin/com/xebia/functional/xef/vectorstores/LocalVectorStore.kt @@ -1,7 +1,6 @@ package com.xebia.functional.xef.vectorstores -import arrow.fx.coroutines.Resource -import arrow.fx.coroutines.continuations.resource +import arrow.fx.coroutines.ResourceScope import arrow.fx.stm.TMap import arrow.fx.stm.TVar import arrow.fx.stm.atomically @@ -11,9 +10,8 @@ import com.xebia.functional.xef.llm.openai.EmbeddingModel import com.xebia.functional.xef.llm.openai.RequestConfig import kotlin.math.sqrt -val LocalVectorStoreBuilder: (Embeddings) -> Resource = { e -> - resource { LocalVectorStore(e) } -} +val LocalVectorStoreBuilder: suspend ResourceScope.(Embeddings) -> LocalVectorStore = + { e -> LocalVectorStore(e) } class LocalVectorStore private constructor( @@ -48,8 +46,8 @@ private constructor( override suspend fun similaritySearchByVector(embedding: Embedding, limit: Int): List = atomically { - documents.read().mapNotNull { doc -> precomputedEmbeddings[doc]?.let { doc to it } } - } + documents.read().mapNotNull { doc -> precomputedEmbeddings[doc]?.let { doc to it } } + } .map { (doc, embedding) -> doc to embedding.cosineSimilarity(embedding) } .sortedByDescending { (_, similarity) -> similarity } .take(limit) From 11214a422a88c97694fc97998e9812acdc052646 Mon Sep 17 00:00:00 2001 From: Simon Vergauwen Date: Mon, 22 May 2023 14:26:49 +0200 Subject: [PATCH 2/2] Update contextScope resourcing --- .../com/xebia/functional/xef/auto/AI.kt | 37 +++++++++---------- .../xef/vectorstores/LocalVectorStore.kt | 9 +++-- 2 files changed, 22 insertions(+), 24 deletions(-) diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/auto/AI.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/auto/AI.kt index aa02beefc..2c66a6819 100644 --- a/core/src/commonMain/kotlin/com/xebia/functional/xef/auto/AI.kt +++ b/core/src/commonMain/kotlin/com/xebia/functional/xef/auto/AI.kt @@ -25,8 +25,7 @@ import kotlinx.serialization.DeserializationStrategy import kotlinx.serialization.descriptors.SerialDescriptor import kotlinx.serialization.json.JsonObject -@DslMarker -annotation class AiDsl +@DslMarker annotation class AiDsl data class SerializationConfig( val jsonSchema: JsonObject, @@ -147,9 +146,7 @@ class AIScope( * } * ``` */ - @AiDsl - @JvmName("invokeAI") - suspend operator fun AI.invoke(): A = invoke(this@AIScope) + @AiDsl @JvmName("invokeAI") suspend operator fun AI.invoke(): A = invoke(this@AIScope) @AiDsl suspend fun extendContext(vararg docs: String) { @@ -157,9 +154,9 @@ class AIScope( } /** - * Creates a new scoped [VectorStore] using [store], which is scoped to the [block] lambda. - * The [block] also runs on a _nested_ [resourceScope], - * meaning that all additional resources created within [block] will be finalized after [block] finishes. + * Creates a new scoped [VectorStore] using [store], which is scoped to the [block] lambda. The + * [block] also runs on a _nested_ [resourceScope], meaning that all additional resources created + * within [block] will be finalized after [block] finishes. */ @AiDsl suspend fun contextScope( @@ -168,13 +165,14 @@ class AIScope( ): A = resourceScope { val newStore = store(this@AIScope.embeddings) AIScope( - this@AIScope.openAIClient, - CombinedVectorStore(newStore, this@AIScope.context), - this@AIScope.embeddings, - this@AIScope.logger, - this, - this@AIScope - ).block() + this@AIScope.openAIClient, + CombinedVectorStore(newStore, this@AIScope.context), + this@AIScope.embeddings, + this@AIScope.logger, + this, + this@AIScope + ) + .block() } @AiDsl @@ -182,9 +180,8 @@ class AIScope( /** Add new [docs] to the [context], and then executes the [block]. */ @AiDsl - suspend fun contextScope(docs: List, block: AI): A = - contextScope { - extendContext(*docs.toTypedArray()) - block(this) - } + suspend fun contextScope(docs: List, block: AI): A = contextScope { + extendContext(*docs.toTypedArray()) + block(this) + } } diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/vectorstores/LocalVectorStore.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/vectorstores/LocalVectorStore.kt index 6004a8b41..2a01c56ef 100644 --- a/core/src/commonMain/kotlin/com/xebia/functional/xef/vectorstores/LocalVectorStore.kt +++ b/core/src/commonMain/kotlin/com/xebia/functional/xef/vectorstores/LocalVectorStore.kt @@ -10,8 +10,9 @@ import com.xebia.functional.xef.llm.openai.EmbeddingModel import com.xebia.functional.xef.llm.openai.RequestConfig import kotlin.math.sqrt -val LocalVectorStoreBuilder: suspend ResourceScope.(Embeddings) -> LocalVectorStore = - { e -> LocalVectorStore(e) } +val LocalVectorStoreBuilder: suspend ResourceScope.(Embeddings) -> LocalVectorStore = { e -> + LocalVectorStore(e) +} class LocalVectorStore private constructor( @@ -46,8 +47,8 @@ private constructor( override suspend fun similaritySearchByVector(embedding: Embedding, limit: Int): List = atomically { - documents.read().mapNotNull { doc -> precomputedEmbeddings[doc]?.let { doc to it } } - } + documents.read().mapNotNull { doc -> precomputedEmbeddings[doc]?.let { doc to it } } + } .map { (doc, embedding) -> doc to embedding.cosineSimilarity(embedding) } .sortedByDescending { (_, similarity) -> similarity } .take(limit)