Skip to content

Add Agent Memory Search Capability with Access Controls (Vector Store Branch) #103

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Copilot wants to merge 6 commits into
base: vector-store
Choose a base branch
from
Draft

Add Agent Memory Search Capability with Access Controls (Vector Store Branch) #103

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e7f5c62
Initial plan
Copilot Sep 1, 2025
1012c18
Add Agent Memory Search API and UI functionality
Copilot Sep 1, 2025
58350e1
Add unit tests for AgentMemorySearchService
Copilot Sep 1, 2025
fd96600
Initial plan
phrocker Aug 29, 2025
d5c4e87
Remove test
phrocker Sep 1, 2025
2a91350
Initial plan
phrocker Aug 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .local.env
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,4 +6,4 @@ SENTRIUS_AI_AGENT_VERSION=1.1.264
LLMPROXY_VERSION=1.0.78
LAUNCHER_VERSION=1.0.82
AGENTPROXY_VERSION=1.0.85
SSHPROXY_VERSION=1.0.87
SSHPROXY_VERSION=1.0.88
2 changes: 1 addition & 1 deletion .local.env.bak
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,4 +6,4 @@ SENTRIUS_AI_AGENT_VERSION=1.1.264
LLMPROXY_VERSION=1.0.78
LAUNCHER_VERSION=1.0.82
AGENTPROXY_VERSION=1.0.85
SSHPROXY_VERSION=1.0.87
SSHPROXY_VERSION=1.0.87
1 change: 1 addition & 0 deletions agent-launcher/src/main/java/io/sentrius/agent/launcher/service/PodLauncherService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,7 @@ public V1Pod launchAgentPod(AgentRegistrationDTO agent) throws Exception {
List<String> argList = new ArrayList<>();
argList.add("--spring.config.location=file:/config/agent.properties");
argList.add("--agent.namePrefix=" + agentId);
argList.add("--agent.type=" + agent.getAgentType());
argList.add("--agent.clientId=" + agent.getClientId());
argList.add("--agent.listen.websocket=true");
argList.add("--agent.callback.url=" + constructedCallbackUrl);
Expand Down
2 changes: 1 addition & 1 deletion agent-launcher/src/main/resources/application.properties
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,4 +8,4 @@ spring.main.web-application-type=servlet
spring.thymeleaf.enabled=true
spring.freemarker.enabled=false

sentrius.agent.registry=local
sentrius.agent.registry=local
70 changes: 70 additions & 0 deletions ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/BaseEnterpriseAgent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
package io.sentrius.agent.analysis.agents.agents;

import java.util.concurrent.TimeUnit;
import com.fasterxml.jackson.databind.node.ArrayNode;
import io.sentrius.agent.analysis.agents.verbs.AgentVerbs;
import io.sentrius.sso.core.dto.agents.AgentExecution;
import io.sentrius.sso.core.dto.agents.AgentExecutionContextDTO;
import io.sentrius.sso.core.dto.ztat.ZtatRequestDTO;
import io.sentrius.sso.core.exceptions.ZtatException;
import io.sentrius.sso.core.services.agents.AgentClientService;
import io.sentrius.sso.core.services.agents.ZeroTrustClientService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.context.event.ApplicationReadyEvent;
import org.springframework.context.ApplicationListener;

@Slf4j
public abstract class BaseEnterpriseAgent implements ApplicationListener<ApplicationReadyEvent> {

@Autowired
protected final AgentVerbs agentVerbs;
@Autowired
protected final ZeroTrustClientService zeroTrustClientService;
@Autowired
protected final AgentClientService agentClientService;
@Autowired
protected final VerbRegistry verbRegistry;

protected BaseEnterpriseAgent(
AgentVerbs agentVerbs, ZeroTrustClientService zeroTrustClientService, AgentClientService agentClientService,
VerbRegistry verbRegistry
) {
this.agentVerbs = agentVerbs;
this.zeroTrustClientService = zeroTrustClientService;
this.agentClientService = agentClientService;
this.verbRegistry = verbRegistry;
}


protected ArrayNode promptAgent(AgentExecution execution) throws ZtatException {
return promptAgent(execution);
}

protected ArrayNode promptAgent(AgentExecution execution, AgentExecutionContextDTO contextDTO) throws ZtatException {
while(true){
try {
log.info("Prompting agent...");
return agentVerbs.promptAgent(execution,contextDTO);
} catch (ZtatException e) {
log.info("Mechanisms {}" , e.getMechanisms());
var endpoint = zeroTrustClientService.createEndPointRequest("prompt_agent", e.getEndpoint());
ZtatRequestDTO ztatRequestDTO = ZtatRequestDTO.builder()
.user(execution.getUser())
.command(endpoint.toString())
.justification("Registered Agent requires ability to prompt LLM endpoints to begin operations")
.summary("Registered Agent requires ability to prompt LLM endpoints to begin operations")
.build();
var request = zeroTrustClientService.requestZtatToken(execution, execution.getUser(),ztatRequestDTO);

var token = zeroTrustClientService.awaitZtatToken(execution, execution.getUser(), request, 60,
TimeUnit.MINUTES);
execution.setZtatToken(token);
} catch (Exception e) {
log.error(e.getMessage());
throw new RuntimeException(e);
}
}
}
}
157 changes: 125 additions & 32 deletions ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/ChatAgent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,77 +1,76 @@
package io.sentrius.agent.analysis.agents.agents;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.node.ArrayNode;
import io.sentrius.agent.analysis.agents.verbs.AgentVerbs;
import io.sentrius.agent.analysis.agents.verbs.ChatVerbs;
import io.sentrius.agent.analysis.api.AgentKeyService;
import io.sentrius.agent.analysis.api.UserCommunicationService;
import io.sentrius.agent.analysis.model.LLMResponse;
import io.sentrius.agent.config.AgentConfigOptions;
import io.sentrius.sso.core.dto.UserDTO;
import io.sentrius.sso.core.dto.agents.AgentExecution;
import io.sentrius.sso.core.dto.ztat.ZtatRequestDTO;
import io.sentrius.sso.core.dto.agents.AgentExecutionContextDTO;
import io.sentrius.sso.core.exceptions.ZtatException;
import io.sentrius.sso.core.model.security.Ztat;
import io.sentrius.sso.core.model.verbs.VerbResponse;
import io.sentrius.sso.core.services.agents.AgentClientService;
import io.sentrius.sso.core.services.agents.AgentExecutionService;
import io.sentrius.sso.core.services.agents.ZeroTrustClientService;
import io.sentrius.sso.core.services.security.KeycloakService;
import io.sentrius.sso.core.utils.JsonUtil;
import io.sentrius.sso.genai.Message;
import jakarta.annotation.PreDestroy;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.event.ApplicationReadyEvent;
import org.springframework.context.ApplicationListener;
import org.springframework.stereotype.Component;

@Slf4j
@Component
@RequiredArgsConstructor
@ConditionalOnProperty(name = "agents.ai.chat.agent.enabled", havingValue = "true", matchIfMissing = false)
public class ChatAgent implements ApplicationListener<ApplicationReadyEvent> {
public class ChatAgent extends BaseEnterpriseAgent {


final ZeroTrustClientService zeroTrustClientService;
final AgentClientService agentClientService;
final VerbRegistry verbRegistry;
final AgentVerbs agentVerbs;
final AgentExecutionService agentExecutionService;
final UserCommunicationService userCommunicationService;
final AgentConfigOptions agentConfigOptions;
final AgentKeyService agentKeyService;
private final KeycloakService keycloakService;
final ChatVerbs chatVerbs;

private volatile boolean running = true;
private Thread workerThread;

private AgentExecution agentExecution;

public ArrayNode promptAgent(AgentExecution execution) throws ZtatException {
while(true){
try {
log.info("Prompting agent...");
return agentVerbs.promptAgent(execution,null);
} catch (ZtatException e) {
log.info("Mechanisms {}" , e.getMechanisms());
var endpoint = zeroTrustClientService.createEndPointRequest("prompt_agent", e.getEndpoint());
ZtatRequestDTO ztatRequestDTO = ZtatRequestDTO.builder()
.user(execution.getUser())
.command(endpoint.toString())
.justification("Registered Agent requires ability to prompt LLM endpoints to begin operations")
.summary("Registered Agent requires ability to prompt LLM endpoints to begin operations")
.build();
var request = zeroTrustClientService.requestZtatToken(execution, execution.getUser(),ztatRequestDTO);

var token = zeroTrustClientService.awaitZtatToken(execution, execution.getUser(), request, 60,
TimeUnit.MINUTES);
execution.setZtatToken(token);
} catch (Exception e) {
log.error(e.getMessage());
throw new RuntimeException(e);
}
}

@Autowired
public ChatAgent(
AgentVerbs agentVerbs, ZeroTrustClientService zeroTrustClientService, AgentClientService agentClientService,
VerbRegistry verbRegistry, AgentExecutionService agentExecutionService, UserCommunicationService userCommunicationService,
AgentConfigOptions agentConfigOptions, AgentKeyService agentKeyService, KeycloakService keycloakService,
ChatVerbs chatVerbs
) {
super(agentVerbs, zeroTrustClientService, agentClientService, verbRegistry);
this.zeroTrustClientService = zeroTrustClientService;
this.agentClientService = agentClientService;
this.verbRegistry = verbRegistry;
this.agentExecutionService = agentExecutionService;
this.userCommunicationService = userCommunicationService;
this.agentConfigOptions = agentConfigOptions;
this.agentKeyService = agentKeyService;
this.keycloakService = keycloakService;
this.chatVerbs = chatVerbs;
}

@Override
Expand Down Expand Up @@ -146,15 +145,109 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {

int allowedFailures = 20;
log.info("Agent Registered...");
AgentExecutionContextDTO agentExecutionContext = AgentExecutionContextDTO.builder().build();
agentExecutionService.setExecutionContextDTO(agentExecution, agentExecutionContext);
LLMResponse response = null;
AgentConfig config = null;
try {
config = chatVerbs.getAgentConfig(agentExecution);
} catch (IOException e) {
throw new RuntimeException(e);
} catch (ZtatException e) {
throw new RuntimeException(e);
}
PromptBuilder promptBuilder = new PromptBuilder(verbRegistry, config);
var prompt = promptBuilder.buildPrompt(false);
try {
if (agentConfigOptions.getType().equalsIgnoreCase("chat-autonomous")) {

response = chatVerbs.promptAgent(agentExecution, agentExecutionContext, prompt);
}
} catch (ZtatException e) {
throw new RuntimeException(e);
} catch (IOException e) {
throw new RuntimeException(e);
}


if (agentConfigOptions.getType().equalsIgnoreCase("chat-autonomous") && response == null) {
log.error("Chat autonomous agent mode enabled but no response received from promptAgent, shutting down...");
throw new RuntimeException("Chat autonomous agent mode enabled but no response received from promptAgent");
}
VerbResponse lastVerbResponse = null;
LLMResponse nextResponse = null;
List<VerbResponse> verbResponses = new ArrayList<>();
while(running) {


try {

Thread.sleep(5_000);
agentClientService.heartbeat(agentExecution, agentExecution.getUser().getUsername());
if (agentConfigOptions.getType().equalsIgnoreCase("chat-autonomous")) {
log.info("Chat autonomous agent mode enabled, executing workload...");
VerbResponse priorResponse = null;
Map<String, Object> args = new HashMap<>();

var arguments = response.getArguments();
if (null != response) {
if (response.getNextOperation() != null && !response.getNextOperation().isEmpty()) {
var executionResponse = verbRegistry.execute(
agentExecution,
agentExecutionContext,
lastVerbResponse,
response.getNextOperation(), arguments
);
verbResponses.add(executionResponse);
lastVerbResponse = executionResponse;


// chatAgent.getAgentExecution().addMessages(Message.builder().role("System")
// .content("System executed operation: " + response.getNextOperation()).build());
var responses = agentExecutionContext.getAgentDataList();
var planResponse =
responses.isEmpty() ? "" :
responses.get(responses.size() - 1).asText();
nextResponse = chatVerbs.interpret_plan_response(
agentExecution,
agentExecutionContext,
verbRegistry.getVerbs().get(response.getNextOperation()),
planResponse
);

var memory = agentExecutionContext.flushPersistentMemory();
if (memory != null) {
for(var memoryEntry : memory.entrySet()){
agentClientService.storeMemory(agentExecution,
agentExecutionContext.getAgentContext().getName(),
io.sentrius.sso.core.dto.agents.AgentMemoryDTO.builder()
.agentName(agentExecutionContext.getAgentContext().getName())
.memoryKey(memoryEntry.getKey())
.memoryValue(memoryEntry.getValue().toString())
.build());
}
}


response = nextResponse;
}

}else {
response = chatVerbs.promptAgent(agentExecution, agentExecutionContext, prompt);

}

continue;
}
allowedFailures = 20; // Reset allowed failures on successful heartbeat
} catch (ZtatException | Exception ex) {
agentExecutionContext.addMessages(Message.builder().role("system").content(
"You caused the following error. Please re-validate you chose the right operations or " +
"endpoints for the context" +
ex.getMessage()).build());


ex.printStackTrace();
if (allowedFailures-- <= 0) {
log.error("Failed to heartbeat agent after multiple attempts, shutting down...");
throw new RuntimeException(ex);
Expand Down
35 changes: 29 additions & 6 deletions ai-agent/src/main/java/io/sentrius/agent/analysis/agents/agents/RegisteredAgent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import io.sentrius.agent.config.AgentConfigOptions;
import io.sentrius.sso.core.dto.agents.AgentExecution;
import io.sentrius.sso.core.dto.agents.AgentExecutionContextDTO;
import io.sentrius.sso.core.dto.agents.AgentMemoryDTO;
import io.sentrius.sso.core.dto.ztat.ZtatRequestDTO;
import io.sentrius.sso.core.exceptions.ZtatException;
import io.sentrius.sso.core.model.security.Ztat;
Expand All @@ -22,22 +23,18 @@
import jakarta.annotation.PreDestroy;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.event.ApplicationReadyEvent;
import org.springframework.context.ApplicationListener;
import org.springframework.stereotype.Component;

@Slf4j
@Component
@RequiredArgsConstructor
@ConditionalOnProperty(name = "agents.ai.registered.agent.enabled", havingValue = "true", matchIfMissing = false)
public class RegisteredAgent implements ApplicationListener<ApplicationReadyEvent> {
public class RegisteredAgent extends BaseEnterpriseAgent {


final ZeroTrustClientService zeroTrustClientService;
final AgentClientService agentClientService;
final VerbRegistry verbRegistry;
final AgentVerbs agentVerbs;
final AgentExecutionService agentExecutionService;
final AgentConfigOptions agentConfigOptions;
final AgentKeyService agentKeyService;
Expand All @@ -46,6 +43,19 @@ public class RegisteredAgent implements ApplicationListener<ApplicationReadyEven
private volatile boolean running = true;
private Thread workerThread;

@Autowired
public RegisteredAgent(
AgentVerbs agentVerbs, ZeroTrustClientService zeroTrustClientService, AgentClientService agentClientService,
VerbRegistry verbRegistry, AgentExecutionService agentExecutionService, AgentConfigOptions agentConfigOptions,
AgentKeyService agentKeyService, KeycloakService keycloakService
) {
super(agentVerbs, zeroTrustClientService, agentClientService, verbRegistry);
this.agentExecutionService = agentExecutionService;
this.agentConfigOptions = agentConfigOptions;
this.agentKeyService = agentKeyService;
this.keycloakService = keycloakService;
}

public ArrayNode promptAgent(AgentExecution execution) throws ZtatException {
while(true){
try {
Expand Down Expand Up @@ -136,6 +146,19 @@ public void onApplicationEvent(final ApplicationReadyEvent event) {
priorResponse = verbRegistry.execute(agentExecution,agentExecutionContext,
priorResponse, verb, args);
}
var memoryList = agentExecutionContext.flushPersistentMemory();
if (memoryList != null) {
for(var memory : memoryList.entrySet()){
AgentMemoryDTO dto = AgentMemoryDTO.builder()
.agentName(agentExecutionContext.getAgentContext().getName())
.memoryKey(memory.getKey())
.memoryValue(memory.getValue().toString())
.build();
agentClientService.storeMemory(agentExecution,
agentExecutionContext.getAgentContext().getName(), dto);
}
}

log.info("Node: {}", node);
}

Expand Down
Loading