openai_messages_digest_session.go 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  1. package service
  2. import (
  3. "encoding/json"
  4. "fmt"
  5. "strings"
  6. "time"
  7. "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
  8. )
  9. type openAICompatAnthropicDigestBinding struct {
  10. PromptCacheKey string
  11. ExpiresAt time.Time
  12. }
  13. func buildOpenAICompatAnthropicDigestChain(req *apicompat.AnthropicRequest) string {
  14. if req == nil {
  15. return ""
  16. }
  17. parts := make([]string, 0, len(req.Messages)+1)
  18. if len(req.System) > 0 && strings.TrimSpace(string(req.System)) != "" && strings.TrimSpace(string(req.System)) != "null" {
  19. parts = append(parts, "s:"+shortHash(req.System))
  20. }
  21. for _, msg := range req.Messages {
  22. content := msg.Content
  23. if len(content) == 0 || strings.TrimSpace(string(content)) == "" {
  24. continue
  25. }
  26. prefix := "u"
  27. if strings.TrimSpace(msg.Role) == "assistant" {
  28. prefix = "a"
  29. }
  30. parts = append(parts, prefix+":"+shortHash(content))
  31. }
  32. return strings.Join(parts, "-")
  33. }
  34. func openAICompatAnthropicDigestNamespace(account *Account, cAPIKeyID int64) string {
  35. if account == nil || account.ID <= 0 {
  36. return ""
  37. }
  38. return fmt.Sprintf("%d|%d|", account.ID, cAPIKeyID)
  39. }
  40. func (s *OpenAIGatewayService) findOpenAICompatAnthropicDigestPromptCacheKey(account *Account, cAPIKeyID int64, digestChain string) (promptCacheKey string, matchedChain string) {
  41. if s == nil || digestChain == "" {
  42. return "", ""
  43. }
  44. ns := openAICompatAnthropicDigestNamespace(account, cAPIKeyID)
  45. if ns == "" {
  46. return "", ""
  47. }
  48. chain := digestChain
  49. for {
  50. if raw, ok := s.openaiCompatAnthropicDigestSessions.Load(ns + chain); ok {
  51. if binding, ok := raw.(openAICompatAnthropicDigestBinding); ok {
  52. if binding.ExpiresAt.IsZero() || time.Now().Before(binding.ExpiresAt) {
  53. if key := strings.TrimSpace(binding.PromptCacheKey); key != "" {
  54. return key, chain
  55. }
  56. }
  57. }
  58. s.openaiCompatAnthropicDigestSessions.Delete(ns + chain)
  59. }
  60. i := strings.LastIndex(chain, "-")
  61. if i < 0 {
  62. return "", ""
  63. }
  64. chain = chain[:i]
  65. }
  66. }
  67. func (s *OpenAIGatewayService) bindOpenAICompatAnthropicDigestPromptCacheKey(account *Account, cAPIKeyID int64, digestChain, promptCacheKey, oldDigestChain string) {
  68. if s == nil || digestChain == "" || strings.TrimSpace(promptCacheKey) == "" {
  69. return
  70. }
  71. ns := openAICompatAnthropicDigestNamespace(account, cAPIKeyID)
  72. if ns == "" {
  73. return
  74. }
  75. binding := openAICompatAnthropicDigestBinding{
  76. PromptCacheKey: strings.TrimSpace(promptCacheKey),
  77. ExpiresAt: time.Now().Add(s.openAIWSResponseStickyTTL()),
  78. }
  79. s.openaiCompatAnthropicDigestSessions.Store(ns+digestChain, binding)
  80. if oldDigestChain != "" && oldDigestChain != digestChain {
  81. s.openaiCompatAnthropicDigestSessions.Delete(ns + oldDigestChain)
  82. }
  83. }
  84. func promptCacheKeyFromAnthropicDigest(digestChain string) string {
  85. if strings.TrimSpace(digestChain) == "" {
  86. return ""
  87. }
  88. return "anthropic-digest-" + hashSensitiveValueForLog(digestChain)
  89. }
  90. func promptCacheKeyFromAnthropicMetadataSession(req *apicompat.AnthropicRequest) string {
  91. if req == nil || len(req.Metadata) == 0 {
  92. return ""
  93. }
  94. var metadata struct {
  95. UserID string `json:"user_id"`
  96. }
  97. if err := json.Unmarshal(req.Metadata, &metadata); err != nil {
  98. return ""
  99. }
  100. parsed := ParseMetadataUserID(metadata.UserID)
  101. if parsed == nil || strings.TrimSpace(parsed.SessionID) == "" {
  102. return ""
  103. }
  104. seed := strings.Join([]string{
  105. "anthropic-metadata",
  106. strings.TrimSpace(parsed.DeviceID),
  107. strings.TrimSpace(parsed.AccountUUID),
  108. strings.TrimSpace(parsed.SessionID),
  109. }, "|")
  110. return "anthropic-metadata-" + hashSensitiveValueForLog(seed)
  111. }
  112. func cloneAnthropicRequestForDigest(req *apicompat.AnthropicRequest) *apicompat.AnthropicRequest {
  113. if req == nil {
  114. return nil
  115. }
  116. cp := *req
  117. if len(req.System) > 0 {
  118. cp.System = append(json.RawMessage(nil), req.System...)
  119. }
  120. if len(req.Messages) > 0 {
  121. cp.Messages = append([]apicompat.AnthropicMessage(nil), req.Messages...)
  122. }
  123. return &cp
  124. }