thread-master/internal/library/websearch/client.go

package websearch

import (
	"context"
	"strings"

	libbrave "haixun-backend/internal/library/brave"
	libexa "haixun-backend/internal/library/exa"
)

type Provider string

const (
	ProviderBrave Provider = "brave"
	ProviderExa   Provider = "exa"
)

func ParseProvider(raw string) Provider {
	switch Provider(strings.ToLower(strings.TrimSpace(raw))) {
	case ProviderExa:
		return ProviderExa
	default:
		return ProviderBrave
	}
}

type Mode string

const (
	ModeKnowledgeExpand Mode = "knowledge_expand"
	ModeThreadsDiscover Mode = "threads_discover"
)

type SearchResult struct {
	Title   string
	Snippet string
	URL     string
}

type SearchResponse struct {
	Results  []SearchResult
	Query    string
	Status   string
	Provider Provider
}

type SearchOptions struct {
	Query              string
	Limit              int
	Mode               Mode
	Country            string
	SearchLang         string
	UserLocation       string
	StartPublishedDate string
}

type Client interface {
	Search(ctx context.Context, opts SearchOptions) (SearchResponse, error)
	Enabled() bool
	Provider() Provider
}

type Config struct {
	Provider     Provider
	BraveKey     string
	ExaKey       string
	Country      string
	SearchLang   string
	UserLocation string
}

func New(cfg Config) Client {
	switch ParseProvider(string(cfg.Provider)) {
	case ProviderExa:
		return &exaAdapter{
			client:       libexa.NewClient(cfg.ExaKey),
			provider:     ProviderExa,
			userLocation: cfg.UserLocation,
		}
	default:
		return &braveAdapter{
			client:     libbrave.NewClient(cfg.BraveKey),
			provider:   ProviderBrave,
			country:    cfg.Country,
			searchLang: cfg.SearchLang,
		}
	}
}

func ConfigFromMember(braveKey, exaKey, provider, country, searchLang, userLocation string) Config {
	return Config{
		Provider:     ParseProvider(provider),
		BraveKey:     braveKey,
		ExaKey:       exaKey,
		Country:      country,
		SearchLang:   searchLang,
		UserLocation: userLocation,
	}
}

type braveAdapter struct {
	client     *libbrave.Client
	provider   Provider
	country    string
	searchLang string
}

func (a *braveAdapter) Provider() Provider { return a.provider }

func (a *braveAdapter) Enabled() bool {
	return a != nil && a.client != nil && a.client.Enabled()
}

func (a *braveAdapter) Search(ctx context.Context, opts SearchOptions) (SearchResponse, error) {
	mode := libbrave.ModeKnowledgeExpand
	if opts.Mode == ModeThreadsDiscover {
		mode = libbrave.ModeThreadsDiscover
	}
	res, err := a.client.Search(ctx, libbrave.SearchOptions{
		Query:      opts.Query,
		Limit:      opts.Limit,
		Mode:       mode,
		Country:    firstNonEmpty(opts.Country, a.country),
		SearchLang: firstNonEmpty(opts.SearchLang, a.searchLang),
	})
	return toBraveResponse(res.Query, res.Status, a.provider, res.Results, err)
}

type exaAdapter struct {
	client       *libexa.Client
	provider     Provider
	userLocation string
}

func (a *exaAdapter) Provider() Provider { return a.provider }

func (a *exaAdapter) Enabled() bool {
	return a != nil && a.client != nil && a.client.Enabled()
}

func (a *exaAdapter) Search(ctx context.Context, opts SearchOptions) (SearchResponse, error) {
	mode := libexa.ModeKnowledgeExpand
	if opts.Mode == ModeThreadsDiscover {
		mode = libexa.ModeThreadsDiscover
	}
	res, err := a.client.Search(ctx, libexa.SearchOptions{
		Query:              opts.Query,
		Limit:              opts.Limit,
		Mode:               mode,
		UserLocation:       firstNonEmpty(opts.UserLocation, a.userLocation),
		StartPublishedDate: opts.StartPublishedDate,
	})
	out := SearchResponse{Query: res.Query, Status: res.Status, Provider: a.provider}
	if err != nil {
		return out, err
	}
	for _, item := range res.Results {
		out.Results = append(out.Results, SearchResult{
			Title:   item.Title,
			Snippet: item.Snippet,
			URL:     item.URL,
		})
	}
	return out, nil
}

func toBraveResponse(query, status string, provider Provider, items []libbrave.SearchResult, err error) (SearchResponse, error) {
	out := SearchResponse{Query: query, Status: status, Provider: provider}
	if err != nil {
		return out, err
	}
	for _, item := range items {
		out.Results = append(out.Results, SearchResult{
			Title:   item.Title,
			Snippet: item.Snippet,
			URL:     item.URL,
		})
	}
	return out, nil
}

func firstNonEmpty(values ...string) string {
	for _, value := range values {
		if trimmed := strings.TrimSpace(value); trimmed != "" {
			return trimmed
		}
	}
	return ""
}

func ActiveAPIKey(cfg Config) string {
	if ParseProvider(string(cfg.Provider)) == ProviderExa {
		return strings.TrimSpace(cfg.ExaKey)
	}
	return strings.TrimSpace(cfg.BraveKey)
}

func ProviderLabel(provider Provider) string {
	switch provider {
	case ProviderExa:
		return "Exa"
	default:
		return "Brave"
	}
}