1b53ffcac7
Inject web search capability for Claude Console (API Key) accounts that don't natively support Anthropic's web_search tool. When a pure web_search request is detected, the gateway calls Brave Search or Tavily API directly and constructs an Anthropic-protocol-compliant SSE/JSON response without forwarding to upstream. Backend: - New `pkg/websearch/` SDK: Brave and Tavily provider implementations with io.LimitReader, proxy support, and Redis-based quota tracking (Lua atomic INCR + TTL, DECR rollback on failure) - Global config via `settings.web_search_emulation_config` (JSON) with in-process cache + singleflight, input validation, API key merge on save, and sanitized API responses - Channel-level toggle via `channels.features_config` JSONB column (DB migration 101) - Account-level toggle via `accounts.extra.web_search_emulation` - Request interception in `Forward()` with SSE streaming response construction using json.Marshal (no manual string concatenation) - Manager hot-reload: `RebuildWebSearchManager()` called on config save and startup via `SetWebSearchRedisClient()` - 70 unit tests covering providers, manager, config validation, sanitization, tool detection, query extraction, and response building Frontend: - Settings → Gateway tab: Web Search Emulation config card with global toggle, provider list (add/remove, API key, priority, quota, proxy) - Channels → Anthropic tab: web search emulation toggle with global state linkage (disabled when global off) - Account Create/Edit modals: web search emulation toggle for API Key type with Toggle component - Full i18n coverage (zh + en)
143 lines
5.2 KiB
Go
143 lines
5.2 KiB
Go
package service
|
|
|
|
import (
|
|
"testing"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/websearch"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// --- isOnlyWebSearchToolInBody ---
|
|
|
|
func TestIsOnlyWebSearchToolInBody_WebSearchType(t *testing.T) {
|
|
require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"type":"web_search"}]}`)))
|
|
}
|
|
|
|
func TestIsOnlyWebSearchToolInBody_WebSearch2025Type(t *testing.T) {
|
|
require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"type":"web_search_20250305"}]}`)))
|
|
}
|
|
|
|
func TestIsOnlyWebSearchToolInBody_GoogleSearchType(t *testing.T) {
|
|
require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"type":"google_search"}]}`)))
|
|
}
|
|
|
|
func TestIsOnlyWebSearchToolInBody_NameWebSearch(t *testing.T) {
|
|
require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"name":"web_search"}]}`)))
|
|
}
|
|
|
|
func TestIsOnlyWebSearchToolInBody_NameWebSearch2025(t *testing.T) {
|
|
require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"name":"web_search_20250305"}]}`)))
|
|
}
|
|
|
|
func TestIsOnlyWebSearchToolInBody_NameGoogleSearch(t *testing.T) {
|
|
require.True(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"name":"google_search"}]}`)))
|
|
}
|
|
|
|
func TestIsOnlyWebSearchToolInBody_MultipleTools(t *testing.T) {
|
|
require.False(t, isOnlyWebSearchToolInBody(
|
|
[]byte(`{"tools":[{"type":"web_search"},{"type":"text_editor"}]}`)))
|
|
}
|
|
|
|
func TestIsOnlyWebSearchToolInBody_NoTools(t *testing.T) {
|
|
require.False(t, isOnlyWebSearchToolInBody([]byte(`{"model":"claude-3"}`)))
|
|
}
|
|
|
|
func TestIsOnlyWebSearchToolInBody_EmptyToolsArray(t *testing.T) {
|
|
require.False(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[]}`)))
|
|
}
|
|
|
|
func TestIsOnlyWebSearchToolInBody_NonWebSearchTool(t *testing.T) {
|
|
require.False(t, isOnlyWebSearchToolInBody([]byte(`{"tools":[{"type":"text_editor"}]}`)))
|
|
}
|
|
|
|
func TestIsOnlyWebSearchToolInBody_ToolsNotArray(t *testing.T) {
|
|
require.False(t, isOnlyWebSearchToolInBody([]byte(`{"tools":"web_search"}`)))
|
|
}
|
|
|
|
// --- extractSearchQueryFromBody ---
|
|
|
|
func TestExtractSearchQueryFromBody_StringContent(t *testing.T) {
|
|
body := `{"messages":[{"role":"user","content":"what is golang"}]}`
|
|
require.Equal(t, "what is golang", extractSearchQueryFromBody([]byte(body)))
|
|
}
|
|
|
|
func TestExtractSearchQueryFromBody_ArrayContent(t *testing.T) {
|
|
body := `{"messages":[{"role":"user","content":[{"type":"text","text":"search this"}]}]}`
|
|
require.Equal(t, "search this", extractSearchQueryFromBody([]byte(body)))
|
|
}
|
|
|
|
func TestExtractSearchQueryFromBody_MultipleMessages(t *testing.T) {
|
|
body := `{"messages":[{"role":"user","content":"first"},{"role":"assistant","content":"ok"},{"role":"user","content":"second"}]}`
|
|
require.Equal(t, "second", extractSearchQueryFromBody([]byte(body)))
|
|
}
|
|
|
|
func TestExtractSearchQueryFromBody_LastMessageNotUser(t *testing.T) {
|
|
body := `{"messages":[{"role":"user","content":"q"},{"role":"assistant","content":"a"}]}`
|
|
require.Equal(t, "", extractSearchQueryFromBody([]byte(body)))
|
|
}
|
|
|
|
func TestExtractSearchQueryFromBody_EmptyMessages(t *testing.T) {
|
|
require.Equal(t, "", extractSearchQueryFromBody([]byte(`{"messages":[]}`)))
|
|
}
|
|
|
|
func TestExtractSearchQueryFromBody_NoMessages(t *testing.T) {
|
|
require.Equal(t, "", extractSearchQueryFromBody([]byte(`{"model":"claude-3"}`)))
|
|
}
|
|
|
|
func TestExtractSearchQueryFromBody_ArrayContentSkipsEmptyText(t *testing.T) {
|
|
body := `{"messages":[{"role":"user","content":[{"type":"image"},{"type":"text","text":""},{"type":"text","text":"real query"}]}]}`
|
|
require.Equal(t, "real query", extractSearchQueryFromBody([]byte(body)))
|
|
}
|
|
|
|
func TestExtractSearchQueryFromBody_ArrayContentNoTextBlock(t *testing.T) {
|
|
body := `{"messages":[{"role":"user","content":[{"type":"image","source":{}}]}]}`
|
|
require.Equal(t, "", extractSearchQueryFromBody([]byte(body)))
|
|
}
|
|
|
|
// --- buildSearchResultBlocks ---
|
|
|
|
func TestBuildSearchResultBlocks_WithResults(t *testing.T) {
|
|
results := []websearch.SearchResult{
|
|
{URL: "https://a.com", Title: "A", Snippet: "snippet a", PageAge: "2 days"},
|
|
{URL: "https://b.com", Title: "B", Snippet: "snippet b"},
|
|
}
|
|
blocks := buildSearchResultBlocks(results)
|
|
require.Len(t, blocks, 2)
|
|
require.Equal(t, "web_search_result", blocks[0]["type"])
|
|
require.Equal(t, "https://a.com", blocks[0]["url"])
|
|
require.Equal(t, "snippet a", blocks[0]["page_content"])
|
|
require.Equal(t, "2 days", blocks[0]["page_age"])
|
|
// Second result has no PageAge
|
|
require.Equal(t, "https://b.com", blocks[1]["url"])
|
|
_, hasPageAge := blocks[1]["page_age"]
|
|
require.False(t, hasPageAge)
|
|
}
|
|
|
|
func TestBuildSearchResultBlocks_Empty(t *testing.T) {
|
|
blocks := buildSearchResultBlocks(nil)
|
|
require.Empty(t, blocks)
|
|
}
|
|
|
|
func TestBuildSearchResultBlocks_SnippetEmpty(t *testing.T) {
|
|
blocks := buildSearchResultBlocks([]websearch.SearchResult{{URL: "https://x.com", Title: "X", Snippet: ""}})
|
|
_, hasContent := blocks[0]["page_content"]
|
|
require.False(t, hasContent)
|
|
}
|
|
|
|
// --- buildTextSummary ---
|
|
|
|
func TestBuildTextSummary_WithResults(t *testing.T) {
|
|
results := []websearch.SearchResult{
|
|
{URL: "https://a.com", Title: "A", Snippet: "desc a"},
|
|
}
|
|
summary := buildTextSummary("test query", results)
|
|
require.Contains(t, summary, "test query")
|
|
require.Contains(t, summary, "1. **A**")
|
|
require.Contains(t, summary, "https://a.com")
|
|
}
|
|
|
|
func TestBuildTextSummary_NoResults(t *testing.T) {
|
|
summary := buildTextSummary("test", nil)
|
|
require.Contains(t, summary, "No search results found for: test")
|
|
}
|