Skip to content

Commit d5f115d

Browse files
phrockerphrocker
committed
move stuff around and what not
1 parent eb5c5b7 commit d5f115d

File tree

124 files changed

+7534
-634
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

124 files changed

+7534
-634
lines changed

‎ai-agent/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,11 @@
3737
<artifactId>sentrius-core</artifactId>
3838
<version>1.0.0-SNAPSHOT</version>
3939
</dependency>
40+
<dependency>
41+
<groupId>io.sentrius</groupId>
42+
<artifactId>sentrius-llm-core</artifactId>
43+
<version>1.0.0-SNAPSHOT</version>
44+
</dependency>
4045
<dependency>
4146
<groupId>io.jsonwebtoken</groupId>
4247
<artifactId>jjwt-api</artifactId>

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ public ArrayNode promptAgent(AgentExecution execution) throws ZtatException {
4646
return agentVerbs.promptAgent(execution,null);
4747
} catch (ZtatException e) {
4848
log.info("Mechanisms {}" , e.getMechanisms());
49-
var endpoint = zeroTrustClientService.createEndPoingRequest("prompt_agent", e.getEndpoint());
49+
var endpoint = zeroTrustClientService.createEndPointRequest("prompt_agent", e.getEndpoint());
5050
ZtatRequestDTO ztatRequestDTO = ZtatRequestDTO.builder()
5151
.user(execution.getUser())
5252
.command(endpoint.toString())
@@ -134,8 +134,8 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
134134
}
135135

136136
// Sleep between prompts
137-
log.info("Sleeping for 60 seconds");
138-
Thread.sleep(10_000);
137+
log.info("Sleeping for 5 seconds");
138+
Thread.sleep(5_000);
139139
}
140140

141141
} catch (Exception e) {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,7 @@ public VerbResponse execute(AgentExecution agentExecution, VerbResponse priorRes
120120
Throwable targetException = e.getTargetException();
121121
if (targetException instanceof ZtatException ztatEx) {
122122
log.info("Mechanisms {}" , ztatEx.getMechanisms());
123-
var endpoint = zeroTrustClientService.createEndPoingRequest("prompt_agent`", ztatEx.getEndpoint());
123+
var endpoint = zeroTrustClientService.createEndPointRequest("prompt_agent`", ztatEx.getEndpoint());
124124
ZtatRequestDTO ztatRequestDTO = ZtatRequestDTO.builder()
125125
.user(agentExecution.getUser())
126126
.command(endpoint.toString())

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

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
package io.sentrius.agent.analysis.agents.interpreters;
2+
3+
import java.util.ArrayList;
4+
import java.util.HashMap;
5+
import java.util.List;
6+
import java.util.Map;
7+
import com.fasterxml.jackson.core.type.TypeReference;
8+
import com.fasterxml.jackson.databind.node.ArrayNode;
9+
import com.fasterxml.jackson.databind.node.ObjectNode;
10+
import io.sentrius.agent.analysis.model.Assessment;
11+
import io.sentrius.sso.core.dto.HostSystemDTO;
12+
import io.sentrius.sso.core.model.verbs.ListInterpreter;
13+
import io.sentrius.sso.core.model.verbs.OutputInterpreterIfc;
14+
import io.sentrius.sso.core.model.verbs.VerbResponse;
15+
import io.sentrius.sso.core.utils.JsonUtil;
16+
import lombok.extern.slf4j.Slf4j;
17+
18+
@Slf4j
19+
public class AsessmentListInterpreter extends ListInterpreter<Assessment> implements OutputInterpreterIfc {
20+
21+
@Override
22+
public Map<String, Object> interpret(VerbResponse input) throws Exception {
23+
log.info("AssessmentListInterpreter: interpret() called with input: {}", input);
24+
// Assuming input.response is a list of strings
25+
Map<String,Object> responseMap = new HashMap<>();
26+
responseMap.put("verb.response.type", "list");
27+
responseMap.put("verb.response.map.key", "assessments");
28+
responseMap.put("verb.response.map.type", Assessment.class.getCanonicalName());
29+
30+
if (input.getResponse() instanceof List<?> list) {
31+
log.info("AssessmentListInterpreter: interpret() called with input list");
32+
if (list.isEmpty() || list.get(0) instanceof Assessment) {
33+
34+
responseMap.put("assessments", list);
35+
} else {
36+
throw new IllegalArgumentException("Input response is not a List of Assessment objects");
37+
}
38+
} else {
39+
40+
var str = input.getResponse().toString();
41+
log.info("AssessmentListInterpreter: interpret() called with input string {} " ,str);
42+
ArrayNode node = (ArrayNode) JsonUtil.MAPPER.readTree(str);
43+
if (node == null) {
44+
throw new IllegalArgumentException("Input response is not a valid JSON array");
45+
}
46+
List<Assessment> list = new ArrayList<>();
47+
for (int i = 0; i < node.size(); i++) {
48+
var item = node.get(i);
49+
if (item.has("sessionId") && item.has("risk") && item.has("description")) {
50+
Assessment hostSystemDTO = new Assessment();
51+
hostSystemDTO.setSessionId(item.get("sessionId").asText());
52+
hostSystemDTO.setDescription(item.get("description").asText());
53+
hostSystemDTO.setRisk(item.get("risk").asText());
54+
list.add(hostSystemDTO);
55+
} else {
56+
throw new IllegalArgumentException("Input response does not contain required fields");
57+
}
58+
}
59+
responseMap.put("assessments",list);
60+
}
61+
62+
return responseMap;
63+
}
64+
65+
}

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

Lines changed: 116 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package io.sentrius.agent.analysis.agents.verbs;
22

3+
import java.io.FileInputStream;
4+
import java.io.FileReader;
35
import java.io.IOException;
46
import java.io.InputStream;
57
import java.util.ArrayList;
@@ -12,15 +14,20 @@
1214
import com.fasterxml.jackson.databind.ObjectMapper;
1315
import com.fasterxml.jackson.databind.node.ArrayNode;
1416
import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
17+
import com.google.common.collect.Maps;
1518
import io.sentrius.agent.analysis.agents.agents.AgentConfig;
1619
import io.sentrius.agent.analysis.agents.agents.PromptBuilder;
1720
import io.sentrius.agent.analysis.agents.agents.VerbRegistry;
21+
import io.sentrius.agent.analysis.agents.interpreters.AsessmentListInterpreter;
1822
import io.sentrius.agent.analysis.agents.interpreters.ObjectListInterpreter;
1923
import io.sentrius.agent.analysis.agents.interpreters.ZtatOutputInterpreter;
24+
import io.sentrius.agent.analysis.model.Assessment;
25+
import io.sentrius.agent.analysis.model.ZtatAsessment;
2026
import io.sentrius.sso.core.dto.JITTrackerDTO;
2127
import io.sentrius.sso.core.dto.ztat.AgentExecution;
2228
import io.sentrius.sso.core.dto.ztat.AtatRequest;
2329
import io.sentrius.sso.core.dto.ztat.TokenDTO;
30+
import io.sentrius.sso.core.dto.ztat.ZtatRequestDTO;
2431
import io.sentrius.sso.core.exceptions.ZtatException;
2532
import io.sentrius.sso.core.model.verbs.Verb;
2633
import io.sentrius.sso.core.services.agents.AgentClientService;
@@ -29,7 +36,7 @@
2936
import io.sentrius.sso.core.utils.JsonUtil;
3037
import io.sentrius.sso.genai.Message;
3138
import io.sentrius.sso.genai.Response;
32-
import io.sentrius.sso.genai.model.ChatRequest;
39+
import io.sentrius.sso.genai.model.LLMRequest;
3340
import lombok.extern.slf4j.Slf4j;
3441
import org.springframework.beans.factory.annotation.Value;
3542
import org.springframework.stereotype.Service;
@@ -98,7 +105,7 @@ public ArrayNode promptAgent(AgentExecution execution, Map<String, Object> args)
98105

99106
messages.add(Message.builder().role("system").content(prompt).build());
100107

101-
ChatRequest chatRequest = ChatRequest.builder().model("gpt-4o").messages(messages).build();
108+
LLMRequest chatRequest = LLMRequest.builder().model("gpt-4o").messages(messages).build();
102109
var resp = llmService.askQuestion(execution, chatRequest);
103110
execution.addMessages( messages );
104111
Response response = JsonUtil.MAPPER.readValue(resp, Response.class);
@@ -132,22 +139,27 @@ public ArrayNode promptAgent(AgentExecution execution, Map<String, Object> args)
132139
*/
133140
@Verb(name = "justify_operations", description = "Chats with an agent to justify operations.", isAiCallable =
134141
false, requiresTokenManagement = true)
135-
public String justifyAgent(AgentExecution execution, String ztatRequest, String reason) throws ZtatException,
136-
IOException {
137-
InputStream is = getClass().getClassLoader().getResourceAsStream(agentConfigFile);
138-
if (is == null) {
139-
throw new RuntimeException("assessor-config.yaml not found on classpath");
140-
}
141-
AgentConfig config = new ObjectMapper(new YAMLFactory()).readValue(is, AgentConfig.class);
142+
public String justifyAgent(AgentExecution execution, ZtatRequestDTO ztatRequest, String reason) throws ZtatException,
143+
IOException, InterruptedException {
142144

143-
log.info("Agent config loaded: {}", config);
144-
PromptBuilder promptBuilder = new PromptBuilder(verbRegistry, config);
145-
var prompt = promptBuilder.buildPrompt();
146-
List<Message> messages = new ArrayList<>();
147145

148-
messages.add(Message.builder().role("system").content(prompt).build());
149146

150-
ChatRequest chatRequest = ChatRequest.builder().model("gpt-4o").messages(messages).build();
147+
var status = zeroTrustClientService.getTokenStatus(execution, execution.getUser(), ztatRequest.getRequestId());
148+
log.info("Status: {} for {} ", status, ztatRequest);
149+
if ("approved".equals(status.get("status").asText())) {
150+
return status.get("ztat_token").asText();
151+
}
152+
153+
while(!status.equals("approved")) {
154+
155+
Thread.sleep(5_000);
156+
157+
status = zeroTrustClientService.getTokenStatus(execution, execution.getUser(), ztatRequest.getRequestId());
158+
log.info("Status: {} for {} ", status, ztatRequest);
159+
if ("approved".equals(status.get("status").asText())) {
160+
return status.get("ztat_token").asText();
161+
}
162+
}
151163

152164
return null;
153165
// return llmService.askQuestion(chatRequest);
@@ -164,9 +176,10 @@ public String justifyAgent(AgentExecution execution, String ztatRequest, String
164176
@Verb(name = "assess_data", returnType = ArrayNode.class, description = "Accepts api server data based on the " +
165177
"context and seeks" +
166178
" to perform the assessment by prompting the LLM. Can be used to assess data or request information from users and/or agents.",
179+
outputInterpreter = AsessmentListInterpreter.class,
167180
inputInterpreter =
168181
ObjectListInterpreter.class, requiresTokenManagement = true)
169-
public ArrayNode assessData(AgentExecution execution, List<?> objectList) throws ZtatException, IOException {
182+
public List<Assessment> assessData(AgentExecution execution, List<?> objectList) throws ZtatException, IOException {
170183
InputStream is = getClass().getClassLoader().getResourceAsStream(agentConfigFile);
171184
if (is == null) {
172185
throw new RuntimeException("assessor-config.yaml not found on classpath");
@@ -175,7 +188,7 @@ public ArrayNode assessData(AgentExecution execution, List<?> objectList) throws
175188

176189
log.info("Agent config loaded: {}", config);
177190

178-
var responses = JsonUtil.MAPPER.createArrayNode();
191+
List<Assessment> responses = new ArrayList<>();
179192
log.info("Object list is {}", objectList);
180193
for (var obj : objectList) {
181194
List<Message> messages = new ArrayList<>();
@@ -184,7 +197,7 @@ public ArrayNode assessData(AgentExecution execution, List<?> objectList) throws
184197
messages.add(Message.builder().role("user").content(obj.toString()).build());
185198
messages.add(Message.builder().role("system").content(context).build());
186199

187-
ChatRequest chatRequest = ChatRequest.builder().model("gpt-4o").messages(messages).build();
200+
LLMRequest chatRequest = LLMRequest.builder().model("gpt-4o").messages(messages).build();
188201
execution.addMessages( messages );
189202
var resp = llmService.askQuestion(execution, chatRequest);
190203
Response response = JsonUtil.MAPPER.readValue(resp, Response.class);
@@ -194,7 +207,10 @@ public ArrayNode assessData(AgentExecution execution, List<?> objectList) throws
194207
if (content.startsWith("```json")) {
195208
content = content.substring(7, content.length() - 3);
196209
}
197-
responses.add(JsonUtil.MAPPER.readTree(content));
210+
211+
212+
responses.add( JsonUtil.MAPPER.readValue(content, Assessment.class) );
213+
//responses.add(JsonUtil.MAPPER.readTree(content));
198214
log.info("content is {}", content);
199215
}
200216
log.info("Object is {}", obj);
@@ -204,7 +220,7 @@ public ArrayNode assessData(AgentExecution execution, List<?> objectList) throws
204220

205221
@Verb(name = "list_ztat_requests", returnType = ArrayNode.class, description = "Lists zero trust access tokens to" +
206222
" review. Does not review access token requests.", outputInterpreter = ZtatOutputInterpreter.class, requiresTokenManagement = true )
207-
public List<AtatRequest> getWork(TokenDTO token, Map<String,Object> args) throws ZtatException, IOException {
223+
public List<AtatRequest> getWork(AgentExecution token, Map<String,Object> args) throws ZtatException, IOException {
208224
List<AtatRequest> requests = new ArrayList<>();
209225

210226
var atatRequests = agentClientService.getAtatRequests(token);
@@ -214,19 +230,90 @@ public List<AtatRequest> getWork(TokenDTO token, Map<String,Object> args) throws
214230
for (var dto : dtos) {
215231
var request = new AtatRequest();
216232
request.setRequestId(dto.getId().toString());
217-
// for each request7
218-
/*request.set
219-
request.setStatus(dto.getStatus());
220-
request.setCreatedAt(dto.getCreatedAt());
221-
request.setUpdatedAt(dto.getUpdatedAt());
222-
request.setAgentId(dto.getAgentId());
223-
requests.add(request);
224-
*
225-
*/
233+
// get messages
234+
request.setRequestedAction( dto.getSummary());
235+
236+
var communications = zeroTrustClientService.callGetOnApi(token,"agent/communications/id",
237+
Maps.immutableEntry("communicationId", List.of(token.getCommunicationId())));
238+
var messages = JsonUtil.MAPPER.readTree(communications);
239+
List<Message> communicationMessages = new ArrayList<>();
240+
for(JsonNode message : messages) {
241+
if (message.has("payload") && message.has("messageType")) {
242+
var type = message.get("messageType").asText();
243+
if (type.equalsIgnoreCase("chat_request")) {
244+
try {
245+
Message msg = JsonUtil.MAPPER.readValue(message.get("payload").asText(), Message.class);
246+
communicationMessages.add(msg);
247+
} catch (JsonProcessingException e) {
248+
throw new RuntimeException(e);
249+
}
250+
}
251+
}
252+
}
253+
request.setMessages(communicationMessages);
226254
}
227255

228256

229257
return requests;
230258
}
231259

260+
@Verb(name = "assess_ztat_requests", returnType = ArrayNode.class, description = "Analyzes ztats according to the" +
261+
" context.",
262+
inputInterpreter = ZtatOutputInterpreter.class, requiresTokenManagement = true )
263+
public List<ZtatAsessment> analyzeAtatRequests(AgentExecution execution, List<AtatRequest> requests) throws ZtatException,
264+
IOException {
265+
// set up context
266+
InputStream is = getClass().getClassLoader().getResourceAsStream(agentConfigFile);
267+
if (is == null) {
268+
throw new RuntimeException("assessor-config.yaml not found on classpath");
269+
270+
}
271+
272+
InputStream assessZtatStream = getClass().getClassLoader().getResourceAsStream("assess-ztat.json");
273+
if (assessZtatStream == null) {
274+
throw new RuntimeException("assessor-config.yaml not found on classpath");
275+
276+
}
277+
String assessZtat = new String(assessZtatStream.readAllBytes());
278+
279+
AgentConfig config = new ObjectMapper(new YAMLFactory()).readValue(is, AgentConfig.class);
280+
log.info("Agent config loaded: {}", config);
281+
List<ZtatAsessment> responses = new ArrayList<>();
282+
for (var request : requests) {
283+
List<Message> messages = new ArrayList<>();
284+
var context = config.getContext();
285+
286+
messages.add(Message.builder().role("system").content(context).build());
287+
messages.add(Message.builder().role("system").content("Ensure your response adheres to the following " +
288+
"json format:" + assessZtat).build());
289+
messages.addAll(execution.getMessages());
290+
messages.addAll(request.getMessages());
291+
292+
LLMRequest chatRequest = LLMRequest.builder().model("gpt-4o").messages(messages).build();
293+
var resp = llmService.askQuestion(execution, chatRequest);
294+
Response response = JsonUtil.MAPPER.readValue(resp, Response.class);
295+
log.info("Response is {}", resp);
296+
for (Response.Choice choice : response.getChoices()) {
297+
var content = choice.getMessage().getContent();
298+
if (content.startsWith("```json")) {
299+
content = content.substring(7, content.length() - 3);
300+
}
301+
log.info("content is {}", content);
302+
var ztat = JsonUtil.MAPPER.readValue(content, ZtatAsessment.class);
303+
if (ztat.isApproved()) {
304+
zeroTrustClientService.approveZtat(execution, request.getRequestId());
305+
}
306+
else {
307+
if (null != ztat.getQuestionToUser() &&
308+
ztat.getQuestionToUser().isEmpty()){
309+
// ask a question of the user
310+
}
311+
}
312+
responses.add(ztat);
313+
314+
}
315+
}
316+
return responses;
317+
}
318+
232319
}

0 commit comments

Comments
 (0)