sift/internal/search/engine.go

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 // static: missing keys
	TavilyCap int
	Now       func() time.Time
	NoRerank  map[string]bool

	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
}

type readySrc interface {
	Name() string
	Ready() bool
}

// Health is missing keys plus sources that still have no pin.
func (e *Engine) Health() []string {
	if e == nil {
		return nil
	}
	out := append([]string(nil), e.Degraded...)
	for _, s := range e.Sources {
		r, ok := s.(readySrc)
		if !ok || r.Ready() {
			continue
		}
		out = append(out, "search."+r.Name()+": unresolved")
	}
	return out
}

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 lists []namedHits
	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)
		lists = append(lists, namedHits{name: o.name, 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 = mergeSearch(e, lists, doRerank)

	var rerankMS int64
	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
}