169 lines
3.4 KiB
Go
169 lines
3.4 KiB
Go
package brave
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
const defaultBaseURL = "https://api.search.brave.com/res/v1/web/search"
|
|
|
|
type Mode string
|
|
|
|
const (
|
|
ModeKnowledgeExpand Mode = "knowledge_expand"
|
|
ModeThreadsDiscover Mode = "threads_discover"
|
|
)
|
|
|
|
type SearchResult struct {
|
|
Title string `json:"title"`
|
|
Snippet string `json:"snippet"`
|
|
URL string `json:"url"`
|
|
}
|
|
|
|
type SearchResponse struct {
|
|
Results []SearchResult
|
|
Query string
|
|
Status string // success | unavailable
|
|
}
|
|
|
|
type Client struct {
|
|
apiKey string
|
|
baseURL string
|
|
http *http.Client
|
|
}
|
|
|
|
func NewClient(apiKey string) *Client {
|
|
return &Client{
|
|
apiKey: strings.TrimSpace(apiKey),
|
|
baseURL: defaultBaseURL,
|
|
http: &http.Client{
|
|
Timeout: 20 * time.Second,
|
|
},
|
|
}
|
|
}
|
|
|
|
func (c *Client) Enabled() bool {
|
|
return c != nil && c.apiKey != ""
|
|
}
|
|
|
|
type SearchOptions struct {
|
|
Query string
|
|
Limit int
|
|
Mode Mode
|
|
Country string
|
|
SearchLang string
|
|
}
|
|
|
|
func (c *Client) Search(ctx context.Context, opts SearchOptions) (SearchResponse, error) {
|
|
out := SearchResponse{Query: strings.TrimSpace(opts.Query), Status: "unavailable"}
|
|
if !c.Enabled() {
|
|
return out, nil
|
|
}
|
|
if BreakerOpen() {
|
|
return out, fmt.Errorf("brave search temporarily paused after repeated failures")
|
|
}
|
|
if out.Query == "" {
|
|
return out, fmt.Errorf("brave search query is required")
|
|
}
|
|
|
|
limit := opts.Limit
|
|
if limit <= 0 {
|
|
limit = 5
|
|
}
|
|
if limit > 20 {
|
|
limit = 20
|
|
}
|
|
|
|
country := strings.TrimSpace(opts.Country)
|
|
if country == "" {
|
|
country = "tw"
|
|
}
|
|
searchLang := strings.TrimSpace(opts.SearchLang)
|
|
if searchLang == "" {
|
|
searchLang = "zh-hant"
|
|
}
|
|
|
|
u, err := url.Parse(c.baseURL)
|
|
if err != nil {
|
|
return out, err
|
|
}
|
|
q := u.Query()
|
|
q.Set("q", out.Query)
|
|
q.Set("count", fmt.Sprintf("%d", limit))
|
|
q.Set("country", country)
|
|
q.Set("search_lang", searchLang)
|
|
u.RawQuery = q.Encode()
|
|
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
|
|
if err != nil {
|
|
return out, err
|
|
}
|
|
req.Header.Set("Accept", "application/json")
|
|
req.Header.Set("X-Subscription-Token", c.apiKey)
|
|
|
|
res, err := c.http.Do(req)
|
|
if err != nil {
|
|
recordBreakerFailure()
|
|
return out, nil
|
|
}
|
|
defer res.Body.Close()
|
|
|
|
if res.StatusCode != http.StatusOK {
|
|
if res.StatusCode == http.StatusTooManyRequests || res.StatusCode >= 500 {
|
|
recordBreakerFailure()
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
body, err := io.ReadAll(io.LimitReader(res.Body, 1<<20))
|
|
if err != nil {
|
|
return out, nil
|
|
}
|
|
|
|
var payload struct {
|
|
Web struct {
|
|
Results []struct {
|
|
Title string `json:"title"`
|
|
Description string `json:"description"`
|
|
URL string `json:"url"`
|
|
} `json:"results"`
|
|
} `json:"web"`
|
|
}
|
|
if err := json.Unmarshal(body, &payload); err != nil {
|
|
return out, nil
|
|
}
|
|
|
|
threadsOnly := opts.Mode == ModeThreadsDiscover
|
|
for _, item := range payload.Web.Results {
|
|
rawURL := strings.TrimSpace(item.URL)
|
|
if rawURL == "" {
|
|
continue
|
|
}
|
|
if threadsOnly && !isThreadsURL(rawURL) {
|
|
continue
|
|
}
|
|
out.Results = append(out.Results, SearchResult{
|
|
Title: strings.TrimSpace(item.Title),
|
|
Snippet: strings.TrimSpace(item.Description),
|
|
URL: rawURL,
|
|
})
|
|
if len(out.Results) >= limit {
|
|
break
|
|
}
|
|
}
|
|
out.Status = "success"
|
|
recordBreakerSuccess()
|
|
return out, nil
|
|
}
|
|
|
|
func isThreadsURL(raw string) bool {
|
|
lower := strings.ToLower(raw)
|
|
return strings.Contains(lower, "threads.com") || strings.Contains(lower, "threads.net")
|
|
}
|