brave search implementation
This commit is contained in:
d1nch8g
@@ -1 +1,122 @@
|
||||
package brave
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
BaseURL = "https://api.search.brave.com/res/v1"
|
||||
httpTimeout = 10 * time.Second
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
apiKey string
|
||||
baseURL string
|
||||
httpClient *http.Client
|
||||
}
|
||||
|
||||
func New(apiKey, baseURL string) *Client {
|
||||
return &Client{
|
||||
apiKey: apiKey,
|
||||
baseURL: baseURL,
|
||||
httpClient: &http.Client{
|
||||
Timeout: httpTimeout,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
type webSearchResult struct {
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
Description string `json:"description"`
|
||||
}
|
||||
|
||||
type webSearchResponse struct {
|
||||
Web struct {
|
||||
Results []webSearchResult `json:"results"`
|
||||
} `json:"web"`
|
||||
Summarizer struct {
|
||||
Key string `json:"key"`
|
||||
} `json:"summarizer"`
|
||||
}
|
||||
|
||||
type summaryMessage struct {
|
||||
Type string `json:"type"`
|
||||
Data string `json:"data"`
|
||||
}
|
||||
|
||||
type summarizerResponse struct {
|
||||
Summary []summaryMessage `json:"summary"`
|
||||
}
|
||||
|
||||
func (c *Client) Search(ctx context.Context, query string) (string, error) {
|
||||
req, _ := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+"/web/search", nil)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("X-Subscription-Token", c.apiKey)
|
||||
|
||||
q := req.URL.Query()
|
||||
q.Add("q", query)
|
||||
q.Add("summary", "1")
|
||||
req.URL.RawQuery = q.Encode()
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return "", fmt.Errorf("status %d: %s", resp.StatusCode, body)
|
||||
}
|
||||
|
||||
var webResp webSearchResponse
|
||||
if err = json.NewDecoder(resp.Body).Decode(&webResp); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if webResp.Summarizer.Key == "" {
|
||||
if len(webResp.Web.Results) > 0 {
|
||||
return webResp.Web.Results[0].Description, nil
|
||||
}
|
||||
return "", errors.New("no results")
|
||||
}
|
||||
|
||||
req, _ = http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+"/summarizer/search", nil)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("X-Subscription-Token", c.apiKey)
|
||||
|
||||
q = req.URL.Query()
|
||||
q.Add("key", webResp.Summarizer.Key)
|
||||
req.URL.RawQuery = q.Encode()
|
||||
|
||||
resp, err = c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
return "", fmt.Errorf("summarizer status %d: %s", resp.StatusCode, body)
|
||||
}
|
||||
|
||||
var sumResp summarizerResponse
|
||||
if err = json.NewDecoder(resp.Body).Decode(&sumResp); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
for _, msg := range sumResp.Summary {
|
||||
if msg.Type == "token" {
|
||||
return msg.Data, nil
|
||||
}
|
||||
}
|
||||
|
||||
return "", errors.New("no token in summary")
|
||||
}
|
||||
|
||||
@@ -1 +1,209 @@
|
||||
package brave
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) {
|
||||
return f(r)
|
||||
}
|
||||
|
||||
func TestNew(t *testing.T) {
|
||||
client := New("key", "url")
|
||||
assert.Equal(t, "key", client.apiKey)
|
||||
assert.Equal(t, "url", client.baseURL)
|
||||
assert.NotNil(t, client.httpClient)
|
||||
}
|
||||
|
||||
func TestSearch_Success_WithSummarizer(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/web/search":
|
||||
assert.Equal(t, "key", r.Header.Get("X-Subscription-Token"))
|
||||
resp := webSearchResponse{}
|
||||
resp.Summarizer.Key = "test-key"
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
case "/summarizer/search":
|
||||
assert.Equal(t, "key", r.Header.Get("X-Subscription-Token"))
|
||||
assert.Equal(t, "test-key", r.URL.Query().Get("key"))
|
||||
resp := summarizerResponse{
|
||||
Summary: []summaryMessage{
|
||||
{Type: "token", Data: "summary text"},
|
||||
},
|
||||
}
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := New("key", server.URL)
|
||||
result, err := client.Search(context.Background(), "query")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "summary text", result)
|
||||
}
|
||||
|
||||
func TestSearch_Success_NoSummarizer_FallbackToDescription(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
resp := webSearchResponse{}
|
||||
resp.Web.Results = []webSearchResult{
|
||||
{Description: "fallback description"},
|
||||
}
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := New("key", server.URL)
|
||||
result, err := client.Search(context.Background(), "query")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "fallback description", result)
|
||||
}
|
||||
|
||||
func TestSearch_NoResults(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
resp := webSearchResponse{}
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := New("key", server.URL)
|
||||
_, err := client.Search(context.Background(), "query")
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "no results")
|
||||
}
|
||||
|
||||
func TestSearch_HTTPError_WebSearch(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := New("key", server.URL)
|
||||
_, err := client.Search(context.Background(), "query")
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "status 500")
|
||||
}
|
||||
|
||||
func TestSearch_InvalidJSON_WebSearch(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Write([]byte("not json"))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := New("key", server.URL)
|
||||
_, err := client.Search(context.Background(), "query")
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestSearch_HTTPError_Summarizer(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/web/search" {
|
||||
resp := webSearchResponse{}
|
||||
resp.Summarizer.Key = "test-key"
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
} else {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := New("key", server.URL)
|
||||
_, err := client.Search(context.Background(), "query")
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "summarizer status 500")
|
||||
}
|
||||
|
||||
func TestSearch_InvalidJSON_Summarizer(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/web/search" {
|
||||
resp := webSearchResponse{}
|
||||
resp.Summarizer.Key = "test-key"
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
} else {
|
||||
w.Write([]byte("not json"))
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := New("key", server.URL)
|
||||
_, err := client.Search(context.Background(), "query")
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestSearch_NoTokenInSummary(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/web/search" {
|
||||
resp := webSearchResponse{}
|
||||
resp.Summarizer.Key = "test-key"
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
} else {
|
||||
resp := summarizerResponse{
|
||||
Summary: []summaryMessage{
|
||||
{Type: "enum_start", Data: "ol"},
|
||||
},
|
||||
}
|
||||
json.NewEncoder(w).Encode(resp)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := New("key", server.URL)
|
||||
_, err := client.Search(context.Background(), "query")
|
||||
assert.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "no token in summary")
|
||||
}
|
||||
|
||||
func TestSearch_DoError_WebSearch(t *testing.T) {
|
||||
client := New("key", "http://example.com")
|
||||
|
||||
client.httpClient.Transport = roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
if r.URL.Path == "/web/search" {
|
||||
return nil, errors.New("web search network error")
|
||||
}
|
||||
t.Fatalf("unexpected call to %s", r.URL.Path)
|
||||
return nil, nil
|
||||
})
|
||||
|
||||
_, err := client.Search(context.Background(), "query")
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "web search network error")
|
||||
}
|
||||
|
||||
func TestSearch_DoError_Summarizer(t *testing.T) {
|
||||
client := New("key", "http://example.com")
|
||||
|
||||
client.httpClient.Transport = roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
switch r.URL.Path {
|
||||
case "/web/search":
|
||||
// successful first response
|
||||
body := `{"summarizer":{"key":"test-key"}}`
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
Header: make(http.Header),
|
||||
}, nil
|
||||
|
||||
case "/summarizer/search":
|
||||
// 💥 fail only here
|
||||
return nil, errors.New("summarizer network error")
|
||||
}
|
||||
|
||||
t.Fatalf("unexpected path %s", r.URL.Path)
|
||||
return nil, nil
|
||||
})
|
||||
|
||||
_, err := client.Search(context.Background(), "query")
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "summarizer network error")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user