Skip to content

Commit fe0ee2f

Browse files
phrockerphrocker
committed
Initial plan
1 parent 6ebce08 commit fe0ee2f

100 files changed

Lines changed: 7875 additions & 293 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: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,4 +6,4 @@ SENTRIUS_AI_AGENT_VERSION=1.1.264
66
LLMPROXY_VERSION=1.0.78
77
LAUNCHER_VERSION=1.0.82
88
AGENTPROXY_VERSION=1.0.85
9-
SSHPROXY_VERSION=1.0.87
9+
SSHPROXY_VERSION=1.0.88

‎.local.env.bak‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,4 +6,4 @@ SENTRIUS_AI_AGENT_VERSION=1.1.264
66
LLMPROXY_VERSION=1.0.78
77
LAUNCHER_VERSION=1.0.82
88
AGENTPROXY_VERSION=1.0.85
9-
SSHPROXY_VERSION=1.0.87
9+
SSHPROXY_VERSION=1.0.87

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,7 @@ 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.type=" + agent.getAgentType());
181182
argList.add("--agent.clientId=" + agent.getClientId());
182183
argList.add("--agent.listen.websocket=true");
183184
argList.add("--agent.callback.url=" + constructedCallbackUrl);

‎agent-launcher/src/main/resources/application.properties‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,4 +8,4 @@ spring.main.web-application-type=servlet
88
spring.thymeleaf.enabled=true
99
spring.freemarker.enabled=false
1010

11-
sentrius.agent.registry=local
11+
sentrius.agent.registry=local

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

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
package io.sentrius.agent.analysis.agents.agents;
2+
3+
import java.util.concurrent.TimeUnit;
4+
import com.fasterxml.jackson.databind.node.ArrayNode;
5+
import io.sentrius.agent.analysis.agents.verbs.AgentVerbs;
6+
import io.sentrius.sso.core.dto.agents.AgentExecution;
7+
import io.sentrius.sso.core.dto.agents.AgentExecutionContextDTO;
8+
import io.sentrius.sso.core.dto.ztat.ZtatRequestDTO;
9+
import io.sentrius.sso.core.exceptions.ZtatException;
10+
import io.sentrius.sso.core.services.agents.AgentClientService;
11+
import io.sentrius.sso.core.services.agents.ZeroTrustClientService;
12+
import lombok.RequiredArgsConstructor;
13+
import lombok.extern.slf4j.Slf4j;
14+
import org.springframework.beans.factory.annotation.Autowired;
15+
import org.springframework.boot.context.event.ApplicationReadyEvent;
16+
import org.springframework.context.ApplicationListener;
17+
18+
@Slf4j
19+
public abstract class BaseEnterpriseAgent implements ApplicationListener<ApplicationReadyEvent> {
20+
21+
@Autowired
22+
protected final AgentVerbs agentVerbs;
23+
@Autowired
24+
protected final ZeroTrustClientService zeroTrustClientService;
25+
@Autowired
26+
protected final AgentClientService agentClientService;
27+
@Autowired
28+
protected final VerbRegistry verbRegistry;
29+
30+
protected BaseEnterpriseAgent(
31+
AgentVerbs agentVerbs, ZeroTrustClientService zeroTrustClientService, AgentClientService agentClientService,
32+
VerbRegistry verbRegistry
33+
) {
34+
this.agentVerbs = agentVerbs;
35+
this.zeroTrustClientService = zeroTrustClientService;
36+
this.agentClientService = agentClientService;
37+
this.verbRegistry = verbRegistry;
38+
}
39+
40+
41+
protected ArrayNode promptAgent(AgentExecution execution) throws ZtatException {
42+
return promptAgent(execution);
43+
}
44+
45+
protected ArrayNode promptAgent(AgentExecution execution, AgentExecutionContextDTO contextDTO) throws ZtatException {
46+
while(true){
47+
try {
48+
log.info("Prompting agent...");
49+
return agentVerbs.promptAgent(execution,contextDTO);
50+
} catch (ZtatException e) {
51+
log.info("Mechanisms {}" , e.getMechanisms());
52+
var endpoint = zeroTrustClientService.createEndPointRequest("prompt_agent", e.getEndpoint());
53+
ZtatRequestDTO ztatRequestDTO = ZtatRequestDTO.builder()
54+
.user(execution.getUser())
55+
.command(endpoint.toString())
56+
.justification("Registered Agent requires ability to prompt LLM endpoints to begin operations")
57+
.summary("Registered Agent requires ability to prompt LLM endpoints to begin operations")
58+
.build();
59+
var request = zeroTrustClientService.requestZtatToken(execution, execution.getUser(),ztatRequestDTO);
60+
61+
var token = zeroTrustClientService.awaitZtatToken(execution, execution.getUser(), request, 60,
62+
TimeUnit.MINUTES);
63+
execution.setZtatToken(token);
64+
} catch (Exception e) {
65+
log.error(e.getMessage());
66+
throw new RuntimeException(e);
67+
}
68+
}
69+
}
70+
}

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

Lines changed: 125 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,77 +1,76 @@
11
package io.sentrius.agent.analysis.agents.agents;
22

3+
import java.io.IOException;
4+
import java.util.ArrayList;
5+
import java.util.HashMap;
6+
import java.util.List;
7+
import java.util.Map;
38
import java.util.UUID;
4-
import java.util.concurrent.TimeUnit;
59
import com.fasterxml.jackson.core.JsonProcessingException;
6-
import com.fasterxml.jackson.databind.node.ArrayNode;
710
import io.sentrius.agent.analysis.agents.verbs.AgentVerbs;
11+
import io.sentrius.agent.analysis.agents.verbs.ChatVerbs;
812
import io.sentrius.agent.analysis.api.AgentKeyService;
913
import io.sentrius.agent.analysis.api.UserCommunicationService;
14+
import io.sentrius.agent.analysis.model.LLMResponse;
1015
import io.sentrius.agent.config.AgentConfigOptions;
1116
import io.sentrius.sso.core.dto.UserDTO;
1217
import io.sentrius.sso.core.dto.agents.AgentExecution;
13-
import io.sentrius.sso.core.dto.ztat.ZtatRequestDTO;
18+
import io.sentrius.sso.core.dto.agents.AgentExecutionContextDTO;
1419
import io.sentrius.sso.core.exceptions.ZtatException;
1520
import io.sentrius.sso.core.model.security.Ztat;
21+
import io.sentrius.sso.core.model.verbs.VerbResponse;
1622
import io.sentrius.sso.core.services.agents.AgentClientService;
1723
import io.sentrius.sso.core.services.agents.AgentExecutionService;
1824
import io.sentrius.sso.core.services.agents.ZeroTrustClientService;
1925
import io.sentrius.sso.core.services.security.KeycloakService;
2026
import io.sentrius.sso.core.utils.JsonUtil;
27+
import io.sentrius.sso.genai.Message;
2128
import jakarta.annotation.PreDestroy;
22-
import lombok.RequiredArgsConstructor;
2329
import lombok.extern.slf4j.Slf4j;
30+
import org.springframework.beans.factory.annotation.Autowired;
2431
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
2532
import org.springframework.boot.context.event.ApplicationReadyEvent;
26-
import org.springframework.context.ApplicationListener;
2733
import org.springframework.stereotype.Component;
2834

2935
@Slf4j
3036
@Component
31-
@RequiredArgsConstructor
3237
@ConditionalOnProperty(name = "agents.ai.chat.agent.enabled", havingValue = "true", matchIfMissing = false)
33-
public class ChatAgent implements ApplicationListener<ApplicationReadyEvent> {
38+
public class ChatAgent extends BaseEnterpriseAgent {
3439

3540

3641
final ZeroTrustClientService zeroTrustClientService;
3742
final AgentClientService agentClientService;
3843
final VerbRegistry verbRegistry;
39-
final AgentVerbs agentVerbs;
4044
final AgentExecutionService agentExecutionService;
4145
final UserCommunicationService userCommunicationService;
4246
final AgentConfigOptions agentConfigOptions;
4347
final AgentKeyService agentKeyService;
4448
private final KeycloakService keycloakService;
49+
final ChatVerbs chatVerbs;
4550

4651
private volatile boolean running = true;
4752
private Thread workerThread;
4853

4954
private AgentExecution agentExecution;
5055

51-
public ArrayNode promptAgent(AgentExecution execution) throws ZtatException {
52-
while(true){
53-
try {
54-
log.info("Prompting agent...");
55-
return agentVerbs.promptAgent(execution,null);
56-
} catch (ZtatException e) {
57-
log.info("Mechanisms {}" , e.getMechanisms());
58-
var endpoint = zeroTrustClientService.createEndPointRequest("prompt_agent", e.getEndpoint());
59-
ZtatRequestDTO ztatRequestDTO = ZtatRequestDTO.builder()
60-
.user(execution.getUser())
61-
.command(endpoint.toString())
62-
.justification("Registered Agent requires ability to prompt LLM endpoints to begin operations")
63-
.summary("Registered Agent requires ability to prompt LLM endpoints to begin operations")
64-
.build();
65-
var request = zeroTrustClientService.requestZtatToken(execution, execution.getUser(),ztatRequestDTO);
66-
67-
var token = zeroTrustClientService.awaitZtatToken(execution, execution.getUser(), request, 60,
68-
TimeUnit.MINUTES);
69-
execution.setZtatToken(token);
70-
} catch (Exception e) {
71-
log.error(e.getMessage());
72-
throw new RuntimeException(e);
73-
}
74-
}
56+
57+
@Autowired
58+
public ChatAgent(
59+
AgentVerbs agentVerbs, ZeroTrustClientService zeroTrustClientService, AgentClientService agentClientService,
60+
VerbRegistry verbRegistry, AgentExecutionService agentExecutionService, UserCommunicationService userCommunicationService,
61+
AgentConfigOptions agentConfigOptions, AgentKeyService agentKeyService, KeycloakService keycloakService,
62+
ChatVerbs chatVerbs
63+
) {
64+
super(agentVerbs, zeroTrustClientService, agentClientService, verbRegistry);
65+
this.zeroTrustClientService = zeroTrustClientService;
66+
this.agentClientService = agentClientService;
67+
this.verbRegistry = verbRegistry;
68+
this.agentExecutionService = agentExecutionService;
69+
this.userCommunicationService = userCommunicationService;
70+
this.agentConfigOptions = agentConfigOptions;
71+
this.agentKeyService = agentKeyService;
72+
this.keycloakService = keycloakService;
73+
this.chatVerbs = chatVerbs;
7574
}
7675

7776
@Override
@@ -146,15 +145,109 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
146145

147146
int allowedFailures = 20;
148147
log.info("Agent Registered...");
148+
AgentExecutionContextDTO agentExecutionContext = AgentExecutionContextDTO.builder().build();
149+
agentExecutionService.setExecutionContextDTO(agentExecution, agentExecutionContext);
150+
LLMResponse response = null;
151+
AgentConfig config = null;
152+
try {
153+
config = chatVerbs.getAgentConfig(agentExecution);
154+
} catch (IOException e) {
155+
throw new RuntimeException(e);
156+
} catch (ZtatException e) {
157+
throw new RuntimeException(e);
158+
}
159+
PromptBuilder promptBuilder = new PromptBuilder(verbRegistry, config);
160+
var prompt = promptBuilder.buildPrompt(false);
161+
try {
162+
if (agentConfigOptions.getType().equalsIgnoreCase("chat-autonomous")) {
163+
164+
response = chatVerbs.promptAgent(agentExecution, agentExecutionContext, prompt);
165+
}
166+
} catch (ZtatException e) {
167+
throw new RuntimeException(e);
168+
} catch (IOException e) {
169+
throw new RuntimeException(e);
170+
}
171+
172+
173+
if (agentConfigOptions.getType().equalsIgnoreCase("chat-autonomous") && response == null) {
174+
log.error("Chat autonomous agent mode enabled but no response received from promptAgent, shutting down...");
175+
throw new RuntimeException("Chat autonomous agent mode enabled but no response received from promptAgent");
176+
}
177+
VerbResponse lastVerbResponse = null;
178+
LLMResponse nextResponse = null;
179+
List<VerbResponse> verbResponses = new ArrayList<>();
149180
while(running) {
150181

151182

152183
try {
153184

154185
Thread.sleep(5_000);
155186
agentClientService.heartbeat(agentExecution, agentExecution.getUser().getUsername());
187+
if (agentConfigOptions.getType().equalsIgnoreCase("chat-autonomous")) {
188+
log.info("Chat autonomous agent mode enabled, executing workload...");
189+
VerbResponse priorResponse = null;
190+
Map<String, Object> args = new HashMap<>();
191+
192+
var arguments = response.getArguments();
193+
if (null != response) {
194+
if (response.getNextOperation() != null && !response.getNextOperation().isEmpty()) {
195+
var executionResponse = verbRegistry.execute(
196+
agentExecution,
197+
agentExecutionContext,
198+
lastVerbResponse,
199+
response.getNextOperation(), arguments
200+
);
201+
verbResponses.add(executionResponse);
202+
lastVerbResponse = executionResponse;
203+
204+
205+
// chatAgent.getAgentExecution().addMessages(Message.builder().role("System")
206+
// .content("System executed operation: " + response.getNextOperation()).build());
207+
var responses = agentExecutionContext.getAgentDataList();
208+
var planResponse =
209+
responses.isEmpty() ? "" :
210+
responses.get(responses.size() - 1).asText();
211+
nextResponse = chatVerbs.interpret_plan_response(
212+
agentExecution,
213+
agentExecutionContext,
214+
verbRegistry.getVerbs().get(response.getNextOperation()),
215+
planResponse
216+
);
217+
218+
var memory = agentExecutionContext.flushPersistentMemory();
219+
if (memory != null) {
220+
for(var memoryEntry : memory.entrySet()){
221+
agentClientService.storeMemory(agentExecution,
222+
agentExecutionContext.getAgentContext().getName(),
223+
io.sentrius.sso.core.dto.agents.AgentMemoryDTO.builder()
224+
.agentName(agentExecutionContext.getAgentContext().getName())
225+
.memoryKey(memoryEntry.getKey())
226+
.memoryValue(memoryEntry.getValue().toString())
227+
.build());
228+
}
229+
}
230+
231+
232+
response = nextResponse;
233+
}
234+
235+
}else {
236+
response = chatVerbs.promptAgent(agentExecution, agentExecutionContext, prompt);
237+
238+
}
239+
240+
continue;
241+
}
156242
allowedFailures = 20; // Reset allowed failures on successful heartbeat
157243
} catch (ZtatException | Exception ex) {
244+
agentExecutionContext.addMessages(Message.builder().role("system").content(
245+
"You caused the following error. Please re-validate you chose the right operations or " +
246+
"endpoints for the context" +
247+
ex.getMessage()).build());
248+
249+
250+
ex.printStackTrace();
158251
if (allowedFailures-- <= 0) {
159252
log.error("Failed to heartbeat agent after multiple attempts, shutting down...");
160253
throw new RuntimeException(ex);

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

Lines changed: 29 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import io.sentrius.agent.config.AgentConfigOptions;
1010
import io.sentrius.sso.core.dto.agents.AgentExecution;
1111
import io.sentrius.sso.core.dto.agents.AgentExecutionContextDTO;
12+
import io.sentrius.sso.core.dto.agents.AgentMemoryDTO;
1213
import io.sentrius.sso.core.dto.ztat.ZtatRequestDTO;
1314
import io.sentrius.sso.core.exceptions.ZtatException;
1415
import io.sentrius.sso.core.model.security.Ztat;
@@ -22,22 +23,18 @@
2223
import jakarta.annotation.PreDestroy;
2324
import lombok.RequiredArgsConstructor;
2425
import lombok.extern.slf4j.Slf4j;
26+
import org.springframework.beans.factory.annotation.Autowired;
2527
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
2628
import org.springframework.boot.context.event.ApplicationReadyEvent;
2729
import org.springframework.context.ApplicationListener;
2830
import org.springframework.stereotype.Component;
2931

3032
@Slf4j
3133
@Component
32-
@RequiredArgsConstructor
3334
@ConditionalOnProperty(name = "agents.ai.registered.agent.enabled", havingValue = "true", matchIfMissing = false)
34-
public class RegisteredAgent implements ApplicationListener<ApplicationReadyEvent> {
35+
public class RegisteredAgent extends BaseEnterpriseAgent {
3536

3637

37-
final ZeroTrustClientService zeroTrustClientService;
38-
final AgentClientService agentClientService;
39-
final VerbRegistry verbRegistry;
40-
final AgentVerbs agentVerbs;
4138
final AgentExecutionService agentExecutionService;
4239
final AgentConfigOptions agentConfigOptions;
4340
final AgentKeyService agentKeyService;
@@ -46,6 +43,19 @@ public class RegisteredAgent implements ApplicationListener<ApplicationReadyEven
4643
private volatile boolean running = true;
4744
private Thread workerThread;
4845

46+
@Autowired
47+
public RegisteredAgent(
48+
AgentVerbs agentVerbs, ZeroTrustClientService zeroTrustClientService, AgentClientService agentClientService,
49+
VerbRegistry verbRegistry, AgentExecutionService agentExecutionService, AgentConfigOptions agentConfigOptions,
50+
AgentKeyService agentKeyService, KeycloakService keycloakService
51+
) {
52+
super(agentVerbs, zeroTrustClientService, agentClientService, verbRegistry);
53+
this.agentExecutionService = agentExecutionService;
54+
this.agentConfigOptions = agentConfigOptions;
55+
this.agentKeyService = agentKeyService;
56+
this.keycloakService = keycloakService;
57+
}
58+
4959
public ArrayNode promptAgent(AgentExecution execution) throws ZtatException {
5060
while(true){
5161
try {
@@ -136,6 +146,19 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
136146
priorResponse = verbRegistry.execute(agentExecution,agentExecutionContext,
137147
priorResponse, verb, args);
138148
}
149+
var memoryList = agentExecutionContext.flushPersistentMemory();
150+
if (memoryList != null) {
151+
for(var memory : memoryList.entrySet()){
152+
AgentMemoryDTO dto = AgentMemoryDTO.builder()
153+
.agentName(agentExecutionContext.getAgentContext().getName())
154+
.memoryKey(memory.getKey())
155+
.memoryValue(memory.getValue().toString())
156+
.build();
157+
agentClientService.storeMemory(agentExecution,
158+
agentExecutionContext.getAgentContext().getName(), dto);
159+
}
160+
}
161+
139162
log.info("Node: {}", node);
140163
}
141164

0 commit comments

Comments
 (0)