初始化提交

mcp/web-search/internal/providers/tavily.go

@@ -0,0 +1,134 @@
+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
+}