105 lines
2.4 KiB
Go
105 lines
2.4 KiB
Go
|
package tavily
|
||
|
|
|
||
|
|
import (
|
||
|
|
"bytes"
|
||
|
|
"encoding/json"
|
||
|
|
"fmt"
|
||
|
|
"net/http"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
const (
|
||
|
|
baseURL = "https://api.tavily.com"
|
||
|
|
searchPath = "/search"
|
||
|
|
timeout = 30 * time.Second
|
||
|
|
)
|
||
|
|
|
||
|
|
// Client calls Tavily Search API.
|
||
|
|
type Client struct {
|
||
|
|
APIKey string
|
||
|
|
HTTPClient *http.Client
|
||
|
|
}
|
||
|
|
|
||
|
|
// NewClient creates a Tavily client. apiKey is required for search.
|
||
|
|
func NewClient(apiKey string) *Client {
|
||
|
|
return &Client{
|
||
|
|
APIKey: apiKey,
|
||
|
|
HTTPClient: &http.Client{
|
||
|
|
Timeout: timeout,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// SearchRequest is the POST body for /search.
|
||
|
|
type SearchRequest struct {
|
||
|
|
Query string `json:"query"`
|
||
|
|
SearchDepth string `json:"search_depth,omitempty"` // basic, advanced, etc.
|
||
|
|
MaxResults int `json:"max_results,omitempty"`
|
||
|
|
}
|
||
|
|
|
||
|
|
// SearchResult is one result item.
|
||
|
|
type SearchResult struct {
|
||
|
|
Title string `json:"title"`
|
||
|
|
URL string `json:"url"`
|
||
|
|
Content string `json:"content"`
|
||
|
|
}
|
||
|
|
|
||
|
|
// SearchResponse is the Tavily search response.
|
||
|
|
type SearchResponse struct {
|
||
|
|
Query string `json:"query"`
|
||
|
|
Answer string `json:"answer,omitempty"`
|
||
|
|
Results []SearchResult `json:"results"`
|
||
|
|
}
|
||
|
|
|
||
|
|
// Search runs a web search and returns a single text suitable for passing to Ollama as tool result.
|
||
|
|
func (c *Client) Search(query string) (string, error) {
|
||
|
|
if c.APIKey == "" {
|
||
|
|
return "", fmt.Errorf("tavily: API key not set")
|
||
|
|
}
|
||
|
|
body, err := json.Marshal(SearchRequest{
|
||
|
|
Query: query,
|
||
|
|
MaxResults: 5,
|
||
|
|
})
|
||
|
|
if err != nil {
|
||
|
|
return "", fmt.Errorf("marshal request: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
url := baseURL + searchPath
|
||
|
|
req, err := http.NewRequest(http.MethodPost, url, bytes.NewReader(body))
|
||
|
|
if err != nil {
|
||
|
|
return "", fmt.Errorf("new request: %w", err)
|
||
|
|
}
|
||
|
|
req.Header.Set("Content-Type", "application/json")
|
||
|
|
req.Header.Set("Authorization", "Bearer "+c.APIKey)
|
||
|
|
|
||
|
|
resp, err := c.HTTPClient.Do(req)
|
||
|
|
if err != nil {
|
||
|
|
return "", fmt.Errorf("do request: %w", err)
|
||
|
|
}
|
||
|
|
defer resp.Body.Close()
|
||
|
|
|
||
|
|
if resp.StatusCode != http.StatusOK {
|
||
|
|
return "", fmt.Errorf("tavily returned %d", resp.StatusCode)
|
||
|
|
}
|
||
|
|
|
||
|
|
var out SearchResponse
|
||
|
|
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
|
||
|
|
return "", fmt.Errorf("decode response: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Build a single text for the model: prefer answer if present, else concatenate results.
|
||
|
|
if out.Answer != "" {
|
||
|
|
return out.Answer, nil
|
||
|
|
}
|
||
|
|
var b bytes.Buffer
|
||
|
|
for i, r := range out.Results {
|
||
|
|
if i > 0 {
|
||
|
|
b.WriteString("\n\n")
|
||
|
|
}
|
||
|
|
b.WriteString(r.Title)
|
||
|
|
b.WriteString(": ")
|
||
|
|
b.WriteString(r.Content)
|
||
|
|
}
|
||
|
|
return b.String(), nil
|
||
|
|
}
|