prompt_scanner.go 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122
  1. package securityaudit
  2. import (
  3. "errors"
  4. "sort"
  5. "strings"
  6. "time"
  7. )
  8. func SplitRunes(value string, limit int) []string {
  9. if limit <= 0 {
  10. return nil
  11. }
  12. segments := strings.Split(value, promptAuditPrioritySeparator)
  13. chunks := make([]string, 0, len(segments))
  14. for _, segment := range segments {
  15. runes := []rune(segment)
  16. for start := 0; start < len(runes); start += limit {
  17. end := start + limit
  18. if end > len(runes) {
  19. end = len(runes)
  20. }
  21. chunks = append(chunks, string(runes[start:end]))
  22. }
  23. }
  24. return chunks
  25. }
  26. func AggregateResults(results []*NormalizedResult, latency time.Duration) (*NormalizedResult, error) {
  27. if len(results) == 0 {
  28. return nil, errors.New("prompt guard produced no complete result")
  29. }
  30. aggregated := &NormalizedResult{
  31. Decision: EventPass, RiskLevel: RiskLow, Action: ActionAllow,
  32. ScannerBackend: "qwen3guard-openai", Categories: []string{}, MatchedScanners: []string{},
  33. ScannerScores: map[string]float64{}, ScannerEvidence: map[string]string{}, ChunkTotal: len(results),
  34. LatencyMS: int(latency.Milliseconds()),
  35. }
  36. categories := map[string]struct{}{}
  37. matched := map[string]struct{}{}
  38. unknown := map[string]struct{}{}
  39. for _, result := range results {
  40. if result == nil {
  41. return nil, errors.New("prompt guard partial result is not allowed")
  42. }
  43. if resultSeverity(result.Decision) > resultSeverity(aggregated.Decision) {
  44. aggregated.Decision = result.Decision
  45. aggregated.RiskLevel = result.RiskLevel
  46. aggregated.Action = result.Action
  47. aggregated.Safety = result.Safety
  48. aggregated.GuardEndpointID = result.GuardEndpointID
  49. aggregated.ScannerVersion = result.ScannerVersion
  50. aggregated.PolicyID = result.PolicyID
  51. aggregated.PolicyVersion = result.PolicyVersion
  52. }
  53. if aggregated.GuardEndpointID == "" {
  54. aggregated.GuardEndpointID = result.GuardEndpointID
  55. aggregated.ScannerVersion = result.ScannerVersion
  56. aggregated.PolicyID = result.PolicyID
  57. aggregated.PolicyVersion = result.PolicyVersion
  58. }
  59. for _, category := range result.Categories {
  60. categories[category] = struct{}{}
  61. }
  62. for _, scanner := range result.MatchedScanners {
  63. matched[scanner] = struct{}{}
  64. }
  65. for scanner, score := range result.ScannerScores {
  66. if score > aggregated.ScannerScores[scanner] {
  67. aggregated.ScannerScores[scanner] = score
  68. }
  69. }
  70. for scanner, evidence := range result.ScannerEvidence {
  71. if _, exists := aggregated.ScannerEvidence[scanner]; !exists {
  72. aggregated.ScannerEvidence[scanner] = RedactPreview(evidence, 160)
  73. }
  74. }
  75. for _, category := range result.UnknownCategories {
  76. unknown[category] = struct{}{}
  77. }
  78. }
  79. aggregated.Categories = orderedScannerKeys(categories)
  80. aggregated.MatchedScanners = orderedScannerKeys(matched)
  81. aggregated.UnknownCategories = sortedKeys(unknown)
  82. return aggregated, nil
  83. }
  84. func resultSeverity(decision EventDecision) int {
  85. switch decision {
  86. case EventCritical:
  87. return 3
  88. case EventFlag:
  89. return 2
  90. default:
  91. return 1
  92. }
  93. }
  94. func sortedKeys(values map[string]struct{}) []string {
  95. result := make([]string, 0, len(values))
  96. for key := range values {
  97. result = append(result, key)
  98. }
  99. sort.Strings(result)
  100. return result
  101. }
  102. func orderedScannerKeys(values map[string]struct{}) []string {
  103. result := make([]string, 0, len(values))
  104. remaining := make(map[string]struct{}, len(values))
  105. for key := range values {
  106. remaining[key] = struct{}{}
  107. }
  108. for _, scannerID := range AllScannerIDs {
  109. if _, ok := remaining[scannerID]; ok {
  110. result = append(result, scannerID)
  111. delete(remaining, scannerID)
  112. }
  113. }
  114. result = append(result, sortedKeys(remaining)...)
  115. return result
  116. }