[LLM] Refactor LLMWithFeedbackCycle and RequestManager

.gitignore

@@ -8,6 +8,8 @@ lib/JUnitRunner.jar
 *.DS_Store
 .intellijPlatform/
 
+.kotlin/
+
 # Test projects
 src/test/resources/project/.idea/
 

build.gradle.kts

@@ -98,6 +98,8 @@ if (spaceCredentialsProvided()) {
 
     // add build of new source set as the part of UI testing
     tasks.prepareTestSandbox.configure {
+        duplicatesStrategy = DuplicatesStrategy.EXCLUDE
+
         dependsOn(hasGrazieAccess.jarTaskName)
         from(
             tasks
@@ -114,6 +116,8 @@ if (spaceCredentialsProvided()) {
     }
     // add build of new source set as the part of pluginBuild process
     tasks.prepareSandbox.configure {
+        duplicatesStrategy = DuplicatesStrategy.EXCLUDE
+
         dependsOn(hasGrazieAccess.jarTaskName)
         from(
             tasks
@@ -212,12 +216,13 @@ dependencies {
     // https://mvnrepository.com/artifact/net.jqwik/jqwik
     testImplementation("net.jqwik:jqwik:1.6.5")
 
+    // https://mvnrepository.com/artifact/org.jetbrains.kotlinx/kotlinx-coroutines-test
+    testImplementation("org.jetbrains.kotlinx:kotlinx-coroutines-test:1.10.1")
+
     // https://mvnrepository.com/artifact/com.github.javaparser/javaparser-symbol-solver-core
     implementation("com.github.javaparser:javaparser-symbol-solver-core:3.24.2")
     // https://mvnrepository.com/artifact/org.jetbrains.kotlin/kotlin-test
     implementation("org.jetbrains.kotlin:kotlin-test:1.8.0")
-
-    implementation("org.jetbrains.kotlinx:kotlinx-serialization-json:1.6.3")
 }
 
 // Configure Gradle IntelliJ Plugin - read more: // Configure Gradle IntelliJ Plugin - read more: https://github.com/JetBrains/gradle-intellij-plugin

core/build.gradle.kts

@@ -1,6 +1,7 @@
 plugins {
-    kotlin("jvm")
+    kotlin("jvm") version "2.1.0"
     `maven-publish`
+    kotlin("plugin.serialization") version "2.1.0"
 }
 
 group = "org.jetbrains.research"
@@ -18,6 +19,13 @@ dependencies {
     compileOnly(kotlin("stdlib"))
 
     implementation("io.github.oshai:kotlin-logging-jvm:6.0.3")
+    implementation("org.jetbrains.kotlinx:kotlinx-coroutines-core:1.10.1")
+    implementation("org.jetbrains.kotlinx:kotlinx-serialization-json:1.6.3")
+
+    val ktorVersion = "2.3.13"
+    implementation("io.ktor:ktor-client-core:$ktorVersion")
+    implementation("io.ktor:ktor-client-cio:$ktorVersion")
+    implementation("io.ktor:ktor-client-logging:$ktorVersion")
 }
 
 tasks.test {

core/src/main/kotlin/org/jetbrains/research/testspark/core/data/ChatMessage.kt

@@ -1,19 +1,20 @@
 package org.jetbrains.research.testspark.core.data
 
-open class ChatMessage protected constructor(
+data class ChatMessage(
     val role: ChatRole,
-    val content: String,
+    val contentBuilder: StringBuilder,
 ) {
     enum class ChatRole {
         User,
         Assistant,
     }
-}
 
-class ChatUserMessage(
-    content: String,
-) : ChatMessage(ChatRole.User, content)
+    val content: String
+        get() = contentBuilder.toString()
+
+    companion object {
+        fun createUserMessage(message: String) = ChatMessage(ChatRole.User, StringBuilder(message))
 
-class ChatAssistantMessage(
-    content: String,
-) : ChatMessage(ChatRole.Assistant, content)
+        fun createAssistantMessage(message: String) = ChatMessage(ChatRole.Assistant, StringBuilder(message))
+    }
+}

core/src/main/kotlin/org/jetbrains/research/testspark/core/error/LlmError.kt

@@ -13,12 +13,12 @@ sealed class LlmError(
 
     data object EmptyLlmResponse : LlmError()
 
-    data object TestSuiteParsingError : LlmError()
+    data class TestSuiteParsingError(
+        override val cause: Throwable?,
+    ) : LlmError(cause)
 
     data object NoCompilableTestCasesGenerated : LlmError()
 
-    data object FailedToSaveTestFiles : LlmError()
-
     data object CompilationError : LlmError()
 
     data object UnsetTokenError : LlmError()

core/src/main/kotlin/org/jetbrains/research/testspark/core/error/Result.kt

@@ -1,27 +1,27 @@
 package org.jetbrains.research.testspark.core.error
 
-sealed interface Result<out D, out E : TestSparkError> {
+sealed interface Result<out D> {
     data class Success<out D>(
         val data: D,
-    ) : Result<D, Nothing>
+    ) : Result<D>
 
-    data class Failure<out E : TestSparkError>(
-        val error: E,
-    ) : Result<Nothing, E>
+    data class Failure(
+        val error: TestSparkError,
+    ) : Result<Nothing>
 
     fun getDataOrNull(): D? = if (this is Success) data else null
 
     fun isSuccess(): Boolean = this is Success
 
     fun isFailure(): Boolean = this is Failure
 
-    fun <R> mapData(transform: (D) -> R): Result<R, E> =
+    fun <R> mapData(transform: (D) -> R): Result<R> =
         when (this) {
             is Success -> Success(transform(data))
             is Failure -> Failure(error)
         }
 
-    fun <R : TestSparkError> mapError(transform: (E) -> R): Result<D, R> =
+    fun <R : TestSparkError> mapError(transform: (TestSparkError) -> R): Result<D> =
         when (this) {
             is Success -> Success(data)
             is Failure -> Failure(transform(error))

core/src/main/kotlin/org/jetbrains/research/testspark/core/exception/CommonException.kt

@@ -11,5 +11,6 @@ sealed class CommonException(
     )
 
 class ProcessCancelledException(
-    module: TestSparkModule,
-) : CommonException(module)
+    module: TestSparkModule = TestSparkModule.Common,
+    cause: Throwable? = null,
+) : CommonException(module, cause)

core/src/main/kotlin/org/jetbrains/research/testspark/core/exception/CompilerException.kt

@@ -43,3 +43,5 @@ class ClassFileNotFoundException(
     val classFilePath: String,
     val filePath: String,
 ) : CompilerException()
+
+class TestSavingFailureException : CompilerException()

core/src/main/kotlin/org/jetbrains/research/testspark/core/generation/llm/ChatSessionManager.kt

@@ -0,0 +1,79 @@
+package org.jetbrains.research.testspark.core.generation.llm
+
+import io.github.oshai.kotlinlogging.KotlinLogging
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.onEach
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
+import org.jetbrains.research.testspark.core.data.ChatMessage
+import org.jetbrains.research.testspark.core.error.LlmError
+import org.jetbrains.research.testspark.core.error.Result
+import org.jetbrains.research.testspark.core.generation.llm.network.RequestManager
+import org.jetbrains.research.testspark.core.generation.llm.network.model.LlmParams
+
+class ChatSessionManager(
+    private val requestManager: RequestManager,
+    private val llmParams: LlmParams,
+) {
+    private val mutex = Mutex()
+    private val chatHistory = mutableListOf<ChatMessage>()
+    private val log = KotlinLogging.logger {}
+
+    suspend fun request(
+        prompt: String,
+        isUserFeedback: Boolean,
+    ): Flow<Result<String>> {
+        log.info { "Sending Request..." }
+
+        recordChatMessage(isUserFeedback, ChatMessage.Companion.createUserMessage(message = prompt))
+
+        val chatHistory =
+            if (isUserFeedback) {
+                chatHistory + ChatMessage.createUserMessage(prompt)
+            } else {
+                chatHistory
+            }
+
+        return requestManager
+            .sendRequest(
+                llmParams,
+                chatHistory,
+            ).map { result ->
+                val responseString = result.getDataOrNull()
+                if (responseString != null && responseString.isEmpty()) {
+                    Result.Failure(error = LlmError.EmptyLlmResponse)
+                } else {
+                    result
+                }
+            }.onEach { result ->
+                val rawText = result.getDataOrNull()
+                if (rawText.isNullOrEmpty().not()) {
+                    recordChatMessage(
+                        isUserFeedback,
+                        ChatMessage.Companion.createAssistantMessage(message = rawText),
+                    )
+                }
+            }
+    }
+
+    private suspend fun recordChatMessage(
+        isUserFeedback: Boolean,
+        message: ChatMessage,
+    ) {
+        if (isUserFeedback) return
+        mutex.withLock {
+            when (message.role) {
+                ChatMessage.ChatRole.User -> chatHistory.add(message)
+                ChatMessage.ChatRole.Assistant -> {
+                    val lastMessage = chatHistory.lastOrNull()
+                    if (lastMessage != null && lastMessage.role == ChatMessage.ChatRole.Assistant) {
+                        lastMessage.contentBuilder.append(message.content)
+                    } else {
+                        chatHistory.add(message)
+                    }
+                }
+            }
+        }
+    }
+}