Merge pull request #1575 from WebFuzzing/llm-mock

basic mock scaffolding for LLM use

Author: arcuri82

core/src/main/kotlin/org/evomaster/core/BaseModule.kt

@@ -3,6 +3,8 @@ package org.evomaster.core
 import com.google.inject.AbstractModule
 import com.google.inject.Provides
 import com.google.inject.Singleton
+import org.evomaster.core.llm.service.DictionaryService
+import org.evomaster.core.llm.service.LlmService
 import org.evomaster.core.output.service.PartialOracles
 import org.evomaster.core.search.service.mutator.genemutation.ArchiveImpactSelector
 import org.evomaster.core.search.service.*
@@ -95,6 +97,12 @@ class BaseModule(val args: Array<String>, val noTests: Boolean = false) : Abstra
         bind(WarningsAggregator::class.java)
             .asEagerSingleton()
 
+        bind(LlmService::class.java)
+            .asEagerSingleton()
+
+        bind(DictionaryService::class.java)
+            .asEagerSingleton()
+
         //no longer needed if TestSuiteWriter is moved out?
 //        if(noTests){
 //            bind(TestCaseWriter::class.java)

core/src/main/kotlin/org/evomaster/core/llm/LlmProvider.kt

@@ -5,5 +5,6 @@ enum class LlmProvider {
     OPENAI,
     DEEPSEEK,   // OpenAI-compatible
     OLLAMA,
-    AZURE_OPENAI
+    AZURE_OPENAI,
+    MOCK
 }

core/src/main/kotlin/org/evomaster/core/llm/LlmSupport.kt

@@ -14,6 +14,8 @@ import dev.langchain4j.model.ollama.OllamaChatModel
 import dev.langchain4j.model.ollama.OllamaStreamingChatModel
 import dev.langchain4j.model.openai.OpenAiChatModel
 import dev.langchain4j.model.openai.OpenAiStreamingChatModel
+import org.evomaster.core.llm.mock.MockChatModel
+import org.evomaster.core.llm.mock.MockStreamingChatModel
 import java.util.concurrent.CompletableFuture
 
 
@@ -35,6 +37,13 @@ object LlmSupport {
 
     }
 
+    fun chat(model: ChatModel, userMessage: String): String {
+
+        val u = UserMessage.userMessage(userMessage)
+
+        return model.chat(listOf(u)).aiMessage().text()
+    }
+
     fun chat(model: ChatModel, systemMessage: String, userMessage: String): String {
 
         val s = SystemMessage.systemMessage(systemMessage)
@@ -76,6 +85,8 @@ object LlmSupport {
          */
 
         return when (provider) {
+            LlmProvider.MOCK -> MockChatModel()
+
             LlmProvider.OPENAI -> OpenAiChatModel.builder()
                 .apiKey(apiKey ?: error("API key required for OpenAI"))
                 .baseUrl(url ?: "https://api.openai.com/v1")
@@ -138,6 +149,8 @@ object LlmSupport {
          */
 
         return when (provider) {
+            LlmProvider.MOCK -> MockStreamingChatModel()
+
             LlmProvider.OPENAI -> OpenAiStreamingChatModel.builder()
                 .apiKey(apiKey ?: error("API key required for OpenAI"))
                 .baseUrl(url ?: "https://api.openai.com/v1")

core/src/main/kotlin/org/evomaster/core/llm/mock/LlmMockResponse.kt

@@ -0,0 +1,14 @@
+package org.evomaster.core.llm.mock
+
+class LlmMockResponse(
+
+    /**
+     * The response producer that will be returned by the LLM, given an input text from user
+     */
+    val responseProducer: (String) -> String,
+
+    /**
+     * Lambda that, given as input a request from the client, decides whether this response should be returned
+     */
+    val matcher: (String) -> Boolean
+)

core/src/main/kotlin/org/evomaster/core/llm/mock/MockChatModel.kt

@@ -0,0 +1,66 @@
+package org.evomaster.core.llm.mock
+
+import dev.langchain4j.data.message.AiMessage
+import dev.langchain4j.model.chat.ChatModel
+import dev.langchain4j.model.chat.ChatRequestOptions
+import dev.langchain4j.model.chat.request.ChatRequest
+import dev.langchain4j.model.chat.response.ChatResponse
+
+class MockChatModel : ChatModel {
+
+    companion object {
+
+        /*
+            WARNING
+            Global mutable state... but only used for testing.
+            But out of the box would prevent parallel tests in same process, which we don't anyway,
+            so shouldn't be a big deal
+         */
+        private val responses : MutableList<LlmMockResponse> = mutableListOf()
+
+        /**
+         * Clear the cache of mocked responses.
+         * Important to avoid dependencies among tests.
+         * Recall this is using mutable static state.
+         */
+        fun reset(){
+            responses.clear()
+        }
+
+        /**
+         * Specify the [response] that will be returned by the LLM,
+         * based on [matcher] that, given as input a request from the client,
+         * decides whether this response should be returned
+         */
+        fun mockResponse(response: String, matcher: (String) -> Boolean){
+            responses.add(LlmMockResponse({req -> response}, matcher))
+        }
+
+        /**
+         * Specify a producer of responses if the text of the query from user does match the matcher.
+         * Input for both lambdas is such text.
+         * This is useful when the response for same query might be different
+         */
+        fun mockResponse(responseProducer: (String) -> String, matcher: (String) -> Boolean){
+            responses.add(LlmMockResponse(responseProducer, matcher))
+        }
+    }
+
+    /**
+     * Main method that needs to be mocked.
+     * All other methods seem to call this one.
+     */
+    override fun chat(chatRequest: ChatRequest, options: ChatRequestOptions) : ChatResponse {
+
+        val text = chatRequest.messages().joinToString { it.toString() }
+
+        val res = responses.find { mock -> mock.matcher(text) }?.responseProducer(text)
+            /*
+               TODO we could also have a boolen parameter to decide what to do in case of no match
+                (eg response as now or throw exception), to make the mock more configurable
+             */
+            ?: "Hei! You forgot to setup the mock for me"
+
+        return ChatResponse.builder().aiMessage(AiMessage.from(res)).build()
+    }
+}

core/src/main/kotlin/org/evomaster/core/llm/mock/MockStreamingChatModel.kt

@@ -0,0 +1,19 @@
+package org.evomaster.core.llm.mock
+
+import dev.langchain4j.model.chat.ChatRequestOptions
+import dev.langchain4j.model.chat.StreamingChatModel
+import dev.langchain4j.model.chat.request.ChatRequest
+import dev.langchain4j.model.chat.response.StreamingChatResponseHandler
+
+/**
+ * WARNING:
+ * we wrote the code for it, but in the end we did not use StreamingChatModel.
+ * TODO If one day we do, then we need to update this class if we want to have tests using it.
+ */
+class MockStreamingChatModel : StreamingChatModel {
+
+    override fun chat(request: ChatRequest?, options: ChatRequestOptions?, handler: StreamingChatResponseHandler?){
+
+        throw NotImplementedError()
+    }
+}

core/src/main/kotlin/org/evomaster/core/problem/enterprise/service/EnterpriseModule.kt

@@ -14,12 +14,6 @@ abstract class EnterpriseModule : AbstractModule() {
         bind(WFCReportWriter::class.java)
             .asEagerSingleton()
 
-        bind(LlmService::class.java)
-            .asEagerSingleton()
-
-        bind(DictionaryService::class.java)
-            .asEagerSingleton()
-
         bind(SSRFAnalyser::class.java)
             .asEagerSingleton()
 

core/src/test/kotlin/org/evomaster/core/llm/LlmSupportMockTest.kt

@@ -0,0 +1,48 @@
+package org.evomaster.core.llm
+
+import org.evomaster.core.llm.mock.MockChatModel
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+
+class LlmSupportMockTest {
+
+    @Test
+    fun testBasePrompt(){
+        val model = LlmSupport.createModel(LlmProvider.MOCK)
+
+        val firstPrompt = """
+            Is A  the first letter in the English alphabet? Answer this question with either a "YES" or a "NO".
+            Do not add anything else in your response.
+            """.trimIndent()
+
+        val secondPrompt = """
+            Is Sweden the capital of Norway? 
+        """.trimIndent()
+
+        MockChatModel.reset()
+        MockChatModel.mockResponse("YES"){ it.contains("English")}
+        MockChatModel.mockResponse("NO"){ it.contains("Norway")}
+
+
+        assertEquals("YES", LlmSupport.chat(model, firstPrompt))
+        assertEquals("NO", LlmSupport.chat(model, secondPrompt))
+    }
+
+    @Test
+    fun testResponseProducer(){
+        val model = LlmSupport.createModel(LlmProvider.MOCK)
+
+        val prompt = "Give me a random value"
+        val x = arrayOf(0)
+
+        MockChatModel.reset()
+        MockChatModel.mockResponse({_ -> "${x[0]++}"}){it.contains("random")}
+
+        val n = 10
+        repeat(n){
+            LlmSupport.chat(model, prompt)
+        }
+
+        assertEquals(n, x[0])
+    }
+}

docs/options.md

@@ -309,7 +309,7 @@ There are 3 types of options:
 |`llm`| __Boolean__. Enable the use of LLMs. *Default value*: `false`.|
 |`llmApiKey`| __String__. API KEY needed to authenticated toward the chosen LLM provider. *Depends on*: `llm=true`. *Default value*: `null`.|
 |`llmName`| __String__. LLM name. If not specified, default will be based on the LLM provider. *Depends on*: `llm=true`. *Default value*: `null`.|
-|`llmProvider`| __Enum__. Provider for the LLM. This could be a local one (e.g., run through Ollama), or a remote one like OpenAI. *Depends on*: `llm=true`. *Valid values*: `OPENAI, DEEPSEEK, OLLAMA, AZURE_OPENAI`. *Default value*: `OLLAMA`.|
+|`llmProvider`| __Enum__. Provider for the LLM. This could be a local one (e.g., run through Ollama), or a remote one like OpenAI. *Depends on*: `llm=true`. *Valid values*: `OPENAI, DEEPSEEK, OLLAMA, AZURE_OPENAI, MOCK`. *Default value*: `OLLAMA`.|
 |`llmTemperature`| __Double__. Temperature parameter for LLM. *Constraints*: `min=0.0, max=2.0`. *Depends on*: `llm=true`. *Default value*: `0.3`.|
 |`llmThreads`| __Int__. The number of threads to use when making calls towards an LLM, in configured. If connecting to Ollama, this value is ignored, and only 1 thread is used. *Depends on*: `llm=true`. *Default value*: `4`.|
 |`llmTimeoutSeconds`| __Long__. How long to wait for LLM's responses. *Constraints*: `min=0.0`. *Depends on*: `llm=true`. *Default value*: `60`.|