144 lines
3.4 KiB
Go
144 lines
3.4 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"sync/atomic"
|
|
)
|
|
|
|
const tavilySearchURL = "https://api.tavily.com/search"
|
|
|
|
type tavilyClient struct {
|
|
keys []string
|
|
idx atomic.Uint64
|
|
http *http.Client
|
|
baseURL string // empty → https://api.tavily.com/search
|
|
}
|
|
|
|
func newTavilyClient(keys []string) *tavilyClient {
|
|
return newTavilyClientWithOptions(keys, nil, "")
|
|
}
|
|
|
|
func newTavilyClientWithOptions(keys []string, httpClient *http.Client, baseURL string) *tavilyClient {
|
|
trimmed := make([]string, 0, len(keys))
|
|
for _, key := range keys {
|
|
if k := strings.TrimSpace(key); k != "" {
|
|
trimmed = append(trimmed, k)
|
|
}
|
|
}
|
|
if httpClient == nil {
|
|
httpClient = &http.Client{}
|
|
}
|
|
return &tavilyClient{
|
|
keys: trimmed,
|
|
http: httpClient,
|
|
baseURL: strings.TrimSpace(baseURL),
|
|
}
|
|
}
|
|
|
|
func (c *tavilyClient) searchEndpoint() string {
|
|
if c != nil && c.baseURL != "" {
|
|
return c.baseURL
|
|
}
|
|
return tavilySearchURL
|
|
}
|
|
|
|
func (c *tavilyClient) available() bool {
|
|
return c != nil && len(c.keys) > 0
|
|
}
|
|
|
|
func (c *tavilyClient) nextKey() string {
|
|
if len(c.keys) == 0 {
|
|
return ""
|
|
}
|
|
n := c.idx.Add(1)
|
|
return c.keys[int(n-1)%len(c.keys)]
|
|
}
|
|
|
|
type tavilySearchRequest struct {
|
|
APIKey string `json:"api_key"`
|
|
Query string `json:"query"`
|
|
SearchDepth string `json:"search_depth,omitempty"`
|
|
MaxResults int `json:"max_results,omitempty"`
|
|
IncludeAnswer bool `json:"include_answer,omitempty"`
|
|
}
|
|
|
|
type tavilySearchResponse struct {
|
|
Answer string `json:"answer"`
|
|
Results []struct {
|
|
Title string `json:"title"`
|
|
URL string `json:"url"`
|
|
Content string `json:"content"`
|
|
} `json:"results"`
|
|
}
|
|
|
|
type claudeWebSearchHit struct {
|
|
Title string
|
|
URL string
|
|
Snippet string
|
|
}
|
|
|
|
func (c *tavilyClient) search(ctx context.Context, query string, maxResults int) ([]claudeWebSearchHit, string, error) {
|
|
if !c.available() {
|
|
return nil, "", fmt.Errorf("tavily_api_keys is empty")
|
|
}
|
|
query = strings.TrimSpace(query)
|
|
if query == "" {
|
|
return nil, "", fmt.Errorf("web search query is empty")
|
|
}
|
|
if maxResults <= 0 {
|
|
maxResults = 5
|
|
}
|
|
payload, errMarshal := json.Marshal(tavilySearchRequest{
|
|
APIKey: c.nextKey(),
|
|
Query: query,
|
|
SearchDepth: "basic",
|
|
MaxResults: maxResults,
|
|
IncludeAnswer: true,
|
|
})
|
|
if errMarshal != nil {
|
|
return nil, "", errMarshal
|
|
}
|
|
req, errNew := http.NewRequestWithContext(ctx, http.MethodPost, c.searchEndpoint(), bytes.NewReader(payload))
|
|
if errNew != nil {
|
|
return nil, "", errNew
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
resp, errDo := c.http.Do(req)
|
|
if errDo != nil {
|
|
return nil, "", errDo
|
|
}
|
|
defer func() { _ = resp.Body.Close() }()
|
|
body, errRead := io.ReadAll(resp.Body)
|
|
if errRead != nil {
|
|
return nil, "", errRead
|
|
}
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
return nil, "", fmt.Errorf("tavily http %d: %s", resp.StatusCode, truncate(string(body), 512))
|
|
}
|
|
var parsed tavilySearchResponse
|
|
if errDecode := json.Unmarshal(body, &parsed); errDecode != nil {
|
|
return nil, "", errDecode
|
|
}
|
|
hits := make([]claudeWebSearchHit, 0, len(parsed.Results))
|
|
for _, r := range parsed.Results {
|
|
hits = append(hits, claudeWebSearchHit{
|
|
Title: strings.TrimSpace(r.Title),
|
|
URL: strings.TrimSpace(r.URL),
|
|
Snippet: strings.TrimSpace(r.Content),
|
|
})
|
|
}
|
|
return hits, strings.TrimSpace(parsed.Answer), nil
|
|
}
|
|
|
|
func truncate(s string, max int) string {
|
|
if len(s) <= max {
|
|
return s
|
|
}
|
|
return s[:max] + "..."
|
|
}
|