kyle
429f6b207e
A 200 with a bad Tavily body is tavily: error, not an empty success that would cache for the hour. Upstream clients stay in the source list if DNS is down at boot and retry the pin on Search; /healthz lists them as unresolved until a pin lands.
109 lines
2.8 KiB
Go
109 lines
2.8 KiB
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() }
|
|
|
|
type tavilyReq struct {
|
|
Query string `json:"query"`
|
|
MaxResults int `json:"max_results"`
|
|
IncludeAnswer bool `json:"include_answer"`
|
|
IncludeRawContent bool `json:"include_raw_content"`
|
|
IncludeImages bool `json:"include_images"`
|
|
SearchDepth string `json:"search_depth"`
|
|
}
|
|
|
|
func (s *Tavily) Search(ctx context.Context, query string, k int) 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(tavilyReq{
|
|
Query: query,
|
|
MaxResults: k,
|
|
IncludeAnswer: false,
|
|
IncludeRawContent: false,
|
|
IncludeImages: false,
|
|
SearchDepth: "basic",
|
|
})
|
|
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 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"`
|
|
} `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})
|
|
}
|
|
credits := env.Usage.Credits
|
|
if credits <= 0 {
|
|
credits = 1
|
|
}
|
|
return hits, credits, nil
|
|
}
|