Files
sub2api/backend/internal/pkg/websearch/brave.go
erio 1b53ffcac7 feat(gateway): add web search emulation for Anthropic API Key accounts
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)
2026-04-14 09:20:39 +08:00

107 lines
2.7 KiB
Go

package websearch
import (
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"net/url"
"strconv"
)
const (
braveSearchEndpoint = "https://api.search.brave.com/res/v1/web/search"
braveMaxCount = 20
braveProviderName = "brave"
)
// braveSearchURL is pre-parsed at init time; url.Parse cannot fail on a constant literal.
var braveSearchURL, _ = url.Parse(braveSearchEndpoint) //nolint:errcheck
// BraveProvider implements web search via the Brave Search API.
type BraveProvider struct {
apiKey string
httpClient *http.Client
}
// NewBraveProvider creates a Brave Search provider.
// The caller is responsible for configuring the http.Client with proxy/timeouts.
func NewBraveProvider(apiKey string, httpClient *http.Client) *BraveProvider {
if httpClient == nil {
httpClient = http.DefaultClient
}
return &BraveProvider{apiKey: apiKey, httpClient: httpClient}
}
func (b *BraveProvider) Name() string { return braveProviderName }
func (b *BraveProvider) Search(ctx context.Context, req SearchRequest) (*SearchResponse, error) {
count := req.MaxResults
if count <= 0 {
count = defaultMaxResults
}
if count > braveMaxCount {
count = braveMaxCount
}
u := *braveSearchURL // copy the pre-parsed URL
q := u.Query()
q.Set("q", req.Query)
q.Set("count", strconv.Itoa(count))
u.RawQuery = q.Encode()
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
if err != nil {
return nil, fmt.Errorf("brave: build request: %w", err)
}
httpReq.Header.Set("X-Subscription-Token", b.apiKey)
httpReq.Header.Set("Accept", "application/json")
resp, err := b.httpClient.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("brave: request failed: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseSize))
if err != nil {
return nil, fmt.Errorf("brave: read body: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("brave: status %d: %s", resp.StatusCode, truncateBody(body))
}
var raw braveResponse
if err := json.Unmarshal(body, &raw); err != nil {
return nil, fmt.Errorf("brave: decode response: %w", err)
}
results := make([]SearchResult, 0, len(raw.Web.Results))
for _, r := range raw.Web.Results {
results = append(results, SearchResult{
URL: r.URL,
Title: r.Title,
Snippet: r.Description,
PageAge: r.Age,
})
}
return &SearchResponse{Results: results, Query: req.Query}, nil
}
// braveResponse is the minimal structure of the Brave Search API response.
type braveResponse struct {
Web struct {
Results []braveResult `json:"results"`
} `json:"web"`
}
type braveResult struct {
URL string `json:"url"`
Title string `json:"title"`
Description string `json:"description"`
Age string `json:"age"`
}