sift/internal/search/pipeline.go

package search

import (
	"context"
	"encoding/json"
	"sync"

	"git.wntrmute.dev/kyle/sift/internal/dcache"
)

type srcOut struct {
	name string
	out  SourceOut
	skip bool
}

type queryFan struct {
	hits      []Hit
	warns     []string
	notes     []string
	okNames   []string
	failNames []string
	att       []string
	degraded  bool
	tavSkip   bool
	tavUsed   bool
}

func (e *Engine) RunFilter(ctx context.Context, queries []string, k int, doRerank bool, caller, class string, f Filter) RunResult {
	start := e.now()
	k = ClampK(k, e.DefaultK)
	if e.DefaultK <= 0 {
		k = ClampK(k, DefaultK)
	}
	if class == "" {
		class = dcache.ClassFleet
	}
	if len(queries) == 0 {
		queries = []string{""}
	}

	attempted := e.attemptedNames(class)
	key := e.searchKey(queries, attempted, k, class, f)
	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()

	var lists [][]Hit
	var warns, notes []string
	okSeen, failSeen := map[string]bool{}, map[string]bool{}
	var okNames, failNames []string
	degraded := false
	capMid := false
	sawTavily := false
	var lastAtt []string

	for _, q := range queries {
		one := e.fanInQuery(ctx, q, k, doRerank, caller, class, f)
		if one.tavUsed {
			sawTavily = true
		}
		if one.tavSkip && sawTavily {
			capMid = true
		}
		if one.degraded {
			degraded = true
		}
		warns = append(warns, one.warns...)
		notes = append(notes, one.notes...)
		for _, n := range one.okNames {
			if !okSeen[n] {
				okSeen[n] = true
				okNames = append(okNames, n)
			}
		}
		for _, n := range one.failNames {
			if !failSeen[n] {
				failSeen[n] = true
				failNames = append(failNames, n)
			}
		}
		if len(one.att) > 0 {
			lastAtt = one.att
		}
		lists = append(lists, one.hits)
	}

	var hits []Hit
	if len(lists) <= 1 {
		if len(lists) == 1 {
			hits = lists[0]
		}
	} else {
		hits = rrfMerge(lists)
		hits = Dedupe(hits)
	}
	hits = clampHits(hits, k)

	warns = append(warns, notes...)
	warns = dedupeStrings(warns)

	att := lastAtt
	if len(att) == 0 {
		att = attempted
	}
	if !degraded && !capMid && e.Cache != nil && len(att) > 0 {
		payload, _ := json.Marshal(hits)
		ck := e.searchKey(queries, att, k, class, f)
		e.Cache.PutSearch(ck, payload, e.ttl())
	}

	return RunResult{
		Hits:          hits,
		Warnings:      warns,
		SourcesOK:     okNames,
		SourcesFailed: failNames,
		SearchMS:      e.now().Sub(start).Milliseconds(),
		Attempted:     att,
	}
}

func (e *Engine) searchKey(queries []string, sources []string, k int, class string, f Filter) string {
	return dcache.EncodeSearchKey(dcache.SearchKeyInput{
		Queries:            queries,
		Sources:            sources,
		K:                  k,
		Hour:               e.now(),
		CallerClass:        class,
		IncludeDomains:     f.IncludeDomains,
		ExcludeDomains:     f.ExcludeDomains,
		IncludeDomainsMode: f.IncludeDomainsMode,
		TimeRange:          f.TimeRange,
		StartDate:          f.StartDate,
		EndDate:            f.EndDate,
		Topic:              f.Topic,
		ExactMatch:         f.ExactMatch,
		MinScore:           f.MinScore,
	})
}

func (e *Engine) fanInQuery(ctx context.Context, query string, k int, doRerank bool, caller, class string, f Filter) queryFan {
	outs := make([]srcOut, len(e.Sources))
	var wg sync.WaitGroup
	tavSkip := false
	for i, src := range e.Sources {
		name := src.Name()
		if name == NameTavily && e.tavilyCapped(class) {
			outs[i] = srcOut{name: name, out: SourceOut{Warn: "tavily: daily cap"}, skip: true}
			e.addCapped(NameTavily)
			tavSkip = true
			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] = srcOut{name: name, out: src.Search(sctx, query, k, f)}
		}(i, src, name)
	}
	wg.Wait()

	var lists []namedHits
	var warns, notes []string
	var okNames, failNames, att []string
	degraded := false
	tavUsed := 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.name == NameTavily {
			tavUsed = true
		}
		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)
		notes = append(notes, o.out.Notes...)
		hits := o.out.Hits
		if o.name == NameTavily && f.MinScore > 0 {
			hits = filterMinScore(hits, f.MinScore)
		}
		lists = append(lists, namedHits{name: o.name, hits: hits})
		if o.out.Credits > 0 {
			e.addCredits(caller, o.name, o.out.Credits)
			if o.name == NameTavily && e.Cache != nil {
				e.Cache.AddSearchCredits(tavilyCapSource(class), e.now().UTC().Format("2006-01-02"), o.out.Credits)
			}
		}
	}
	merged := mergeSearch(e, lists, doRerank)
	merged = filterDomains(merged, f)
	return queryFan{
		hits:      merged,
		warns:     warns,
		notes:     notes,
		okNames:   okNames,
		failNames: failNames,
		att:       att,
		degraded:  degraded,
		tavSkip:   tavSkip,
		tavUsed:   tavUsed,
	}
}