[WIP] Add asynchronous AI service infrastructure for Kotlin (#128)

# [WIP] Add asynchronous AI service infrastructure for Kotlin

Implemented `AsyncAiServices` and supporting components to enable
asynchronous AI service handling in Kotlin. Added tests, a service
factory, and _InvocationHandler_ to integrate and validate the
functionality.

Author: kpavlov

Makefile

@@ -8,7 +8,7 @@ apidocs:
 
 lint:prepare
 		ktlint && \
-		mvn spotless:check
+		mvn spotless:check detekt:check
 
 # https://docs.openrewrite.org/recipes/maven/bestpractices
 format:prepare

langchain4j-kotlin/src/main/kotlin/dev/langchain4j/service/DefaultAiServicesOpener.kt

@@ -0,0 +1,52 @@
+package dev.langchain4j.service
+
+import me.kpavlov.langchain4j.kotlin.ChatMemoryId
+import java.lang.reflect.Method
+import java.util.Optional
+
+/**
+ * This is a hack to access package-private methods in [DefaultAiServices].
+ * It is not supposed to be used directly.
+ */
+internal object DefaultAiServicesOpener {
+    /**
+     * This class is used to open package-private methods in [DefaultAiServices].
+     * It is not supposed to be used directly.
+     */
+    @Suppress("UNCHECKED_CAST")
+    internal fun findMemoryId(
+        method: Method,
+        args: Array<Any?>,
+    ): Optional<ChatMemoryId> {
+        val findMemoryId =
+            DefaultAiServices::class.java.getDeclaredMethod(
+                "findMemoryId",
+                Method::class.java,
+                Array<Any?>::class.java,
+            )
+        findMemoryId.isAccessible = true
+        @Suppress("UNCHECKED_CAST")
+        return findMemoryId.invoke(null, method, args) as Optional<Any>
+    }
+
+    @Suppress("TooGenericExceptionCaught")
+    internal fun validateParameters(method: Method) {
+        val validateParameters =
+            DefaultAiServices::class.java.getDeclaredMethod(
+                "validateParameters",
+                Method::class.java,
+            )
+        validateParameters.isAccessible = true
+        try {
+            validateParameters.invoke(null, method)
+        } catch (e: Exception) {
+            if (e.cause is dev.langchain4j.service.IllegalConfigurationException) {
+                val illegalConfigurationException =
+                    e.cause as dev.langchain4j.service.IllegalConfigurationException
+                illegalConfigurationException.printStackTrace()
+                return
+            }
+            throw e
+        }
+    }
+}

langchain4j-kotlin/src/main/kotlin/me/kpavlov/langchain4j/kotlin/rag/RetrievalAugmentorExtensions.kt

@@ -0,0 +1,10 @@
+package me.kpavlov.langchain4j.kotlin.rag
+
+import dev.langchain4j.rag.AugmentationRequest
+import dev.langchain4j.rag.AugmentationResult
+import dev.langchain4j.rag.RetrievalAugmentor
+import kotlinx.coroutines.coroutineScope
+
+public suspend fun RetrievalAugmentor.augmentAsync(
+    request: AugmentationRequest,
+): AugmentationResult = coroutineScope { this@augmentAsync.augment(request) }

langchain4j-kotlin/src/main/kotlin/me/kpavlov/langchain4j/kotlin/service/AiServicesExtensions.kt

@@ -0,0 +1,13 @@
+package me.kpavlov.langchain4j.kotlin.service
+
+import dev.langchain4j.service.AiServiceContext
+import dev.langchain4j.service.AiServices
+import dev.langchain4j.spi.services.AiServicesFactory
+
+/**
+ * Creates an [AiServices] instance using the provided [AiServicesFactory].
+ */
+public fun <T> createAiService(
+    serviceClass: Class<T>,
+    factory: AiServicesFactory,
+): AiServices<T> = AiServiceContext(serviceClass).let { context -> factory.create<T>(context) }

langchain4j-kotlin/src/main/kotlin/me/kpavlov/langchain4j/kotlin/service/AsyncAiServices.kt

@@ -0,0 +1,79 @@
+package me.kpavlov.langchain4j.kotlin.service
+
+import dev.langchain4j.service.AiServiceContext
+import dev.langchain4j.service.AiServices
+import dev.langchain4j.service.ChatMemoryAccess
+import dev.langchain4j.service.IllegalConfigurationException.illegalConfiguration
+import dev.langchain4j.service.MemoryId
+import dev.langchain4j.service.Moderate
+import dev.langchain4j.service.Result
+import dev.langchain4j.service.TypeUtils
+import dev.langchain4j.service.output.ServiceOutputParser
+import dev.langchain4j.spi.ServiceHelper
+import dev.langchain4j.spi.services.TokenStreamAdapter
+
+public class AsyncAiServices<T : Any>(
+    context: AiServiceContext,
+) : AiServices<T>(context) {
+    private val serviceOutputParser = ServiceOutputParser()
+    private val tokenStreamAdapters =
+        ServiceHelper.loadFactories<TokenStreamAdapter>(TokenStreamAdapter::class.java)
+
+    @Suppress("NestedBlockDepth")
+    override fun build(): T {
+        performBasicValidation()
+
+        if (!context.hasChatMemory() &&
+            ChatMemoryAccess::class.java.isAssignableFrom(context.aiServiceClass)
+        ) {
+            throw illegalConfiguration(
+                "In order to have a service implementing ChatMemoryAccess, " +
+                    "please configure the ChatMemoryProvider on the '%s'.",
+                context.aiServiceClass.name,
+            )
+        }
+
+        for (method in context.aiServiceClass.methods) {
+            if (method.isAnnotationPresent(Moderate::class.java) &&
+                context.moderationModel == null
+            ) {
+                throw illegalConfiguration(
+                    "The @Moderate annotation is present, but the moderationModel is not set up. " +
+                        "Please ensure a valid moderationModel is configured " +
+                        "before using the @Moderate annotation.",
+                )
+            }
+            if (method.returnType in
+                arrayOf(
+                    // supported collection types
+                    Result::class.java,
+                    MutableList::class.java,
+                    MutableSet::class.java,
+                )
+            ) {
+                TypeUtils.validateReturnTypesAreProperlyParametrized(
+                    method.name,
+                    method.genericReturnType,
+                )
+            }
+
+            if (!context.hasChatMemory()) {
+                for (parameter in method.parameters) {
+                    if (parameter.isAnnotationPresent(MemoryId::class.java)) {
+                        throw illegalConfiguration(
+                            "In order to use @MemoryId, please configure " +
+                                "ChatMemoryProvider on the '%s'.",
+                            context.aiServiceClass.name,
+                        )
+                    }
+                }
+            }
+        }
+
+        val handler = ServiceInvocationHandler<T>(context, serviceOutputParser, tokenStreamAdapters)
+        @Suppress("UNCHECKED_CAST", "unused")
+        return ReflectionHelper.createSuspendProxy(context.aiServiceClass) { method, args ->
+            return@createSuspendProxy handler.invoke(method, args)
+        } as T
+    }
+}

langchain4j-kotlin/src/main/kotlin/me/kpavlov/langchain4j/kotlin/service/AsyncAiServicesFactory.kt

@@ -0,0 +1,10 @@
+package me.kpavlov.langchain4j.kotlin.service
+
+import dev.langchain4j.service.AiServiceContext
+import dev.langchain4j.service.AiServices
+import dev.langchain4j.spi.services.AiServicesFactory
+
+public class AsyncAiServicesFactory : AiServicesFactory {
+    override fun <T : Any> create(context: AiServiceContext): AiServices<T> =
+        AsyncAiServices(context)
+}

langchain4j-kotlin/src/main/kotlin/me/kpavlov/langchain4j/kotlin/service/ReflectionHelper.kt

@@ -0,0 +1,116 @@
+package me.kpavlov.langchain4j.kotlin.service
+
+import kotlinx.coroutines.DelicateCoroutinesApi
+import kotlinx.coroutines.GlobalScope
+import kotlinx.coroutines.asCoroutineDispatcher
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
+import java.lang.reflect.InvocationHandler
+import java.lang.reflect.Method
+import java.lang.reflect.ParameterizedType
+import java.lang.reflect.Proxy
+import java.lang.reflect.Type
+import java.lang.reflect.WildcardType
+import java.util.concurrent.Executors
+import kotlin.coroutines.Continuation
+import kotlin.coroutines.intrinsics.COROUTINE_SUSPENDED
+import kotlin.coroutines.resume
+import kotlin.coroutines.resumeWithException
+
+@OptIn(DelicateCoroutinesApi::class)
+internal object ReflectionHelper {
+    private val vtDispatcher = Executors.newVirtualThreadPerTaskExecutor().asCoroutineDispatcher()
+
+    @Throws(kotlin.IllegalStateException::class)
+    private fun getReturnType(method: Method): Type {
+        val continuationParam = method.parameterTypes.findLast { it.kotlin is Continuation<*> }
+        if (continuationParam != null) {
+            @Suppress("UseCheckOrError")
+            return continuationParam.genericInterfaces[0]
+                ?: throw IllegalStateException(
+                    "Can't find generic interface of continuation parameter",
+                )
+        }
+        return method.getGenericReturnType()
+    }
+
+    fun getSuspendReturnType(method: Method): java.lang.reflect.Type {
+        val parameters = method.genericParameterTypes
+        if (parameters.isEmpty()) return getReturnType(method)
+        val lastParameter = parameters.last()
+        // Check if the last parameter is Continuation<T>
+        return (
+            if (lastParameter is ParameterizedType &&
+                (lastParameter.rawType as? Class<*>)?.name == Continuation::class.java.name
+            ) {
+                // T is the first (and only) type argument
+                val type = lastParameter.actualTypeArguments.first()
+                if (type is WildcardType) {
+                    type.lowerBounds.first()
+                } else {
+                    type
+                }
+            } else {
+                getReturnType(method) // Not a suspend function, or not detectably so
+            }
+        )
+    }
+
+    @Suppress("UNCHECKED_CAST")
+    fun <T : Any> createSuspendProxy(
+        iface: Class<T>,
+        handler: suspend (method: java.lang.reflect.Method, args: Array<Any?>) -> Any?,
+    ): T {
+        return Proxy.newProxyInstance(
+            iface.classLoader,
+            arrayOf(iface),
+            InvocationHandler { _, method, args ->
+                // If not a suspend method, optionally fall back
+                val cont =
+                    args.lastOrNull() as? Continuation<Any?>
+                        ?: return@InvocationHandler method.invoke(this, args)
+
+                // Remove Continuation for our handler
+                val argsForSuspend = args.dropLast(1).toTypedArray()
+
+                // Launch coroutine for the suspend implementation
+                // (here, for demonstration, using a helper)
+                // Use coroutine machinery to start the suspend block
+                // If using Kotlin 1.3+, this is the correct way
+                // Uses GlobalScope (be sure that's okay for your use-case!)
+                GlobalScope.launch(vtDispatcher) {
+                    @Suppress("TooGenericExceptionCaught")
+                    try {
+                        val result = handler(method, argsForSuspend)
+                        cont.resume(result)
+                    } catch (e: Throwable) {
+                        cont.resumeWithException(e)
+                    }
+                }
+                COROUTINE_SUSPENDED
+            },
+        ) as T
+    }
+
+    @FunctionalInterface
+    interface MyApi {
+        suspend fun greet(name: String): String
+    }
+
+    internal fun dropContinuationArg(args: Array<Any?>): Array<Any?> =
+        args
+            .dropLastWhile {
+                it is Continuation<*>
+            }.toTypedArray()
+}
+
+public fun main() {
+    val proxy =
+        ReflectionHelper.createSuspendProxy(ReflectionHelper.MyApi::class.java) { method, args ->
+            "${Thread.currentThread()}: Hello, $method(${args[0]} )"
+        }
+
+    runBlocking {
+        println(proxy.greet("world")) // Prints: Hello, world
+    }
+}