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:
+17
-3
@@ -23,6 +23,7 @@ import (
|
||||
"git.wntrmute.dev/kyle/sift/internal/fleetctx"
|
||||
"git.wntrmute.dev/kyle/sift/internal/guard"
|
||||
"git.wntrmute.dev/kyle/sift/internal/recover"
|
||||
"git.wntrmute.dev/kyle/sift/internal/search"
|
||||
"git.wntrmute.dev/kyle/sift/internal/selfip"
|
||||
)
|
||||
|
||||
@@ -42,6 +43,7 @@ type Server struct {
|
||||
OutCap int
|
||||
SelfIP *selfip.Discoverer
|
||||
Digest *digest.Engine
|
||||
Search *search.Engine
|
||||
Ctx *fleetctx.Client
|
||||
// PermitRestricted is context.permit_restricted: those callers never
|
||||
// have a fetch host written to the log (host=-). Independent of Honcho.
|
||||
@@ -74,6 +76,8 @@ func New(version string, g *guard.Guard, f *fetch.Fetcher, keys *auth.Store, log
|
||||
s.Mux.HandleFunc("GET /metrics", s.metrics)
|
||||
s.Mux.HandleFunc("POST /v2/scrape", s.scrape)
|
||||
s.Mux.HandleFunc("POST /v1/fetch", s.v1fetch)
|
||||
s.Mux.HandleFunc("POST /v1/search", s.v1search)
|
||||
s.Mux.HandleFunc("POST /v2/search", s.v2search)
|
||||
return s
|
||||
}
|
||||
|
||||
@@ -92,7 +96,7 @@ func (s *Server) HTTPServer(addr string) *http.Server {
|
||||
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
r.Body = http.MaxBytesReader(w, r.Body, maxReqBody)
|
||||
switch r.URL.Path {
|
||||
case "/v1/fetch", "/v2/scrape":
|
||||
case "/v1/fetch", "/v2/scrape", "/v1/search", "/v2/search":
|
||||
ctx, cancel := context.WithTimeout(r.Context(), requestDeadline)
|
||||
defer cancel()
|
||||
r = r.WithContext(ctx)
|
||||
@@ -106,6 +110,10 @@ func (s *Server) healthz(w http.ResponseWriter, _ *http.Request) {
|
||||
if s.SelfIP != nil {
|
||||
out["self_ip"] = s.SelfIP.Status()
|
||||
}
|
||||
if s.Search != nil && len(s.Search.Degraded) > 0 {
|
||||
out["status"] = "degraded"
|
||||
out["search"] = s.Search.Degraded
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(out)
|
||||
}
|
||||
|
||||
@@ -115,12 +123,18 @@ func (s *Server) metrics(w http.ResponseWriter, _ *http.Request) {
|
||||
if s.SelfIP != nil {
|
||||
fails = s.SelfIP.Failures.Load()
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
m := map[string]any{
|
||||
"requests": s.reqN.Load(),
|
||||
"blocks": s.blockN.Load(),
|
||||
"scrapes": s.scrapeN.Load(),
|
||||
"self_ip_failures": fails,
|
||||
})
|
||||
}
|
||||
if s.Search != nil {
|
||||
for k, v := range s.Search.Metrics() {
|
||||
m[k] = v
|
||||
}
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(m)
|
||||
}
|
||||
|
||||
type scrapeReq struct {
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/search"
|
||||
)
|
||||
|
||||
type v1SearchReq struct {
|
||||
Query string `json:"query"`
|
||||
Question string `json:"question"`
|
||||
K int `json:"k"`
|
||||
Rerank *bool `json:"rerank"`
|
||||
Digest bool `json:"digest"`
|
||||
URL string `json:"url"`
|
||||
Limit int `json:"limit"`
|
||||
}
|
||||
|
||||
type v2SearchReq struct {
|
||||
Query string `json:"query"`
|
||||
Limit int `json:"limit"`
|
||||
URL string `json:"url"`
|
||||
Digest bool `json:"digest"`
|
||||
}
|
||||
|
||||
func (s *Server) v1search(w http.ResponseWriter, r *http.Request) {
|
||||
s.handleSearch(w, r, false)
|
||||
}
|
||||
|
||||
func (s *Server) v2search(w http.ResponseWriter, r *http.Request) {
|
||||
s.handleSearch(w, r, true)
|
||||
}
|
||||
|
||||
func (s *Server) handleSearch(w http.ResponseWriter, r *http.Request, v2 bool) {
|
||||
s.reqN.Add(1)
|
||||
c := s.caller(w, r)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
dec := json.NewDecoder(r.Body)
|
||||
query := ""
|
||||
k := 0
|
||||
doRerank := true
|
||||
digest := false
|
||||
if v2 {
|
||||
var req v2SearchReq
|
||||
if err := dec.Decode(&req); err != nil && err.Error() != "EOF" {
|
||||
writeErr(w, 400, "invalid query", "")
|
||||
return
|
||||
}
|
||||
query = req.Query
|
||||
k = req.Limit
|
||||
digest = req.Digest
|
||||
_ = req.URL
|
||||
} else {
|
||||
var req v1SearchReq
|
||||
if err := dec.Decode(&req); err != nil && err.Error() != "EOF" {
|
||||
writeErr(w, 400, "invalid query", "")
|
||||
return
|
||||
}
|
||||
query = req.Query
|
||||
k = req.K
|
||||
if k == 0 {
|
||||
k = req.Limit
|
||||
}
|
||||
if req.Rerank != nil {
|
||||
doRerank = *req.Rerank
|
||||
}
|
||||
digest = req.Digest
|
||||
_ = req.URL
|
||||
_ = req.Question
|
||||
}
|
||||
if msg := search.ValidateQuery(query); msg != "" {
|
||||
writeErr(w, 400, msg, "")
|
||||
return
|
||||
}
|
||||
defK := search.DefaultK
|
||||
if v2 {
|
||||
defK = search.DefaultV2K
|
||||
}
|
||||
k = search.ClampK(k, defK)
|
||||
if s.Search == nil {
|
||||
writeErr(w, 422, "no search results", "")
|
||||
return
|
||||
}
|
||||
res := s.Search.Run(r.Context(), strings.TrimSpace(query), k, doRerank, c.Name, s.callerClass(c))
|
||||
warns := append([]string(nil), res.Warnings...)
|
||||
if digest {
|
||||
warns = append(warns, "digest ignored")
|
||||
}
|
||||
if !res.Cached && len(res.SourcesOK) == 0 && len(res.Hits) == 0 {
|
||||
writeErr(w, 422, "no search results", "")
|
||||
return
|
||||
}
|
||||
s.logSearch(c, query, res)
|
||||
if v2 {
|
||||
web := make([]map[string]any, 0, len(res.Hits))
|
||||
for i, h := range res.Hits {
|
||||
web = append(web, map[string]any{
|
||||
"url": h.URL,
|
||||
"title": h.Title,
|
||||
"description": h.Snippet,
|
||||
"position": i + 1,
|
||||
})
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{"success": true, "data": map[string]any{"web": web}})
|
||||
return
|
||||
}
|
||||
results := make([]map[string]any, 0, len(res.Hits))
|
||||
for _, h := range res.Hits {
|
||||
results = append(results, map[string]any{
|
||||
"url": h.URL,
|
||||
"title": h.Title,
|
||||
"snippet": h.Snippet,
|
||||
"score": h.Score,
|
||||
"source": h.Source,
|
||||
})
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{
|
||||
"results": results,
|
||||
"warnings": warns,
|
||||
"sizes": map[string]any{
|
||||
"results": len(results),
|
||||
"sources_ok": res.SourcesOK,
|
||||
"sources_failed": res.SourcesFailed,
|
||||
},
|
||||
"timings": map[string]any{"search_ms": res.SearchMS, "rerank_ms": res.RerankMS},
|
||||
"cached": res.Cached,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) logSearch(c *auth.Caller, query string, res search.RunResult) {
|
||||
qLen := len([]rune(query))
|
||||
cache := "miss"
|
||||
if res.Cached {
|
||||
cache = "hit"
|
||||
}
|
||||
src := strings.Join(res.Attempted, ",")
|
||||
if s.restrictedLog(c) {
|
||||
s.Log.Info("search", "caller", c.Name, "q_len", qLen, "sources", src, "n", len(res.Hits), "cache", cache)
|
||||
return
|
||||
}
|
||||
s.Log.Info("search", "caller", c.Name, "q_len", qLen, "sources", src, "n", len(res.Hits), "cache", cache, "rerank_ms", res.RerankMS)
|
||||
}
|
||||
@@ -0,0 +1,175 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
"git.wntrmute.dev/kyle/sift/internal/fetch"
|
||||
"git.wntrmute.dev/kyle/sift/internal/guard"
|
||||
"git.wntrmute.dev/kyle/sift/internal/search"
|
||||
"git.wntrmute.dev/kyle/sift/internal/search/dial"
|
||||
)
|
||||
|
||||
func searchServer(t *testing.T) (*Server, string, string, *recHandler, *httptest.Server) {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
keys := filepath.Join(dir, "keys.toml")
|
||||
fleetOut := filepath.Join(dir, "fleet")
|
||||
gsOut := filepath.Join(dir, "gs")
|
||||
if err := auth.Add(keys, "test", "hermes-fleet", fleetOut); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := auth.Add(keys, "straylight-gs", "guardian-star", gsOut); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
fleetKey, _ := os.ReadFile(fleetOut)
|
||||
gsKey, _ := os.ReadFile(gsOut)
|
||||
store, err := auth.Load(keys)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if strings.Contains(r.URL.Path, "search") {
|
||||
_, _ = w.Write([]byte(`{"results":[{"url":"https://example.com/hit","title":"Hit","content":"snippet about q","score":1}]}`))
|
||||
return
|
||||
}
|
||||
w.WriteHeader(404)
|
||||
}))
|
||||
t.Cleanup(ts.Close)
|
||||
c, err := dial.New(ts.URL, dial.Options{Loopback: true, Timeout: time.Second})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
db, err := dcache.Open(filepath.Join(dir, "c.db"), 1<<20, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
h := &recHandler{}
|
||||
s := New("test", &guard.Guard{}, &fetch.Fetcher{Guard: &guard.Guard{}, UA: "sift-test"}, store, slog.New(h))
|
||||
s.PermitRestricted = []string{"straylight-gs"}
|
||||
s.Search = &search.Engine{
|
||||
Sources: []search.Source{&search.SearxNG{Client: c}},
|
||||
DefaultK: 8,
|
||||
RerankOn: false,
|
||||
Cache: db,
|
||||
}
|
||||
return s, strings.TrimSpace(string(fleetKey)), strings.TrimSpace(string(gsKey)), h, ts
|
||||
}
|
||||
|
||||
func TestV2SearchGolden(t *testing.T) {
|
||||
s, key, _, _, _ := searchServer(t)
|
||||
rec := postJSON(t, s, "/v2/search", key, map[string]any{"query": "contract", "limit": 3, "url": "http://127.0.0.1:11434"})
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("%d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var env map[string]any
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &env); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if env["success"] != true {
|
||||
t.Fatalf("%v", env)
|
||||
}
|
||||
data, _ := env["data"].(map[string]any)
|
||||
web, _ := data["web"].([]any)
|
||||
if len(web) < 1 {
|
||||
t.Fatalf("web %v", env)
|
||||
}
|
||||
row, _ := web[0].(map[string]any)
|
||||
if row["url"] == "" || row["title"] == "" || row["description"] == nil {
|
||||
t.Fatalf("%v", row)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchCachePartitionAndRestrictedLogs(t *testing.T) {
|
||||
s, fleet, gs, h, _ := searchServer(t)
|
||||
rec := postJSON(t, s, "/v1/search", gs, map[string]any{"query": "secret topic", "k": 5, "rerank": false})
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("gs %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
rec = postJSON(t, s, "/v1/search", fleet, map[string]any{"query": "secret topic", "k": 5, "rerank": false})
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("fleet %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var env map[string]any
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &env); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if env["cached"] == true {
|
||||
t.Fatal("fleet observed restricted cache")
|
||||
}
|
||||
dump := h.dump()
|
||||
if strings.Contains(dump, "secret topic") || strings.Contains(dump, "example.com") || strings.Contains(dump, "127.0.0.1") {
|
||||
t.Fatalf("restricted log leak:\n%s", dump)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchInvalidQuery(t *testing.T) {
|
||||
s, key, _, _, _ := searchServer(t)
|
||||
rec := postJSON(t, s, "/v1/search", key, map[string]any{"query": ""})
|
||||
if rec.Code != 400 {
|
||||
t.Fatalf("%d", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHealthzMissingKey(t *testing.T) {
|
||||
s, _ := setupAPI(t, nil)
|
||||
s.Search = &search.Engine{Degraded: []string{"search.tavily: missing key"}}
|
||||
req := httptest.NewRequest(http.MethodGet, "/healthz", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
s.ServeHTTP(rec, req)
|
||||
if rec.Code != 200 || !strings.Contains(rec.Body.String(), "missing key") {
|
||||
t.Fatalf("%s", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestMetricsCacheHitMiss(t *testing.T) {
|
||||
s, key, _, _, _ := searchServer(t)
|
||||
postJSON(t, s, "/v1/search", key, map[string]any{"query": "same q", "k": 5, "rerank": false})
|
||||
postJSON(t, s, "/v1/search", key, map[string]any{"query": "same q", "k": 5, "rerank": false})
|
||||
req := httptest.NewRequest(http.MethodGet, "/metrics", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
s.ServeHTTP(rec, req)
|
||||
var m map[string]any
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &m)
|
||||
if m["search_cache_hit"] == 0 && m["search_cache_miss"] == 0 {
|
||||
t.Fatalf("%v", m)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNoSearchResults422(t *testing.T) {
|
||||
s, key := setupAPI(t, nil)
|
||||
s.Search = &search.Engine{DefaultK: 8}
|
||||
rec := postJSON(t, s, "/v1/search", key, map[string]any{"query": "none"})
|
||||
if rec.Code != 422 {
|
||||
t.Fatalf("%d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func (h *recHandler) dump() string {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
var b bytes.Buffer
|
||||
for _, r := range h.recs {
|
||||
b.WriteString(recMsg(r))
|
||||
r.Attrs(func(a slog.Attr) bool {
|
||||
b.WriteByte(' ')
|
||||
b.WriteString(a.Key)
|
||||
b.WriteByte('=')
|
||||
b.WriteString(a.Value.String())
|
||||
return true
|
||||
})
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
@@ -16,6 +16,7 @@ type File struct {
|
||||
Dixie Dixie `toml:"dixie"`
|
||||
Digest Digest `toml:"digest"`
|
||||
Context Context `toml:"context"`
|
||||
Search Search `toml:"search"`
|
||||
}
|
||||
|
||||
type Dixie struct {
|
||||
@@ -39,6 +40,40 @@ type Context struct {
|
||||
PermitRestricted []string `toml:"permit_restricted"`
|
||||
}
|
||||
|
||||
// Search is the [search] block. Sources default off; bake-off picks defaults.
|
||||
type Search struct {
|
||||
K int `toml:"k"`
|
||||
Timeout string `toml:"timeout"`
|
||||
Rerank *bool `toml:"rerank"`
|
||||
Digest bool `toml:"digest"`
|
||||
CacheTTL string `toml:"cache_ttl"`
|
||||
SearxNG SearchSource `toml:"searxng"`
|
||||
BraveLLM SearchSource `toml:"brave_llm"`
|
||||
Tavily TavilySource `toml:"tavily"`
|
||||
}
|
||||
|
||||
type SearchSource struct {
|
||||
Enabled bool `toml:"enabled"`
|
||||
URL string `toml:"url"`
|
||||
}
|
||||
|
||||
type TavilySource struct {
|
||||
Enabled bool `toml:"enabled"`
|
||||
URL string `toml:"url"`
|
||||
DailyCreditCap int `toml:"daily_credit_cap"`
|
||||
}
|
||||
|
||||
func (s Search) RerankOn() bool {
|
||||
if s.Rerank == nil {
|
||||
return true
|
||||
}
|
||||
return *s.Rerank
|
||||
}
|
||||
|
||||
func (s Search) AnyEnabled() bool {
|
||||
return s.SearxNG.Enabled || s.BraveLLM.Enabled || s.Tavily.Enabled
|
||||
}
|
||||
|
||||
// SelfIP is the [self_ip] block.
|
||||
type SelfIP struct {
|
||||
Sources []string `toml:"sources"`
|
||||
|
||||
+104
-5
@@ -11,6 +11,8 @@ import (
|
||||
"log/slog"
|
||||
"math"
|
||||
"os"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -41,6 +43,8 @@ type DB struct {
|
||||
DigestM uint64
|
||||
EmbedH uint64
|
||||
EmbedM uint64
|
||||
SearchH uint64
|
||||
SearchM uint64
|
||||
Evictions uint64
|
||||
}
|
||||
|
||||
@@ -75,6 +79,19 @@ func Open(path string, maxBytes int64, log *slog.Logger) (*DB, error) {
|
||||
last_hit INTEGER NOT NULL,
|
||||
nbytes INTEGER NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS search (
|
||||
k TEXT PRIMARY KEY,
|
||||
payload TEXT NOT NULL,
|
||||
expires INTEGER NOT NULL,
|
||||
last_hit INTEGER NOT NULL,
|
||||
nbytes INTEGER NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS search_cap (
|
||||
source TEXT NOT NULL,
|
||||
day TEXT NOT NULL,
|
||||
credits INTEGER NOT NULL,
|
||||
PRIMARY KEY (source, day)
|
||||
);
|
||||
`); err != nil {
|
||||
_ = sq.Close()
|
||||
return nil, err
|
||||
@@ -118,6 +135,17 @@ func QuestionEmbedKey(question, embedModel, callerClass string) string {
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func SearchKey(query string, sources []string, k int, hour time.Time, callerClass string) string {
|
||||
if callerClass == "" {
|
||||
callerClass = ClassFleet
|
||||
}
|
||||
src := append([]string(nil), sources...)
|
||||
sort.Strings(src)
|
||||
bucket := hour.UTC().Format("2006-01-02T15")
|
||||
sum := sha256.Sum256([]byte(NormalizeQuestion(query) + "\x00" + strings.Join(src, ",") + "\x00" + strconv.Itoa(k) + "\x00" + bucket + "\x00" + callerClass))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
type DigestHit struct {
|
||||
Text string
|
||||
ChunksUsed int
|
||||
@@ -193,11 +221,79 @@ func (d *DB) Size() int64 {
|
||||
if d == nil {
|
||||
return 0
|
||||
}
|
||||
var n sql.NullInt64
|
||||
var n, m, s sql.NullInt64
|
||||
_ = d.sql.QueryRow(`SELECT IFNULL(SUM(nbytes),0) FROM digest`).Scan(&n)
|
||||
var m sql.NullInt64
|
||||
_ = d.sql.QueryRow(`SELECT IFNULL(SUM(nbytes),0) FROM embed`).Scan(&m)
|
||||
return n.Int64 + m.Int64
|
||||
_ = d.sql.QueryRow(`SELECT IFNULL(SUM(nbytes),0) FROM search`).Scan(&s)
|
||||
return n.Int64 + m.Int64 + s.Int64
|
||||
}
|
||||
|
||||
func (d *DB) GetSearch(key string) (payload []byte, ok bool) {
|
||||
if d == nil {
|
||||
return nil, false
|
||||
}
|
||||
var b string
|
||||
var exp int64
|
||||
err := d.sql.QueryRow(`SELECT payload, expires FROM search WHERE k = ?`, key).Scan(&b, &exp)
|
||||
if err != nil {
|
||||
d.mu.Lock()
|
||||
d.SearchM++
|
||||
d.mu.Unlock()
|
||||
return nil, false
|
||||
}
|
||||
if exp > 0 && time.Now().Unix() > exp {
|
||||
_, _ = d.sql.Exec(`DELETE FROM search WHERE k = ?`, key)
|
||||
d.mu.Lock()
|
||||
d.SearchM++
|
||||
d.mu.Unlock()
|
||||
return nil, false
|
||||
}
|
||||
_, _ = d.sql.Exec(`UPDATE search SET last_hit = ? WHERE k = ?`, time.Now().Unix(), key)
|
||||
d.mu.Lock()
|
||||
d.SearchH++
|
||||
d.mu.Unlock()
|
||||
return []byte(b), true
|
||||
}
|
||||
|
||||
func (d *DB) PutSearch(key string, payload []byte, ttl time.Duration) {
|
||||
if d == nil {
|
||||
return
|
||||
}
|
||||
exp := time.Now().Add(ttl).Unix()
|
||||
n := len(key) + len(payload)
|
||||
_, err := d.sql.Exec(`INSERT OR REPLACE INTO search(k, payload, expires, last_hit, nbytes) VALUES(?,?,?,?,?)`,
|
||||
key, string(payload), exp, time.Now().Unix(), n)
|
||||
if err != nil && d.log != nil {
|
||||
d.log.Warn("search cache write", "err", err)
|
||||
}
|
||||
d.evict()
|
||||
}
|
||||
|
||||
func (d *DB) GetSearchCredits(source, day string) int {
|
||||
if d == nil {
|
||||
return 0
|
||||
}
|
||||
var n int
|
||||
err := d.sql.QueryRow(`SELECT credits FROM search_cap WHERE source = ? AND day = ?`, source, day).Scan(&n)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
func (d *DB) AddSearchCredits(source, day string, delta int) int {
|
||||
if d == nil || delta == 0 {
|
||||
if d == nil {
|
||||
return 0
|
||||
}
|
||||
return d.GetSearchCredits(source, day)
|
||||
}
|
||||
_, err := d.sql.Exec(`INSERT INTO search_cap(source, day, credits) VALUES(?,?,?)
|
||||
ON CONFLICT(source, day) DO UPDATE SET credits = credits + excluded.credits`, source, day, delta)
|
||||
if err != nil && d.log != nil {
|
||||
d.log.Warn("search cap write", "err", err)
|
||||
}
|
||||
return d.GetSearchCredits(source, day)
|
||||
}
|
||||
|
||||
func (d *DB) evict() {
|
||||
@@ -208,15 +304,18 @@ func (d *DB) evict() {
|
||||
var k string
|
||||
err := d.sql.QueryRow(`
|
||||
SELECT k FROM (
|
||||
SELECT k, last_hit, 0 AS kind FROM digest
|
||||
SELECT k, last_hit FROM digest
|
||||
UNION ALL
|
||||
SELECT k, last_hit, 1 FROM embed
|
||||
SELECT k, last_hit FROM embed
|
||||
UNION ALL
|
||||
SELECT k, last_hit FROM search
|
||||
) ORDER BY last_hit ASC LIMIT 1`).Scan(&k)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
_, _ = d.sql.Exec(`DELETE FROM digest WHERE k = ?`, k)
|
||||
_, _ = d.sql.Exec(`DELETE FROM embed WHERE k = ?`, k)
|
||||
_, _ = d.sql.Exec(`DELETE FROM search WHERE k = ?`, k)
|
||||
d.mu.Lock()
|
||||
d.Evictions++
|
||||
d.mu.Unlock()
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestHitMissAndNormalise(t *testing.T) {
|
||||
@@ -95,6 +96,21 @@ func TestEvictUnderCap(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchKeyPartitions(t *testing.T) {
|
||||
now := time.Date(2026, 9, 23, 10, 0, 0, 0, time.UTC)
|
||||
a := SearchKey("Hello", []string{"tavily", "searxng"}, 8, now, ClassFleet)
|
||||
b := SearchKey("hello", []string{"searxng", "tavily"}, 8, now, ClassFleet)
|
||||
if a != b {
|
||||
t.Fatal("norm/sources order")
|
||||
}
|
||||
if SearchKey("hello", []string{"searxng"}, 8, now, ClassRestricted) == a {
|
||||
t.Fatal("class")
|
||||
}
|
||||
if SearchKey("hello", []string{"searxng", "tavily"}, 8, now.Add(time.Hour), ClassFleet) == a {
|
||||
t.Fatal("hour")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCorruptDBBypass(t *testing.T) {
|
||||
p := filepath.Join(t.TempDir(), "c.db")
|
||||
if err := os.WriteFile(p, []byte("not a database"), 0o600); err != nil {
|
||||
|
||||
@@ -210,6 +210,27 @@ func (e *Engine) embedModel() string {
|
||||
return "honcho-embed"
|
||||
}
|
||||
|
||||
// Embed is the search-rerank entry. Same gate as digest.
|
||||
func (e *Engine) Embed(ctx context.Context, inputs []string) ([][]float64, error) {
|
||||
return e.embed(ctx, inputs)
|
||||
}
|
||||
|
||||
func (e *Engine) Acquire(ctx context.Context) (bool, string) {
|
||||
if e == nil {
|
||||
return false, "no embed"
|
||||
}
|
||||
if e.EmbedURL == "" {
|
||||
return false, "no embed"
|
||||
}
|
||||
return e.Gate.Acquire(ctx)
|
||||
}
|
||||
|
||||
func (e *Engine) Release() {
|
||||
if e != nil {
|
||||
e.Gate.Release()
|
||||
}
|
||||
}
|
||||
|
||||
func (e *Engine) logFallback(reason string, err error) {
|
||||
e.noteFallback(reason)
|
||||
if e == nil || e.Log == nil {
|
||||
|
||||
@@ -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