Files
sub2api/backend/internal/pkg/websearch/tavily.go
erio 6a08efeef9 fix: resolve upstream CI failures (lint, test, gofmt)
- Fix errcheck: handle Write/Encode return values in brave_test.go
- Fix errcheck: defer resp.Body.Close() with _ assignment in tavily.go
- Fix gofmt: payment.go, channel.go, payment_config_providers.go
- Fix unused: remove dead decodeURLValue in easypay.go
- Restore shouldFallbackGeminiModel function (deleted during cherry-pick)
- Add missing balanceNotifyService param to NewGatewayService in test
- Fix platform default test expectation (empty stays empty)
- Fix wildcard pricing test (longest prefix wins, not config order)
- Fix subscription group test (SUBSCRIPTION_REPOSITORY_UNAVAILABLE)
2026-04-14 12:11:08 +08:00

108 lines
2.7 KiB
Go

package websearch
import (
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
)
const (
tavilySearchEndpoint = "https://api.tavily.com/search"
tavilyProviderName = "tavily"
tavilySearchDepthBasic = "basic"
)
// TavilyProvider implements web search via the Tavily Search API.
type TavilyProvider struct {
apiKey string
httpClient *http.Client
}
// NewTavilyProvider creates a Tavily Search provider.
// The caller is responsible for configuring the http.Client with proxy/timeouts.
func NewTavilyProvider(apiKey string, httpClient *http.Client) *TavilyProvider {
if httpClient == nil {
httpClient = http.DefaultClient
}
return &TavilyProvider{apiKey: apiKey, httpClient: httpClient}
}
func (t *TavilyProvider) Name() string { return tavilyProviderName }
func (t *TavilyProvider) Search(ctx context.Context, req SearchRequest) (*SearchResponse, error) {
maxResults := req.MaxResults
if maxResults <= 0 {
maxResults = defaultMaxResults
}
payload := tavilyRequest{
APIKey: t.apiKey,
Query: req.Query,
MaxResults: maxResults,
SearchDepth: tavilySearchDepthBasic,
}
bodyBytes, err := json.Marshal(payload)
if err != nil {
return nil, fmt.Errorf("tavily: encode request: %w", err)
}
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, tavilySearchEndpoint, bytes.NewReader(bodyBytes))
if err != nil {
return nil, fmt.Errorf("tavily: build request: %w", err)
}
httpReq.Header.Set("Content-Type", "application/json")
resp, err := t.httpClient.Do(httpReq)
if err != nil {
return nil, fmt.Errorf("tavily: request failed: %w", err)
}
defer func() { _ = resp.Body.Close() }()
body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseSize))
if err != nil {
return nil, fmt.Errorf("tavily: read body: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("tavily: status %d: %s", resp.StatusCode, truncateBody(body))
}
var raw tavilyResponse
if err := json.Unmarshal(body, &raw); err != nil {
return nil, fmt.Errorf("tavily: decode response: %w", err)
}
results := make([]SearchResult, 0, len(raw.Results))
for _, r := range raw.Results {
results = append(results, SearchResult{
URL: r.URL,
Title: r.Title,
Snippet: r.Content,
})
}
return &SearchResponse{Results: results, Query: req.Query}, nil
}
type tavilyRequest struct {
APIKey string `json:"api_key"`
Query string `json:"query"`
MaxResults int `json:"max_results"`
SearchDepth string `json:"search_depth"`
}
type tavilyResponse struct {
Results []tavilyResult `json:"results"`
}
type tavilyResult struct {
URL string `json:"url"`
Title string `json:"title"`
Content string `json:"content"`
Score float64 `json:"score"`
}