Skip to content

Commit 6e21da1

Browse files
Copilotphrocker
andauthored
Add memoryLookup field to LLMResponse for prescriptive memory retrieval (#111)
* Initial plan * Add memoryLookup field to LLMResponse with automated lookup execution Co-authored-by: phrocker <[email protected]> --------- Co-authored-by: copilot-swe-agent[bot] <[email protected]> Co-authored-by: phrocker <[email protected]>
1 parent a617100 commit 6e21da1

File tree

5 files changed

+208
-2
lines changed

5 files changed

+208
-2
lines changed

enterprise-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/ChatAgent.java

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -215,6 +215,48 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
215215

216216
var arguments = response.getArguments();
217217
if (null != response) {
218+
// Handle memory lookup if specified
219+
if (response.getMemoryLookup() != null && !response.getMemoryLookup().isEmpty()) {
220+
log.info("Memory lookup requested: {}", response.getMemoryLookup());
221+
try {
222+
// Set up memory lookup arguments
223+
Map<String, Object> memoryArgs = new HashMap<>();
224+
memoryArgs.put("query", response.getMemoryLookup());
225+
226+
// Execute memory lookup
227+
var memoryResponse = verbRegistry.execute(
228+
agentExecution,
229+
agentExecutionContext,
230+
lastVerbResponse,
231+
"lookupAgentMemory",
232+
memoryArgs
233+
);
234+
235+
// Add memory results to context for LLM
236+
if (memoryResponse != null && memoryResponse.getReturnName() != null) {
237+
var memoryResult = agentExecutionContext.getAgentShortTermMemory()
238+
.get(memoryResponse.getReturnName());
239+
if (memoryResult != null) {
240+
agentExecutionContext.addMessages(
241+
Message.builder()
242+
.role("system")
243+
.content("Memory lookup results: " + memoryResult.toString())
244+
.build()
245+
);
246+
log.info("Memory lookup completed, results added to context");
247+
}
248+
}
249+
} catch (Exception e) {
250+
log.warn("Memory lookup failed: {}", e.getMessage());
251+
agentExecutionContext.addMessages(
252+
Message.builder()
253+
.role("system")
254+
.content("Memory lookup failed: " + e.getMessage())
255+
.build()
256+
);
257+
}
258+
}
259+
218260
if (response.getNextOperation() != null && !response.getNextOperation().isEmpty()) {
219261
var executionResponse = verbRegistry.execute(
220262
agentExecution,

enterprise-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/ChatVerbs.java

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,13 @@ public LLMResponse interpretUserData(
104104
". Please summarize prior terminal " +
105105
"sessions, using " +
106106
"terminal output if needed " +
107-
"for clarity of the next LLM request and for the user. Ensure your all future responses meets this " +
107+
"for clarity of the next LLM request and for the user. " +
108+
"IMPORTANT: Use the 'memoryLookup' field to search for information that is NOT in the current context. " +
109+
"For example, if the user mentions something from a previous conversation that is not in the recent messages, " +
110+
"populate memoryLookup with a query string to retrieve that information. " +
111+
"Leave memoryLookup empty if the information is already available in the context or not needed. " +
112+
"The memory lookup will be executed BEFORE the nextOperation. " +
113+
"Ensure your all future responses meets this " +
108114
"json format (LLMResponse format): " + terminalResponse).build());
109115
messages.add(Message.builder().role("user").content(userMessage.getContentAsString()).build());
110116
LLMRequest chatRequest = LLMRequest.builder().model("gpt-4o-mini").messages(messages).build();
@@ -222,7 +228,13 @@ public LLMResponse promptAgent(
222228
"." +
223229
". " +
224230
". Please summarize messaging, using " +
225-
"for clarity of the next LLM request and for the user. Ensure your all future responses meets this " +
231+
"for clarity of the next LLM request and for the user. " +
232+
"IMPORTANT: Use the 'memoryLookup' field to search for information that is NOT in the current context. " +
233+
"For example, if you need information from previous sessions that is not in recent messages, " +
234+
"populate memoryLookup with a query string to retrieve that information. " +
235+
"Leave memoryLookup empty if the information is already available in the context or not needed. " +
236+
"The memory lookup will be executed BEFORE the nextOperation. " +
237+
"Ensure your all future responses meets this " +
226238
"json format (LLMResponse format): " + terminalResponse).build());
227239
LLMRequest chatRequest = LLMRequest.builder().model("gpt-4o-mini").messages(messages).build();
228240
var resp = llmService.askQuestion(execution, chatRequest);

enterprise-agent/src/main/java/io/sentrius/agent/analysis/model/LLMResponse.java

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
public class LLMResponse {
2121
String previousOperation;
2222
String nextOperation;
23+
String memoryLookup;
2324
String summaryForLLM;
2425
String responseForUser;
2526
@Builder.Default

enterprise-agent/src/main/resources/terminal-helper.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
{
22
"previousOperation": "<previousOperation>",
33
"nextOperation": "<nextOperation>",
4+
"memoryLookup": "<Query string to lookup in memory before nextOperation. Only populate this if the information you need is NOT already in the current context. Leave empty if the information is already available in context or if no memory lookup is needed>",
45
"arguments": {
56
"argumentname": "<argumentvalue for nextOperation>",
67
"argumentname2": "<argumentvalue>"
Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
package io.sentrius.sentrius.analysis.model;
2+
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.junit.jupiter.api.Assertions.assertNotNull;
5+
import static org.junit.jupiter.api.Assertions.assertNull;
6+
7+
import com.fasterxml.jackson.core.JsonProcessingException;
8+
import io.sentrius.agent.analysis.model.LLMResponse;
9+
import io.sentrius.sso.core.utils.JsonUtil;
10+
import java.util.HashMap;
11+
import java.util.Map;
12+
import org.junit.jupiter.api.Test;
13+
14+
class LLMResponseTest {
15+
16+
@Test
17+
void testLLMResponseWithMemoryLookup() {
18+
// Given
19+
Map<String, Object> args = new HashMap<>();
20+
args.put("arg1", "value1");
21+
22+
// When
23+
LLMResponse response = LLMResponse.builder()
24+
.previousOperation("getPlan")
25+
.nextOperation("executeOperation")
26+
.memoryLookup("search for user preferences")
27+
.summaryForLLM("Summary of operations")
28+
.responseForUser("Task completed successfully")
29+
.arguments(args)
30+
.build();
31+
32+
// Then
33+
assertEquals("getPlan", response.getPreviousOperation());
34+
assertEquals("executeOperation", response.getNextOperation());
35+
assertEquals("search for user preferences", response.getMemoryLookup());
36+
assertEquals("Summary of operations", response.getSummaryForLLM());
37+
assertEquals("Task completed successfully", response.getResponseForUser());
38+
assertNotNull(response.getArguments());
39+
assertEquals("value1", response.getArguments().get("arg1"));
40+
}
41+
42+
@Test
43+
void testLLMResponseWithoutMemoryLookup() {
44+
// Given & When
45+
LLMResponse response = LLMResponse.builder()
46+
.previousOperation("getPlan")
47+
.nextOperation("executeOperation")
48+
.summaryForLLM("Summary of operations")
49+
.responseForUser("Task completed successfully")
50+
.build();
51+
52+
// Then
53+
assertEquals("getPlan", response.getPreviousOperation());
54+
assertEquals("executeOperation", response.getNextOperation());
55+
assertNull(response.getMemoryLookup());
56+
assertEquals("Summary of operations", response.getSummaryForLLM());
57+
assertEquals("Task completed successfully", response.getResponseForUser());
58+
}
59+
60+
@Test
61+
void testLLMResponseJsonSerialization() throws JsonProcessingException {
62+
// Given
63+
Map<String, Object> args = new HashMap<>();
64+
args.put("query", "test query");
65+
66+
LLMResponse response = LLMResponse.builder()
67+
.previousOperation("lookupData")
68+
.nextOperation("processData")
69+
.memoryLookup("find previous user interactions")
70+
.summaryForLLM("Looking up and processing data")
71+
.responseForUser("Processing your request")
72+
.arguments(args)
73+
.build();
74+
75+
// When
76+
String json = JsonUtil.MAPPER.writeValueAsString(response);
77+
LLMResponse deserialized = JsonUtil.MAPPER.readValue(json, LLMResponse.class);
78+
79+
// Then
80+
assertNotNull(deserialized);
81+
assertEquals(response.getPreviousOperation(), deserialized.getPreviousOperation());
82+
assertEquals(response.getNextOperation(), deserialized.getNextOperation());
83+
assertEquals(response.getMemoryLookup(), deserialized.getMemoryLookup());
84+
assertEquals(response.getSummaryForLLM(), deserialized.getSummaryForLLM());
85+
assertEquals(response.getResponseForUser(), deserialized.getResponseForUser());
86+
}
87+
88+
@Test
89+
void testLLMResponseJsonDeserializationWithMemoryLookup() throws JsonProcessingException {
90+
// Given
91+
String json = "{"
92+
+ "\"previousOperation\":\"getStatus\","
93+
+ "\"nextOperation\":\"updateStatus\","
94+
+ "\"memoryLookup\":\"search for system configuration\","
95+
+ "\"summaryForLLM\":\"Status check and update\","
96+
+ "\"responseForUser\":\"Status updated\","
97+
+ "\"arguments\":{\"status\":\"active\"}"
98+
+ "}";
99+
100+
// When
101+
LLMResponse response = JsonUtil.MAPPER.readValue(json, LLMResponse.class);
102+
103+
// Then
104+
assertNotNull(response);
105+
assertEquals("getStatus", response.getPreviousOperation());
106+
assertEquals("updateStatus", response.getNextOperation());
107+
assertEquals("search for system configuration", response.getMemoryLookup());
108+
assertEquals("Status check and update", response.getSummaryForLLM());
109+
assertEquals("Status updated", response.getResponseForUser());
110+
assertNotNull(response.getArguments());
111+
}
112+
113+
@Test
114+
void testLLMResponseJsonDeserializationWithoutMemoryLookup() throws JsonProcessingException {
115+
// Given
116+
String json = "{"
117+
+ "\"previousOperation\":\"getStatus\","
118+
+ "\"nextOperation\":\"updateStatus\","
119+
+ "\"summaryForLLM\":\"Status check and update\","
120+
+ "\"responseForUser\":\"Status updated\""
121+
+ "}";
122+
123+
// When
124+
LLMResponse response = JsonUtil.MAPPER.readValue(json, LLMResponse.class);
125+
126+
// Then
127+
assertNotNull(response);
128+
assertEquals("getStatus", response.getPreviousOperation());
129+
assertEquals("updateStatus", response.getNextOperation());
130+
assertNull(response.getMemoryLookup());
131+
assertEquals("Status check and update", response.getSummaryForLLM());
132+
assertEquals("Status updated", response.getResponseForUser());
133+
}
134+
135+
@Test
136+
void testLLMResponseWithEmptyMemoryLookup() {
137+
// Given & When
138+
LLMResponse response = LLMResponse.builder()
139+
.previousOperation("operation1")
140+
.nextOperation("operation2")
141+
.memoryLookup("")
142+
.summaryForLLM("Summary")
143+
.responseForUser("Response")
144+
.build();
145+
146+
// Then
147+
assertEquals("", response.getMemoryLookup());
148+
assertNotNull(response.getMemoryLookup());
149+
}
150+
}

0 commit comments

Comments
 (0)