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 }