|
10 | 10 | import com.dbaagent.service.SemanticModelService; |
11 | 11 | import com.dbaagent.service.TrainingJobService; |
12 | 12 | import com.dbaagent.service.TrainingService; |
| 13 | +import com.dbaagent.service.UserDataAccessPolicyService; |
13 | 14 | import com.dbaagent.service.security.AccessControlService; |
14 | 15 | import lombok.Data; |
15 | 16 | import lombok.RequiredArgsConstructor; |
@@ -37,6 +38,7 @@ public class TrainingController { |
37 | 38 | private final CredentialRepository credentialRepository; |
38 | 39 | private final SemanticModelService semanticModelService; |
39 | 40 | private final AccessControlService accessControlService; |
| 41 | + private final UserDataAccessPolicyService userDataAccessPolicyService; |
40 | 42 |
|
41 | 43 | /** |
42 | 44 | * Train with schema DDL |
@@ -186,7 +188,9 @@ public ResponseEntity<Map<String, Object>> debugRetrieve( |
186 | 188 | if (question == null || question.isBlank()) { |
187 | 189 | return ResponseEntity.badRequest().body(Map.of("error", "Query parameter 'q' is required")); |
188 | 190 | } |
189 | | - return ResponseEntity.ok(trainingService.debugRetrieve(connectionId, question, topK)); |
| 191 | + Map<String, Object> payload = trainingService.debugRetrieve(connectionId, question, topK); |
| 192 | + filterDebugRetrieval(connectionId, payload); |
| 193 | + return ResponseEntity.ok(payload); |
190 | 194 | } catch (org.springframework.web.server.ResponseStatusException e) { |
191 | 195 | throw e; |
192 | 196 | } catch (Exception e) { |
@@ -346,6 +350,34 @@ private void rebuildAndReindexConnection(String connectionId) { |
346 | 350 | trainingService.reindexConnection(connectionId); |
347 | 351 | } |
348 | 352 |
|
| 353 | + @SuppressWarnings("unchecked") |
| 354 | + private void filterDebugRetrieval(String connectionId, Map<String, Object> payload) { |
| 355 | + if (payload == null) { |
| 356 | + return; |
| 357 | + } |
| 358 | + Object results = payload.get("results"); |
| 359 | + if (!(results instanceof List<?> rows)) { |
| 360 | + return; |
| 361 | + } |
| 362 | + List<Map<String, Object>> filtered = new java.util.ArrayList<>(); |
| 363 | + for (Object row : rows) { |
| 364 | + if (!(row instanceof Map<?, ?> map)) { |
| 365 | + continue; |
| 366 | + } |
| 367 | + Object metadata = map.get("metadata"); |
| 368 | + if (userDataAccessPolicyService.isRagMetadataInScope( |
| 369 | + connectionId, |
| 370 | + accessControlService.getCurrentUsername(), |
| 371 | + accessControlService.isCurrentUserAdmin(), |
| 372 | + metadata == null ? null : String.valueOf(metadata) |
| 373 | + )) { |
| 374 | + filtered.add((Map<String, Object>) map); |
| 375 | + } |
| 376 | + } |
| 377 | + payload.put("results", filtered); |
| 378 | + payload.put("resultCount", filtered.size()); |
| 379 | + } |
| 380 | + |
349 | 381 | @Data |
350 | 382 | public static class DocumentationRequest { |
351 | 383 | private String connectionId; |
|
0 commit comments