sift/internal/search/tavily.go

package search

import (
	"bytes"
	"context"
	"encoding/json"
	"io"
	"net/http"

	"git.wntrmute.dev/kyle/sift/internal/search/dial"
)

// Tavily is POST {url} JSON. Never requests or returns answer.
type Tavily struct {
	Client *dial.Client
	Key    string
}

func (s *Tavily) Name() string { return NameTavily }
func (s *Tavily) Ready() bool  { return s != nil && s.Client.Ready() }

func (s *Tavily) Search(ctx context.Context, query string, k int, f Filter) SourceOut {
	if s == nil || s.Client == nil {
		return SourceOut{Warn: NameTavily + ": error"}
	}
	if w := pinErr(ctx, s.Client, NameTavily); w != "" {
		return SourceOut{Warn: w}
	}
	if k <= 0 {
		k = DefaultK
	}
	if k > 20 {
		k = 20
	}
	payload, _ := json.Marshal(tavilyBody(query, k, f))
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.Client.URL.String(), bytes.NewReader(payload))
	if err != nil {
		return SourceOut{Warn: warnErr(NameTavily, err)}
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+s.Key)
	resp, err := s.Client.HTTP.Do(req)
	if err != nil {
		return SourceOut{Warn: warnErr(NameTavily, err)}
	}
	if resp.StatusCode != 200 {
		code := resp.StatusCode
		resp.Body.Close()
		return SourceOut{Warn: warnStatus(NameTavily, code)}
	}
	raw, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	resp.Body.Close()
	if err != nil {
		return SourceOut{Warn: warnErr(NameTavily, err)}
	}
	hits, credits, err := mapTavily(raw)
	if err != nil {
		return SourceOut{Warn: NameTavily + ": error"}
	}
	return SourceOut{Hits: hits, Credits: credits}
}

func tavilyBody(query string, k int, f Filter) map[string]any {
	body := map[string]any{
		"query":               query,
		"max_results":         k,
		"include_answer":      false,
		"include_raw_content": false,
		"include_images":      false,
		"search_depth":        "basic",
		"auto_parameters":     false,
	}
	if len(f.IncludeDomains) > 0 {
		body["include_domains"] = f.IncludeDomains
		mode := f.IncludeDomainsMode
		if mode == "" {
			mode = "restrict"
		}
		body["include_domains_mode"] = mode
	}
	if len(f.ExcludeDomains) > 0 {
		body["exclude_domains"] = f.ExcludeDomains
	}
	if f.TimeRange != "" {
		body["time_range"] = f.TimeRange
	} else {
		if f.StartDate != "" {
			body["start_date"] = f.StartDate
		}
		if f.EndDate != "" {
			body["end_date"] = f.EndDate
		}
	}
	if f.Topic != "" && f.Topic != "general" {
		body["topic"] = f.Topic
	}
	if f.ExactMatch {
		body["exact_match"] = true
	}
	if f.Topic == "news" || f.TimeRange != "" || f.StartDate != "" || f.EndDate != "" {
		body["include_published_date"] = true
	}
	return body
}

func mapTavily(raw []byte) ([]Hit, int, error) {
	var env struct {
		Answer  string `json:"answer"`
		Results []struct {
			URL           string  `json:"url"`
			Title         string  `json:"title"`
			Content       string  `json:"content"`
			Score         float64 `json:"score"`
			PublishedDate string  `json:"published_date"`
		} `json:"results"`
		Usage struct {
			Credits int `json:"credits"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(raw, &env); err != nil {
		return nil, 0, err
	}
	_ = env.Answer // dropped; never returned
	hits := make([]Hit, 0, len(env.Results))
	for _, r := range env.Results {
		if r.URL == "" {
			continue
		}
		hits = append(hits, Hit{
			URL: r.URL, Title: r.Title, Snippet: r.Content,
			Score: r.Score, Source: NameTavily, PublishedDate: r.PublishedDate,
		})
	}
	credits := env.Usage.Credits
	if credits <= 0 {
		credits = 1
	}
	return hits, credits, nil
}