Skip to content

Commit fcabbd3

Browse files
phrockerphrocker
committed
commit
1 parent 5bf2cf8 commit fcabbd3

31 files changed

Lines changed: 697 additions & 184 deletions

File tree

‎ai-agent/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,11 @@
8989
<artifactId>snowflake-ingest-sdk</artifactId>
9090
<version>${snowflake-ingest-version}</version>
9191
</dependency>
92+
<dependency>
93+
<groupId>io.github.classgraph</groupId>
94+
<artifactId>classgraph</artifactId>
95+
<version>${classgraph.version}</version>
96+
</dependency>
9297
</dependencies>
9398

9499
<build>

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

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
package io.sentrius.agent.analysis.agents.agents;
2+
3+
import java.util.ArrayList;
4+
import java.util.List;
5+
import com.fasterxml.jackson.annotation.JsonProperty;
6+
import lombok.AllArgsConstructor;
7+
import lombok.Builder;
8+
import lombok.Data;
9+
import lombok.Getter;
10+
import lombok.NoArgsConstructor;
11+
import lombok.ToString;
12+
13+
@Builder
14+
@Getter
15+
@Data
16+
@ToString
17+
@AllArgsConstructor
18+
@NoArgsConstructor
19+
public class AgentConfig {
20+
@JsonProperty(required=false)
21+
@Builder.Default
22+
private List<String> roles = new ArrayList<>();
23+
@JsonProperty(required=false)
24+
private String description;
25+
@JsonProperty(required=true)
26+
private String context;
27+
28+
// getters and setters
29+
}

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
package io.sentrius.agent.analysis.agents.agents;
2+
3+
import java.util.Map;
4+
5+
public interface AgentVerb {
6+
String name();
7+
String description();
8+
Object execute(Map<String, Object> params) throws Exception;
9+
}

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

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
package io.sentrius.agent.analysis.agents.agents;
2+
3+
import java.lang.reflect.Method;
4+
import java.util.Arrays;
5+
import java.util.stream.Collectors;
6+
7+
public class PromptBuilder {
8+
9+
private final VerbRegistry verbRegistry;
10+
private final AgentConfig agentConfig;
11+
12+
public PromptBuilder(VerbRegistry verbRegistry, AgentConfig agentConfig) {
13+
this.verbRegistry = verbRegistry;
14+
this.agentConfig = agentConfig;
15+
}
16+
17+
public String buildPrompt() {
18+
StringBuilder prompt = new StringBuilder();
19+
20+
//prompt.append("Agent Name: ").append(agentConfig.getAgentName()).append("\n");
21+
//prompt.append("Purpose: ").append(agentConfig.getPurpose()).append("\n");
22+
prompt.append("Roles: ").append(String.join(", ", agentConfig.getRoles())).append("\n");
23+
prompt.append("Context: ").append(agentConfig.getContext()).append("\n\n");
24+
prompt.append("Instructions: ").append("Respond using this JSON format. Only use verbs provided in " +
25+
"Available Verbs:\n" +
26+
"\n" +
27+
"{\n" +
28+
" \"plan\": [\n" +
29+
" {\n" +
30+
" \"verb\": \"list_open_terminals\",\n" +
31+
" \"params\": {}\n" +
32+
" },\n" +
33+
" {\n" +
34+
" \"verb\": \"send_terminal_command\",\n" +
35+
" \"params\": {}\n" +
36+
" }\n" +
37+
" ]\n" +
38+
"}\n" );
39+
prompt.append("Available Verbs:\n");
40+
41+
verbRegistry.getVerbs().forEach((name, method) -> {
42+
prompt.append("- ").append(name);
43+
prompt.append(" (").append(buildMethodSignature(method)).append(")\n");
44+
});
45+
46+
return prompt.toString();
47+
}
48+
49+
private String buildMethodSignature(Method method) {
50+
return Arrays.stream(method.getParameters())
51+
.map(p -> p.getName() + ": " + p.getType().getSimpleName())
52+
.collect(Collectors.joining(", "));
53+
}
54+
}

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

Lines changed: 63 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,15 @@
11
package io.sentrius.agent.analysis.agents.agents;
22

3-
import io.sentrius.agent.services.ZeroTrustClientService;
3+
import java.io.IOException;
4+
import java.util.List;
5+
import com.fasterxml.jackson.core.JsonProcessingException;
6+
import io.sentrius.agent.analysis.agents.verbs.AgentVerbs;
7+
import io.sentrius.sso.core.exceptions.ZtatException;
8+
import io.sentrius.sso.core.services.agents.LLMService;
9+
import io.sentrius.sso.core.services.agents.ZeroTrustClientService;
410
import io.sentrius.sso.core.dto.UserDTO;
511
import lombok.RequiredArgsConstructor;
612
import lombok.extern.slf4j.Slf4j;
7-
import org.springframework.beans.factory.annotation.Value;
813
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
914
import org.springframework.boot.context.event.ApplicationReadyEvent;
1015
import org.springframework.context.ApplicationListener;
@@ -19,23 +24,68 @@ public class RegisteredAgent implements ApplicationListener<ApplicationReadyEven
1924

2025
final ZeroTrustClientService zeroTrustClientService;
2126

27+
final VerbRegistry verbRegistry;
28+
final AgentVerbs agentVerbs;
29+
30+
31+
32+
33+
2234
@Override
2335
public void onApplicationEvent(final ApplicationReadyEvent event) {
2436

25-
// get username from token
26-
UserDTO user = UserDTO.builder()
27-
.username(zeroTrustClientService.getUsername())
28-
.build();
37+
verbRegistry.scanClasspath();
38+
39+
// Load agent config
40+
41+
42+
43+
try {
44+
45+
46+
47+
// get username from token
48+
UserDTO user = UserDTO.builder()
49+
.username(zeroTrustClientService.getUsername())
50+
.build();
51+
52+
log.info(zeroTrustClientService.getUsername());
53+
54+
log.info("Registering v1.0.2 agent...");
55+
56+
// register
57+
var register = zeroTrustClientService.registerAgent(user);
58+
log.info("Registered agent is running {} ", register);
59+
60+
61+
//while(true){
62+
63+
try {
64+
// this phase is called "prompting"
65+
var response = agentVerbs.promptAgent(null);
66+
log.info("got " + response);
67+
} catch (ZtatException e) {
68+
e.printStackTrace();
69+
//zeroTrustClientService.requestZtatToken()
70+
// we have been requested to get a ztat
71+
} catch (JsonProcessingException e) {
72+
throw new RuntimeException(e);
73+
} catch (IOException e) {
74+
throw new RuntimeException(e);
75+
}
76+
// execute command
77+
// var result = verbRegistry.execute(command, null);
78+
//log.info("Command executed: {}", result);
79+
2980

30-
log.info(zeroTrustClientService.getUsername());
31-
String command = "ssh connect host123";
3281

33-
log.info("Registering v1.0.2 agent...");
82+
// sleep for 5 seconds
83+
Thread.sleep(5000);
84+
// }
3485

35-
// register
36-
var register = zeroTrustClientService.registerAgent(user);
37-
log.info("Registered agent is running {} ", register);
38-
return;
86+
} catch (InterruptedException e) {
87+
throw new RuntimeException(e);
88+
}
3989
}
4090

4191
}

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
package io.sentrius.agent.analysis.agents.agents;
2+
3+
import java.util.Map;
4+
5+
public interface VerbHandler {
6+
Object execute(Map<String, Object> args);
7+
}

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

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
package io.sentrius.agent.analysis.agents.agents;
2+
3+
import io.github.classgraph.ClassGraph;
4+
import io.github.classgraph.ScanResult;
5+
import io.sentrius.sso.core.model.verbs.Verb;
6+
import lombok.RequiredArgsConstructor;
7+
import lombok.extern.slf4j.Slf4j;
8+
import org.springframework.context.ApplicationContext;
9+
import org.springframework.stereotype.Component;
10+
import org.springframework.stereotype.Service;
11+
12+
13+
import java.lang.reflect.Method;
14+
import java.util.HashMap;
15+
import java.util.Map;
16+
17+
@Slf4j
18+
@Component
19+
@RequiredArgsConstructor
20+
public class VerbRegistry {
21+
22+
private final ApplicationContext applicationContext;
23+
24+
private final Map<String, Method> verbs = new HashMap<>();
25+
private final Map<String, Object> instances = new HashMap<>();
26+
27+
public void scanClasspath() {
28+
// Scan the classpath for classes with the @Verb annotation
29+
synchronized (this) {
30+
try (
31+
ScanResult scanResult = new ClassGraph()
32+
.enableAllInfo()
33+
.acceptPackages("io.sentrius.agent.analysis.agents.verbs")
34+
.scan()
35+
) {
36+
37+
scanResult.getClassesWithMethodAnnotation(Verb.class.getName()).forEach(classInfo -> {
38+
try {
39+
Class<?> clazz = classInfo.loadClass();
40+
Object instance = applicationContext.getBean(clazz); // <-- Get from Spring
41+
42+
for (Method method : clazz.getDeclaredMethods()) {
43+
if (method.isAnnotationPresent(Verb.class)) {
44+
Verb annotation = method.getAnnotation(Verb.class);
45+
String name = annotation.name();
46+
log.info("Found verb: {} in class: {}", name, clazz.getName());
47+
if (annotation.isAiCallable()) {
48+
log.info("Registering verb: {} in class: {}", name, clazz.getName());
49+
verbs.put(name, method);
50+
instances.put(name, instance);
51+
}
52+
}
53+
}
54+
} catch (Exception e) {
55+
throw new RuntimeException("Failed to load verb class", e);
56+
}
57+
});
58+
}
59+
}
60+
}
61+
62+
public Object execute(String verb, Map<String, Object> args) {
63+
synchronized (this) {
64+
Method method = verbs.get(verb);
65+
Object instance = instances.get(verb);
66+
if (method == null) {
67+
throw new IllegalArgumentException("Unknown verb: " + verb);
68+
}
69+
70+
try {
71+
return method.invoke(instance, args);
72+
} catch (Exception e) {
73+
throw new RuntimeException("Failed to execute verb: " + verb, e);
74+
}
75+
}
76+
}
77+
78+
public Map<String, Method> getVerbs() {
79+
return new HashMap<>(verbs);
80+
}
81+
}

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

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
package io.sentrius.agent.analysis.agents.verbs;
2+
3+
import java.io.IOException;
4+
import java.io.InputStream;
5+
import java.util.ArrayList;
6+
import java.util.List;
7+
import java.util.Map;
8+
import com.fasterxml.jackson.core.JsonProcessingException;
9+
import com.fasterxml.jackson.databind.ObjectMapper;
10+
import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
11+
import io.sentrius.agent.analysis.agents.agents.AgentConfig;
12+
import io.sentrius.agent.analysis.agents.agents.PromptBuilder;
13+
import io.sentrius.agent.analysis.agents.agents.VerbRegistry;
14+
import io.sentrius.sso.core.exceptions.ZtatException;
15+
import io.sentrius.sso.core.model.verbs.Verb;
16+
import io.sentrius.sso.core.services.agents.LLMService;
17+
import io.sentrius.sso.core.services.agents.ZeroTrustClientService;
18+
import io.sentrius.sso.genai.Message;
19+
import io.sentrius.sso.genai.model.ChatRequest;
20+
import lombok.extern.slf4j.Slf4j;
21+
import org.springframework.beans.factory.annotation.Value;
22+
import org.springframework.stereotype.Service;
23+
24+
@Service
25+
@Slf4j
26+
public class AgentVerbs {
27+
28+
29+
final ZeroTrustClientService zeroTrustClientService;
30+
final LLMService llmService;
31+
final VerbRegistry verbRegistry;
32+
@Value("${agent.ai.config:agent-config.yaml}")
33+
private String agentConfigFile;
34+
35+
final ObjectMapper mapper = new ObjectMapper(new YAMLFactory()); // jackson databind
36+
37+
public AgentVerbs(ZeroTrustClientService zeroTrustClientService, LLMService llmService, VerbRegistry verbRegistry) throws JsonProcessingException {
38+
this.zeroTrustClientService = zeroTrustClientService;
39+
this.llmService = llmService;
40+
this.verbRegistry = verbRegistry;
41+
42+
log.info("Loading agent config from {}", agentConfigFile);
43+
44+
45+
46+
}
47+
48+
49+
@Verb(name = "prompt_agent", description = "Prompts for agent workload.", isAiCallable = false)
50+
public String promptAgent(Map<String, Object> args) throws ZtatException, IOException {
51+
InputStream is = getClass().getClassLoader().getResourceAsStream("agent-config.yaml");
52+
if (is == null) {
53+
throw new RuntimeException("agent-config.yaml not found on classpath");
54+
}
55+
AgentConfig config = new ObjectMapper(new YAMLFactory()).readValue(is, AgentConfig.class);
56+
57+
log.info("Agent config loaded: {}", config);
58+
PromptBuilder promptBuilder = new PromptBuilder(verbRegistry, config);
59+
var prompt = promptBuilder.buildPrompt();
60+
List<Message> messages = new ArrayList<>();
61+
62+
messages.add(Message.builder().role("system").content(prompt).build());
63+
64+
ChatRequest chatRequest = ChatRequest.builder().model("gpt-3.5-turbo").messages(messages).build();
65+
66+
return llmService.askQuestion(chatRequest);
67+
}
68+
}

0 commit comments

Comments
 (0)