Skip to content

Commit dda969d

Browse files
phrockerphrocker
committed
minor update
1 parent 1f5871b commit dda969d

File tree

16 files changed

+422
-125
lines changed

16 files changed

+422
-125
lines changed

‎ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/VerbRegistry.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ public VerbResponse execute(AgentExecution agentExecution, VerbResponse priorRes
8585
if (null != args) {
8686
newArgs.putAll(args);
8787
}
88+
log.info("Executing verb: {}", verb);
8889
var returnType = verbs.get(verb).getReturnType();
8990
if (null != priorResponse ) {
9091
Class<? extends OutputInterpreterIfc> interpreter = priorResponse.getOutputInterpreter();
@@ -134,8 +135,12 @@ public VerbResponse execute(AgentExecution agentExecution, VerbResponse priorRes
134135
request, 60, TimeUnit.MINUTES);
135136
agentExecution.setZtatToken(token);
136137

138+
log.info("Re-attempting verb execution after Ztat token acquisition: {}", verb);
139+
137140
var interpretedInput = interpreterInstance.interpret(newArgs);
138141

142+
log.info("Re-attempting verb execution after Ztat token acquisition: {}", verb);
143+
139144
return VerbResponse.builder()
140145
.response(verbs.get(verb).isRequiresTokenManagement() ?
141146
method.invoke(instance,agentExecution, interpretedInput) :

‎ai-agent/src/main/java/io/sentrius/agent/analysis/agents/verbs/AgentVerbs.java‎

Lines changed: 0 additions & 109 deletions
Original file line numberDiff line numberDiff line change
@@ -141,116 +141,7 @@ public ArrayNode promptAgent(AgentExecution execution, Map<String, Object> args)
141141
return JsonUtil.MAPPER.createArrayNode();
142142
}
143143

144-
/**
145-
* Prompts the agent for workload based on the provided arguments.
146-
*
147-
* @return An `ArrayNode` containing the plan generated by the agent.
148-
* @throws ZtatException If there is an error during the operation.
149-
* @throws IOException If there is an error reading the configuration file.
150-
*/
151-
@Verb(name = "interpret_user_request", returnType = ArrayNode.class, description = "Queries the LLM using the " +
152-
"user input.",
153-
isAiCallable = false, requiresTokenManagement = true)
154-
public TerminalResponse interpretUserData(AgentExecution execution, @NonNull WebSocky socketConnection,
155-
@NonNull Message userMessage) throws ZtatException,
156-
IOException {
157-
158-
var lastMessage = socketConnection.getMessages().stream().reduce((prev, next) -> next).orElse(null);
159-
if (socketConnection.getMessages().isEmpty()) {
160-
161-
InputStream terminalHelperStream = getClass().getClassLoader().getResourceAsStream("terminal-helper.json");
162-
if (terminalHelperStream == null) {
163-
throw new RuntimeException("assessor-config.yaml not found on classpath");
164-
165-
}
166-
String terminalResponse = new String(terminalHelperStream.readAllBytes());
167-
InputStream is = getClass().getClassLoader().getResourceAsStream(agentConfigFile);
168-
if (is == null) {
169-
throw new RuntimeException(agentConfigFile + " not found on classpath");
170-
}
171-
AgentConfig config = new ObjectMapper(new YAMLFactory()).readValue(is, AgentConfig.class);
172144

173-
log.info("Agent config loaded: {}", config);
174-
PromptBuilder promptBuilder = new PromptBuilder(verbRegistry, config);
175-
var prompt = promptBuilder.buildPrompt();
176-
List<Message> messages = new ArrayList<>();
177-
178-
messages.add(Message.builder().role("system").content(prompt).build());
179-
messages.add(Message.builder().role("system").content("Please ensure your nextOperation abides by the " +
180-
"following json format and leave it empty if user's request doesn't require explicit use of system " +
181-
"operations" +
182-
". Please summarize prior terminal " +
183-
"sessions, using " +
184-
"terminal output if needed " +
185-
"for clarity of the next LLM request and for the user: " + terminalResponse).build());
186-
messages.add(Message.builder().role("user").content(userMessage.getContent()).build());
187-
LLMRequest chatRequest = LLMRequest.builder().model("gpt-4o").messages(messages).build();
188-
var resp = llmService.askQuestion(execution, chatRequest);
189-
execution.addMessages( messages );
190-
Response response = JsonUtil.MAPPER.readValue(resp, Response.class);
191-
log.info("Response is {}", resp);
192-
for (Response.Choice choice : response.getChoices()) {
193-
var content = choice.getMessage().getContent();
194-
if (content.startsWith("```json")) {
195-
content = content.substring(7, content.length() - 3);
196-
} else if (content.startsWith("```")) {
197-
content = content.substring(3, content.length() - 3);
198-
}
199-
log.info("content is {}", content);
200-
if (null != content && !content.isEmpty()) {
201-
var newResponse = JsonUtil.MAPPER.enable(JsonParser.Feature.ALLOW_COMMENTS).readValue(content,
202-
TerminalResponse.class);
203-
return newResponse;
204-
}
205-
}
206-
} else {
207-
InputStream terminalHelperStream = getClass().getClassLoader().getResourceAsStream("terminal-helper.json");
208-
if (terminalHelperStream == null) {
209-
throw new RuntimeException("assessor-config.yaml not found on classpath");
210-
211-
}
212-
String terminalResponse = new String(terminalHelperStream.readAllBytes());
213-
InputStream is = getClass().getClassLoader().getResourceAsStream(agentConfigFile);
214-
if (is == null) {
215-
throw new RuntimeException(agentConfigFile + " not found on classpath");
216-
}
217-
AgentConfig config = new ObjectMapper(new YAMLFactory()).readValue(is, AgentConfig.class);
218-
219-
log.info("Agent config loaded: {}", config);
220-
PromptBuilder promptBuilder = new PromptBuilder(verbRegistry, config);
221-
var prompt = promptBuilder.buildPrompt();
222-
List<Message> messages = new ArrayList<>();
223-
224-
messages.add(Message.builder().role("system").content(prompt).build());
225-
messages.add(Message.builder().role("system").content("Please ensure your nextOperation abide by the " +
226-
"following json format. Please summarize prior terminal sessions, using terminal output if needed " +
227-
"for clarity of the next LLM request and for the user: " + terminalResponse).build());
228-
messages.add(Message.builder().role("assistant").content("prior response: " + lastMessage.getTerminalSummaryForLLM()).build());
229-
messages.add(Message.builder().role("user").content(userMessage.getContent()).build());
230-
231-
LLMRequest chatRequest = LLMRequest.builder().model("gpt-4o").messages(messages).build();
232-
var resp = llmService.askQuestion(execution, chatRequest);
233-
execution.addMessages( messages );
234-
Response response = JsonUtil.MAPPER.readValue(resp, Response.class);
235-
log.info("Response is {}", resp);
236-
for (Response.Choice choice : response.getChoices()) {
237-
var content = choice.getMessage().getContent();
238-
if (content.startsWith("```json")) {
239-
content = content.substring(7, content.length() - 3);
240-
} else if (content.startsWith("```")) {
241-
content = content.substring(3, content.length() - 3);
242-
}
243-
log.info("content is {}", content);
244-
if (null != content && !content.isEmpty()) {
245-
var newResponse = JsonUtil.MAPPER.enable(JsonParser.Feature.ALLOW_COMMENTS).readValue(content,
246-
TerminalResponse.class);
247-
return newResponse;
248-
}
249-
}
250-
}
251-
252-
return null;
253-
}
254145

255146
/**
256147
* Chats with an agent to justify operations based on the provided arguments.

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

Lines changed: 217 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,217 @@
1+
package io.sentrius.agent.analysis.agents.verbs;
2+
3+
import java.io.IOException;
4+
import java.io.InputStream;
5+
import java.util.ArrayList;
6+
import java.util.List;
7+
import com.fasterxml.jackson.core.JsonParser;
8+
import com.fasterxml.jackson.databind.ObjectMapper;
9+
import com.fasterxml.jackson.databind.node.ArrayNode;
10+
import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
11+
import io.sentrius.agent.analysis.agents.agents.AgentConfig;
12+
import io.sentrius.agent.analysis.agents.agents.AgentVerb;
13+
import io.sentrius.agent.analysis.agents.agents.PromptBuilder;
14+
import io.sentrius.agent.analysis.agents.agents.VerbRegistry;
15+
import io.sentrius.agent.analysis.model.TerminalResponse;
16+
import io.sentrius.agent.analysis.model.WebSocky;
17+
import io.sentrius.sso.core.dto.ztat.AgentExecution;
18+
import io.sentrius.sso.core.exceptions.ZtatException;
19+
import io.sentrius.sso.core.model.verbs.Verb;
20+
import io.sentrius.sso.core.services.agents.AgentClientService;
21+
import io.sentrius.sso.core.services.agents.LLMService;
22+
import io.sentrius.sso.core.services.agents.ZeroTrustClientService;
23+
import io.sentrius.sso.core.utils.JsonUtil;
24+
import io.sentrius.sso.genai.Message;
25+
import io.sentrius.sso.genai.Response;
26+
import io.sentrius.sso.genai.model.LLMRequest;
27+
import lombok.NonNull;
28+
import lombok.RequiredArgsConstructor;
29+
import lombok.extern.slf4j.Slf4j;
30+
import org.springframework.beans.factory.annotation.Value;
31+
import org.springframework.stereotype.Service;
32+
33+
@Service
34+
@Slf4j
35+
@RequiredArgsConstructor
36+
public class ChatVerbs {
37+
38+
@Value("${agent.ai.config}")
39+
private String agentConfigFile;
40+
41+
final ZeroTrustClientService zeroTrustClientService;
42+
final LLMService llmService;
43+
final VerbRegistry verbRegistry;
44+
final AgentClientService agentClientService;
45+
46+
/**
47+
* Prompts the agent for workload based on the provided arguments.
48+
*
49+
* @return An `ArrayNode` containing the plan generated by the agent.
50+
* @throws io.sentrius.sso.core.exceptions.ZtatException If there is an error during the operation.
51+
* @throws java.io.IOException If there is an error reading the configuration file.
52+
*/
53+
@Verb(name = "interpret_user_request", returnType = ArrayNode.class, description = "Queries the LLM using the " +
54+
"user input.",
55+
isAiCallable = false, requiresTokenManagement = true)
56+
public TerminalResponse interpretUserData(
57+
AgentExecution execution, @NonNull WebSocky socketConnection,
58+
@NonNull Message userMessage) throws ZtatException,
59+
IOException {
60+
61+
var lastMessage = socketConnection.getMessages().stream().reduce((prev, next) -> next).orElse(null);
62+
if (socketConnection.getMessages().isEmpty()) {
63+
64+
InputStream terminalHelperStream = getClass().getClassLoader().getResourceAsStream("terminal-helper.json");
65+
if (terminalHelperStream == null) {
66+
throw new RuntimeException("assessor-config.yaml not found on classpath");
67+
68+
}
69+
String terminalResponse = new String(terminalHelperStream.readAllBytes());
70+
InputStream is = getClass().getClassLoader().getResourceAsStream(agentConfigFile);
71+
if (is == null) {
72+
throw new RuntimeException(agentConfigFile + " not found on classpath");
73+
}
74+
AgentConfig config = new ObjectMapper(new YAMLFactory()).readValue(is, AgentConfig.class);
75+
76+
log.info("Agent config loaded: {}", config);
77+
PromptBuilder promptBuilder = new PromptBuilder(verbRegistry, config);
78+
var prompt = promptBuilder.buildPrompt();
79+
List<Message> messages = new ArrayList<>();
80+
81+
messages.add(Message.builder().role("system").content(prompt).build());
82+
messages.add(Message.builder().role("system").content("Please ensure your nextOperation abides by the " +
83+
"following json format and leave it empty if user's request doesn't require explicit use of system " +
84+
"operations" +
85+
". Please summarize prior terminal " +
86+
"sessions, using " +
87+
"terminal output if needed " +
88+
"for clarity of the next LLM request and for the user: " + terminalResponse).build());
89+
messages.add(Message.builder().role("user").content(userMessage.getContent()).build());
90+
LLMRequest chatRequest = LLMRequest.builder().model("gpt-4o").messages(messages).build();
91+
var resp = llmService.askQuestion(execution, chatRequest);
92+
execution.addMessages( messages );
93+
Response response = JsonUtil.MAPPER.readValue(resp, Response.class);
94+
log.info("Response is {}", resp);
95+
for (Response.Choice choice : response.getChoices()) {
96+
var content = choice.getMessage().getContent();
97+
if (content.startsWith("```json")) {
98+
content = content.substring(7, content.length() - 3);
99+
} else if (content.startsWith("```")) {
100+
content = content.substring(3, content.length() - 3);
101+
}
102+
log.info("content is {}", content);
103+
if (null != content && !content.isEmpty()) {
104+
var newResponse = JsonUtil.MAPPER.enable(JsonParser.Feature.ALLOW_COMMENTS).readValue(content,
105+
TerminalResponse.class);
106+
return newResponse;
107+
}
108+
}
109+
} else {
110+
InputStream terminalHelperStream = getClass().getClassLoader().getResourceAsStream("terminal-helper.json");
111+
if (terminalHelperStream == null) {
112+
throw new RuntimeException("assessor-config.yaml not found on classpath");
113+
114+
}
115+
String terminalResponse = new String(terminalHelperStream.readAllBytes());
116+
InputStream is = getClass().getClassLoader().getResourceAsStream(agentConfigFile);
117+
if (is == null) {
118+
throw new RuntimeException(agentConfigFile + " not found on classpath");
119+
}
120+
AgentConfig config = new ObjectMapper(new YAMLFactory()).readValue(is, AgentConfig.class);
121+
122+
log.info("Agent config loaded: {}", config);
123+
PromptBuilder promptBuilder = new PromptBuilder(verbRegistry, config);
124+
var prompt = promptBuilder.buildPrompt();
125+
List<Message> messages = new ArrayList<>();
126+
127+
messages.add(Message.builder().role("system").content(prompt).build());
128+
messages.add(Message.builder().role("system").content("Please ensure your nextOperation abide by the " +
129+
"following json format. Please summarize prior terminal sessions, using terminal output if needed " +
130+
"for clarity of the next LLM request and for the user: " + terminalResponse).build());
131+
messages.add(Message.builder().role("assistant").content("prior response: " + lastMessage.getTerminalSummaryForLLM()).build());
132+
messages.add(Message.builder().role("user").content(userMessage.getContent()).build());
133+
134+
LLMRequest chatRequest = LLMRequest.builder().model("gpt-4o").messages(messages).build();
135+
var resp = llmService.askQuestion(execution, chatRequest);
136+
execution.addMessages( messages );
137+
Response response = JsonUtil.MAPPER.readValue(resp, Response.class);
138+
log.info("Response is {}", resp);
139+
for (Response.Choice choice : response.getChoices()) {
140+
var content = choice.getMessage().getContent();
141+
if (content.startsWith("```json")) {
142+
content = content.substring(7, content.length() - 3);
143+
} else if (content.startsWith("```")) {
144+
content = content.substring(3, content.length() - 3);
145+
}
146+
log.info("content is {}", content);
147+
if (null != content && !content.isEmpty()) {
148+
var newResponse = JsonUtil.MAPPER.enable(JsonParser.Feature.ALLOW_COMMENTS).readValue(content,
149+
TerminalResponse.class);
150+
return newResponse;
151+
}
152+
}
153+
}
154+
155+
return null;
156+
}
157+
158+
159+
public TerminalResponse interpret_plan_response(
160+
AgentExecution execution, @NonNull WebSocky socketConnection,
161+
AgentVerb agentVerb, String planExecutionOutput) throws ZtatException,
162+
IOException {
163+
164+
165+
log.info("interpret_plan_response");
166+
167+
var lastMessage = socketConnection.getMessages().stream().reduce((prev, next) -> next).orElse(null);
168+
169+
InputStream terminalHelperStream = getClass().getClassLoader().getResourceAsStream("terminal-helper.json");
170+
if (terminalHelperStream == null) {
171+
throw new RuntimeException("assessor-config.yaml not found on classpath");
172+
173+
}
174+
String terminalResponse = new String(terminalHelperStream.readAllBytes());
175+
InputStream is = getClass().getClassLoader().getResourceAsStream(agentConfigFile);
176+
if (is == null) {
177+
throw new RuntimeException(agentConfigFile + " not found on classpath");
178+
}
179+
AgentConfig config = new ObjectMapper(new YAMLFactory()).readValue(is, AgentConfig.class);
180+
181+
log.info("Agent config loaded: {}", config);
182+
PromptBuilder promptBuilder = new PromptBuilder(verbRegistry, config);
183+
var prompt = promptBuilder.buildPrompt();
184+
List<Message> messages = new ArrayList<>();
185+
186+
messages.add(Message.builder().role("system").content(prompt).build());
187+
messages.add(Message.builder().role("system").content("You have executed verbs for the previous user " +
188+
"messages. Please generate a user response that summarizes the last message.").build());
189+
if (null != agentVerb ){
190+
messages.add(Message.builder().role("system").content("You have executed verb: " + agentVerb.getName() +
191+
" with the following description: " + agentVerb.getDescription()).build());
192+
}
193+
messages.add(Message.builder().role("assistant").content("prior response: " + lastMessage.getTerminalSummaryForLLM()).build());
194+
messages.add(Message.builder().role("assistant").content(planExecutionOutput).build());
195+
196+
LLMRequest chatRequest = LLMRequest.builder().model("gpt-4o").messages(messages).build();
197+
var resp = llmService.askQuestion(execution, chatRequest);
198+
execution.addMessages( messages );
199+
Response response = JsonUtil.MAPPER.readValue(resp, Response.class);
200+
log.info("Response is {}", resp);
201+
for (Response.Choice choice : response.getChoices()) {
202+
var content = choice.getMessage().getContent();
203+
if (content.startsWith("```json")) {
204+
content = content.substring(7, content.length() - 3);
205+
} else if (content.startsWith("```")) {
206+
content = content.substring(3, content.length() - 3);
207+
}
208+
log.info("content is {}", content);
209+
if (null != content && !content.isEmpty()) {
210+
var newResponse = JsonUtil.MAPPER.enable(JsonParser.Feature.ALLOW_COMMENTS).readValue(content,
211+
TerminalResponse.class);
212+
return newResponse;
213+
}
214+
}
215+
return null;
216+
}
217+
}

0 commit comments

Comments
 (0)