gemini_session_integration_test.go 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145
  1. package service
  2. import (
  3. "testing"
  4. "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
  5. )
  6. // TestGeminiSessionContinuousConversation 测试连续会话的摘要链匹配
  7. func TestGeminiSessionContinuousConversation(t *testing.T) {
  8. store := NewDigestSessionStore()
  9. groupID := int64(1)
  10. prefixHash := "test_prefix_hash"
  11. sessionUUID := "session-uuid-12345"
  12. accountID := int64(100)
  13. // 模拟第一轮对话
  14. req1 := &antigravity.GeminiRequest{
  15. SystemInstruction: &antigravity.GeminiContent{
  16. Parts: []antigravity.GeminiPart{{Text: "You are a helpful assistant"}},
  17. },
  18. Contents: []antigravity.GeminiContent{
  19. {Role: "user", Parts: []antigravity.GeminiPart{{Text: "Hello, what's your name?"}}},
  20. },
  21. }
  22. chain1 := BuildGeminiDigestChain(req1)
  23. t.Logf("Round 1 chain: %s", chain1)
  24. // 第一轮:没有找到会话,创建新会话
  25. _, _, _, found := store.Find(groupID, prefixHash, chain1)
  26. if found {
  27. t.Error("Round 1: should not find existing session")
  28. }
  29. // 保存第一轮会话(首轮无旧 chain)
  30. store.Save(groupID, prefixHash, chain1, sessionUUID, accountID, "")
  31. // 模拟第二轮对话(用户继续对话)
  32. req2 := &antigravity.GeminiRequest{
  33. SystemInstruction: &antigravity.GeminiContent{
  34. Parts: []antigravity.GeminiPart{{Text: "You are a helpful assistant"}},
  35. },
  36. Contents: []antigravity.GeminiContent{
  37. {Role: "user", Parts: []antigravity.GeminiPart{{Text: "Hello, what's your name?"}}},
  38. {Role: "model", Parts: []antigravity.GeminiPart{{Text: "I'm Claude, nice to meet you!"}}},
  39. {Role: "user", Parts: []antigravity.GeminiPart{{Text: "What can you do?"}}},
  40. },
  41. }
  42. chain2 := BuildGeminiDigestChain(req2)
  43. t.Logf("Round 2 chain: %s", chain2)
  44. // 第二轮:应该能找到会话(通过前缀匹配)
  45. foundUUID, foundAccID, matchedChain, found := store.Find(groupID, prefixHash, chain2)
  46. if !found {
  47. t.Error("Round 2: should find session via prefix matching")
  48. }
  49. if foundUUID != sessionUUID {
  50. t.Errorf("Round 2: expected UUID %s, got %s", sessionUUID, foundUUID)
  51. }
  52. if foundAccID != accountID {
  53. t.Errorf("Round 2: expected accountID %d, got %d", accountID, foundAccID)
  54. }
  55. // 保存第二轮会话,传入 Find 返回的 matchedChain 以删旧 key
  56. store.Save(groupID, prefixHash, chain2, sessionUUID, accountID, matchedChain)
  57. // 模拟第三轮对话
  58. req3 := &antigravity.GeminiRequest{
  59. SystemInstruction: &antigravity.GeminiContent{
  60. Parts: []antigravity.GeminiPart{{Text: "You are a helpful assistant"}},
  61. },
  62. Contents: []antigravity.GeminiContent{
  63. {Role: "user", Parts: []antigravity.GeminiPart{{Text: "Hello, what's your name?"}}},
  64. {Role: "model", Parts: []antigravity.GeminiPart{{Text: "I'm Claude, nice to meet you!"}}},
  65. {Role: "user", Parts: []antigravity.GeminiPart{{Text: "What can you do?"}}},
  66. {Role: "model", Parts: []antigravity.GeminiPart{{Text: "I can help with coding, writing, and more!"}}},
  67. {Role: "user", Parts: []antigravity.GeminiPart{{Text: "Great, help me write some Go code"}}},
  68. },
  69. }
  70. chain3 := BuildGeminiDigestChain(req3)
  71. t.Logf("Round 3 chain: %s", chain3)
  72. // 第三轮:应该能找到会话(通过第二轮的前缀匹配)
  73. foundUUID, foundAccID, _, found = store.Find(groupID, prefixHash, chain3)
  74. if !found {
  75. t.Error("Round 3: should find session via prefix matching")
  76. }
  77. if foundUUID != sessionUUID {
  78. t.Errorf("Round 3: expected UUID %s, got %s", sessionUUID, foundUUID)
  79. }
  80. if foundAccID != accountID {
  81. t.Errorf("Round 3: expected accountID %d, got %d", accountID, foundAccID)
  82. }
  83. }
  84. // TestGeminiSessionDifferentConversations 测试不同会话不会错误匹配
  85. func TestGeminiSessionDifferentConversations(t *testing.T) {
  86. store := NewDigestSessionStore()
  87. groupID := int64(1)
  88. prefixHash := "test_prefix_hash"
  89. // 第一个会话
  90. req1 := &antigravity.GeminiRequest{
  91. Contents: []antigravity.GeminiContent{
  92. {Role: "user", Parts: []antigravity.GeminiPart{{Text: "Tell me about Go programming"}}},
  93. },
  94. }
  95. chain1 := BuildGeminiDigestChain(req1)
  96. store.Save(groupID, prefixHash, chain1, "session-1", 100, "")
  97. // 第二个完全不同的会话
  98. req2 := &antigravity.GeminiRequest{
  99. Contents: []antigravity.GeminiContent{
  100. {Role: "user", Parts: []antigravity.GeminiPart{{Text: "What's the weather today?"}}},
  101. },
  102. }
  103. chain2 := BuildGeminiDigestChain(req2)
  104. // 不同会话不应该匹配
  105. _, _, _, found := store.Find(groupID, prefixHash, chain2)
  106. if found {
  107. t.Error("Different conversations should not match")
  108. }
  109. }
  110. // TestGeminiSessionPrefixMatchingOrder 测试前缀匹配的优先级(最长匹配优先)
  111. func TestGeminiSessionPrefixMatchingOrder(t *testing.T) {
  112. store := NewDigestSessionStore()
  113. groupID := int64(1)
  114. prefixHash := "test_prefix_hash"
  115. // 保存不同轮次的会话到不同账号
  116. store.Save(groupID, prefixHash, "s:sys-u:q1", "session-round1", 1, "")
  117. store.Save(groupID, prefixHash, "s:sys-u:q1-m:a1", "session-round2", 2, "")
  118. store.Save(groupID, prefixHash, "s:sys-u:q1-m:a1-u:q2", "session-round3", 3, "")
  119. // 查找更长的链,应该返回最长匹配(账号 3)
  120. _, accID, _, found := store.Find(groupID, prefixHash, "s:sys-u:q1-m:a1-u:q2-m:a2")
  121. if !found {
  122. t.Error("Should find session")
  123. }
  124. if accID != 3 {
  125. t.Errorf("Should match longest prefix (account 3), got account %d", accID)
  126. }
  127. }