Skip to content

Commit f4a4116

Browse files
phrockerphrocker
committed
Update
1 parent e117d08 commit f4a4116

76 files changed

Lines changed: 1640 additions & 965 deletions

File tree

Some content is hidden

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

‎.local.env‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1-
SENTRIUS_VERSION=1.1.261
1+
SENTRIUS_VERSION=1.1.320
22
SENTRIUS_SSH_VERSION=1.1.40
33
SENTRIUS_KEYCLOAK_VERSION=1.1.52
4-
SENTRIUS_AGENT_VERSION=1.1.39
5-
SENTRIUS_AI_AGENT_VERSION=1.1.148
6-
LLMPROXY_VERSION=1.0.53
7-
LAUNCHER_VERSION=1.0.73
4+
SENTRIUS_AGENT_VERSION=1.1.41
5+
SENTRIUS_AI_AGENT_VERSION=1.1.255
6+
LLMPROXY_VERSION=1.0.74
7+
LAUNCHER_VERSION=1.0.81
88
AGENTPROXY_VERSION=1.0.74

‎.local.env.bak‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
1-
SENTRIUS_VERSION=1.1.261
1+
SENTRIUS_VERSION=1.1.319
22
SENTRIUS_SSH_VERSION=1.1.40
33
SENTRIUS_KEYCLOAK_VERSION=1.1.52
4-
SENTRIUS_AGENT_VERSION=1.1.39
5-
SENTRIUS_AI_AGENT_VERSION=1.1.148
6-
LLMPROXY_VERSION=1.0.53
7-
LAUNCHER_VERSION=1.0.73
4+
SENTRIUS_AGENT_VERSION=1.1.41
5+
SENTRIUS_AI_AGENT_VERSION=1.1.254
6+
LLMPROXY_VERSION=1.0.73
7+
LAUNCHER_VERSION=1.0.81
88
AGENTPROXY_VERSION=1.0.74

‎agent-launcher/src/main/java/io/sentrius/agent/launcher/api/AgentLauncherController.java‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,16 +23,19 @@
2323
@Slf4j
2424
@RestController
2525
@RequestMapping(ApiPaths.API_V1 + "/agent/launcher")
26-
public class AgentLauncherController {
26+
public class AgentLauncherController {
2727
private final PodLauncherService podLauncherService;
2828
private final KeycloakService keycloakService;
2929

30-
public AgentLauncherController(
31-
PodLauncherService podLauncherService, KeycloakService keycloakService) {
30+
protected AgentLauncherController(
31+
PodLauncherService podLauncherService, KeycloakService keycloakService
32+
) {
33+
3234
this.podLauncherService = podLauncherService;
3335
this.keycloakService = keycloakService;
3436
}
3537

38+
3639
@PostMapping("/create")
3740
@LimitAccess(applicationAccess = {ApplicationAccessEnum.CAN_MANAGE_APPLICATION})
3841
public ResponseEntity<?> createPod(
@@ -49,6 +52,9 @@ public ResponseEntity<?> createPod(
4952
return ResponseEntity.status(HttpStatus.SC_UNAUTHORIZED).body("Invalid Keycloak token");
5053
}
5154

55+
// create the user, assign the policy
56+
57+
5258
podLauncherService.launchAgentPod(agent);
5359

5460
return ResponseEntity.ok(Map.of("status", "success"));
@@ -65,9 +71,9 @@ public ResponseEntity<String> deleteAgent(@RequestParam(name="agentId") String a
6571
}
6672

6773
@GetMapping("/status")
68-
public ResponseEntity<String> getAgentStatus(@RequestParam(name="agentId") String agentId) {
74+
public ResponseEntity<?> getAgentStatus(@RequestParam(name="agentId") String agentId) {
6975
try {
70-
return ResponseEntity.ok(podLauncherService.statusById(agentId) );
76+
return ResponseEntity.ok(Map.of("status", podLauncherService.statusById(agentId)) );
7177
} catch (Exception e) {
7278
log.error("Status failed", e);
7379
return ResponseEntity.status(500).body("Status retrieval failed");

‎agent-launcher/src/main/java/io/sentrius/agent/launcher/service/PodLauncherService.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -178,8 +178,12 @@ public V1Pod launchAgentPod(AgentRegistrationDTO agent) throws Exception {
178178
List<String> argList = new ArrayList<>();
179179
argList.add("--spring.config.location=file:/config/agent.properties");
180180
argList.add("--agent.namePrefix=" + agentId);
181+
argList.add("--agent.clientId=" + agent.getClientId());
181182
argList.add("--agent.listen.websocket=true");
182183
argList.add("--agent.callback.url=" + constructedCallbackUrl);
184+
if (agent.getAgentPolicyId() != null && !agent.getAgentPolicyId().isEmpty()) {
185+
argList.add("--agent.ai.policy.id=" + agent.getAgentPolicyId());
186+
}
183187
if (agent.getAgentContextId() != null && !agent.getAgentContextId().isEmpty()) {
184188
argList.add("--agent.ai.context.db.id=" + agent.getAgentContextId());
185189
}else {
@@ -215,8 +219,8 @@ public V1Pod launchAgentPod(AgentRegistrationDTO agent) throws Exception {
215219
.args(argList)
216220
.resources(new V1ResourceRequirements()
217221
.limits(Map.of(
218-
"cpu", Quantity.fromString("1000m"),
219-
"memory", Quantity.fromString("1Gi")
222+
"cpu", Quantity.fromString("2000m"),
223+
"memory", Quantity.fromString("2Gi")
220224
)))
221225
.volumeMounts(List.of(
222226
new V1VolumeMount()

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

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,6 @@
33
import java.lang.reflect.Method;
44
import java.util.List;
55
import io.sentrius.sso.core.dto.capabilities.ParameterDescriptor;
6-
import io.sentrius.sso.core.model.verbs.DefaultInterpreter;
7-
import io.sentrius.sso.core.model.verbs.InputInterpreterIfc;
8-
import io.sentrius.sso.core.model.verbs.OutputInterpreterIfc;
96
import lombok.Builder;
107
import lombok.Getter;
118

@@ -21,9 +18,11 @@ public class AgentVerb {
2118
private boolean requiresTokenManagement = false;
2219
@Builder.Default
2320
private Class<?> returnType = String.class;
21+
2422
@Builder.Default
25-
Class<? extends OutputInterpreterIfc> outputInterpreter = DefaultInterpreter.class;
26-
Class<? extends InputInterpreterIfc> inputInterpreter = DefaultInterpreter.class;
23+
private String returnName = "";
2724

2825
private String exampleJson = "";
26+
@Builder.Default
27+
private String argName = "arg1";
2928
}

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

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
import io.sentrius.agent.analysis.api.UserCommunicationService;
1010
import io.sentrius.agent.config.AgentConfigOptions;
1111
import io.sentrius.sso.core.dto.UserDTO;
12-
import io.sentrius.sso.core.dto.ztat.AgentExecution;
12+
import io.sentrius.sso.core.dto.agents.AgentExecution;
1313
import io.sentrius.sso.core.dto.ztat.ZtatRequestDTO;
1414
import io.sentrius.sso.core.exceptions.ZtatException;
1515
import io.sentrius.sso.core.model.security.Ztat;
@@ -85,7 +85,8 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
8585
try {
8686
var agentName = agentConfigOptions.getNamePrefix() + "-" + UUID.randomUUID().toString();
8787
var base64PublicKey = agentKeyService.getBase64PublicKey(keyPair.getPublic());
88-
var agentRegistrationDTO = agentClientService.bootstrap(agentName, base64PublicKey
88+
var agentRegistrationDTO = agentClientService.bootstrap(agentConfigOptions.getClientId(), agentName,
89+
base64PublicKey
8990
, keyPair.getPublic().getAlgorithm());
9091

9192
var encryptedSecret = agentRegistrationDTO.getClientSecret();
@@ -143,15 +144,24 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
143144
throw new RuntimeException(e);
144145
}
145146

147+
int allowedFailures = 20;
148+
log.info("Agent Registered...");
146149
while(running) {
147150

148-
log.info("Agent Registered...");
151+
149152
try {
150153

151154
Thread.sleep(5_000);
152155
agentClientService.heartbeat(agentExecution, agentExecution.getUser().getUsername());
153-
} catch (InterruptedException | ZtatException ex) {
154-
throw new RuntimeException(ex);
156+
allowedFailures = 20; // Reset allowed failures on successful heartbeat
157+
} catch (ZtatException | Exception ex) {
158+
if (allowedFailures-- <= 0) {
159+
log.error("Failed to heartbeat agent after multiple attempts, shutting down...");
160+
throw new RuntimeException(ex);
161+
} else {
162+
log.warn("Heartbeat failed, retrying... Remaining attempts: {}", allowedFailures);
163+
}
164+
155165
}
156166

157167
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -81,11 +81,11 @@ public String buildPrompt(boolean applyInstructions)
8181
// your DTO
8282
try {
8383
if (verb.getExampleJson() != null && !verb.getExampleJson().isEmpty()) {
84-
prompt.append(" Example arg1: ").append(verb.getExampleJson()).append("\n");
84+
prompt.append(" Example \"" + verb.getArgName() + "\": ").append(verb.getExampleJson()).append("\n");
8585
} else if (example != null) {
8686
// Serialize the example object to JSON
8787
String exampleJson = JsonUtil.MAPPER.writeValueAsString(example);
88-
prompt.append(" Example arg1: ").append(exampleJson).append("\n");
88+
prompt.append(" Example " + verb.getArgName() + ": ").append(exampleJson).append("\n");
8989
}
9090

9191
} catch (Exception e) {

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77
import io.sentrius.agent.analysis.agents.verbs.AgentVerbs;
88
import io.sentrius.agent.analysis.api.AgentKeyService;
99
import io.sentrius.agent.config.AgentConfigOptions;
10-
import io.sentrius.sso.core.dto.ztat.AgentExecution;
10+
import io.sentrius.sso.core.dto.agents.AgentExecution;
11+
import io.sentrius.sso.core.dto.agents.AgentExecutionContextDTO;
1112
import io.sentrius.sso.core.dto.ztat.ZtatRequestDTO;
1213
import io.sentrius.sso.core.exceptions.ZtatException;
1314
import io.sentrius.sso.core.model.security.Ztat;
@@ -120,6 +121,7 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
120121

121122
var agentExecution = agentExecutionService.getAgentExecution(finalUser);
122123
var response = promptAgent(agentExecution);
124+
AgentExecutionContextDTO agentExecutionContext = AgentExecutionContextDTO.builder().build();
123125
while (running) {
124126
try {
125127
log.info("Got response: {}", response);
@@ -131,7 +133,8 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
131133
if (node.get("verb") != null) {
132134
var verb = node.get("verb").asText();
133135
log.info("Executing verb: {}", verb);
134-
priorResponse = verbRegistry.execute(agentExecution, priorResponse, verb, args);
136+
priorResponse = verbRegistry.execute(agentExecution,agentExecutionContext,
137+
priorResponse, verb, args);
135138
}
136139
log.info("Node: {}", node);
137140
}

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

Lines changed: 71 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,22 @@
11
package io.sentrius.agent.analysis.agents.agents;
22

33
import com.fasterxml.jackson.core.JsonProcessingException;
4+
import com.fasterxml.jackson.databind.JsonNode;
5+
import com.fasterxml.jackson.databind.node.ObjectNode;
46
import io.github.classgraph.ClassGraph;
57
import io.github.classgraph.ScanResult;
8+
import io.sentrius.sso.core.dto.agents.AgentExecutionContextDTO;
69
import io.sentrius.sso.core.dto.capabilities.EndpointDescriptor;
710
import io.sentrius.agent.discovery.AgentEndpointDiscoveryService;
8-
import io.sentrius.sso.core.dto.ztat.AgentExecution;
11+
import io.sentrius.sso.core.dto.agents.AgentExecution;
912
import io.sentrius.sso.core.dto.ztat.ZtatRequestDTO;
1013
import io.sentrius.sso.core.exceptions.ZtatException;
11-
import io.sentrius.sso.core.model.verbs.OutputInterpreterIfc;
1214
import io.sentrius.sso.core.model.verbs.Verb;
1315
import io.sentrius.sso.core.model.verbs.VerbResponse;
1416
import io.sentrius.sso.core.services.agents.AgentClientService;
1517
import io.sentrius.sso.core.services.agents.ZeroTrustClientService;
1618
import io.sentrius.sso.core.services.capabilities.EndpointScanningService;
19+
import io.sentrius.sso.core.utils.JsonUtil;
1720
import lombok.RequiredArgsConstructor;
1821
import lombok.extern.slf4j.Slf4j;
1922
import org.springframework.context.ApplicationContext;
@@ -93,11 +96,11 @@ public void scanClasspath() {
9396
.name(name)
9497
.description(annotation.description())
9598
.method(method)
99+
.argName(annotation.argName())
100+
.returnName(annotation.returnName())
96101
.exampleJson(annotation.exampleJson())
97102
.requiresTokenManagement(annotation.requiresTokenManagement())
98103
.returnType(annotation.returnType())
99-
.outputInterpreter(annotation.outputInterpreter())
100-
.inputInterpreter(annotation.inputInterpreter())
101104
.build();
102105
verbs.put(name, verb);
103106
instances.put(name, instance);
@@ -119,46 +122,75 @@ public boolean isVerbRegistered(String verb) {
119122
}
120123
}
121124

122-
public VerbResponse execute(AgentExecution agentExecution, VerbResponse priorResponse, String verb,
125+
public VerbResponse execute(AgentExecution agentExecution,
126+
AgentExecutionContextDTO contextDTO, VerbResponse priorResponse,
127+
String verb,
123128
Map<String, Object> args)
124129
throws Exception {
130+
131+
log.info("Executing {}", contextDTO);
125132
synchronized (this) {
126-
Map<String, Object> newArgs = new HashMap<>();
133+
var agentVerb = verbs.get(verb);
134+
if (null == agentVerb) {
135+
throw new IllegalArgumentException("Unknown verb: " + verb);
136+
}
137+
127138
if (null != args) {
128-
newArgs.putAll(args);
139+
140+
for (Map.Entry<String, Object> entry : args.entrySet()) {
141+
contextDTO.getExecutionArgs().put(entry.getKey(), entry.getValue().toString());
142+
}
143+
144+
if (agentVerb.getArgName() != null && !agentVerb.getArgName().isEmpty() && !agentVerb.getArgName().equals("arg1")) {
145+
ObjectNode object = JsonUtil.MAPPER.createObjectNode();
146+
for (Map.Entry<String, Object> entry : args.entrySet()) {
147+
var node = JsonUtil.MAPPER.valueToTree(entry.getValue());
148+
object.put(entry.getKey(), node);
149+
log.info("Interpreting input " +
150+
"for AgentExecutionContextDTO: {}", entry.getValue());
151+
}
152+
153+
154+
contextDTO.getExecutionArgs().put(agentVerb.getArgName(), object);
155+
156+
}
129157
}
158+
130159
log.info("Executing verb: {}", verb);
131-
var returnType = verbs.get(verb).getReturnType();
160+
var returnType = agentVerb.getReturnType();
132161
if (null != priorResponse ) {
133-
Class<? extends OutputInterpreterIfc> interpreter = priorResponse.getOutputInterpreter();
134162
log.info("Interpreting prior response for verb: {}", verb);
135-
log.info("Interpreting prior response with: {}", interpreter.getName());
136-
log.info("Interpreting prior response: {}", priorResponse.getReturnType());
163+
log.info("Interpreting prior response: {}", priorResponse.getReturnType());
137164
log.info("New response return type: {}", returnType);
138-
newArgs.putAll(interpreter.getConstructor().newInstance().interpret(priorResponse));
139165
}
140-
Method method = verbs.get(verb).getMethod();
141-
var inputInterpreter = verbs.get(verb).getInputInterpreter();
166+
Method method = agentVerb.getMethod();
142167
Object instance = instances.get(verb);
143168
if (method == null) {
144169
throw new IllegalArgumentException("Unknown verb: " + verb);
145170
}
146-
var interpreterInstance = inputInterpreter.getConstructor().newInstance();
147171
try {
148172

149173

150-
log.info("Interpreting input with: {}", interpreterInstance.getClass().getName());
151-
log.info("Interpreting input: {}", newArgs.getClass().getName());
152-
log.info("Interpreting input: {}", newArgs);
153-
var interpretedInput = interpreterInstance.interpret(newArgs);
174+
175+
var thisVerb = agentVerb;
176+
var exec = thisVerb.isRequiresTokenManagement() ?
177+
method.invoke(instance, agentExecution, contextDTO) :
178+
method.invoke(instance, contextDTO);
179+
180+
JsonNode execNode = JsonUtil.MAPPER.valueToTree(exec);
181+
log.info("Interpreting output for AgentExecutionContextDTO: {}", execNode);
182+
// add the output
183+
if (null != thisVerb.getReturnName() && !thisVerb.getReturnName().isEmpty()) {
184+
contextDTO.addToMemory(thisVerb.getReturnName(), execNode);
185+
} else {
186+
contextDTO.addToMemory(verb, execNode);
187+
}
188+
154189

155190
return VerbResponse.builder()
156-
.response(verbs.get(verb).isRequiresTokenManagement() ?
157-
method.invoke(instance,agentExecution, interpretedInput) :
158-
method.invoke(instance, interpretedInput))
159191
.returnType(returnType)
160-
.outputInterpreter(verbs.get(verb).getOutputInterpreter())
161192
.build();
193+
162194
} catch (InvocationTargetException e) {
163195
Throwable targetException = e.getTargetException();
164196
if (targetException instanceof ZtatException ztatEx) {
@@ -179,22 +211,31 @@ public VerbResponse execute(AgentExecution agentExecution, VerbResponse priorRes
179211

180212
log.info("Re-attempting verb execution after Ztat token acquisition: {}", verb);
181213

182-
var interpretedInput = interpreterInstance.interpret(newArgs);
183-
184-
log.info("Re-attempting verb execution after Ztat token acquisition: {}", verb);
214+
var thisVerb = agentVerb;
215+
var exec = thisVerb.isRequiresTokenManagement() ?
216+
method.invoke(instance, agentExecution, contextDTO) :
217+
method.invoke(instance, contextDTO);
218+
JsonNode execNode = JsonUtil.MAPPER.valueToTree(exec);
219+
// add the output
220+
if (null != thisVerb.getReturnName() && !thisVerb.getReturnName().isEmpty()) {
221+
contextDTO.addToMemory(thisVerb.getReturnName(), execNode);
222+
} else {
223+
contextDTO.addToMemory(verb, execNode);
224+
}
185225

186226
return VerbResponse.builder()
187-
.response(verbs.get(verb).isRequiresTokenManagement() ?
188-
method.invoke(instance,agentExecution, interpretedInput) :
189-
method.invoke(instance, interpretedInput))
190227
.returnType(returnType)
191-
.outputInterpreter(verbs.get(verb).getOutputInterpreter())
192228
.build();
229+
230+
231+
193232
// re-attempt
194233
} else {
195234
throw new RuntimeException(targetException);
196235
}
197236
} catch (Exception e) {
237+
log.info(method.getName() + " failed", e);
238+
e.printStackTrace();
198239
throw new RuntimeException("Failed to execute verb: " + verb, e);
199240
}
200241
}

0 commit comments

Comments
 (0)