1+ package org.evomaster.core.llm.service
2+
3+ import com.fasterxml.jackson.core.type.TypeReference
4+ import com.fasterxml.jackson.databind.ObjectMapper
5+ import org.evomaster.core.EMConfig
6+ import org.evomaster.core.llm.DictionarySearchResult
7+ import org.evomaster.core.llm.FieldInfo
8+ import org.evomaster.core.logging.LoggingUtil
9+ import org.evomaster.core.search.service.DataPool
10+ import org.slf4j.Logger
11+ import org.slf4j.LoggerFactory
12+ import javax.inject.Inject
13+
14+
15+ class DictionaryService {
16+
17+ companion object {
18+ private val log: Logger = LoggerFactory .getLogger(DictionaryService ::class .java)
19+ private const val location = " /llm_dictionary.jsonl"
20+
21+ /* *
22+ * WARNING: possibly quite expensive... should only be used for debugging/testing/statistics
23+ */
24+ fun loadAll () : Map <String , Set <String >> {
25+
26+ val data = mutableMapOf<String , Set <String >>()
27+ val mapper = ObjectMapper ()
28+
29+ val reader = DictionaryService ::class .java.getResourceAsStream(location)!! .bufferedReader()
30+
31+ reader.forEachLine { line ->
32+ addLineEntry(mapper, line, data)
33+ }
34+
35+ return data
36+ }
37+
38+
39+ /* *
40+ * Load the dictionary from file, and scan it for the given name entries.
41+ * Considering the large size of the dictionary, this functions tries to be efficient.
42+ * Still, it can be quite expensive, and should be called ideally only once.
43+ */
44+ fun searchForNames (names : Collection <String >): DictionarySearchResult {
45+
46+ if (names.isEmpty()) {
47+ throw IllegalArgumentException (" No name to find is specified"
48+ }
49+
50+ val found = mutableMapOf<String , Set <String >>()
51+ val missing = mutableSetOf<String >()
52+ val result = DictionarySearchResult (found, missing)
53+ val mapper = ObjectMapper ()
54+
55+ val reader = DictionaryService ::class .java.getResourceAsStream(location)!! .bufferedReader()
56+
57+ val toFind = names.toList().sorted()
58+ var index = 0
59+
60+ reader.lineSequence()
61+ .takeWhile { index < toFind.size }
62+ .forEach { line ->
63+
64+ // let's avoid parsing whole line with Jackson if we do not have a name-match
65+ val startQuote = line.indexOf(' "'
66+ val endQuote = line.indexOf(' "' + 1 )
67+ val x = line.substring(startQuote+ 1 , endQuote)
68+
69+ var handled = false
70+
71+ while (! handled && index < toFind.size) {
72+ val f = toFind[index]
73+
74+ if (f == x) {
75+ addLineEntry(mapper, line, found)
76+ index++
77+ handled = true
78+ } else if (x < f) {
79+ // nothing to do. f is not found, but could be found in the next lines
80+ // recall names and dictionary is alphabetically sorted
81+ handled = true
82+ } else {
83+ missing.add(f)
84+ index++
85+ }
86+ }
87+ }
88+
89+ // in case any asked name come alphabetically after the last element in the dictionary
90+ while (index < toFind.size) {
91+ missing.add(toFind[index])
92+ index++
93+ }
94+
95+ // would not hold if there are errors in parsing... but then if so we should remove those from dictionary!!!
96+ assert (names.size == result.data.size + result.missing.size)
97+ return result
98+ }
99+
100+ private fun addLineEntry (
101+ mapper : ObjectMapper ,
102+ line : String ,
103+ found : MutableMap <String , Set <String >>
104+ ) {
105+ val node = mapper.readTree(line)
106+ if (! node.isObject) {
107+ log.warn(" Not an object: $line "
108+ } else {
109+ node.fields().forEach { field ->
110+ if (! field.value.isArray) {
111+ log.warn(" Not containing an array: $line "
112+ } else {
113+ found[field.key] = mapper.convertValue(field.value, object : TypeReference <Set <String >>() {})
114+ }
115+ }
116+ }
117+ }
118+ }
119+
120+
121+
122+ @Inject
123+ private lateinit var config: EMConfig
124+
125+ @Inject
126+ private lateinit var llmService: LlmService
127+
128+ @Inject
129+ private lateinit var dataPool : DataPool
130+
131+
132+ /* *
133+ * Update the data pool with all the info from dictionary based on the input field's names.
134+ * If some names are not in the dictionary, if LLM is enabled then it will be queried,
135+ * including using the names' descriptions (if any is provided)
136+ */
137+ fun updatePoolFromDictionary (fields : Collection <FieldInfo >) {
138+ if (! isActive()){
139+ throw IllegalStateException (" Dictionary service is not active"
140+ }
141+
142+ // TODO check how expensive this is... put on thread if too expensive
143+
144+ val result = searchForNames(fields.map { it.name })
145+
146+ result.data.entries.forEach { entry ->
147+ entry.value.forEach { exm -> dataPool.addValue(entry.key, exm) }
148+ }
149+
150+ if (config.llm && result.missing.isNotEmpty()) {
151+ LoggingUtil .getInfoLogger().info(" Going to make ${result.missing.size} calls to LLM to obtain new input data" +
152+ " based on field names and descriptions"
153+
154+ val toInfer = fields.filter{ result.missing.contains(it.name)}
155+
156+ llmService.askForNewExamples(toInfer){ name, examples ->
157+ examples.forEach { ex -> dataPool.addValue(name, ex) }
158+ }
159+ }
160+ }
161+
162+
163+ fun isActive () : Boolean {
164+ return config.useDictionaryDataPool
165+ }
166+ }
0 commit comments