fetching new data if missing in Dictionary

Author: arcuri82

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

@@ -3015,6 +3015,12 @@ class EMConfig {
     @Cfg("Enable the use of LLMs.")
     var llm = false
 
+    @Experimental
+    @DependsOnTrueFor("llm")
+    @Cfg("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.")
+    var llmThreads = 4
+
     @Experimental
     @DependsOnTrueFor("llm")
     @Cfg("LLM external service URL. If not specified, default will be based on the LLM provider.")

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

@@ -13,6 +13,7 @@ import org.evomaster.core.AnsiColor.Companion.inYellow
 import org.evomaster.core.DocumentationLinks.EM_DOCKER_LINK
 import org.evomaster.core.DocumentationLinks.EM_ISSUES_LINK
 import org.evomaster.core.config.ConfigProblemException
+import org.evomaster.core.llm.service.LlmService
 import org.evomaster.core.logging.LoggingUtil
 import org.evomaster.core.output.TestSuiteCode
 import org.evomaster.core.output.TestSuiteSplitter
@@ -902,10 +903,17 @@ class Main {
                                 snapshotTimestamp: String ->
                 writeTestsAsSnapshots(injector, solution, controllerInfo, snapshotTimestamp)
             }.also {
+                /*
+                    TODO should have a better way to specify that some services need to be shutdown
+                 */
                 if (config.isEnabledHarvestingActualResponse()) {
                     val hp = injector.getInstance(HarvestActualHttpWsResponseHandler::class.java)
                     hp.shutdown()
                 }
+                if(config.llm){
+                    val llm = injector.getInstance(LlmService::class.java)
+                    llm.shutdown()
+                }
             }
         }
 

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

@@ -17,11 +17,6 @@ import kotlin.io.path.Path
 
 object DictionaryCreator {
 
-    private data class FieldInfo(
-        val name: String,
-        val description: String?
-    )
-
     @JvmStatic
     fun main(args: Array<String>) {
 

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

@@ -0,0 +1,6 @@
+package org.evomaster.core.llm
+
+data class FieldInfo(
+    val name: String,
+    val description: String?
+)

core/src/main/kotlin/org/evomaster/core/llm/service/DictionaryService.kt

@@ -4,10 +4,11 @@ import com.fasterxml.jackson.core.type.TypeReference
 import com.fasterxml.jackson.databind.ObjectMapper
 import org.evomaster.core.EMConfig
 import org.evomaster.core.llm.DictionarySearchResult
+import org.evomaster.core.llm.FieldInfo
+import org.evomaster.core.logging.LoggingUtil
 import org.evomaster.core.search.service.DataPool
 import org.slf4j.Logger
 import org.slf4j.LoggerFactory
-import javax.annotation.PostConstruct
 import javax.inject.Inject
 
 
@@ -128,20 +129,33 @@ class DictionaryService {
     private lateinit var dataPool : DataPool
 
 
-    fun updatePoolFromDictionary(names: Collection<String>) {
+    /**
+     * Update the data pool with all the info from dictionary based on the input field's names.
+     * If some names are not in the dictionary, if LLM is enabled then it will be queried,
+     * including using the names' descriptions (if any is provided)
+     */
+    fun updatePoolFromDictionary(fields: Collection<FieldInfo>) {
         if(!isActive()){
             throw IllegalStateException("Dictionary service is not active")
         }
 
         //TODO check how expensive this is... put on thread if too expensive
 
-        val result = searchForNames(names)
+        val result = searchForNames(fields.map { it.name })
+
         result.data.entries.forEach { entry ->
             entry.value.forEach { exm ->  dataPool.addValue(entry.key, exm) }
         }
 
         if(config.llm && result.missing.isNotEmpty()) {
-            //TODO defenetively this needs to be on thread
+            LoggingUtil.getInfoLogger().info("Going to make ${result.missing.size} calls to LLM to obtain new input data" +
+                    " based on field names and descriptions")
+
+            val toInfer = fields.filter{ result.missing.contains(it.name)}
+
+            llmService.askForNewExamples(toInfer){ name, examples ->
+                examples.forEach { ex -> dataPool.addValue(name, ex) }
+            }
         }
     }
 

core/src/main/kotlin/org/evomaster/core/llm/service/LlmService.kt

@@ -1,10 +1,17 @@
 package org.evomaster.core.llm.service
 
+import com.fasterxml.jackson.core.type.TypeReference
+import com.fasterxml.jackson.databind.ObjectMapper
 import dev.langchain4j.model.chat.ChatModel
 import dev.langchain4j.model.chat.StreamingChatModel
 import org.evomaster.core.EMConfig
 import org.evomaster.core.config.ConfigProblemException
+import org.evomaster.core.llm.FieldInfo
+import org.evomaster.core.llm.LlmProvider
 import org.evomaster.core.llm.LlmSupport
+import org.evomaster.core.llm.Prompts
+import java.util.concurrent.ExecutorService
+import java.util.concurrent.Executors
 import java.util.concurrent.Future
 import javax.annotation.PostConstruct
 import javax.inject.Inject
@@ -18,6 +25,8 @@ class LlmService {
 
     private lateinit var syncModel: ChatModel
 
+    private lateinit var executor: ExecutorService
+
     @PostConstruct
     private fun initService() {
 
@@ -47,6 +56,13 @@ class LlmService {
         }catch (e: Exception){
             throw ConfigProblemException("Failed to connect and initialize the chosen LLM: ${e.message}")
         }
+
+        val n = if(config.llmProvider == LlmProvider.OLLAMA) 1 else config.llmThreads
+        executor = Executors.newFixedThreadPool(n)
+    }
+
+    fun shutdown() {
+        executor.shutdown()
     }
 
     private fun checkUsingLLM(){
@@ -64,4 +80,36 @@ class LlmService {
     fun chat(userMessage: String) : String{
         return LlmSupport.chat(syncModel, "", userMessage)
     }
+
+    fun askForNewExamples(fields: Collection<FieldInfo>, callback: (name: String, examples: Collection<String>) -> Unit){
+
+        fields.toList()
+            .sortedBy { it.name }
+            .forEach { entry ->
+                executor.submit { askForExamplesTask(entry.name, entry.description, callback) }
+            }
+        //here we do not wait... updates are done asynchronously
+    }
+
+    private fun askForExamplesTask(
+        name: String,
+        description: String?,
+        callback: (name: String, examples: Collection<String>) -> Unit
+    ){
+        val mapper = ObjectMapper()
+        val prompt = Prompts.getPromptForNameDescription(name, description)
+        var result = LlmSupport.chat(syncModel, prompt.first, prompt.second)
+        val list = try {
+            mapper.readValue(result, object : TypeReference<List<String>>() {})
+        } catch (e: Exception) {
+            try {
+                val failed = Prompts.getPromptForFailedName(e.toString())
+                result = LlmSupport.chat(syncModel, failed.first, failed.second)
+                mapper.readValue(result, object : TypeReference<List<String>>() {})
+            } catch (e: Exception) {
+                throw e
+            }
+        }
+        callback(name, list)
+    }
 }

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

@@ -1,6 +1,7 @@
 package org.evomaster.core.problem.enterprise.service
 
 import com.google.inject.AbstractModule
+import org.evomaster.core.llm.service.DictionaryService
 import org.evomaster.core.llm.service.LlmService
 import org.evomaster.core.problem.security.service.SSRFAnalyser
 import org.evomaster.core.problem.security.service.HttpCallbackVerifier
@@ -16,6 +17,9 @@ abstract class EnterpriseModule : AbstractModule() {
         bind(LlmService::class.java)
             .asEagerSingleton()
 
+        bind(DictionaryService::class.java)
+            .asEagerSingleton()
+
         bind(SSRFAnalyser::class.java)
             .asEagerSingleton()