content_moderation_helper.go 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134
  1. package handler
  2. import (
  3. "context"
  4. "net/http"
  5. "strings"
  6. "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
  7. middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
  8. "github.com/Wei-Shaw/sub2api/internal/service"
  9. "github.com/gin-gonic/gin"
  10. "go.uber.org/zap"
  11. )
  12. func contentModerationStatus(decision *service.ContentModerationDecision) int {
  13. if decision == nil || decision.StatusCode < 400 || decision.StatusCode > 599 {
  14. return http.StatusForbidden
  15. }
  16. return decision.StatusCode
  17. }
  18. func contentModerationErrorCode(decision *service.ContentModerationDecision) string {
  19. return "content_policy_violation"
  20. }
  21. func clientRequestedModel(c *gin.Context, fallback string) string {
  22. fallback = strings.TrimSpace(fallback)
  23. if c == nil || c.Request == nil {
  24. return fallback
  25. }
  26. if model, ok := service.RequestedPublicModelFromContext(c.Request.Context()); ok {
  27. return model
  28. }
  29. return fallback
  30. }
  31. func clientRequestedUsageFields(c *gin.Context, mapping service.ChannelMappingResult, fallbackModel, upstreamModel string) service.ChannelUsageFields {
  32. return mapping.ToUsageFields(clientRequestedModel(c, fallbackModel), upstreamModel)
  33. }
  34. func runContentModeration(c *gin.Context, reqLog *zap.Logger, svc *service.ContentModerationService, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol string, model string, body []byte) *service.ContentModerationDecision {
  35. if svc == nil || c == nil || c.Request == nil {
  36. return nil
  37. }
  38. input := buildContentModerationInput(c, apiKey, subject, protocol, model, body)
  39. if reqLog != nil {
  40. reqLog.Info("content_moderation.gateway_check_start",
  41. zap.String("request_id", input.RequestID),
  42. zap.Int64("user_id", input.UserID),
  43. zap.Int64("api_key_id", input.APIKeyID),
  44. zap.String("api_key_name", input.APIKeyName),
  45. zap.Int64p("group_id", input.GroupID),
  46. zap.String("group_name", input.GroupName),
  47. zap.String("endpoint", input.Endpoint),
  48. zap.String("provider", input.Provider),
  49. zap.String("protocol", input.Protocol),
  50. zap.String("model", input.Model),
  51. zap.Int("body_bytes", len(body)),
  52. )
  53. }
  54. decision, err := svc.Check(c.Request.Context(), input)
  55. if err != nil {
  56. if reqLog != nil {
  57. reqLog.Warn("content_moderation.check_failed", zap.Error(err))
  58. }
  59. return nil
  60. }
  61. if reqLog != nil && decision != nil {
  62. reqLog.Info("content_moderation.gateway_check_done",
  63. zap.String("request_id", input.RequestID),
  64. zap.Bool("allowed", decision.Allowed),
  65. zap.Bool("blocked", decision.Blocked),
  66. zap.Bool("flagged", decision.Flagged),
  67. zap.String("action", decision.Action),
  68. zap.Int("status_code", decision.StatusCode),
  69. zap.String("highest_category", decision.HighestCategory),
  70. zap.Float64("highest_score", decision.HighestScore),
  71. )
  72. }
  73. return decision
  74. }
  75. func buildContentModerationInput(c *gin.Context, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol string, model string, body []byte) service.ContentModerationCheckInput {
  76. input := service.ContentModerationCheckInput{
  77. RequestID: contentModerationRequestID(c.Request.Context()),
  78. UserID: subject.UserID,
  79. Endpoint: GetInboundEndpoint(c),
  80. Provider: contentModerationProvider(apiKey),
  81. Model: clientRequestedModel(c, model),
  82. Protocol: protocol,
  83. Body: body,
  84. }
  85. if resolvedPlatform, ok := service.ResolvedTargetPlatformFromContext(c.Request.Context()); ok {
  86. input.Provider = resolvedPlatform
  87. }
  88. if forcedPlatform, ok := middleware2.GetForcePlatformFromContext(c); ok {
  89. input.Provider = strings.TrimSpace(forcedPlatform)
  90. }
  91. if apiKey != nil {
  92. input.APIKeyID = apiKey.ID
  93. input.APIKeyName = apiKey.Name
  94. if apiKey.User != nil {
  95. input.UserEmail = apiKey.User.Email
  96. }
  97. if apiKey.GroupID != nil {
  98. groupID := *apiKey.GroupID
  99. input.GroupID = &groupID
  100. }
  101. if apiKey.Group != nil {
  102. input.GroupName = apiKey.Group.Name
  103. }
  104. }
  105. if input.Endpoint == "" && c.Request != nil && c.Request.URL != nil {
  106. input.Endpoint = c.Request.URL.Path
  107. }
  108. return input
  109. }
  110. func contentModerationProvider(apiKey *service.APIKey) string {
  111. if apiKey == nil || apiKey.Group == nil {
  112. return ""
  113. }
  114. return strings.TrimSpace(apiKey.Group.Platform)
  115. }
  116. func contentModerationRequestID(ctx context.Context) string {
  117. if ctx == nil {
  118. return ""
  119. }
  120. if requestID, ok := ctx.Value(ctxkey.RequestID).(string); ok {
  121. return strings.TrimSpace(requestID)
  122. }
  123. return ""
  124. }