Skip to content

Commit 1f5871b

Browse files
phrockerphrocker
committed
update
1 parent 06315bf commit 1f5871b

File tree

11 files changed

+128
-20
lines changed

11 files changed

+128
-20
lines changed

‎ai-agent/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,11 @@
3737
<artifactId>sentrius-core</artifactId>
3838
<version>1.0.0-SNAPSHOT</version>
3939
</dependency>
40+
<dependency>
41+
<groupId>io.sentrius</groupId>
42+
<artifactId>provenance-core</artifactId>
43+
<version>1.0.0-SNAPSHOT</version>
44+
</dependency>
4045
<dependency>
4146
<groupId>io.sentrius</groupId>
4247
<artifactId>sentrius-llm-core</artifactId>

‎ai-agent/src/main/java/io/sentrius/agent/analysis/api/websocket/ChatWSHandler.java‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,15 @@
1515
import io.sentrius.agent.analysis.agents.verbs.TerminalVerbs;
1616
import io.sentrius.agent.analysis.api.UserCommunicationService;
1717
import io.sentrius.sso.core.exceptions.ZtatException;
18+
import io.sentrius.sso.core.services.agents.AgentClientService;
1819
import io.sentrius.sso.core.services.agents.ZeroTrustClientService;
1920
import io.sentrius.sso.genai.Message;
2021
import io.sentrius.sso.protobuf.Session;
22+
import io.sentrius.sso.provenance.ProvenanceEvent;
2123
import lombok.RequiredArgsConstructor;
2224
import lombok.extern.slf4j.Slf4j;
2325
import org.springframework.beans.factory.annotation.Autowired;
26+
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
2427
import org.springframework.stereotype.Component;
2528
import org.springframework.web.socket.CloseStatus;
2629
import org.springframework.web.socket.TextMessage;
@@ -29,6 +32,7 @@
2932

3033
@Slf4j
3134
@Component
35+
@ConditionalOnProperty(name = "agents.ai.chat.agent.enabled", havingValue = "true", matchIfMissing = false)
3236
public class ChatWSHandler extends TextWebSocketHandler {
3337

3438
final UserCommunicationService userCommunicationService;
@@ -39,15 +43,19 @@ public class ChatWSHandler extends TextWebSocketHandler {
3943

4044

4145
private final ChatAgent chatAgent;
46+
private final AgentClientService agentClientService;
4247

4348
@Autowired
4449
public ChatWSHandler(UserCommunicationService userCommunicationService, ZeroTrustClientService zeroTrustClientService,
45-
TerminalVerbs terminalVerbs, AgentVerbs agentVerbs, ChatAgent chatAgent) {
50+
TerminalVerbs terminalVerbs, AgentVerbs agentVerbs, ChatAgent chatAgent,
51+
AgentClientService agentClientService
52+
) {
4653
this.userCommunicationService = userCommunicationService;
4754
this.zeroTrustClientService = zeroTrustClientService;
4855
this.terminalVerbs = terminalVerbs;
4956
this.agentVerbs = agentVerbs;
5057
this.chatAgent = chatAgent;
58+
this.agentClientService = agentClientService;
5159
}
5260

5361
@Override
@@ -96,6 +104,17 @@ public void afterConnectionEstablished(WebSocketSession session) throws Exceptio
96104
));
97105

98106
userCommunicationService.createSession(queryParams.get("sessionId"), session);
107+
108+
109+
ProvenanceEvent provenanceEvent = ProvenanceEvent.builder()
110+
.eventType(ProvenanceEvent.EventType.USER_CHAT_AGENT)
111+
.actor(session.getPrincipal().getName())
112+
.triggeringUser(chatAgent.getAgentExecution().getUser().getName())
113+
.outputSummary("New chat session established")
114+
.sessionId(session.getId())
115+
.build();
116+
117+
agentClientService.submitProvenance(chatAgent.getAgentExecution(), provenanceEvent);
99118
}
100119

101120

‎ai-agent/src/main/java/io/sentrius/agent/analysis/api/websocket/JwtHandshakeInterceptor.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import lombok.RequiredArgsConstructor;
44
import lombok.extern.slf4j.Slf4j;
55
import org.springframework.beans.factory.annotation.Autowired;
6+
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
67
import org.springframework.http.HttpHeaders;
78
import org.springframework.http.server.ServerHttpRequest;
89
import org.springframework.http.server.ServerHttpResponse;
@@ -19,6 +20,7 @@
1920

2021
@Slf4j
2122
@Component
23+
@ConditionalOnProperty(name = "agents.ai.chat.agent.enabled", havingValue = "true", matchIfMissing = false)
2224
public class JwtHandshakeInterceptor implements HandshakeInterceptor {
2325

2426
private final JwtDecoder jwtDecoder;

‎ai-agent/src/main/java/io/sentrius/agent/analysis/api/websocket/WebSocketConfig.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import lombok.RequiredArgsConstructor;
44
import lombok.extern.slf4j.Slf4j;
55
import org.springframework.beans.factory.annotation.Value;
6+
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
67
import org.springframework.context.annotation.Configuration;
78
import org.springframework.web.socket.config.annotation.EnableWebSocket;
89
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
@@ -13,6 +14,7 @@
1314
@EnableWebSocket
1415
@RequiredArgsConstructor
1516
@Slf4j
17+
@ConditionalOnProperty(name = "agents.ai.chat.agent.enabled", havingValue = "true", matchIfMissing = false)
1618
public class WebSocketConfig implements WebSocketConfigurer {
1719

1820
@Value("${agent.listen.websocket:false}") // Default is false

‎api/src/main/java/io/sentrius/sso/controllers/api/AgentApiController.java‎

Lines changed: 52 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55
import java.security.GeneralSecurityException;
66
import java.sql.SQLException;
77
import java.time.LocalDateTime;
8-
import java.util.ArrayList;
98
import java.util.HashSet;
109
import java.util.List;
1110
import java.util.Map;
@@ -37,7 +36,6 @@
3736
import io.sentrius.sso.core.services.security.ZeroTrustAccessTokenService;
3837
import io.sentrius.sso.core.services.security.ZeroTrustRequestService;
3938
import io.sentrius.sso.core.services.terminal.SessionTrackingService;
40-
import io.sentrius.sso.protobuf.Session;
4139
import io.sentrius.sso.provenance.ProvenanceEvent;
4240
import io.sentrius.sso.provenance.kafka.ProvenanceKafkaProducer;
4341
import jakarta.servlet.http.HttpServletRequest;
@@ -134,16 +132,6 @@ public ResponseEntity<?> heartbeat(
134132
}
135133
agentService.recordHeartbeat(operatingUser.getUserId(),status.getName(), status);
136134
log.info("Heartbeat status recorded for agent: {} {}", agentId, status);
137-
ProvenanceEvent event = ProvenanceEvent.builder()
138-
.eventId(UUID.randomUUID().toString())
139-
.sessionId(status.getAgentId())
140-
.actor(operatingUser.getUsername())
141-
.triggeringUser(operatingUser.getUsername())
142-
.eventType(ProvenanceEvent.EventType.AGENT_RESPONSE)
143-
.outputSummary("Heartbeat received from agent: " + status.getName() + " with status: " + status.getStatus())
144-
.timestamp(LocalDateTime.now().toInstant(java.time.ZoneOffset.UTC))
145-
.build();
146-
provenanceKafkaProducer.send(event);
147135
return ResponseEntity.ok(Map.of("status", "success"));
148136
}
149137

@@ -198,6 +186,43 @@ public ResponseEntity<?> requestRegistration(
198186

199187

200188

189+
}
190+
191+
@PostMapping("/provenance/submit")
192+
// no LimitAccess
193+
public ResponseEntity<?> submitProvenance(
194+
@RequestHeader("Authorization") String token,
195+
HttpServletRequest request, HttpServletResponse response, @RequestBody ProvenanceEvent event) throws SQLException,
196+
GeneralSecurityException {
197+
198+
String compactJwt = token.startsWith("Bearer ") ? token.substring(7) : token;
199+
200+
201+
if (!keycloakService.validateJwt(compactJwt)) {
202+
log.warn("Invalid Keycloak token");
203+
return ResponseEntity.status(HttpStatus.SC_UNAUTHORIZED).body("Invalid Keycloak token");
204+
}
205+
206+
var operatingUser = getOperatingUser(request, response );
207+
208+
// Extract agent identity from the JWT
209+
String agentId = keycloakService.extractAgentId(compactJwt);
210+
211+
if (null == operatingUser) {
212+
log.warn("No operating user found for agent: {}", agentId);
213+
var username = keycloakService.extractUsername(compactJwt);
214+
operatingUser = userService.getUserByUsername(username);
215+
216+
}
217+
218+
provenanceKafkaProducer.send(event);
219+
220+
return ResponseEntity.ok(Map.of("status", "success", "message", "Provenance event submitted successfully"));
221+
222+
223+
224+
225+
201226
}
202227

203228

@@ -549,6 +574,21 @@ public ResponseEntity<?> getNextMessage(
549574
.build())
550575
.toList();
551576

577+
for (var comm : comms) {
578+
log.info("Found communication: {} {} {} {}", comm.getId(), comm.getSourceAgent(), comm.getTargetAgent(), comm.getPayload());
579+
ProvenanceEvent event = ProvenanceEvent.builder()
580+
.eventId(communicationId)
581+
.sessionId(communicationId)
582+
.actor(operatingUser.getUsername())
583+
.triggeringUser(comm.getTargetAgent())
584+
.eventType(ProvenanceEvent.EventType.KNOWLEDGE_GENERATED)
585+
.outputSummary("Ask agent " + comm.getPayload())
586+
.timestamp(LocalDateTime.now().toInstant(java.time.ZoneOffset.UTC))
587+
.build();
588+
provenanceKafkaProducer.send(event);
589+
}
590+
591+
552592
return ResponseEntity.ok(commsDto);
553593

554594
}

‎core/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,11 @@
2020

2121

2222
<dependencies>
23+
<dependency>
24+
<groupId>io.sentrius</groupId>
25+
<artifactId>provenance-core</artifactId>
26+
<version>1.0.0-SNAPSHOT</version>
27+
</dependency>
2328
<dependency>
2429
<groupId>org.apache.commons</groupId>
2530
<artifactId>commons-lang3</artifactId>

‎core/src/main/java/io/sentrius/sso/core/services/agents/AgentClientService.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import io.sentrius.sso.core.exceptions.ZtatException;
1919
import io.sentrius.sso.core.services.security.KeycloakService;
2020
import io.sentrius.sso.core.utils.JsonUtil;
21+
import io.sentrius.sso.provenance.ProvenanceEvent;
2122
import lombok.extern.slf4j.Slf4j;
2223
import org.springframework.beans.factory.annotation.Value;
2324
import org.springframework.stereotype.Service;
@@ -86,6 +87,16 @@ public Set<String> getCommunicationIds(AgentExecution execution, ZtatRequestDTO
8687
return Set.of();
8788
}
8889

90+
public void submitProvenance(AgentExecution execution, ProvenanceEvent event){
91+
String url = "/agent/provenance/submit";
92+
93+
try {
94+
zeroTrustClientService.callPostOnApi(execution, url, event);
95+
} catch (ZtatException e) {
96+
log.error("Failed to submit provenance event: {}", e.getMessage());
97+
}
98+
}
99+
89100
public List<AgentCommunicationDTO> getResponse(AgentExecution execution, AtatRequest atatRequest,
90101
AgentCommunicationDTO lastCommunication,
91102
long timeToWait, TimeUnit timeUnit)

‎dataplane/src/main/java/io/sentrius/sso/core/services/agents/AgentService.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -166,8 +166,8 @@ public CompletableFuture<AgentCommunication> saveCommunication(String communicat
166166
try {
167167
var eventType = switch(messageType){
168168
case "intercept" -> ProvenanceEvent.EventType.ENDPOINT_ACCESS;
169-
case "chat_request" -> ProvenanceEvent.EventType.AGENT_RESPONSE;
170-
case "interpretation_response" -> ProvenanceEvent.EventType.AGENT_RESPONSE;
169+
case "chat_request" -> ProvenanceEvent.EventType.AGENT_RESPOND;
170+
case "interpretation_response" -> ProvenanceEvent.EventType.AGENT_RESPOND;
171171
default -> ProvenanceEvent.EventType.UNKNOWN;
172172
};
173173
ProvenanceEvent event = ProvenanceEvent.builder()
@@ -203,8 +203,8 @@ public CompletableFuture<AgentCommunication> saveCommunication(AgentCommunicatio
203203
try {
204204
var eventType = switch(communication.getMessageType()){
205205
case "intercept" -> ProvenanceEvent.EventType.ENDPOINT_ACCESS;
206-
case "chat_request" -> ProvenanceEvent.EventType.AGENT_RESPONSE;
207-
case "interpretation_response" -> ProvenanceEvent.EventType.AGENT_RESPONSE;
206+
case "chat_request" -> ProvenanceEvent.EventType.AGENT_RESPOND;
207+
case "interpretation_response" -> ProvenanceEvent.EventType.INTERPRET_MESSAGE;
208208
default -> ProvenanceEvent.EventType.UNKNOWN;
209209
};
210210
ProvenanceEvent event = ProvenanceEvent.builder()

‎llm-proxy/src/main/java/io/sentrius/sso/controllers/api/OpenAIProxyController.java‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package io.sentrius.sso.controllers.api;
22

3+
import java.time.LocalDateTime;
34
import java.util.ArrayList;
45
import java.util.UUID;
56
import java.util.concurrent.ExecutionException;
@@ -153,6 +154,27 @@ public ResponseEntity<?> chat(@RequestHeader("Authorization") String token,
153154
);
154155

155156

157+
ProvenanceEvent event = ProvenanceEvent.builder()
158+
.eventId(communicationId)
159+
.sessionId(communicationId)
160+
.actor(operatingUser.getUsername())
161+
.triggeringUser("LLM")
162+
.eventType(ProvenanceEvent.EventType.KNOWLEDGE_REQUESTED)
163+
.outputSummary("prompt LLM" + chatRequest.getMessages().get(0).getContent())
164+
.timestamp(LocalDateTime.now().toInstant(java.time.ZoneOffset.UTC))
165+
.build();
166+
provenanceKafkaProducer.send(event);
167+
168+
event = ProvenanceEvent.builder()
169+
.eventId(communicationId)
170+
.sessionId(communicationId)
171+
.actor("LLM")
172+
.triggeringUser(operatingUser.getUsername())
173+
.eventType(ProvenanceEvent.EventType.KNOWLEDGE_GENERATED)
174+
.outputSummary("prompt LLM")
175+
.timestamp(LocalDateTime.now().toInstant(java.time.ZoneOffset.UTC))
176+
.build();
177+
provenanceKafkaProducer.send(event);
156178

157179

158180

@@ -248,7 +270,7 @@ public ResponseEntity<?> justify(@RequestHeader("Authorization") String token,
248270
"chat_request",
249271
rawBody
250272
);
251-
273+
252274
Span span = tracer.spanBuilder("AgentToAgentCommunication").startSpan();
253275
try (Scope scope = span.makeCurrent()) {
254276
var resp = endpoint.sample(RawConversationRequest.builder().request(chatRequest).build());

‎pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,8 @@
77
<name>sentrius</name>
88
<description>Multi-module project for sentrius</description>
99
<modules>
10-
<module>core</module>
1110
<module>provenance-core</module>
11+
<module>core</module>
1212
<module>dataplane</module>
1313
<module>sentrius-llm-core</module>
1414
<module>sentrius-llm-dataplane</module>

0 commit comments

Comments
 (0)