135 lines
3.1 KiB
Go
135 lines
3.1 KiB
Go
package providers
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"net/http"
|
|
"time"
|
|
|
|
"github.com/chenxin/hunjing/mcp/web-search/internal/models"
|
|
)
|
|
|
|
const tavilyAPIBase = "https://api.tavily.com"
|
|
|
|
type tavilyRequest struct {
|
|
APIKey string `json:"api_key"`
|
|
Query string `json:"query"`
|
|
MaxResults int `json:"max_results"`
|
|
}
|
|
|
|
type tavilyResponse struct {
|
|
Answer string `json:"answer"`
|
|
Query string `json:"query"`
|
|
ResponseTime float64 `json:"response_time"`
|
|
Results []tavilyResult `json:"results"`
|
|
}
|
|
|
|
type tavilyResult struct {
|
|
Title string `json:"title"`
|
|
URL string `json:"url"`
|
|
Content string `json:"content"`
|
|
Score float64 `json:"score"`
|
|
PublishedAt string `json:"published_date"`
|
|
}
|
|
|
|
// TavilyProvider 实现 SearchProvider。
|
|
type TavilyProvider struct {
|
|
client *http.Client
|
|
apiBase string
|
|
maxResults int
|
|
}
|
|
|
|
func NewTavilyProvider(timeout time.Duration, maxResults int) *TavilyProvider {
|
|
return &TavilyProvider{
|
|
client: &http.Client{
|
|
Timeout: timeout,
|
|
},
|
|
apiBase: tavilyAPIBase,
|
|
maxResults: maxResults,
|
|
}
|
|
}
|
|
|
|
func (t *TavilyProvider) Name() models.ProviderKind {
|
|
return models.ProviderTavily
|
|
}
|
|
|
|
func (t *TavilyProvider) MaxResults() int {
|
|
return t.maxResults
|
|
}
|
|
|
|
func (t *TavilyProvider) Search(ctx context.Context, apiKey string, req *models.SearchRequest) (*models.SearchResult, error) {
|
|
maxRes := req.MaxResults
|
|
if maxRes <= 0 || maxRes > t.maxResults {
|
|
maxRes = t.maxResults
|
|
}
|
|
|
|
body := tavilyRequest{
|
|
APIKey: apiKey,
|
|
Query: req.Query,
|
|
MaxResults: maxRes,
|
|
}
|
|
payload, _ := json.Marshal(body)
|
|
|
|
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, t.apiBase+"/search", bytes.NewReader(payload))
|
|
if err != nil {
|
|
return nil, fmt.Errorf("tavily: create request: %w", err)
|
|
}
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
|
|
start := time.Now()
|
|
resp, err := t.client.Do(httpReq)
|
|
duration := time.Since(start)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("tavily: do request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
raw, err := io.ReadAll(resp.Body)
|
|
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, string(raw))
|
|
}
|
|
|
|
var tr tavilyResponse
|
|
if err := json.Unmarshal(raw, &tr); err != nil {
|
|
return nil, fmt.Errorf("tavily: decode response: %w", err)
|
|
}
|
|
|
|
items := make([]models.ResultItem, 0, len(tr.Results))
|
|
for _, r := range tr.Results {
|
|
items = append(items, models.ResultItem{
|
|
Title: r.Title,
|
|
URL: r.URL,
|
|
Snippet: r.Content,
|
|
Score: r.Score,
|
|
PublishedAt: r.PublishedAt,
|
|
})
|
|
}
|
|
|
|
return &models.SearchResult{
|
|
Provider: models.ProviderTavily,
|
|
Query: req.Query,
|
|
Results: items,
|
|
TotalCount: len(items),
|
|
DurationMs: duration.Milliseconds(),
|
|
Cached: false,
|
|
}, nil
|
|
}
|
|
|
|
func (t *TavilyProvider) HealthCheck(ctx context.Context) error {
|
|
var d net.Dialer
|
|
conn, err := d.DialContext(ctx, "tcp", "api.tavily.com:443")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
conn.Close()
|
|
return nil
|
|
}
|