126 lines
2.8 KiB
Go
126 lines
2.8 KiB
Go
package providers
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"net/http"
|
|
"net/url"
|
|
"time"
|
|
|
|
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
|
)
|
|
|
|
const braveAPIBase = "https://api.search.brave.com"
|
|
|
|
// BraveProvider 实现 SearchProvider。
|
|
type BraveProvider struct {
|
|
client *http.Client
|
|
apiBase string
|
|
maxResults int
|
|
}
|
|
|
|
func NewBraveProvider(timeout time.Duration, maxResults int) *BraveProvider {
|
|
return &BraveProvider{
|
|
client: &http.Client{
|
|
Timeout: timeout,
|
|
},
|
|
apiBase: braveAPIBase,
|
|
maxResults: maxResults,
|
|
}
|
|
}
|
|
|
|
func (b *BraveProvider) Name() models.ProviderKind {
|
|
return models.ProviderBrave
|
|
}
|
|
|
|
func (b *BraveProvider) MaxResults() int {
|
|
return b.maxResults
|
|
}
|
|
|
|
type braveWebResult struct {
|
|
Title string `json:"title"`
|
|
URL string `json:"url"`
|
|
Description string `json:"description"`
|
|
Age string `json:"age"`
|
|
}
|
|
|
|
type braveResponse struct {
|
|
Web struct {
|
|
Results []braveWebResult `json:"results"`
|
|
} `json:"web"`
|
|
}
|
|
|
|
func (b *BraveProvider) Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error) {
|
|
maxRes := req.MaxResults
|
|
if maxRes <= 0 || maxRes > b.maxResults {
|
|
maxRes = b.maxResults
|
|
}
|
|
|
|
u, _ := url.Parse(b.apiBase + "/res/v1/web/search")
|
|
q := u.Query()
|
|
q.Set("q", req.Query)
|
|
q.Set("count", fmt.Sprintf("%d", maxRes))
|
|
u.RawQuery = q.Encode()
|
|
|
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("brave: create request: %w", err)
|
|
}
|
|
httpReq.Header.Set("Accept", "application/json")
|
|
httpReq.Header.Set("X-Subscription-Token", apiKey)
|
|
|
|
start := time.Now()
|
|
resp, err := b.client.Do(httpReq)
|
|
duration := time.Since(start)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("brave: do request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
raw, err := io.ReadAll(resp.Body)
|
|
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, string(raw))
|
|
}
|
|
|
|
var br braveResponse
|
|
if err := json.Unmarshal(raw, &br); err != nil {
|
|
return nil, fmt.Errorf("brave: decode: %w", err)
|
|
}
|
|
|
|
items := make([]models.ResultItem, 0, len(br.Web.Results))
|
|
for _, r := range br.Web.Results {
|
|
items = append(items, models.ResultItem{
|
|
Title: r.Title,
|
|
URL: r.URL,
|
|
Snippet: r.Description,
|
|
PublishedAt: r.Age,
|
|
})
|
|
}
|
|
|
|
return &models.SearchResult{
|
|
Provider: models.ProviderBrave,
|
|
Query: req.Query,
|
|
Results: items,
|
|
TotalCount: len(items),
|
|
DurationMs: duration.Milliseconds(),
|
|
Cached: false,
|
|
}, nil
|
|
}
|
|
|
|
func (b *BraveProvider) HealthCheck(ctx context.Context) error {
|
|
var d net.Dialer
|
|
conn, err := d.DialContext(ctx, "tcp", "api.search.brave.com:443")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
conn.Close()
|
|
return nil
|
|
}
|