p3: search API, pinned sources, cache, CLI
POST /v1/search and /v2/search fan in SearXNG, Brave LLM Context, and Tavily from config-fixed endpoints. Public sources refuse private A records; SearXNG is a loopback-only client, not the fetch guard. Cache keys include caller class. Restricted callers search like the fleet; logs omit query, hosts, URLs, and upstream errors. Degraded results are not stored. Tavily sends search_depth=basic and drops answer. Missing keys surface in /healthz. Default-on sources stay off until the bake-off. Query harvest is not in this commit.
This commit is contained in:
@@ -0,0 +1,100 @@
|
||||
package search
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/config"
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
"git.wntrmute.dev/kyle/sift/internal/search/dial"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultSearx = "http://127.0.0.1:8888"
|
||||
defaultBrave = "https://api.search.brave.com/res/v1/llm/context"
|
||||
defaultTav = "https://api.tavily.com/search"
|
||||
)
|
||||
|
||||
// Build constructs sources from config. Missing keys are listed on Engine.Degraded
|
||||
// and logged once; those sources are not attempted.
|
||||
func Build(cfg config.Search, cache *dcache.DB, emb Embedder, log *slog.Logger) *Engine {
|
||||
timeout := config.ParseDuration(cfg.Timeout, 8*time.Second)
|
||||
ttl := config.ParseDuration(cfg.CacheTTL, time.Hour)
|
||||
k := cfg.K
|
||||
if k <= 0 {
|
||||
k = DefaultK
|
||||
}
|
||||
e := &Engine{
|
||||
Timeout: timeout,
|
||||
DefaultK: k,
|
||||
RerankOn: cfg.RerankOn(),
|
||||
Cache: cache,
|
||||
CacheTTL: ttl,
|
||||
Embedder: emb,
|
||||
Log: log,
|
||||
TavilyCap: cfg.Tavily.DailyCreditCap,
|
||||
}
|
||||
if cfg.SearxNG.Enabled {
|
||||
u := cfg.SearxNG.URL
|
||||
if u == "" {
|
||||
u = defaultSearx
|
||||
}
|
||||
c, err := dial.New(u, dial.Options{Loopback: true, Timeout: timeout})
|
||||
if err != nil {
|
||||
e.Degraded = append(e.Degraded, "search.searxng: "+err.Error())
|
||||
if log != nil {
|
||||
log.Error("search source", "source", NameSearxNG, "err", err)
|
||||
}
|
||||
} else {
|
||||
e.Sources = append(e.Sources, &SearxNG{Client: c})
|
||||
}
|
||||
}
|
||||
if cfg.BraveLLM.Enabled {
|
||||
key := os.Getenv("BRAVE_API_KEY")
|
||||
if key == "" {
|
||||
e.Degraded = append(e.Degraded, "search.brave_llm: missing key")
|
||||
if log != nil {
|
||||
log.Error("search source", "source", NameBraveLLM, "err", "missing key")
|
||||
}
|
||||
} else {
|
||||
u := cfg.BraveLLM.URL
|
||||
if u == "" {
|
||||
u = defaultBrave
|
||||
}
|
||||
c, err := dial.New(u, dial.Options{Timeout: timeout})
|
||||
if err != nil {
|
||||
e.Degraded = append(e.Degraded, "search.brave_llm: "+err.Error())
|
||||
if log != nil {
|
||||
log.Error("search source", "source", NameBraveLLM, "err", err)
|
||||
}
|
||||
} else {
|
||||
e.Sources = append(e.Sources, &BraveLLM{Client: c, Key: key})
|
||||
}
|
||||
}
|
||||
}
|
||||
if cfg.Tavily.Enabled {
|
||||
key := os.Getenv("TAVILY_API_KEY")
|
||||
if key == "" {
|
||||
e.Degraded = append(e.Degraded, "search.tavily: missing key")
|
||||
if log != nil {
|
||||
log.Error("search source", "source", NameTavily, "err", "missing key")
|
||||
}
|
||||
} else {
|
||||
u := cfg.Tavily.URL
|
||||
if u == "" {
|
||||
u = defaultTav
|
||||
}
|
||||
c, err := dial.New(u, dial.Options{Timeout: timeout})
|
||||
if err != nil {
|
||||
e.Degraded = append(e.Degraded, "search.tavily: "+err.Error())
|
||||
if log != nil {
|
||||
log.Error("search source", "source", NameTavily, "err", err)
|
||||
}
|
||||
} else {
|
||||
e.Sources = append(e.Sources, &Tavily{Client: c, Key: key})
|
||||
}
|
||||
}
|
||||
}
|
||||
return e
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
package search
|
||||
|
||||
import (
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// CanonicalURL lowercases scheme+host, strips default ports and fragments,
|
||||
// and drops a trailing slash on the path. www. is left as-is.
|
||||
func CanonicalURL(raw string) string {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return ""
|
||||
}
|
||||
u, err := url.Parse(raw)
|
||||
if err != nil || u.Host == "" && u.Opaque == "" {
|
||||
return strings.TrimRight(raw, "/")
|
||||
}
|
||||
u.Fragment = ""
|
||||
u.Scheme = strings.ToLower(u.Scheme)
|
||||
host := strings.ToLower(u.Hostname())
|
||||
port := u.Port()
|
||||
switch {
|
||||
case port == "80" && u.Scheme == "http":
|
||||
port = ""
|
||||
case port == "443" && u.Scheme == "https":
|
||||
port = ""
|
||||
}
|
||||
if port != "" {
|
||||
u.Host = host + ":" + port
|
||||
} else {
|
||||
u.Host = host
|
||||
}
|
||||
if u.Path == "/" {
|
||||
u.Path = ""
|
||||
} else if len(u.Path) > 1 && strings.HasSuffix(u.Path, "/") {
|
||||
u.Path = strings.TrimSuffix(u.Path, "/")
|
||||
}
|
||||
return u.String()
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
package search
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestCanonicalURL(t *testing.T) {
|
||||
cases := []struct {
|
||||
in, want string
|
||||
}{
|
||||
{"https://Example.COM/a/b#frag", "https://example.com/a/b"},
|
||||
{"http://example.com:80/x", "http://example.com/x"},
|
||||
{"https://example.com:443/x", "https://example.com/x"},
|
||||
{"https://example.com/x/", "https://example.com/x"},
|
||||
{"https://www.example.com/x", "https://www.example.com/x"},
|
||||
{"https://example.com/", "https://example.com"},
|
||||
{"HTTP://Ex.com:80/a/?q=1", "http://ex.com/a?q=1"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got := CanonicalURL(c.in)
|
||||
if got != c.want {
|
||||
t.Errorf("CanonicalURL(%q)=%q want %q", c.in, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,232 @@
|
||||
// Package dial builds an HTTP client pinned to one config URL.
|
||||
// The agent cannot name the target; only the operator-configured host:port is reachable.
|
||||
package dial
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/guard"
|
||||
)
|
||||
|
||||
const DefaultRefresh = 10 * time.Minute
|
||||
|
||||
// Lookup resolves a hostname. Tests inject a fake.
|
||||
type Lookup func(ctx context.Context, host string) ([]netip.Addr, error)
|
||||
|
||||
type Options struct {
|
||||
// Loopback permits loopback (and only loopback) addresses. Public sources
|
||||
// must leave this false: every record must pass guard.CheckAddr.
|
||||
Loopback bool
|
||||
Lookup Lookup
|
||||
Refresh time.Duration
|
||||
Timeout time.Duration
|
||||
}
|
||||
|
||||
// Client is an http.Client that may only connect to the config URL's host:port.
|
||||
type Client struct {
|
||||
URL *url.URL
|
||||
HTTP *http.Client
|
||||
Port string
|
||||
|
||||
loopback bool
|
||||
lookup Lookup
|
||||
refresh time.Duration
|
||||
|
||||
mu sync.Mutex
|
||||
ips []netip.Addr
|
||||
resolved time.Time
|
||||
}
|
||||
|
||||
func New(raw string, opt Options) (*Client, error) {
|
||||
u, err := parseConfigURL(raw)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
port := u.Port()
|
||||
if port == "" {
|
||||
if u.Scheme == "https" {
|
||||
port = "443"
|
||||
} else {
|
||||
port = "80"
|
||||
}
|
||||
}
|
||||
if opt.Refresh <= 0 {
|
||||
opt.Refresh = DefaultRefresh
|
||||
}
|
||||
if opt.Timeout <= 0 {
|
||||
opt.Timeout = 8 * time.Second
|
||||
}
|
||||
lookup := opt.Lookup
|
||||
if lookup == nil {
|
||||
lookup = defaultLookup
|
||||
}
|
||||
c := &Client{
|
||||
URL: u,
|
||||
Port: port,
|
||||
loopback: opt.Loopback,
|
||||
lookup: lookup,
|
||||
refresh: opt.Refresh,
|
||||
}
|
||||
if err := c.resolve(context.Background()); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
d := &net.Dialer{
|
||||
Timeout: opt.Timeout,
|
||||
Control: c.control,
|
||||
}
|
||||
tr := &http.Transport{
|
||||
Proxy: nil,
|
||||
DialContext: d.DialContext,
|
||||
ForceAttemptHTTP2: true,
|
||||
TLSHandshakeTimeout: 10 * time.Second,
|
||||
ResponseHeaderTimeout: opt.Timeout,
|
||||
TLSClientConfig: &tls.Config{MinVersion: tls.VersionTLS12}, // system roots; InsecureSkipVerify stays false
|
||||
DisableCompression: false,
|
||||
}
|
||||
c.HTTP = &http.Client{
|
||||
Timeout: opt.Timeout + time.Second,
|
||||
Transport: tr,
|
||||
CheckRedirect: c.checkRedirect,
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
func parseConfigURL(raw string) (*url.URL, error) {
|
||||
u, err := url.Parse(raw)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if u.Scheme != "http" && u.Scheme != "https" {
|
||||
return nil, fmt.Errorf("search upstream: scheme %q", u.Scheme)
|
||||
}
|
||||
if u.User != nil {
|
||||
return nil, fmt.Errorf("search upstream: userinfo not allowed")
|
||||
}
|
||||
if u.Hostname() == "" {
|
||||
return nil, fmt.Errorf("search upstream: empty host")
|
||||
}
|
||||
return u, nil
|
||||
}
|
||||
|
||||
func defaultLookup(ctx context.Context, host string) ([]netip.Addr, error) {
|
||||
if ip, err := netip.ParseAddr(host); err == nil {
|
||||
return []netip.Addr{ip}, nil
|
||||
}
|
||||
return net.DefaultResolver.LookupNetIP(ctx, "ip", host)
|
||||
}
|
||||
|
||||
func (c *Client) resolve(ctx context.Context) error {
|
||||
addrs, err := c.lookup(ctx, c.URL.Hostname())
|
||||
if err != nil {
|
||||
return fmt.Errorf("search resolve: %w", err)
|
||||
}
|
||||
if len(addrs) == 0 {
|
||||
return fmt.Errorf("search resolve: no addresses")
|
||||
}
|
||||
for i, a := range addrs {
|
||||
addrs[i] = a.Unmap()
|
||||
if err := c.checkAddr(addrs[i]); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
c.mu.Lock()
|
||||
c.ips = addrs
|
||||
c.resolved = time.Now()
|
||||
c.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) checkAddr(a netip.Addr) error {
|
||||
if c.loopback {
|
||||
if !a.IsLoopback() {
|
||||
return fmt.Errorf("search loopback source got non-loopback %s", a)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if err := guard.CheckAddr(a); err != nil {
|
||||
return fmt.Errorf("search public source: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) allow(ip netip.Addr, port string) bool {
|
||||
if port != c.Port {
|
||||
return false
|
||||
}
|
||||
ip = ip.Unmap()
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
for _, a := range c.ips {
|
||||
if a == ip {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (c *Client) control(network, address string, _ syscall.RawConn) error {
|
||||
switch network {
|
||||
case "tcp", "tcp4", "tcp6":
|
||||
default:
|
||||
return fmt.Errorf("search dial: network %s", network)
|
||||
}
|
||||
host, port, err := net.SplitHostPort(address)
|
||||
if err != nil {
|
||||
return fmt.Errorf("search dial: %w", err)
|
||||
}
|
||||
ip, err := netip.ParseAddr(host)
|
||||
if err != nil {
|
||||
return fmt.Errorf("search dial: not an IP")
|
||||
}
|
||||
if !c.allow(ip, port) {
|
||||
return fmt.Errorf("search dial refused")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Client) checkRedirect(req *http.Request, via []*http.Request) error {
|
||||
if len(via) > 1 {
|
||||
return fmt.Errorf("search redirect: extra hop")
|
||||
}
|
||||
if len(via) == 1 && !strings.EqualFold(req.URL.Hostname(), via[0].URL.Hostname()) {
|
||||
return fmt.Errorf("search redirect: different host")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RefreshIfStale re-resolves when the pin is older than Refresh.
|
||||
func (c *Client) RefreshIfStale(ctx context.Context) error {
|
||||
c.mu.Lock()
|
||||
stale := time.Since(c.resolved) >= c.refresh
|
||||
c.mu.Unlock()
|
||||
if !stale {
|
||||
return nil
|
||||
}
|
||||
return c.resolve(ctx)
|
||||
}
|
||||
|
||||
func (c *Client) IPs() []netip.Addr {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
out := make([]netip.Addr, len(c.ips))
|
||||
copy(out, c.ips)
|
||||
return out
|
||||
}
|
||||
|
||||
func (c *Client) TLSConfig() *tls.Config {
|
||||
tr, _ := c.HTTP.Transport.(*http.Transport)
|
||||
if tr == nil {
|
||||
return nil
|
||||
}
|
||||
return tr.TLSClientConfig
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
package dial_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/netip"
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/search/dial"
|
||||
)
|
||||
|
||||
func TestLoopbackAllowsConfiguredPort(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(200)
|
||||
_, _ = w.Write([]byte("ok"))
|
||||
}))
|
||||
defer ts.Close()
|
||||
c, err := dial.New(ts.URL, dial.Options{Loopback: true, Timeout: 0})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp, err := c.HTTP.Get(ts.URL + "/")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != 200 {
|
||||
t.Fatalf("status %d", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoopbackRefusesOtherPort(t *testing.T) {
|
||||
c, err := dial.New("http://127.0.0.1:8888", dial.Options{
|
||||
Loopback: true,
|
||||
Lookup: func(_ context.Context, _ string) ([]netip.Addr, error) {
|
||||
return []netip.Addr{netip.MustParseAddr("127.0.0.1")}, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp, err := c.HTTP.Get("http://127.0.0.1:11434/")
|
||||
if err == nil {
|
||||
resp.Body.Close()
|
||||
t.Fatal("dialed other loopback port")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "refused") && !strings.Contains(err.Error(), "search dial") {
|
||||
t.Fatalf("err %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPublicPoisonedPrivateRefused(t *testing.T) {
|
||||
_, err := dial.New("https://api.tavily.com/search", dial.Options{
|
||||
Loopback: false,
|
||||
Lookup: func(_ context.Context, host string) ([]netip.Addr, error) {
|
||||
return []netip.Addr{netip.MustParseAddr("127.0.0.1")}, nil
|
||||
},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("poisoned A record allowed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPublicPrivateRFC1918Refused(t *testing.T) {
|
||||
_, err := dial.New("https://api.search.brave.com/res/v1/llm/context", dial.Options{
|
||||
Lookup: func(_ context.Context, _ string) ([]netip.Addr, error) {
|
||||
return []netip.Addr{netip.MustParseAddr("10.0.0.5")}, nil
|
||||
},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("private A allowed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestTLSSystemRoots(t *testing.T) {
|
||||
c, err := dial.New("https://api.tavily.com/search", dial.Options{
|
||||
Lookup: func(_ context.Context, _ string) ([]netip.Addr, error) {
|
||||
return []netip.Addr{netip.MustParseAddr("1.1.1.1")}, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg := c.TLSConfig()
|
||||
if cfg == nil {
|
||||
t.Fatal("nil tls")
|
||||
}
|
||||
if cfg.InsecureSkipVerify {
|
||||
t.Fatal("InsecureSkipVerify")
|
||||
}
|
||||
if cfg.RootCAs != nil {
|
||||
t.Fatal("want system roots (nil RootCAs)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoUserinfo(t *testing.T) {
|
||||
_, err := dial.New("http://u:p@127.0.0.1:8888", dial.Options{Loopback: true})
|
||||
if err == nil {
|
||||
t.Fatal("userinfo allowed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseRejectsNonHTTP(t *testing.T) {
|
||||
_, err := dial.New("ftp://example.com/", dial.Options{})
|
||||
if err == nil {
|
||||
t.Fatal("ftp allowed")
|
||||
}
|
||||
_ = url.URL{}
|
||||
}
|
||||
@@ -0,0 +1,331 @@
|
||||
package search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
)
|
||||
|
||||
type Engine struct {
|
||||
Sources []Source
|
||||
Timeout time.Duration
|
||||
DefaultK int
|
||||
RerankOn bool
|
||||
Cache *dcache.DB
|
||||
CacheTTL time.Duration
|
||||
Embedder Embedder
|
||||
Log *slog.Logger
|
||||
Degraded []string // healthz: "search.brave_llm: missing key"
|
||||
TavilyCap int
|
||||
Now func() time.Time
|
||||
|
||||
mu sync.Mutex
|
||||
Queries map[string]map[string]uint64
|
||||
Credits map[string]map[string]uint64
|
||||
Errors map[string]map[string]uint64
|
||||
CacheHit uint64
|
||||
CacheMiss uint64
|
||||
RerankFallback uint64
|
||||
DailyCredits map[string]uint64
|
||||
Capped map[string]uint64
|
||||
}
|
||||
|
||||
type RunResult struct {
|
||||
Hits []Hit
|
||||
Warnings []string
|
||||
SourcesOK []string
|
||||
SourcesFailed []string
|
||||
SearchMS int64
|
||||
RerankMS int64
|
||||
Cached bool
|
||||
Attempted []string
|
||||
}
|
||||
|
||||
func (e *Engine) now() time.Time {
|
||||
if e != nil && e.Now != nil {
|
||||
return e.Now()
|
||||
}
|
||||
return time.Now()
|
||||
}
|
||||
|
||||
func (e *Engine) ttl() time.Duration {
|
||||
if e != nil && e.CacheTTL > 0 {
|
||||
return e.CacheTTL
|
||||
}
|
||||
return time.Hour
|
||||
}
|
||||
|
||||
func (e *Engine) timeout() time.Duration {
|
||||
if e != nil && e.Timeout > 0 {
|
||||
return e.Timeout
|
||||
}
|
||||
return 8 * time.Second
|
||||
}
|
||||
|
||||
func (e *Engine) Run(ctx context.Context, query string, k int, doRerank bool, caller, class string) RunResult {
|
||||
start := e.now()
|
||||
k = ClampK(k, e.DefaultK)
|
||||
if e.DefaultK <= 0 {
|
||||
k = ClampK(k, DefaultK)
|
||||
}
|
||||
if class == "" {
|
||||
class = dcache.ClassFleet
|
||||
}
|
||||
|
||||
attempted := e.attemptedNames()
|
||||
key := dcache.SearchKey(query, attempted, k, e.now(), class)
|
||||
if b, ok := e.Cache.GetSearch(key); ok {
|
||||
var hits []Hit
|
||||
if json.Unmarshal(b, &hits) == nil {
|
||||
e.mu.Lock()
|
||||
e.CacheHit++
|
||||
e.mu.Unlock()
|
||||
return RunResult{Hits: clampHits(hits, k), Cached: true, Attempted: attempted, SearchMS: e.now().Sub(start).Milliseconds()}
|
||||
}
|
||||
}
|
||||
e.mu.Lock()
|
||||
e.CacheMiss++
|
||||
e.mu.Unlock()
|
||||
|
||||
type one struct {
|
||||
name string
|
||||
out SourceOut
|
||||
skip bool
|
||||
}
|
||||
outs := make([]one, len(e.Sources))
|
||||
var wg sync.WaitGroup
|
||||
for i, src := range e.Sources {
|
||||
name := src.Name()
|
||||
if name == NameTavily && e.tavilyCapped() {
|
||||
outs[i] = one{name: name, out: SourceOut{Warn: "tavily: daily cap"}, skip: true}
|
||||
e.addCapped(NameTavily)
|
||||
continue
|
||||
}
|
||||
wg.Add(1)
|
||||
go func(i int, src Source, name string) {
|
||||
defer wg.Done()
|
||||
sctx, cancel := context.WithTimeout(ctx, e.timeout())
|
||||
defer cancel()
|
||||
outs[i] = one{name: name, out: src.Search(sctx, query, k)}
|
||||
}(i, src, name)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
var hits []Hit
|
||||
var warns []string
|
||||
var okNames, failNames, att []string
|
||||
degraded := false
|
||||
for _, o := range outs {
|
||||
if o.skip {
|
||||
if o.out.Warn != "" {
|
||||
warns = append(warns, o.out.Warn)
|
||||
}
|
||||
continue
|
||||
}
|
||||
att = append(att, o.name)
|
||||
e.addQuery(caller, o.name)
|
||||
if o.out.Warn != "" {
|
||||
warns = append(warns, o.out.Warn)
|
||||
failNames = append(failNames, o.name)
|
||||
degraded = true
|
||||
e.addErr(o.name, o.out.Warn)
|
||||
continue
|
||||
}
|
||||
okNames = append(okNames, o.name)
|
||||
hits = append(hits, o.out.Hits...)
|
||||
if o.out.Credits > 0 {
|
||||
e.addCredits(caller, o.name, o.out.Credits)
|
||||
if o.name == NameTavily && e.Cache != nil {
|
||||
e.Cache.AddSearchCredits(NameTavily, e.now().UTC().Format("2006-01-02"), o.out.Credits)
|
||||
}
|
||||
}
|
||||
}
|
||||
hits = Dedupe(hits)
|
||||
sortBySourceOrder(hits)
|
||||
|
||||
var rerankMS int64
|
||||
if doRerank && e.RerankOn {
|
||||
rs := e.now()
|
||||
ranked, fb := rankHits(ctx, query, hits, e.Embedder)
|
||||
rerankMS = e.now().Sub(rs).Milliseconds()
|
||||
if fb != "" {
|
||||
warns = append(warns, "rerank fell back: "+fb)
|
||||
degraded = true
|
||||
e.mu.Lock()
|
||||
e.RerankFallback++
|
||||
e.mu.Unlock()
|
||||
} else {
|
||||
hits = ranked
|
||||
}
|
||||
}
|
||||
hits = clampHits(hits, k)
|
||||
|
||||
if !degraded && e.Cache != nil && len(att) > 0 {
|
||||
payload, _ := json.Marshal(hits)
|
||||
ck := dcache.SearchKey(query, att, k, e.now(), class)
|
||||
e.Cache.PutSearch(ck, payload, e.ttl())
|
||||
}
|
||||
|
||||
return RunResult{
|
||||
Hits: hits,
|
||||
Warnings: warns,
|
||||
SourcesOK: okNames,
|
||||
SourcesFailed: failNames,
|
||||
SearchMS: e.now().Sub(start).Milliseconds(),
|
||||
RerankMS: rerankMS,
|
||||
Attempted: att,
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Engine) attemptedNames() []string {
|
||||
var n []string
|
||||
for _, s := range e.Sources {
|
||||
if s.Name() == NameTavily && e.tavilyCapped() {
|
||||
continue
|
||||
}
|
||||
n = append(n, s.Name())
|
||||
}
|
||||
sort.Strings(n)
|
||||
return n
|
||||
}
|
||||
|
||||
func (e *Engine) tavilyCapped() bool {
|
||||
if e == nil || e.TavilyCap <= 0 {
|
||||
return false
|
||||
}
|
||||
n := 0
|
||||
if e.Cache != nil {
|
||||
n = e.Cache.GetSearchCredits(NameTavily, e.now().UTC().Format("2006-01-02"))
|
||||
}
|
||||
return n >= e.TavilyCap
|
||||
}
|
||||
|
||||
func clampHits(hits []Hit, k int) []Hit {
|
||||
if k < 0 {
|
||||
k = 0
|
||||
}
|
||||
if len(hits) > k {
|
||||
return hits[:k]
|
||||
}
|
||||
return hits
|
||||
}
|
||||
|
||||
func (e *Engine) addQuery(caller, source string) {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
if e.Queries == nil {
|
||||
e.Queries = map[string]map[string]uint64{}
|
||||
}
|
||||
if e.Queries[caller] == nil {
|
||||
e.Queries[caller] = map[string]uint64{}
|
||||
}
|
||||
e.Queries[caller][source]++
|
||||
}
|
||||
|
||||
func (e *Engine) addCredits(caller, source string, n int) {
|
||||
if n <= 0 {
|
||||
return
|
||||
}
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
if e.Credits == nil {
|
||||
e.Credits = map[string]map[string]uint64{}
|
||||
}
|
||||
if e.Credits[caller] == nil {
|
||||
e.Credits[caller] = map[string]uint64{}
|
||||
}
|
||||
e.Credits[caller][source] += uint64(n)
|
||||
if e.DailyCredits == nil {
|
||||
e.DailyCredits = map[string]uint64{}
|
||||
}
|
||||
e.DailyCredits[source] += uint64(n)
|
||||
}
|
||||
|
||||
func (e *Engine) addErr(source, warn string) {
|
||||
kind := "error"
|
||||
if len(warn) > 0 {
|
||||
kind = warn
|
||||
if i := lastColon(warn); i >= 0 && i+1 < len(warn) {
|
||||
kind = warn[i+1:]
|
||||
if len(kind) > 0 && kind[0] == ' ' {
|
||||
kind = kind[1:]
|
||||
}
|
||||
}
|
||||
}
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
if e.Errors == nil {
|
||||
e.Errors = map[string]map[string]uint64{}
|
||||
}
|
||||
if e.Errors[source] == nil {
|
||||
e.Errors[source] = map[string]uint64{}
|
||||
}
|
||||
e.Errors[source][kind]++
|
||||
}
|
||||
|
||||
func (e *Engine) addCapped(source string) {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
if e.Capped == nil {
|
||||
e.Capped = map[string]uint64{}
|
||||
}
|
||||
e.Capped[source]++
|
||||
}
|
||||
|
||||
func lastColon(s string) int {
|
||||
for i := len(s) - 1; i >= 0; i-- {
|
||||
if s[i] == ':' {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
func (e *Engine) Metrics() map[string]any {
|
||||
if e == nil {
|
||||
return nil
|
||||
}
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
return map[string]any{
|
||||
"search_queries": clone2(e.Queries),
|
||||
"search_credits": clone2(e.Credits),
|
||||
"search_errors": clone2(e.Errors),
|
||||
"search_cache_hit": e.CacheHit,
|
||||
"search_cache_miss": e.CacheMiss,
|
||||
"search_rerank_fallback": e.RerankFallback,
|
||||
"search_daily_credits": clone1(e.DailyCredits),
|
||||
"search_capped": clone1(e.Capped),
|
||||
}
|
||||
}
|
||||
|
||||
func clone2(m map[string]map[string]uint64) map[string]map[string]uint64 {
|
||||
if m == nil {
|
||||
return map[string]map[string]uint64{}
|
||||
}
|
||||
out := make(map[string]map[string]uint64, len(m))
|
||||
for k, v := range m {
|
||||
inner := make(map[string]uint64, len(v))
|
||||
for k2, v2 := range v {
|
||||
inner[k2] = v2
|
||||
}
|
||||
out[k] = inner
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func clone1(m map[string]uint64) map[string]uint64 {
|
||||
if m == nil {
|
||||
return map[string]uint64{}
|
||||
}
|
||||
out := make(map[string]uint64, len(m))
|
||||
for k, v := range m {
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
package search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
)
|
||||
|
||||
type fakeSrc struct {
|
||||
name string
|
||||
hits []Hit
|
||||
warn string
|
||||
delay time.Duration
|
||||
credits int
|
||||
}
|
||||
|
||||
func (f fakeSrc) Name() string { return f.name }
|
||||
func (f fakeSrc) Search(ctx context.Context, query string, k int) SourceOut {
|
||||
if f.delay > 0 {
|
||||
select {
|
||||
case <-time.After(f.delay):
|
||||
case <-ctx.Done():
|
||||
return SourceOut{Warn: f.name + ": timeout"}
|
||||
}
|
||||
}
|
||||
return SourceOut{Hits: f.hits, Warn: f.warn, Credits: f.credits}
|
||||
}
|
||||
|
||||
type stubEmb struct {
|
||||
ok bool
|
||||
reason string
|
||||
}
|
||||
|
||||
func (s stubEmb) Embed(ctx context.Context, inputs []string) ([][]float64, error) {
|
||||
return nil, context.Canceled
|
||||
}
|
||||
func (s stubEmb) Acquire(ctx context.Context) (bool, string) { return s.ok, s.reason }
|
||||
func (s stubEmb) Release() {}
|
||||
|
||||
func TestFanInPartial(t *testing.T) {
|
||||
e := &Engine{
|
||||
Sources: []Source{
|
||||
fakeSrc{name: NameSearxNG, hits: []Hit{{URL: "https://ok.example/a", Title: "ok", Snippet: "x", Source: NameSearxNG}}},
|
||||
fakeSrc{name: NameBraveLLM, warn: "brave_llm: 500"},
|
||||
fakeSrc{name: NameTavily, delay: 2 * time.Second},
|
||||
},
|
||||
Timeout: 50 * time.Millisecond,
|
||||
DefaultK: 8,
|
||||
RerankOn: false,
|
||||
}
|
||||
r := e.Run(context.Background(), "q", 8, false, "titan", dcache.ClassFleet)
|
||||
if len(r.Hits) != 1 || r.Hits[0].URL != "https://ok.example/a" {
|
||||
t.Fatalf("hits %+v", r.Hits)
|
||||
}
|
||||
joined := strings.Join(r.Warnings, " ")
|
||||
if !strings.Contains(joined, "500") || !strings.Contains(joined, "timeout") {
|
||||
t.Fatalf("warnings %v", r.Warnings)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCacheClassAndHour(t *testing.T) {
|
||||
db, err := dcache.Open(filepath.Join(t.TempDir(), "c.db"), 1<<20, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
now := time.Date(2026, 9, 23, 10, 0, 0, 0, time.UTC)
|
||||
src := fakeSrc{name: NameSearxNG, hits: []Hit{{URL: "https://ok.example/a", Title: "ok", Snippet: "x", Source: NameSearxNG}}}
|
||||
e := &Engine{Sources: []Source{src}, DefaultK: 8, Cache: db, RerankOn: false, Now: func() time.Time { return now }}
|
||||
r1 := e.Run(context.Background(), "Hello World", 8, false, "titan", dcache.ClassFleet)
|
||||
if r1.Cached {
|
||||
t.Fatal("first miss")
|
||||
}
|
||||
r2 := e.Run(context.Background(), "hello world", 8, false, "titan", dcache.ClassFleet)
|
||||
if !r2.Cached {
|
||||
t.Fatal("norm/hour hit")
|
||||
}
|
||||
r3 := e.Run(context.Background(), "hello world", 8, false, "gs", dcache.ClassRestricted)
|
||||
if r3.Cached {
|
||||
t.Fatal("restricted must miss fleet cache")
|
||||
}
|
||||
e.Now = func() time.Time { return now.Add(time.Hour) }
|
||||
r4 := e.Run(context.Background(), "hello world", 8, false, "titan", dcache.ClassFleet)
|
||||
if r4.Cached {
|
||||
t.Fatal("next hour must miss")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoCacheOnDegraded(t *testing.T) {
|
||||
db, err := dcache.Open(filepath.Join(t.TempDir(), "c.db"), 1<<20, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
e := &Engine{
|
||||
Sources: []Source{
|
||||
fakeSrc{name: NameSearxNG, hits: []Hit{{URL: "https://ok.example/a", Source: NameSearxNG}}},
|
||||
fakeSrc{name: NameBraveLLM, warn: "brave_llm: 500"},
|
||||
},
|
||||
DefaultK: 8, Cache: db, RerankOn: false,
|
||||
}
|
||||
e.Run(context.Background(), "q", 8, false, "titan", dcache.ClassFleet)
|
||||
r := e.Run(context.Background(), "q", 8, false, "titan", dcache.ClassFleet)
|
||||
if r.Cached {
|
||||
t.Fatal("degraded cached")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRerankFallbackStill200(t *testing.T) {
|
||||
e := &Engine{
|
||||
Sources: []Source{fakeSrc{name: NameSearxNG, hits: []Hit{{URL: "https://ok.example/a", Source: NameSearxNG}}}},
|
||||
DefaultK: 8,
|
||||
RerankOn: true,
|
||||
Embedder: stubEmb{ok: false, reason: "queue_full"},
|
||||
}
|
||||
r := e.Run(context.Background(), "q", 8, true, "titan", dcache.ClassFleet)
|
||||
if len(r.Hits) != 1 {
|
||||
t.Fatal(r.Hits)
|
||||
}
|
||||
if !strings.Contains(strings.Join(r.Warnings, " "), "rerank fell back") {
|
||||
t.Fatalf("%v", r.Warnings)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDailyCapSkipsTavily(t *testing.T) {
|
||||
db, err := dcache.Open(filepath.Join(t.TempDir(), "c.db"), 1<<20, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
var n int
|
||||
tav := fakeSrc{name: NameTavily, hits: []Hit{{URL: "https://t.example/a", Source: NameTavily}}, credits: 1}
|
||||
e := &Engine{Sources: []Source{tav}, DefaultK: 8, Cache: db, TavilyCap: 1, RerankOn: false}
|
||||
e.Run(context.Background(), "q1", 8, false, "titan", dcache.ClassFleet)
|
||||
r := e.Run(context.Background(), "q2", 8, false, "titan", dcache.ClassFleet)
|
||||
for _, w := range r.Warnings {
|
||||
if strings.Contains(w, "daily cap") {
|
||||
n++
|
||||
}
|
||||
}
|
||||
if n != 1 {
|
||||
t.Fatalf("warnings %v", r.Warnings)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDedupeLongerSnippet(t *testing.T) {
|
||||
in := []Hit{
|
||||
{URL: "https://Ex.com/a#x", Snippet: "aa", Source: NameSearxNG},
|
||||
{URL: "https://ex.com/a", Snippet: "aabbcc", Source: NameTavily},
|
||||
}
|
||||
out := Dedupe(in)
|
||||
if len(out) != 1 || out[0].Source != NameSearxNG || out[0].Snippet != "aabbcc" {
|
||||
t.Fatalf("%+v", out)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
package search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"math"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type Embedder interface {
|
||||
Embed(ctx context.Context, inputs []string) ([][]float64, error)
|
||||
Acquire(ctx context.Context) (ok bool, reason string)
|
||||
Release()
|
||||
}
|
||||
|
||||
func Cosine(a, b []float64) float64 {
|
||||
if len(a) == 0 || len(a) != len(b) {
|
||||
return 0
|
||||
}
|
||||
var dot, na, nb float64
|
||||
for i := range a {
|
||||
dot += a[i] * b[i]
|
||||
na += a[i] * a[i]
|
||||
nb += b[i] * b[i]
|
||||
}
|
||||
if na == 0 || nb == 0 {
|
||||
return 0
|
||||
}
|
||||
return dot / (math.Sqrt(na) * math.Sqrt(nb))
|
||||
}
|
||||
|
||||
func rankHits(ctx context.Context, query string, hits []Hit, emb Embedder) (out []Hit, fallback string) {
|
||||
if emb == nil || len(hits) == 0 {
|
||||
if emb == nil {
|
||||
return hits, "no embed"
|
||||
}
|
||||
return hits, ""
|
||||
}
|
||||
ok, reason := emb.Acquire(ctx)
|
||||
if !ok {
|
||||
if reason == "" {
|
||||
reason = "queue_full"
|
||||
}
|
||||
return hits, reason
|
||||
}
|
||||
defer emb.Release()
|
||||
inputs := make([]string, 0, 1+len(hits))
|
||||
inputs = append(inputs, query)
|
||||
for _, h := range hits {
|
||||
s := strings.TrimSpace(h.Title + " " + h.Snippet)
|
||||
if s == "" {
|
||||
s = h.URL
|
||||
}
|
||||
inputs = append(inputs, s)
|
||||
}
|
||||
vecs, err := emb.Embed(ctx, inputs)
|
||||
if err != nil || len(vecs) != len(inputs) {
|
||||
return hits, "embed error"
|
||||
}
|
||||
q := vecs[0]
|
||||
type pair struct {
|
||||
h Hit
|
||||
s float64
|
||||
}
|
||||
ps := make([]pair, len(hits))
|
||||
for i := range hits {
|
||||
ps[i] = pair{h: hits[i], s: Cosine(q, vecs[i+1])}
|
||||
ps[i].h.Score = ps[i].s
|
||||
}
|
||||
for i := 1; i < len(ps); i++ {
|
||||
j := i
|
||||
for j > 0 && ps[j].s > ps[j-1].s {
|
||||
ps[j], ps[j-1] = ps[j-1], ps[j]
|
||||
j--
|
||||
}
|
||||
}
|
||||
out = make([]Hit, len(ps))
|
||||
for i := range ps {
|
||||
out[i] = ps[i].h
|
||||
}
|
||||
return out, ""
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
package search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/search/dial"
|
||||
)
|
||||
|
||||
func warnStatus(name string, status int) string {
|
||||
if status == 0 {
|
||||
return name + ": error"
|
||||
}
|
||||
return fmt.Sprintf("%s: %d", name, status)
|
||||
}
|
||||
|
||||
func warnErr(name string, err error) string {
|
||||
if err == nil {
|
||||
return ""
|
||||
}
|
||||
if strings.Contains(strings.ToLower(err.Error()), "timeout") || strings.Contains(err.Error(), "deadline") {
|
||||
return name + ": timeout"
|
||||
}
|
||||
return name + ": error"
|
||||
}
|
||||
|
||||
func readJSON(resp *http.Response, dest any) error {
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != 200 {
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 512))
|
||||
return fmt.Errorf("status %d", resp.StatusCode)
|
||||
}
|
||||
dec := json.NewDecoder(io.LimitReader(resp.Body, 1<<20))
|
||||
return dec.Decode(dest)
|
||||
}
|
||||
|
||||
// SearxNG is GET {url}/search?q=&format=json&pageno=1
|
||||
type SearxNG struct {
|
||||
Client *dial.Client
|
||||
}
|
||||
|
||||
func (s *SearxNG) Name() string { return NameSearxNG }
|
||||
|
||||
func (s *SearxNG) Search(ctx context.Context, query string, k int) SourceOut {
|
||||
if s == nil || s.Client == nil {
|
||||
return SourceOut{Warn: NameSearxNG + ": error"}
|
||||
}
|
||||
_ = s.Client.RefreshIfStale(ctx)
|
||||
base := strings.TrimRight(s.Client.URL.String(), "/")
|
||||
u, err := url.Parse(base + "/search")
|
||||
if err != nil {
|
||||
return SourceOut{Warn: warnErr(NameSearxNG, err)}
|
||||
}
|
||||
q := u.Query()
|
||||
q.Set("q", query)
|
||||
q.Set("format", "json")
|
||||
q.Set("pageno", "1")
|
||||
u.RawQuery = q.Encode()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
|
||||
if err != nil {
|
||||
return SourceOut{Warn: warnErr(NameSearxNG, err)}
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
resp, err := s.Client.HTTP.Do(req)
|
||||
if err != nil {
|
||||
return SourceOut{Warn: warnErr(NameSearxNG, err)}
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 256))
|
||||
resp.Body.Close()
|
||||
return SourceOut{Warn: warnStatus(NameSearxNG, resp.StatusCode)}
|
||||
}
|
||||
var env struct {
|
||||
Results []struct {
|
||||
URL string `json:"url"`
|
||||
Title string `json:"title"`
|
||||
Content string `json:"content"`
|
||||
Score float64 `json:"score"`
|
||||
} `json:"results"`
|
||||
}
|
||||
if err := readJSON(resp, &env); err != nil {
|
||||
return SourceOut{Warn: warnErr(NameSearxNG, err)}
|
||||
}
|
||||
sort.SliceStable(env.Results, func(i, j int) bool { return env.Results[i].Score > env.Results[j].Score })
|
||||
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: NameSearxNG})
|
||||
}
|
||||
return SourceOut{Hits: hits}
|
||||
}
|
||||
|
||||
// BraveLLM is GET {url}?q= with X-Subscription-Token.
|
||||
type BraveLLM struct {
|
||||
Client *dial.Client
|
||||
Key string
|
||||
}
|
||||
|
||||
func (s *BraveLLM) Name() string { return NameBraveLLM }
|
||||
|
||||
func (s *BraveLLM) Search(ctx context.Context, query string, k int) SourceOut {
|
||||
if s == nil || s.Client == nil {
|
||||
return SourceOut{Warn: NameBraveLLM + ": error"}
|
||||
}
|
||||
_ = s.Client.RefreshIfStale(ctx)
|
||||
u, err := url.Parse(s.Client.URL.String())
|
||||
if err != nil {
|
||||
return SourceOut{Warn: warnErr(NameBraveLLM, err)}
|
||||
}
|
||||
q := u.Query()
|
||||
q.Set("q", query)
|
||||
u.RawQuery = q.Encode()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil)
|
||||
if err != nil {
|
||||
return SourceOut{Warn: warnErr(NameBraveLLM, err)}
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("X-Subscription-Token", s.Key)
|
||||
resp, err := s.Client.HTTP.Do(req)
|
||||
if err != nil {
|
||||
return SourceOut{Warn: warnErr(NameBraveLLM, err)}
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
code := resp.StatusCode
|
||||
resp.Body.Close()
|
||||
return SourceOut{Warn: warnStatus(NameBraveLLM, code)}
|
||||
}
|
||||
raw, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
resp.Body.Close()
|
||||
if err != nil {
|
||||
return SourceOut{Warn: warnErr(NameBraveLLM, err)}
|
||||
}
|
||||
hits := mapBrave(raw)
|
||||
credits := 0
|
||||
if len(hits) >= 0 {
|
||||
credits = 1
|
||||
}
|
||||
return SourceOut{Hits: hits, Credits: credits}
|
||||
}
|
||||
|
||||
func mapBrave(raw []byte) []Hit {
|
||||
var top map[string]any
|
||||
if err := json.Unmarshal(raw, &top); err != nil {
|
||||
return nil
|
||||
}
|
||||
var buckets []any
|
||||
for _, key := range []string{"sources", "results", "web"} {
|
||||
if v, ok := top[key]; ok {
|
||||
buckets = append(buckets, v)
|
||||
}
|
||||
}
|
||||
var hits []Hit
|
||||
seen := map[string]bool{}
|
||||
var walk func(any)
|
||||
walk = func(v any) {
|
||||
switch t := v.(type) {
|
||||
case []any:
|
||||
for _, x := range t {
|
||||
walk(x)
|
||||
}
|
||||
case map[string]any:
|
||||
u := strField(t, "url", "link", "href")
|
||||
if u == "" {
|
||||
for _, x := range t {
|
||||
walk(x)
|
||||
}
|
||||
return
|
||||
}
|
||||
if seen[u] {
|
||||
return
|
||||
}
|
||||
seen[u] = true
|
||||
hits = append(hits, Hit{
|
||||
URL: u,
|
||||
Title: strField(t, "title", "name"),
|
||||
Snippet: strField(t, "snippet", "description", "content", "text"),
|
||||
Source: NameBraveLLM,
|
||||
})
|
||||
}
|
||||
}
|
||||
for _, b := range buckets {
|
||||
walk(b)
|
||||
}
|
||||
return hits
|
||||
}
|
||||
|
||||
func strField(m map[string]any, keys ...string) string {
|
||||
for _, k := range keys {
|
||||
if v, ok := m[k]; ok {
|
||||
if s, ok := v.(string); ok && s != "" {
|
||||
return s
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
package search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/search/dial"
|
||||
)
|
||||
|
||||
func TestTavilyOutboundAndDropsAnswer(t *testing.T) {
|
||||
var gotBody []byte
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
t.Errorf("method %s", r.Method)
|
||||
}
|
||||
if !strings.HasPrefix(r.Header.Get("Authorization"), "Bearer ") {
|
||||
t.Error("missing bearer")
|
||||
}
|
||||
gotBody, _ = io.ReadAll(r.Body)
|
||||
_, _ = w.Write([]byte(`{"answer":"SECRET","results":[{"url":"https://ex.com/a","title":"A","content":"snip","score":0.9}],"usage":{"credits":1}}`))
|
||||
}))
|
||||
defer ts.Close()
|
||||
c, err := dial.New(ts.URL, dial.Options{Loopback: true, Timeout: time.Second})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
out := (&Tavily{Client: c, Key: "test-key"}).Search(context.Background(), "q", 5)
|
||||
if out.Warn != "" {
|
||||
t.Fatal(out.Warn)
|
||||
}
|
||||
var req tavilyReq
|
||||
if err := json.Unmarshal(gotBody, &req); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if req.IncludeAnswer || req.IncludeRawContent || req.IncludeImages || req.SearchDepth != "basic" {
|
||||
t.Fatalf("%+v", req)
|
||||
}
|
||||
if strings.Contains(string(gotBody), "extract") || strings.Contains(string(gotBody), "crawl") {
|
||||
t.Fatal(string(gotBody))
|
||||
}
|
||||
if len(out.Hits) != 1 || out.Hits[0].URL != "https://ex.com/a" {
|
||||
t.Fatalf("%+v", out.Hits)
|
||||
}
|
||||
for _, h := range out.Hits {
|
||||
if strings.Contains(h.Snippet, "SECRET") || strings.Contains(h.Title, "SECRET") {
|
||||
t.Fatal("answer leaked into hit")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearxNGMapsAndSorts(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/search" || r.URL.Query().Get("format") != "json" {
|
||||
t.Errorf("path %s", r.URL.String())
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"results":[
|
||||
{"url":"https://b.example/","title":"b","content":"bb","score":0.1},
|
||||
{"url":"https://a.example/","title":"a","content":"aa","score":0.9}
|
||||
]}`))
|
||||
}))
|
||||
defer ts.Close()
|
||||
c, err := dial.New(ts.URL, dial.Options{Loopback: true, Timeout: time.Second})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
out := (&SearxNG{Client: c}).Search(context.Background(), "q", 8)
|
||||
if out.Warn != "" || len(out.Hits) != 2 || out.Hits[0].Title != "a" {
|
||||
t.Fatalf("%+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBraveMapsSources(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Header.Get("X-Subscription-Token") == "" {
|
||||
t.Error("token")
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"grounding":"ignore me","sources":[{"url":"https://s.example/x","title":"S","snippet":"hello"}]}`))
|
||||
}))
|
||||
defer ts.Close()
|
||||
c, err := dial.New(ts.URL, dial.Options{Loopback: true, Timeout: time.Second})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
out := (&BraveLLM{Client: c, Key: "k"}).Search(context.Background(), "q", 5)
|
||||
if len(out.Hits) != 1 || out.Hits[0].URL != "https://s.example/x" {
|
||||
t.Fatalf("%+v", out)
|
||||
}
|
||||
if strings.Contains(out.Hits[0].Snippet, "ignore") {
|
||||
t.Fatal("grounding used as snippet")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,102 @@
|
||||
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 }
|
||||
|
||||
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"}
|
||||
}
|
||||
_ = s.Client.RefreshIfStale(ctx)
|
||||
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 := mapTavily(raw)
|
||||
return SourceOut{Hits: hits, Credits: credits}
|
||||
}
|
||||
|
||||
func mapTavily(raw []byte) ([]Hit, int) {
|
||||
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
|
||||
}
|
||||
_ = 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
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sort"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
const (
|
||||
MaxQueryRunes = 1500
|
||||
MaxK = 20
|
||||
DefaultK = 8
|
||||
DefaultV2K = 5
|
||||
NameSearxNG = "searxng"
|
||||
NameBraveLLM = "brave_llm"
|
||||
NameTavily = "tavily"
|
||||
)
|
||||
|
||||
var sourceRank = map[string]int{
|
||||
NameSearxNG: 0,
|
||||
NameBraveLLM: 1,
|
||||
NameTavily: 2,
|
||||
}
|
||||
|
||||
// Hit is one search result after mapping. No page body, no vendor answers.
|
||||
type Hit struct {
|
||||
URL string `json:"url"`
|
||||
Title string `json:"title"`
|
||||
Snippet string `json:"snippet"`
|
||||
Score float64 `json:"score"`
|
||||
Source string `json:"source"`
|
||||
}
|
||||
|
||||
// Source is one upstream. Config-fixed endpoint; never a request URL.
|
||||
type Source interface {
|
||||
Name() string
|
||||
Search(ctx context.Context, query string, k int) SourceOut
|
||||
}
|
||||
|
||||
type SourceOut struct {
|
||||
Hits []Hit
|
||||
Warn string
|
||||
Credits int
|
||||
}
|
||||
|
||||
func ValidateQuery(q string) string {
|
||||
q = strings.TrimSpace(q)
|
||||
if q == "" {
|
||||
return "invalid query"
|
||||
}
|
||||
if utf8.RuneCountInString(q) > MaxQueryRunes {
|
||||
return "invalid query"
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func ClampK(k, def int) int {
|
||||
if k <= 0 {
|
||||
k = def
|
||||
}
|
||||
if k > MaxK {
|
||||
return MaxK
|
||||
}
|
||||
return k
|
||||
}
|
||||
|
||||
// Dedupe by canonical URL. First source (in fan-in order) wins the Source
|
||||
// field; a later longer snippet replaces the text.
|
||||
func Dedupe(in []Hit) []Hit {
|
||||
type slot struct {
|
||||
h Hit
|
||||
i int
|
||||
}
|
||||
order := make([]string, 0)
|
||||
seen := map[string]*slot{}
|
||||
for _, h := range in {
|
||||
canon := CanonicalURL(h.URL)
|
||||
if canon == "" {
|
||||
continue
|
||||
}
|
||||
h.URL = canon
|
||||
if s, ok := seen[canon]; ok {
|
||||
if len(h.Snippet) > len(s.h.Snippet) {
|
||||
src := s.h.Source
|
||||
sc := s.h.Score
|
||||
s.h = h
|
||||
s.h.Source = src
|
||||
if sc > s.h.Score {
|
||||
s.h.Score = sc
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
cp := h
|
||||
seen[canon] = &slot{h: cp, i: len(order)}
|
||||
order = append(order, canon)
|
||||
}
|
||||
out := make([]Hit, 0, len(order))
|
||||
for _, c := range order {
|
||||
out = append(out, seen[c].h)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func sortBySourceOrder(hits []Hit) {
|
||||
sort.SliceStable(hits, func(i, j int) bool {
|
||||
ri, rj := sourceRank[hits[i].Source], sourceRank[hits[j].Source]
|
||||
if ri != rj {
|
||||
return ri < rj
|
||||
}
|
||||
return false
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user