Skip to content

Commit dd9ee06

Browse files
phrockerphrocker
committed
fix document
1 parent fd20ebe commit dd9ee06

File tree

8 files changed

+270
-43
lines changed

8 files changed

+270
-43
lines changed

‎.azure.env‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,14 @@
1-
SENTRIUS_VERSION=1.1.100
2-
SENTRIUS_SSH_VERSION=1.1.11
3-
SENTRIUS_KEYCLOAK_VERSION=1.1.14
4-
SENTRIUS_AGENT_VERSION=1.1.23
5-
SENTRIUS_AI_AGENT_VERSION=1.1.4
6-
LLMPROXY_VERSION=1.1.4
7-
LAUNCHER_VERSION=1.1.5
8-
AGENTPROXY_VERSION=1.1.5
9-
SSHPROXY_VERSION=1.1.4
10-
RDPPROXY_VERSION=1.1.4
11-
GITHUB_MCP_VERSION=1.1.4
12-
PROMPT_ADVISOR_VERSION=1.1.7
13-
MONITORING_AGENT_VERSION=1.1.22
14-
SSH_AGENT_VERSION=1.1.4
1+
SENTRIUS_VERSION=1.1.104
2+
SENTRIUS_SSH_VERSION=1.1.12
3+
SENTRIUS_KEYCLOAK_VERSION=1.1.15
4+
SENTRIUS_AGENT_VERSION=1.1.24
5+
SENTRIUS_AI_AGENT_VERSION=1.1.5
6+
LLMPROXY_VERSION=1.1.5
7+
LAUNCHER_VERSION=1.1.6
8+
AGENTPROXY_VERSION=1.1.6
9+
SSHPROXY_VERSION=1.1.5
10+
RDPPROXY_VERSION=1.1.5
11+
GITHUB_MCP_VERSION=1.1.5
12+
PROMPT_ADVISOR_VERSION=1.1.8
13+
MONITORING_AGENT_VERSION=1.1.23
14+
SSH_AGENT_VERSION=1.1.5

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

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import java.util.Map;
88
import java.util.Optional;
99
import java.util.Set;
10+
import java.util.UUID;
1011
import java.util.stream.Collectors;
1112
import io.sentrius.sso.core.annotations.LimitAccess;
1213
import io.sentrius.sso.core.config.SystemOptions;
@@ -319,7 +320,7 @@ public ResponseEntity<Map<String, Object>> generateRule(
319320
// Create token for LLM service
320321
// Note: Using empty builder as LLM service will use system authentication
321322
// If LLM requires user-specific tokens, this would need to be enhanced
322-
TokenDTO token = TokenDTO.builder().build();
323+
TokenDTO token = TokenDTO.builder().communicationId(UUID.randomUUID().toString()).build();
323324

324325
var authToken = integrationSecurityTokenService
325326
.selectToken(systemOptions.getDefaultLlmProvider())
@@ -334,12 +335,27 @@ public ResponseEntity<Map<String, Object>> generateRule(
334335

335336
// Parse the response
336337
JsonNode responseJson = objectMapper.readTree(llmResponse);
337-
assistantMessage = responseJson.get("choices")
338-
.get(0)
339-
.get("message")
340-
.get("content")
341-
.asText();
342-
338+
339+
// Handle both old format (choices) and new format (output)
340+
JsonNode outputNode = responseJson.get("output");
341+
if (outputNode != null && outputNode.isArray() && outputNode.size() > 0) {
342+
// New format: output array
343+
JsonNode messageNode = outputNode.get(0);
344+
JsonNode contentArray = messageNode.get("content");
345+
if (contentArray != null && contentArray.isArray() && contentArray.size() > 0) {
346+
assistantMessage = contentArray.get(0).get("text").asText();
347+
}
348+
} else {
349+
// Old format: choices array
350+
JsonNode choicesNode = responseJson.get("choices");
351+
if (choicesNode != null && choicesNode.isArray() && choicesNode.size() > 0) {
352+
assistantMessage = choicesNode.get(0)
353+
.get("message")
354+
.get("content")
355+
.asText();
356+
}
357+
}
358+
343359
// Try to extract rule configuration from the response
344360
ruleConfig = extractRuleConfig(assistantMessage);
345361
} catch (io.sentrius.sso.core.exceptions.ZtatException ztatEx) {

‎api/src/main/resources/static/js/rules.js‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -118,7 +118,7 @@ function editRule(ruleId) {
118118

119119
function deleteRule(ruleId) {
120120
const csrfToken = document.getElementById("assignCsrf").value;
121-
var url = '/api/v1/zerotrust/rules/delete/' + ruleId
121+
const url = '/api/v1/zerotrust/rules/delete/' + ruleId;
122122
fetch(url, {
123123
method: "DELETE",
124124
headers: {
@@ -254,7 +254,7 @@ $(document).ready(function () {
254254

255255

256256
(async () => {
257-
url = "/api/v1/zerotrust/rules/save";
257+
let url = "/api/v1/zerotrust/rules/save";
258258
const csrfToken = document.getElementById('csrf-token').value; // Get CSRF token value
259259
const payload = {
260260
ruleName: ruleName,
@@ -277,7 +277,7 @@ $(document).ready(function () {
277277
return;
278278
} else {
279279
(async () => {
280-
url = "/api/v1/zerotrust/rules/save";
280+
let url = "/api/v1/zerotrust/rules/save";
281281
const csrfToken = document.getElementById('csrf-token').value; // Get CSRF token value
282282
const payload = {
283283
ruleName: ruleName,

‎api/src/main/resources/templates/fragments/header.html‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -340,6 +340,7 @@
340340
const submenuPatterns = {
341341
'infrastructureSubmenu': ['/sso/v1/enclaves/', '/sso/v1/integrations'],
342342
'securitySubmenu': ['/sso/v1/zerotrust/', '/sso/v1/atpl/', '/sso/trust-scores', '/sso/v1/attributes/'],
343+
'dataSubmenu': ['/sso/v1/data/sources/'],
343344
'aiSubmenu': ['/sso/v1/users/', '/sso/v1/agent/'],
344345
'systemSubmenu': ['/sso/v1/system/', '/sso/v1/telemetry', '/sso/v1/automation/', '/sso/v1/k8s/']
345346
};

‎api/src/main/resources/templates/sso/data/sources.html‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -66,13 +66,13 @@ <h2 class="page-header">
6666
<!-- Search and Filters -->
6767
<div class="search-container">
6868
<div class="row">
69-
<div class="col-md-8 filter-section">
69+
<div class="col-md-6 filter-section">
7070
<label for="searchQuery" class="form-label">
7171
<i class="fas fa-search"></i> Search Data Sources
7272
</label>
7373
<input type="text" class="form-control" id="searchQuery" placeholder="Search by name, content, or tags...">
7474
</div>
75-
<div class="col-md-4 filter-section">
75+
<div class="col-md-3 filter-section">
7676
<label for="sourceTypeFilter" class="form-label">
7777
<i class="fas fa-filter"></i> Source Type
7878
</label>
@@ -88,6 +88,12 @@ <h2 class="page-header">
8888
<option value="OTHER">Other</option>
8989
</select>
9090
</div>
91+
<div class="col-md-3 filter-section">
92+
<label for="markingsFilter" class="form-label">
93+
<i class="fas fa-shield-alt"></i> Markings
94+
</label>
95+
<input type="text" class="form-control" id="markingsFilter" placeholder="Filter by markings...">
96+
</div>
9197
</div>
9298
<div class="row mt-3">
9399
<div class="col-md-12">
@@ -218,6 +224,7 @@ <h5 class="modal-title" id="retrieveExternalModalLabel">Retrieve External Data S
218224
document.getElementById('clearBtn').addEventListener('click', function() {
219225
document.getElementById('searchQuery').value = '';
220226
document.getElementById('sourceTypeFilter').value = '';
227+
document.getElementById('markingsFilter').value = '';
221228
searchSources();
222229
});
223230

@@ -260,10 +267,12 @@ <h5 class="modal-title" id="retrieveExternalModalLabel">Retrieve External Data S
260267
function searchSources() {
261268
const query = document.getElementById('searchQuery').value;
262269
const sourceType = document.getElementById('sourceTypeFilter').value;
270+
const markings = document.getElementById('markingsFilter').value;
263271

264272
const searchRequest = {
265273
query: query || null,
266-
documentType: sourceType || null
274+
documentType: sourceType || null,
275+
markings: markings || null
267276
};
268277

269278
fetch('/api/v1/documents/search', {

‎dataplane/src/main/java/io/sentrius/sso/core/services/documents/DocumentService.java‎

Lines changed: 52 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -126,10 +126,11 @@ public Optional<Document> getDocumentByName(String documentName) {
126126
* Search documents using hybrid text and vector search
127127
*/
128128
public List<Document> searchDocuments(DocumentSearchDTO searchDTO) {
129-
log.info("Searching documents with query: {}", searchDTO.getQuery());
129+
log.info("Searching documents with query: {}, type: {}, markings: {}",
130+
searchDTO.getQuery(), searchDTO.getDocumentType(), searchDTO.getMarkings());
130131

131132
if (searchDTO.getQuery() == null || searchDTO.getQuery().trim().isEmpty()) {
132-
return getAllDocuments(searchDTO.getPage(), searchDTO.getSize());
133+
return getAllDocuments(searchDTO);
133134
}
134135

135136
if (!searchDTO.isUseSemanticSearch() || embeddingService == null || !embeddingService.isAvailable()) {
@@ -333,17 +334,7 @@ private List<Document> textSearchDocuments(DocumentSearchDTO searchDTO) {
333334
List<Document> results = documentRepository.searchByContent(searchDTO.getQuery());
334335

335336
// Apply filters
336-
if (searchDTO.getDocumentType() != null) {
337-
results = results.stream()
338-
.filter(d -> d.getDocumentType().equals(searchDTO.getDocumentType()))
339-
.collect(Collectors.toList());
340-
}
341-
342-
if (searchDTO.getTags() != null && searchDTO.getTags().length > 0) {
343-
results = results.stream()
344-
.filter(d -> containsAnyTag(d, searchDTO.getTags()))
345-
.collect(Collectors.toList());
346-
}
337+
results = applySearchFilters(results, searchDTO);
347338

348339
// Apply limit
349340
if (searchDTO.getLimit() != null && searchDTO.getLimit() > 0) {
@@ -380,6 +371,10 @@ private List<Document> hybridSearchDocuments(DocumentSearchDTO searchDTO) {
380371
vectorResults = documentRepository.findSimilarDocuments(embeddingString, limit * 2);
381372
}
382373

374+
// Apply filters to both text and vector results
375+
textResults = applySearchFilters(textResults, searchDTO);
376+
vectorResults = applySearchFilters(vectorResults, searchDTO);
377+
383378
// Score and combine results
384379
Map<Long, Double> scores = new HashMap<>();
385380

@@ -415,10 +410,51 @@ private List<Document> hybridSearchDocuments(DocumentSearchDTO searchDTO) {
415410
}
416411
}
417412

418-
private List<Document> getAllDocuments(int page, int size) {
419-
Pageable pageable = PageRequest.of(page, size, Sort.by(Sort.Direction.DESC, "createdAt"));
413+
private List<Document> getAllDocuments(DocumentSearchDTO searchDTO) {
414+
Pageable pageable = PageRequest.of(searchDTO.getPage(), searchDTO.getSize(),
415+
Sort.by(Sort.Direction.DESC, "createdAt"));
420416
Page<Document> documentPage = documentRepository.findAll(pageable);
421-
return documentPage.getContent();
417+
List<Document> results = documentPage.getContent();
418+
419+
// Apply filters even when no query is provided
420+
return applySearchFilters(results, searchDTO);
421+
}
422+
423+
/**
424+
* Apply common filters to search results
425+
*/
426+
private List<Document> applySearchFilters(List<Document> results, DocumentSearchDTO searchDTO) {
427+
// Filter by document type
428+
if (searchDTO.getDocumentType() != null && !searchDTO.getDocumentType().trim().isEmpty()) {
429+
results = results.stream()
430+
.filter(d -> d.getDocumentType().equals(searchDTO.getDocumentType()))
431+
.collect(Collectors.toList());
432+
}
433+
434+
// Filter by tags
435+
if (searchDTO.getTags() != null && searchDTO.getTags().length > 0) {
436+
results = results.stream()
437+
.filter(d -> containsAnyTag(d, searchDTO.getTags()))
438+
.collect(Collectors.toList());
439+
}
440+
441+
// Filter by classification
442+
if (searchDTO.getClassification() != null && !searchDTO.getClassification().trim().isEmpty()) {
443+
results = results.stream()
444+
.filter(d -> d.getClassification() != null &&
445+
d.getClassification().equals(searchDTO.getClassification()))
446+
.collect(Collectors.toList());
447+
}
448+
449+
// Filter by markings
450+
if (searchDTO.getMarkings() != null && !searchDTO.getMarkings().trim().isEmpty()) {
451+
results = results.stream()
452+
.filter(d -> d.getMarkings() != null &&
453+
d.getMarkings().contains(searchDTO.getMarkings()))
454+
.collect(Collectors.toList());
455+
}
456+
457+
return results;
422458
}
423459

424460
private boolean containsAnyTag(Document document, String[] tags) {

‎dataplane/src/test/java/io/sentrius/sso/core/services/documents/DocumentServiceTest.java‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import org.junit.jupiter.api.extension.ExtendWith;
1212
import org.mockito.Mock;
1313
import org.mockito.junit.jupiter.MockitoExtension;
14+
import org.springframework.data.domain.Pageable;
1415

1516
import java.util.*;
1617

@@ -181,6 +182,111 @@ void testSearchDocuments_TextOnly() {
181182
verify(documentRepository).searchByContent("test query");
182183
}
183184

185+
@Test
186+
void testSearchDocuments_WithMarkingsFilter() {
187+
// Arrange
188+
DocumentSearchDTO searchDTO = DocumentSearchDTO.builder()
189+
.query("test query")
190+
.markings("ABC")
191+
.useSemanticSearch(false)
192+
.limit(10)
193+
.build();
194+
195+
Document doc1 = Document.builder().id(1L).documentName("Doc1").markings("ABC//DEF").build();
196+
Document doc2 = Document.builder().id(2L).documentName("Doc2").markings("XYZ").build();
197+
Document doc3 = Document.builder().id(3L).documentName("Doc3").markings("ABC").build();
198+
List<Document> allResults = Arrays.asList(doc1, doc2, doc3);
199+
200+
when(documentRepository.searchByContent("test query")).thenReturn(allResults);
201+
202+
// Act
203+
List<Document> results = documentService.searchDocuments(searchDTO);
204+
205+
// Assert
206+
assertEquals(2, results.size());
207+
assertTrue(results.stream().anyMatch(d -> d.getId().equals(1L)));
208+
assertTrue(results.stream().anyMatch(d -> d.getId().equals(3L)));
209+
assertFalse(results.stream().anyMatch(d -> d.getId().equals(2L)));
210+
verify(documentRepository).searchByContent("test query");
211+
}
212+
213+
@Test
214+
void testSearchDocuments_WithDocumentTypeFilter() {
215+
// Arrange
216+
DocumentSearchDTO searchDTO = DocumentSearchDTO.builder()
217+
.query("test query")
218+
.documentType("TSG")
219+
.useSemanticSearch(false)
220+
.limit(10)
221+
.build();
222+
223+
Document doc1 = Document.builder().id(1L).documentName("Doc1").documentType("TSG").build();
224+
Document doc2 = Document.builder().id(2L).documentName("Doc2").documentType("MANUAL").build();
225+
List<Document> allResults = Arrays.asList(doc1, doc2);
226+
227+
when(documentRepository.searchByContent("test query")).thenReturn(allResults);
228+
229+
// Act
230+
List<Document> results = documentService.searchDocuments(searchDTO);
231+
232+
// Assert
233+
assertEquals(1, results.size());
234+
assertEquals("TSG", results.get(0).getDocumentType());
235+
verify(documentRepository).searchByContent("test query");
236+
}
237+
238+
@Test
239+
void testSearchDocuments_EmptyQueryWithFilters() {
240+
// Arrange
241+
DocumentSearchDTO searchDTO = DocumentSearchDTO.builder()
242+
.query("")
243+
.documentType("TSG")
244+
.markings("ABC")
245+
.page(0)
246+
.size(20)
247+
.build();
248+
249+
Document doc1 = Document.builder().id(1L).documentType("TSG").markings("ABC").build();
250+
Document doc2 = Document.builder().id(2L).documentType("TSG").markings("XYZ").build();
251+
Document doc3 = Document.builder().id(3L).documentType("MANUAL").markings("ABC").build();
252+
253+
when(documentRepository.findAll(any(Pageable.class)))
254+
.thenReturn(new org.springframework.data.domain.PageImpl<>(Arrays.asList(doc1, doc2, doc3)));
255+
256+
// Act
257+
List<Document> results = documentService.searchDocuments(searchDTO);
258+
259+
// Assert
260+
assertEquals(1, results.size());
261+
assertEquals(1L, results.get(0).getId());
262+
assertEquals("TSG", results.get(0).getDocumentType());
263+
assertEquals("ABC", results.get(0).getMarkings());
264+
}
265+
266+
@Test
267+
void testSearchDocuments_WithClassificationFilter() {
268+
// Arrange
269+
DocumentSearchDTO searchDTO = DocumentSearchDTO.builder()
270+
.query("test query")
271+
.classification("UNCLASSIFIED")
272+
.useSemanticSearch(false)
273+
.build();
274+
275+
Document doc1 = Document.builder().id(1L).documentName("Doc1").classification("UNCLASSIFIED").build();
276+
Document doc2 = Document.builder().id(2L).documentName("Doc2").classification("SECRET").build();
277+
List<Document> allResults = Arrays.asList(doc1, doc2);
278+
279+
when(documentRepository.searchByContent("test query")).thenReturn(allResults);
280+
281+
// Act
282+
List<Document> results = documentService.searchDocuments(searchDTO);
283+
284+
// Assert
285+
assertEquals(1, results.size());
286+
assertEquals("UNCLASSIFIED", results.get(0).getClassification());
287+
verify(documentRepository).searchByContent("test query");
288+
}
289+
184290
@Test
185291
void testGetDocumentsByType() {
186292
// Arrange

0 commit comments

Comments
 (0)