Skip to content

Commit 6fa1712

Browse files
phrockerphrocker
committed
update
1 parent 634f22c commit 6fa1712

209 files changed

Lines changed: 1456 additions & 6327 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.

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
11
package io.sentrius.agent.analysis;
22

3+
import io.sentrius.agent.config.AgentConfigOptions;
34
import org.springframework.boot.SpringApplication;
45
import org.springframework.boot.autoconfigure.SpringBootApplication;
6+
import org.springframework.boot.context.properties.EnableConfigurationProperties;
57
import org.springframework.scheduling.annotation.EnableScheduling;
68

79

@@ -11,6 +13,7 @@
1113
org.springframework.boot.autoconfigure.orm.jpa.HibernateJpaAutoConfiguration.class,
1214
org.springframework.boot.autoconfigure.data.jpa.JpaRepositoriesAutoConfiguration.class
1315
})
16+
@EnableConfigurationProperties(AgentConfigOptions.class)
1417
@EnableScheduling
1518
public class AiAgent {
1619
public static void main(String[] args) {

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

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
package io.sentrius.agent.analysis.agents.agents;
2+
3+
import java.util.concurrent.TimeUnit;
4+
import com.fasterxml.jackson.core.JsonProcessingException;
5+
import com.fasterxml.jackson.databind.node.ArrayNode;
6+
import io.sentrius.agent.analysis.agents.verbs.AgentVerbs;
7+
import io.sentrius.agent.analysis.api.AgentKeyService;
8+
import io.sentrius.agent.analysis.api.UserCommunicationService;
9+
import io.sentrius.agent.config.AgentConfigOptions;
10+
import io.sentrius.sso.core.dto.UserDTO;
11+
import io.sentrius.sso.core.dto.ztat.AgentExecution;
12+
import io.sentrius.sso.core.dto.ztat.ZtatRequestDTO;
13+
import io.sentrius.sso.core.exceptions.ZtatException;
14+
import io.sentrius.sso.core.model.security.Ztat;
15+
import io.sentrius.sso.core.services.agents.AgentClientService;
16+
import io.sentrius.sso.core.services.agents.AgentExecutionService;
17+
import io.sentrius.sso.core.services.agents.ZeroTrustClientService;
18+
import io.sentrius.sso.core.services.security.KeycloakService;
19+
import io.sentrius.sso.core.utils.JsonUtil;
20+
import jakarta.annotation.PreDestroy;
21+
import lombok.RequiredArgsConstructor;
22+
import lombok.extern.slf4j.Slf4j;
23+
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
24+
import org.springframework.boot.context.event.ApplicationReadyEvent;
25+
import org.springframework.context.ApplicationListener;
26+
import org.springframework.stereotype.Component;
27+
28+
@Slf4j
29+
@Component
30+
@RequiredArgsConstructor
31+
@ConditionalOnProperty(name = "agents.ai.chat.agent.enabled", havingValue = "true", matchIfMissing = false)
32+
public class ChatAgent implements ApplicationListener<ApplicationReadyEvent> {
33+
34+
35+
final ZeroTrustClientService zeroTrustClientService;
36+
final AgentClientService agentClientService;
37+
final VerbRegistry verbRegistry;
38+
final AgentVerbs agentVerbs;
39+
final AgentExecutionService agentExecutionService;
40+
final UserCommunicationService userCommunicationService;
41+
final AgentConfigOptions agentConfigOptions;
42+
final AgentKeyService agentKeyService;
43+
private final KeycloakService keycloakService;
44+
45+
private volatile boolean running = true;
46+
private Thread workerThread;
47+
48+
public ArrayNode promptAgent(AgentExecution execution) throws ZtatException {
49+
while(true){
50+
try {
51+
log.info("Prompting agent...");
52+
return agentVerbs.promptAgent(execution,null);
53+
} catch (ZtatException e) {
54+
log.info("Mechanisms {}" , e.getMechanisms());
55+
var endpoint = zeroTrustClientService.createEndPointRequest("prompt_agent", e.getEndpoint());
56+
ZtatRequestDTO ztatRequestDTO = ZtatRequestDTO.builder()
57+
.user(execution.getUser())
58+
.command(endpoint.toString())
59+
.justification("Registered Agent requires ability to prompt LLM endpoints to begin operations")
60+
.summary("Registered Agent requires ability to prompt LLM endpoints to begin operations")
61+
.build();
62+
var request = zeroTrustClientService.requestZtatToken(execution, execution.getUser(),ztatRequestDTO);
63+
64+
var token = zeroTrustClientService.awaitZtatToken(execution, execution.getUser(), request, 60,
65+
TimeUnit.MINUTES);
66+
execution.setZtatToken(token);
67+
} catch (Exception e) {
68+
log.error(e.getMessage());
69+
throw new RuntimeException(e);
70+
}
71+
}
72+
}
73+
74+
@Override
75+
public void onApplicationEvent(final ApplicationReadyEvent event) {
76+
77+
verbRegistry.scanClasspath();
78+
79+
var keyPair = agentKeyService.getKeyPair();
80+
81+
try {
82+
var base64PublicKey = agentKeyService.getBase64PublicKey(keyPair.getPublic());
83+
var agentRegistrationDTO = agentClientService.bootstrap(agentConfigOptions.getName(), base64PublicKey
84+
, keyPair.getPublic().getAlgorithm());
85+
86+
var encryptedSecret = agentRegistrationDTO.getClientSecret();
87+
var decryptedSecret = agentKeyService.
88+
decryptWithPrivateKey(encryptedSecret, keyPair.getPrivate());
89+
keycloakService.createKeycloakClient(agentRegistrationDTO.getAgentName(),
90+
decryptedSecret);
91+
92+
93+
} catch (ZtatException e) {
94+
throw new RuntimeException(e);
95+
} catch (JsonProcessingException e) {
96+
throw new RuntimeException(e);
97+
}
98+
99+
100+
final UserDTO user = UserDTO.builder()
101+
.username(zeroTrustClientService.getUsername())
102+
.build();
103+
var execution = agentExecutionService.getAgentExecution(user);
104+
try {
105+
agentClientService.heartbeat(execution, execution.getUser().getUsername());
106+
} catch (ZtatException e) {
107+
throw new RuntimeException(e);
108+
}
109+
110+
while(running) {
111+
112+
try {
113+
var register = zeroTrustClientService.registerAgent(execution);
114+
log.info("Registered agent response: {}", register);
115+
116+
var ztat = JsonUtil.MAPPER.readValue(register, Ztat.class);
117+
execution.setZtatToken(ztat.getZtatToken());
118+
execution.setCommunicationId(ztat.getCommunicationId());
119+
break;
120+
}catch (Exception | ZtatException e) {
121+
122+
log.error(e.getMessage());
123+
log.info("Registering v1.0.2 agent failed. Retrying in 10 seconds...");
124+
try {
125+
Thread.sleep(10_000);
126+
} catch (InterruptedException ex) {
127+
throw new RuntimeException(ex);
128+
}
129+
}
130+
}
131+
132+
while(running) {
133+
134+
log.info("Registering v1.0.2 agent failed. Retrying in 10 seconds...");
135+
try {
136+
137+
Thread.sleep(5_000);
138+
} catch (InterruptedException ex) {
139+
throw new RuntimeException(ex);
140+
}
141+
142+
}
143+
144+
}
145+
146+
@PreDestroy
147+
public void shutdown() {
148+
log.info("Shutting down ChatAgent...");
149+
running = false;
150+
if (workerThread != null) {
151+
workerThread.interrupt();
152+
}
153+
}
154+
155+
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ public String buildPrompt() {
6363
verbRegistry.getVerbs().forEach((name, verb) -> {
6464
prompt.append("- ").append(name);
6565
prompt.append(" (").append(buildMethodSignature(verb.getMethod())).append(") - ");
66-
prompt.append(verb.getDescription()).append(")\n");
66+
prompt.append(verb.getDescription()).append("\n");
6767
});
6868

6969
return prompt.toString();

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

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -109,10 +109,10 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
109109

110110

111111

112+
var agentExecution = agentExecutionService.getAgentExecution(user);
113+
var response = promptAgent(agentExecution);
112114
while (running) {
113115
try {
114-
var agentExecution = agentExecutionService.getAgentExecution(user);
115-
var response = promptAgent(agentExecution);
116116
log.info("Got response: {}", response);
117117

118118
VerbResponse priorResponse = null;
@@ -129,16 +129,14 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
129129

130130
} catch (Exception e) {
131131
log.error("Exception in agent loop", e);
132-
} catch (ZtatException e) {
133-
throw new RuntimeException(e);
134132
}
135133

136134
// Sleep between prompts
137135
log.info("Sleeping for 5 seconds");
138-
Thread.sleep(1_000);
136+
Thread.sleep(5_000);
139137
}
140138

141-
} catch (Exception e) {
139+
} catch (Exception | ZtatException e) {
142140
log.error("Fatal error in RegisteredAgent", e);
143141
}
144142
});

0 commit comments

Comments
 (0)