sift/internal/search/engine.go

package search

import (
	"context"
	"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 {
	return e.RunFilter(ctx, []string{query}, k, doRerank, caller, class, Filter{})
}

func (e *Engine) attemptedNames(class string) []string {
	var n []string
	for _, s := range e.Sources {
		if s.Name() == NameTavily && e.tavilyCapped(class) {
			continue
		}
		n = append(n, s.Name())
	}
	sort.Strings(n)
	return n
}

func tavilyCapSource(class string) string {
	if class == dcache.ClassRestricted {
		return NameTavily + ":restricted"
	}
	return NameTavily
}

func (e *Engine) tavilyCapped(class string) bool {
	if e == nil || e.TavilyCap <= 0 {
		return false
	}
	if class == dcache.ClassRestricted {
		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
}