openai_messages_replay_guard.go 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990
  1. package service
  2. import (
  3. "encoding/json"
  4. "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
  5. )
  6. const openAICompatAnthropicReplayMaxTailMessages = 12
  7. func applyAnthropicCompatFullReplayGuard(req *apicompat.AnthropicRequest) bool {
  8. if req == nil || len(req.Messages) <= openAICompatAnthropicReplayMaxTailMessages {
  9. return false
  10. }
  11. start := len(req.Messages) - openAICompatAnthropicReplayMaxTailMessages
  12. start = expandAnthropicCompatTrimBoundary(req.Messages, start)
  13. if start <= 0 {
  14. return false
  15. }
  16. req.Messages = append([]apicompat.AnthropicMessage(nil), req.Messages[start:]...)
  17. return true
  18. }
  19. func expandAnthropicCompatTrimBoundary(messages []apicompat.AnthropicMessage, start int) int {
  20. if start <= 0 || start >= len(messages) {
  21. return start
  22. }
  23. toolUseIndex := make(map[string]int)
  24. toolResultIndex := make(map[string]int)
  25. for i, msg := range messages {
  26. uses, results := anthropicCompatMessageToolIDs(msg)
  27. for _, id := range uses {
  28. if _, exists := toolUseIndex[id]; !exists {
  29. toolUseIndex[id] = i
  30. }
  31. }
  32. for _, id := range results {
  33. if _, exists := toolResultIndex[id]; !exists {
  34. toolResultIndex[id] = i
  35. }
  36. }
  37. }
  38. for {
  39. next := start
  40. for i := start; i < len(messages); i++ {
  41. uses, results := anthropicCompatMessageToolIDs(messages[i])
  42. for _, id := range results {
  43. if useIdx, ok := toolUseIndex[id]; ok && useIdx < next {
  44. next = useIdx
  45. }
  46. }
  47. for _, id := range uses {
  48. if resultIdx, ok := toolResultIndex[id]; ok && resultIdx < next {
  49. next = resultIdx
  50. }
  51. }
  52. }
  53. if next == start {
  54. return start
  55. }
  56. start = next
  57. }
  58. }
  59. func anthropicCompatMessageToolIDs(msg apicompat.AnthropicMessage) ([]string, []string) {
  60. var blocks []apicompat.AnthropicContentBlock
  61. if err := json.Unmarshal(msg.Content, &blocks); err != nil {
  62. return nil, nil
  63. }
  64. uses := make([]string, 0, 1)
  65. results := make([]string, 0, 1)
  66. for _, block := range blocks {
  67. switch block.Type {
  68. case "tool_use":
  69. if block.ID != "" {
  70. uses = append(uses, block.ID)
  71. }
  72. case "tool_result":
  73. if block.ToolUseID != "" {
  74. results = append(results, block.ToolUseID)
  75. }
  76. }
  77. }
  78. return uses, results
  79. }