gateway_websearch_emulation_test.go 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380
  1. //go:build unit
  2. package service
  3. import (
  4. "context"
  5. "encoding/json"
  6. "testing"
  7. "time"
  8. "github.com/Wei-Shaw/sub2api/internal/config"
  9. "github.com/Wei-Shaw/sub2api/internal/pkg/websearch"
  10. "github.com/stretchr/testify/require"
  11. )
  12. // --- isOnlyWebSearchToolInBody ---
  13. func TestIsOnlyWebSearchToolInBody_WebSearchType(t *testing.T) {
  14. require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"type":"web_search"}]}`)))
  15. }
  16. func TestIsOnlyWebSearchToolInBody_WebSearch2025Type(t *testing.T) {
  17. require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"type":"web_search_20250305"}]}`)))
  18. }
  19. func TestIsOnlyWebSearchToolInBody_GoogleSearchType(t *testing.T) {
  20. require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"type":"google_search"}]}`)))
  21. }
  22. func TestIsOnlyWebSearchToolInBody_NameWebSearch(t *testing.T) {
  23. require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"name":"web_search"}]}`)))
  24. }
  25. func TestIsOnlyWebSearchToolInBody_NameWebSearch2025(t *testing.T) {
  26. require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"name":"web_search_20250305"}]}`)))
  27. }
  28. func TestIsOnlyWebSearchToolInBody_NameGoogleSearch(t *testing.T) {
  29. require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"name":"google_search"}]}`)))
  30. }
  31. func TestIsOnlyWebSearchToolInBody_MultipleTools(t *testing.T) {
  32. require.False(t, isOnlyWebSearchToolInBody(
  33. []byte(`{"tools":[{"type":"web_search"},{"type":"text_editor"}]}`)))
  34. }
  35. func TestIsOnlyWebSearchToolInBody_NoTools(t *testing.T) {
  36. require.False(t, isOnlyWebSearchToolInBody([]byte(`{"model":"claude-3"}`)))
  37. }
  38. func TestIsOnlyWebSearchToolInBody_EmptyToolsArray(t *testing.T) {
  39. require.False(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[]}`)))
  40. }
  41. func TestIsOnlyWebSearchToolInBody_NonWebSearchTool(t *testing.T) {
  42. require.False(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"type":"text_editor"}]}`)))
  43. }
  44. func TestIsOnlyWebSearchToolInBody_ToolsNotArray(t *testing.T) {
  45. require.False(t, isOnlyWebSearchToolInBody([]byte(`{"tools":"web_search"}`)))
  46. }
  47. // --- extractSearchQueryFromBody ---
  48. func TestExtractSearchQueryFromBody_StringContent(t *testing.T) {
  49. body := `{"messages":[{"role":"user","content":"what is golang"}]}`
  50. require.Equal(t, "what is golang", extractSearchQueryFromBody([]byte(body)))
  51. }
  52. func TestExtractSearchQueryFromBody_ArrayContent(t *testing.T) {
  53. body := `{"messages":[{"role":"user","content":[{"type":"text","text":"search this"}]}]}`
  54. require.Equal(t, "search this", extractSearchQueryFromBody([]byte(body)))
  55. }
  56. func TestExtractSearchQueryFromBody_MultipleMessages(t *testing.T) {
  57. body := `{"messages":[{"role":"user","content":"first"},{"role":"assistant","content":"ok"},{"role":"user","content":"second"}]}`
  58. require.Equal(t, "second", extractSearchQueryFromBody([]byte(body)))
  59. }
  60. func TestExtractSearchQueryFromBody_LastMessageNotUser(t *testing.T) {
  61. body := `{"messages":[{"role":"user","content":"q"},{"role":"assistant","content":"a"}]}`
  62. require.Equal(t, "", extractSearchQueryFromBody([]byte(body)))
  63. }
  64. func TestExtractSearchQueryFromBody_EmptyMessages(t *testing.T) {
  65. require.Equal(t, "", extractSearchQueryFromBody([]byte(`{"messages":[]}`)))
  66. }
  67. func TestExtractSearchQueryFromBody_NoMessages(t *testing.T) {
  68. require.Equal(t, "", extractSearchQueryFromBody([]byte(`{"model":"claude-3"}`)))
  69. }
  70. func TestExtractSearchQueryFromBody_ArrayContentSkipsEmptyText(t *testing.T) {
  71. body := `{"messages":[{"role":"user","content":[{"type":"image"},{"type":"text","text":""},{"type":"text","text":"real query"}]}]}`
  72. require.Equal(t, "real query", extractSearchQueryFromBody([]byte(body)))
  73. }
  74. func TestExtractSearchQueryFromBody_ArrayContentNoTextBlock(t *testing.T) {
  75. body := `{"messages":[{"role":"user","content":[{"type":"image","source":{}}]}]}`
  76. require.Equal(t, "", extractSearchQueryFromBody([]byte(body)))
  77. }
  78. // --- buildSearchResultBlocks ---
  79. func TestBuildSearchResultBlocks_WithResults(t *testing.T) {
  80. results := []websearch.SearchResult{
  81. {URL: "https://a.com", Title: "A", Snippet: "snippet a", PageAge: "2 days"},
  82. {URL: "https://b.com", Title: "B", Snippet: "snippet b"},
  83. }
  84. blocks := buildSearchResultBlocks(results)
  85. require.Len(t, blocks, 2)
  86. require.Equal(t, "web_search_result", blocks[0]["type"])
  87. require.Equal(t, "https://a.com", blocks[0]["url"])
  88. require.Equal(t, "snippet a", blocks[0]["page_content"])
  89. require.Equal(t, "2 days", blocks[0]["page_age"])
  90. // Second result has no PageAge
  91. require.Equal(t, "https://b.com", blocks[1]["url"])
  92. _, hasPageAge := blocks[1]["page_age"]
  93. require.False(t, hasPageAge)
  94. }
  95. func TestBuildSearchResultBlocks_Empty(t *testing.T) {
  96. blocks := buildSearchResultBlocks(nil)
  97. require.Empty(t, blocks)
  98. }
  99. func TestBuildSearchResultBlocks_SnippetEmpty(t *testing.T) {
  100. blocks := buildSearchResultBlocks([]websearch.SearchResult{{URL: "https://x.com", Title: "X", Snippet: ""}})
  101. _, hasContent := blocks[0]["page_content"]
  102. require.False(t, hasContent)
  103. }
  104. // --- buildTextSummary ---
  105. func TestBuildTextSummary_WithResults(t *testing.T) {
  106. results := []websearch.SearchResult{
  107. {URL: "https://a.com", Title: "A", Snippet: "desc a"},
  108. }
  109. summary := buildTextSummary("test query", results)
  110. require.Contains(t, summary, "test query")
  111. require.Contains(t, summary, "1. **A**")
  112. require.Contains(t, summary, "https://a.com")
  113. }
  114. func TestBuildTextSummary_NoResults(t *testing.T) {
  115. summary := buildTextSummary("test", nil)
  116. require.Contains(t, summary, "No search results found for: test")
  117. }
  118. // --- shouldEmulateWebSearch ---
  119. // webSearchToolBody is a valid request body with exactly one web_search tool.
  120. var webSearchToolBody = []byte(`{"tools":[{"type":"web_search"}],"messages":[{"role":"user","content":"test"}]}`)
  121. // nonWebSearchToolBody is a request body without web_search tool.
  122. var nonWebSearchToolBody = []byte(`{"tools":[{"type":"text_editor"}],"messages":[{"role":"user","content":"test"}]}`)
  123. // newAnthropicAPIKeyAccount creates a test Account with the given web search emulation mode.
  124. func newAnthropicAPIKeyAccount(mode string) *Account {
  125. return &Account{
  126. ID: 1,
  127. Platform: PlatformAnthropic,
  128. Type: AccountTypeAPIKey,
  129. Extra: map[string]any{featureKeyWebSearchEmulation: mode},
  130. }
  131. }
  132. // setGlobalWebSearchConfig stores a config in the global cache used by SettingService.IsWebSearchEmulationEnabled.
  133. func setGlobalWebSearchConfig(cfg *WebSearchEmulationConfig) {
  134. webSearchEmulationCache.Store(&cachedWebSearchEmulationConfig{
  135. config: cfg,
  136. expiresAt: time.Now().Add(10 * time.Minute).UnixNano(),
  137. })
  138. }
  139. // clearGlobalWebSearchConfig resets the global cache to force re-read.
  140. func clearGlobalWebSearchConfig() {
  141. webSearchEmulationCache.Store((*cachedWebSearchEmulationConfig)(nil))
  142. }
  143. // newSettingServiceForWebSearchTest creates a SettingService with a mock repo pre-loaded with config.
  144. func newSettingServiceForWebSearchTest(enabled bool) *SettingService {
  145. repo := newMockSettingRepo()
  146. cfg := &WebSearchEmulationConfig{
  147. Enabled: enabled,
  148. Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "sk-test"}},
  149. }
  150. data, _ := json.Marshal(cfg)
  151. repo.data[SettingKeyWebSearchEmulationConfig] = string(data)
  152. return NewSettingService(repo, &config.Config{})
  153. }
  154. // newChannelServiceWithCache creates a ChannelService with a pre-built cache containing the channel.
  155. func newChannelServiceWithCache(groupID int64, ch *Channel) *ChannelService {
  156. svc := &ChannelService{}
  157. cache := &channelCache{
  158. channelByGroupID: map[int64]*Channel{groupID: ch},
  159. byID: map[int64]*Channel{ch.ID: ch},
  160. groupPlatform: map[int64]string{},
  161. loadedAt: time.Now(),
  162. }
  163. svc.cache.Store(cache)
  164. return svc
  165. }
  166. func TestShouldEmulateWebSearch_NilManager(t *testing.T) {
  167. SetWebSearchManager(nil)
  168. defer SetWebSearchManager(nil)
  169. settingSvc := newSettingServiceForWebSearchTest(true)
  170. setGlobalWebSearchConfig(&WebSearchEmulationConfig{
  171. Enabled: true,
  172. Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
  173. })
  174. defer clearGlobalWebSearchConfig()
  175. svc := &GatewayService{settingService: settingSvc}
  176. account := newAnthropicAPIKeyAccount(WebSearchModeEnabled)
  177. require.False(t, svc.shouldEmulateWebSearch(context.Background(), account, nil, webSearchToolBody))
  178. }
  179. func TestShouldEmulateWebSearch_NotOnlyWebSearchTool(t *testing.T) {
  180. mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
  181. SetWebSearchManager(mgr)
  182. defer SetWebSearchManager(nil)
  183. settingSvc := newSettingServiceForWebSearchTest(true)
  184. setGlobalWebSearchConfig(&WebSearchEmulationConfig{
  185. Enabled: true,
  186. Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
  187. })
  188. defer clearGlobalWebSearchConfig()
  189. svc := &GatewayService{settingService: settingSvc}
  190. account := newAnthropicAPIKeyAccount(WebSearchModeEnabled)
  191. require.False(t, svc.shouldEmulateWebSearch(context.Background(), account, nil, nonWebSearchToolBody))
  192. }
  193. func TestShouldEmulateWebSearch_GlobalDisabled(t *testing.T) {
  194. mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
  195. SetWebSearchManager(mgr)
  196. defer SetWebSearchManager(nil)
  197. // Global config disabled
  198. setGlobalWebSearchConfig(&WebSearchEmulationConfig{
  199. Enabled: false,
  200. Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
  201. })
  202. defer clearGlobalWebSearchConfig()
  203. settingSvc := newSettingServiceForWebSearchTest(false)
  204. svc := &GatewayService{settingService: settingSvc}
  205. account := newAnthropicAPIKeyAccount(WebSearchModeEnabled)
  206. require.False(t, svc.shouldEmulateWebSearch(context.Background(), account, nil, webSearchToolBody))
  207. }
  208. func TestShouldEmulateWebSearch_AccountDisabled(t *testing.T) {
  209. mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
  210. SetWebSearchManager(mgr)
  211. defer SetWebSearchManager(nil)
  212. setGlobalWebSearchConfig(&WebSearchEmulationConfig{
  213. Enabled: true,
  214. Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
  215. })
  216. defer clearGlobalWebSearchConfig()
  217. settingSvc := newSettingServiceForWebSearchTest(true)
  218. svc := &GatewayService{settingService: settingSvc}
  219. account := newAnthropicAPIKeyAccount(WebSearchModeDisabled)
  220. require.False(t, svc.shouldEmulateWebSearch(context.Background(), account, nil, webSearchToolBody))
  221. }
  222. func TestShouldEmulateWebSearch_AccountEnabled(t *testing.T) {
  223. mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
  224. SetWebSearchManager(mgr)
  225. defer SetWebSearchManager(nil)
  226. setGlobalWebSearchConfig(&WebSearchEmulationConfig{
  227. Enabled: true,
  228. Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
  229. })
  230. defer clearGlobalWebSearchConfig()
  231. settingSvc := newSettingServiceForWebSearchTest(true)
  232. svc := &GatewayService{settingService: settingSvc}
  233. account := newAnthropicAPIKeyAccount(WebSearchModeEnabled)
  234. require.True(t, svc.shouldEmulateWebSearch(context.Background(), account, nil, webSearchToolBody))
  235. }
  236. func TestShouldEmulateWebSearch_DefaultMode_ChannelEnabled(t *testing.T) {
  237. mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
  238. SetWebSearchManager(mgr)
  239. defer SetWebSearchManager(nil)
  240. setGlobalWebSearchConfig(&WebSearchEmulationConfig{
  241. Enabled: true,
  242. Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
  243. })
  244. defer clearGlobalWebSearchConfig()
  245. settingSvc := newSettingServiceForWebSearchTest(true)
  246. ch := &Channel{
  247. ID: 10,
  248. Status: StatusActive,
  249. FeaturesConfig: map[string]any{
  250. featureKeyWebSearchEmulation: map[string]any{PlatformAnthropic: true},
  251. },
  252. }
  253. channelSvc := newChannelServiceWithCache(42, ch)
  254. svc := &GatewayService{settingService: settingSvc, channelService: channelSvc}
  255. account := newAnthropicAPIKeyAccount(WebSearchModeDefault)
  256. groupID := int64(42)
  257. require.True(t, svc.shouldEmulateWebSearch(context.Background(), account, &groupID, webSearchToolBody))
  258. }
  259. func TestShouldEmulateWebSearch_DefaultMode_ChannelDisabled(t *testing.T) {
  260. mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
  261. SetWebSearchManager(mgr)
  262. defer SetWebSearchManager(nil)
  263. setGlobalWebSearchConfig(&WebSearchEmulationConfig{
  264. Enabled: true,
  265. Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
  266. })
  267. defer clearGlobalWebSearchConfig()
  268. settingSvc := newSettingServiceForWebSearchTest(true)
  269. ch := &Channel{
  270. ID: 10,
  271. Status: StatusActive,
  272. FeaturesConfig: map[string]any{
  273. featureKeyWebSearchEmulation: map[string]any{PlatformAnthropic: false},
  274. },
  275. }
  276. channelSvc := newChannelServiceWithCache(42, ch)
  277. svc := &GatewayService{settingService: settingSvc, channelService: channelSvc}
  278. account := newAnthropicAPIKeyAccount(WebSearchModeDefault)
  279. groupID := int64(42)
  280. require.False(t, svc.shouldEmulateWebSearch(context.Background(), account, &groupID, webSearchToolBody))
  281. }
  282. func TestShouldEmulateWebSearch_DefaultMode_NilGroupID(t *testing.T) {
  283. mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
  284. SetWebSearchManager(mgr)
  285. defer SetWebSearchManager(nil)
  286. setGlobalWebSearchConfig(&WebSearchEmulationConfig{
  287. Enabled: true,
  288. Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
  289. })
  290. defer clearGlobalWebSearchConfig()
  291. settingSvc := newSettingServiceForWebSearchTest(true)
  292. svc := &GatewayService{settingService: settingSvc}
  293. account := newAnthropicAPIKeyAccount(WebSearchModeDefault)
  294. // nil groupID + default mode → falls through to channel check → returns false
  295. require.False(t, svc.shouldEmulateWebSearch(context.Background(), account, nil, webSearchToolBody))
  296. }
  297. func TestShouldEmulateWebSearch_DefaultMode_NilChannelService(t *testing.T) {
  298. mgr := websearch.NewManager([]websearch.ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
  299. SetWebSearchManager(mgr)
  300. defer SetWebSearchManager(nil)
  301. setGlobalWebSearchConfig(&WebSearchEmulationConfig{
  302. Enabled: true,
  303. Providers: []WebSearchProviderConfig{{Type: "brave", APIKey: "k"}},
  304. })
  305. defer clearGlobalWebSearchConfig()
  306. settingSvc := newSettingServiceForWebSearchTest(true)
  307. svc := &GatewayService{settingService: settingSvc, channelService: nil}
  308. account := newAnthropicAPIKeyAccount(WebSearchModeDefault)
  309. groupID := int64(42)
  310. // nil channelService + default mode → returns false
  311. require.False(t, svc.shouldEmulateWebSearch(context.Background(), account, &groupID, webSearchToolBody))
  312. }