sift/cmd/sift/searchcmd.go

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
)

type stringList []string

func (s *stringList) String() string { return strings.Join(*s, ",") }
func (s *stringList) Set(v string) error {
	*s = append(*s, v)
	return nil
}

func searchMain() {
	fs := flag.NewFlagSet("search", flag.ContinueOnError)
	fs.SetOutput(io.Discard)
	k := fs.Int("k", 0, "result cap")
	asJSON := fs.Bool("json", false, "native envelope")
	var include, exclude stringList
	fs.Var(&include, "include-domain", "restrict or prefer this host")
	fs.Var(&exclude, "exclude-domain", "drop this host")
	domainsMode := fs.String("domains-mode", "", "restrict|prefer")
	timeRange := fs.String("time-range", "", "day|week|month|year")
	startDate := fs.String("start-date", "", "YYYY-MM-DD")
	endDate := fs.String("end-date", "", "YYYY-MM-DD")
	topic := fs.String("topic", "", "general|news|finance")
	exact := fs.Bool("exact-match", false, "require quoted phrases")
	minScore := fs.Float64("min-score", 0, "Tavily score cut")
	queriesFile := fs.String("queries-file", "", "JSON array of 1-4 queries")
	args := os.Args[1:]
	flagArgs, pos, err := splitAroundPositional(args, searchValueFlags)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(2)
	}
	if err := fs.Parse(flagArgs); err != nil {
		fmt.Fprintln(os.Stderr, searchUsage)
		os.Exit(2)
	}
	body := map[string]any{}
	if *queriesFile != "" {
		if len(pos) != 0 {
			fmt.Fprintln(os.Stderr, searchUsage)
			os.Exit(2)
		}
		raw, err := os.ReadFile(*queriesFile)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(2)
		}
		var qs []string
		if err := json.Unmarshal(raw, &qs); err != nil {
			fmt.Fprintln(os.Stderr, "queries-file: want JSON array of strings")
			os.Exit(2)
		}
		body["queries"] = qs
	} else {
		if len(pos) != 1 {
			fmt.Fprintln(os.Stderr, searchUsage)
			os.Exit(2)
		}
		body["query"] = pos[0]
	}
	if *k > 0 {
		body["k"] = *k
	}
	if len(include) > 0 {
		body["include_domains"] = []string(include)
	}
	if len(exclude) > 0 {
		body["exclude_domains"] = []string(exclude)
	}
	if *domainsMode != "" {
		body["include_domains_mode"] = *domainsMode
	}
	if *timeRange != "" {
		body["time_range"] = *timeRange
	}
	if *startDate != "" {
		body["start_date"] = *startDate
	}
	if *endDate != "" {
		body["end_date"] = *endDate
	}
	if *topic != "" {
		body["topic"] = *topic
	}
	if *exact {
		body["exact_match"] = true
	}
	if *minScore != 0 {
		body["min_score"] = *minScore
	}
	base, key, err := loadClient()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(5)
	}
	raw, _ := json.Marshal(body)
	req, err := http.NewRequest(http.MethodPost, base+"/v1/search", bytes.NewReader(raw))
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(5)
	}
	req.Header.Set("Authorization", "Bearer "+key)
	req.Header.Set("Content-Type", "application/json")
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(5)
	}
	defer resp.Body.Close()
	var env map[string]any
	if err := json.NewDecoder(resp.Body).Decode(&env); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(5)
	}
	if resp.StatusCode != 200 {
		msg, _ := env["error"].(string)
		fmt.Fprintln(os.Stderr, msg)
		os.Exit(exitFor(resp.StatusCode, msg))
	}
	if *asJSON {
		enc := json.NewEncoder(os.Stdout)
		enc.SetIndent("", "  ")
		_ = enc.Encode(env)
		return
	}
	results, _ := env["results"].([]any)
	for i, it := range results {
		m, _ := it.(map[string]any)
		if i > 0 {
			fmt.Println()
		}
		fmt.Println(m["title"])
		fmt.Println(m["url"])
		if sn, _ := m["snippet"].(string); sn != "" {
			fmt.Println(sn)
		}
		if d, _ := m["published_date"].(string); d != "" {
			fmt.Println(d)
		}
	}
}

var searchValueFlags = map[string]bool{
	"k": true, "include-domain": true, "exclude-domain": true,
	"domains-mode": true, "time-range": true, "start-date": true, "end-date": true,
	"topic": true, "min-score": true, "queries-file": true,
}

const searchUsage = `usage: sift search "<query>" [--k n] [--json]
            [--include-domain host]... [--exclude-domain host]...
            [--domains-mode restrict|prefer]
            [--time-range day|week|month|year]
            [--start-date YYYY-MM-DD] [--end-date YYYY-MM-DD]
            [--topic general|news|finance]
            [--exact-match]
            [--min-score 0.5]
            [--queries-file path]`