thread-master/backend/internal/library/knowledge/brave_collect.go

package knowledge

import (
	"context"
	"strings"
	"sync"

	"haixun-backend/internal/library/websearch"
)

type BraveSearchLocale struct {
	Country      string
	SearchLang   string
	UserLocation string
}

type BraveCollectConfig struct {
	ResultsPerQuery      int
	MinSourcesBeforeStop int
	MaxSourcesCap        int
	Concurrency          int
}

const minUsefulSourceRunes = 24

func BraveCollectConfigFromQueryCfg(cfg queryConfig) BraveCollectConfig {
	out := BraveCollectConfig{
		ResultsPerQuery:      cfg.ResultsPerQuery,
		MinSourcesBeforeStop: cfg.MinSourcesBeforeStop,
		MaxSourcesCap:        cfg.MaxSourcesCap,
		Concurrency:          cfg.BraveCollectConcurrency,
	}
	if out.ResultsPerQuery <= 0 {
		out.ResultsPerQuery = 8
	}
	if out.MinSourcesBeforeStop <= 0 {
		out.MinSourcesBeforeStop = 18
	}
	if out.MaxSourcesCap <= 0 {
		out.MaxSourcesCap = 32
	}
	if out.Concurrency <= 0 {
		out.Concurrency = 4
	}
	return out
}

func DefaultBraveCollectConfig() BraveCollectConfig {
	cfg, err := loadQueryConfig()
	if err != nil {
		return BraveCollectConfig{ResultsPerQuery: 8, MinSourcesBeforeStop: 18, MaxSourcesCap: 32, Concurrency: 4}
	}
	return BraveCollectConfigFromQueryCfg(cfg)
}

func CollectBraveSources(
	ctx context.Context,
	client websearch.Client,
	locale BraveSearchLocale,
	queries []string,
	cfg BraveCollectConfig,
	onProgress func(i, total int),
	heartbeat func() error,
) []BraveSource {
	return CollectWebSources(ctx, client, locale, queries, cfg, onProgress, heartbeat)
}

func CollectWebSources(
	ctx context.Context,
	client websearch.Client,
	locale BraveSearchLocale,
	queries []string,
	cfg BraveCollectConfig,
	onProgress func(i, total int),
	heartbeat func() error,
) []BraveSource {
	if client == nil || !client.Enabled() || len(queries) == 0 {
		return nil
	}
	if cfg.Concurrency <= 1 {
		return collectWebSourcesSequential(ctx, client, locale, queries, cfg, onProgress, heartbeat)
	}
	return collectWebSourcesParallel(ctx, client, locale, queries, cfg, onProgress, heartbeat)
}

func collectWebSourcesSequential(
	ctx context.Context,
	client websearch.Client,
	locale BraveSearchLocale,
	queries []string,
	cfg BraveCollectConfig,
	onProgress func(i, total int),
	heartbeat func() error,
) []BraveSource {
	capHint := cfg.MaxSourcesCap
	if est := len(queries) * cfg.ResultsPerQuery; est < capHint {
		capHint = est
	}
	out := make([]BraveSource, 0, capHint)
	seenURL := map[string]struct{}{}

	for i, query := range queries {
		if shouldStopCollect(out, queries, cfg) {
			break
		}
		if heartbeat != nil {
			if err := heartbeat(); err != nil {
				return out
			}
		}
		appendBraveResults(&out, seenURL, query, searchWebQuery(ctx, client, locale, query, cfg.ResultsPerQuery), cfg.MaxSourcesCap)
		if onProgress != nil {
			onProgress(i, len(queries))
		}
	}
	return out
}

type braveCollectState struct {
	cfg     BraveCollectConfig
	mu      sync.Mutex
	out     []BraveSource
	seenURL map[string]struct{}
	stop    bool
}

func (s *braveCollectState) shouldStop(queries []string, cfg BraveCollectConfig) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.stop {
		return true
	}
	if shouldStopCollect(s.out, queries, cfg) {
		s.stop = true
		return true
	}
	return false
}

func (s *braveCollectState) appendResults(query string, items []BraveSource, queries []string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.stop {
		return
	}
	appendBraveResults(&s.out, s.seenURL, query, items, s.cfg.MaxSourcesCap)
	if shouldStopCollect(s.out, queries, s.cfg) {
		s.stop = true
	}
}

func (s *braveCollectState) snapshot() []BraveSource {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]BraveSource(nil), s.out...)
}

func collectWebSourcesParallel(
	ctx context.Context,
	client websearch.Client,
	locale BraveSearchLocale,
	queries []string,
	cfg BraveCollectConfig,
	onProgress func(i, total int),
	heartbeat func() error,
) []BraveSource {
	state := &braveCollectState{
		cfg:     cfg,
		out:     make([]BraveSource, 0, cfg.MaxSourcesCap),
		seenURL: map[string]struct{}{},
	}

	workers := cfg.Concurrency
	if workers > len(queries) {
		workers = len(queries)
	}
	if workers <= 0 {
		workers = 1
	}

	completed := 0
	for next := 0; next < len(queries); {
		if state.shouldStop(queries, cfg) {
			break
		}
		batchEnd := next + workers
		if batchEnd > len(queries) {
			batchEnd = len(queries)
		}

		type queryResult struct {
			index int
			items []BraveSource
		}
		results := make(chan queryResult, batchEnd-next)
		var wg sync.WaitGroup
		for i := next; i < batchEnd; i++ {
			if heartbeat != nil {
				if err := heartbeat(); err != nil {
					return state.snapshot()
				}
			}
			query := queries[i]
			wg.Add(1)
			go func(index int, q string) {
				defer wg.Done()
				results <- queryResult{
					index: index,
					items: searchWebQuery(ctx, client, locale, q, cfg.ResultsPerQuery),
				}
			}(i, query)
		}
		wg.Wait()
		close(results)

		ordered := make([]queryResult, batchEnd-next)
		for result := range results {
			ordered[result.index-next] = result
		}
		for _, result := range ordered {
			state.appendResults(queries[result.index], result.items, queries)
			completed++
			if onProgress != nil {
				onProgress(completed-1, len(queries))
			}
		}
		next = batchEnd
	}
	return state.snapshot()
}

func shouldStopCollect(out []BraveSource, queries []string, cfg BraveCollectConfig) bool {
	if len(out) >= cfg.MaxSourcesCap {
		return true
	}
	return sourcesAreSufficient(out, queries, cfg)
}

func sourcesAreSufficient(out []BraveSource, queries []string, cfg BraveCollectConfig) bool {
	if cfg.MinSourcesBeforeStop <= 0 {
		return len(out) > 0
	}
	stats := sourceStats(out)
	if stats.UniqueURLs < cfg.MinSourcesBeforeStop {
		return false
	}
	minQueryCoverage := 1
	if len(queries) > 1 {
		minQueryCoverage = 2
	}
	if stats.QueryCoverage < minQueryCoverage {
		return false
	}
	minUseful := (cfg.MinSourcesBeforeStop * 2) / 3
	if minUseful < 1 {
		minUseful = 1
	}
	if stats.UsefulSources < minUseful {
		return false
	}
	return stats.TextRunes >= cfg.MinSourcesBeforeStop*minUsefulSourceRunes
}

type collectSourceStats struct {
	UniqueURLs    int
	QueryCoverage int
	UsefulSources int
	TextRunes     int
}

func sourceStats(items []BraveSource) collectSourceStats {
	seenURL := map[string]struct{}{}
	seenQuery := map[string]struct{}{}
	stats := collectSourceStats{}
	for _, item := range items {
		url := strings.TrimSpace(item.URL)
		if url == "" {
			continue
		}
		if _, ok := seenURL[url]; ok {
			continue
		}
		seenURL[url] = struct{}{}
		query := strings.TrimSpace(item.Query)
		if query != "" {
			seenQuery[query] = struct{}{}
		}
		textRunes := len([]rune(strings.TrimSpace(item.Title))) + len([]rune(strings.TrimSpace(item.Snippet)))
		stats.TextRunes += textRunes
		if textRunes >= minUsefulSourceRunes {
			stats.UsefulSources++
		}
	}
	stats.UniqueURLs = len(seenURL)
	stats.QueryCoverage = len(seenQuery)
	return stats
}

func searchWebQuery(
	ctx context.Context,
	client websearch.Client,
	locale BraveSearchLocale,
	query string,
	limit int,
) []BraveSource {
	res, _ := client.Search(ctx, websearch.SearchOptions{
		Query:        query,
		Limit:        limit,
		Mode:         websearch.ModeKnowledgeExpand,
		Country:      locale.Country,
		SearchLang:   locale.SearchLang,
		UserLocation: locale.UserLocation,
	})
	items := make([]BraveSource, 0, len(res.Results))
	for _, item := range res.Results {
		url := strings.TrimSpace(item.URL)
		if url == "" {
			continue
		}
		items = append(items, BraveSource{
			Query:   query,
			Snippet: item.Snippet,
			URL:     url,
			Title:   item.Title,
		})
	}
	return items
}

func appendBraveResults(out *[]BraveSource, seenURL map[string]struct{}, query string, items []BraveSource, max int) {
	for _, item := range items {
		if max > 0 && len(*out) >= max {
			return
		}
		url := strings.TrimSpace(item.URL)
		if url == "" {
			continue
		}
		if _, ok := seenURL[url]; ok {
			continue
		}
		seenURL[url] = struct{}{}
		if item.Query == "" {
			item.Query = query
		}
		*out = append(*out, item)
	}
}

func MergeBraveSources(chunks ...[]BraveSource) []BraveSource {
	seen := map[string]struct{}{}
	out := make([]BraveSource, 0)
	for _, chunk := range chunks {
		for _, item := range chunk {
			url := strings.TrimSpace(item.URL)
			if url == "" {
				continue
			}
			if _, ok := seen[url]; ok {
				continue
			}
			seen[url] = struct{}{}
			out = append(out, item)
		}
	}
	return out
}

func uniqueSourceCount(items []BraveSource) int {
	seen := map[string]struct{}{}
	for _, item := range items {
		url := strings.TrimSpace(item.URL)
		if url == "" {
			continue
		}
		seen[url] = struct{}{}
	}
	return len(seen)
}