Skip to content

Commit d5c4e79

Browse files
phrockerphrocker
authored
Merge pull request #43 from SentriusLLC/feature.chat_agent
add chat agent to dashboard
2 parents 634f22c + 613a67e commit d5c4e79

File tree

240 files changed

+2748
-6414
lines changed

Some content is hidden

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

240 files changed

+2748
-6414
lines changed

‎.gitignore‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,10 @@ llm-proxy/target/**
4646
llm-proxy/target/
4747
llm-dataplane/target/**
4848
sentrius-llm-dataplane/target/**
49+
sentrius-llm-core/target/**
4950
llm-dataplane/target/
5051
sentrius-llm-dataplane/target/
52+
sentrius-llm-core/target/
5153
node/*
5254
node_modules/*
5355
api/node_modules/*

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

‎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)