init, working

1 files changed, 130 insertions, 0 deletions

diff --git a/main.go b/main.go
new file mode 100644
index 0000000..9290c6c
--- /dev/null
+++ b/main.go
@@ -0,0 +1,130 @@
+package main
+
+import (
+	"bufio"
+	"bytes"
+	"encoding/json"
+	"flag"
+	"io"
+	"log"
+	"net/http"
+	"os"
+	"strings"
+	"text/template"
+)
+
+var (
+	token     = flag.String("token", os.Getenv("FASTGPT_API_TOKEN"), "API-Token, default from FASTGPT_API_TOKEN environment")
+	stdin     = flag.Bool("stdin", true, "Use stdin instead of command-line arguments for search")
+	cache     = flag.Bool("cache", false, "Whether to allow cached requests & responses.")
+	websearch = flag.Bool("websearch", true, "Whether to perform web searches to enrich answers. MUST   be true for now.")
+	raw       = flag.Bool("raw", false, "Output raw json result")
+)
+
+type Query struct {
+	Query     string `json:"query"`      // A query to be answered.
+	Cache     bool   `json:"cache"`      // Whether to allow cached requests & responses. (default true)
+	WebSearch bool   `json:"web_search"` // Whether to perform web searches to enrich answers. MUST be true for now.
+}
+
+type Result struct {
+	Meta struct {
+		Id   string `json:"id"`
+		Node string `json:"node"`
+		MS   int    `json:"ms"`
+	} `json:"meta"`
+	Data struct {
+		Query      string      `json:"-"`
+		Output     string      `json:"output"`     // Answer output
+		References []Reference `json:"references"` // The search results that are referred in the answer.
+		Tokens     int         `json:"tokens"`     // Amount of tokens processed
+	} `json:"data"`
+}
+
+type Reference struct {
+	Title   string `json:"title"`   // Title of the referenced search result.
+	Snippet string `json:"snippet"` // Snippet of the referenced search result.
+	URL     string `json:"url"`     // URL of the referenced search result.
+}
+
+var fm = template.FuncMap{
+	"quote": func(s string) string {
+		buf := &bytes.Buffer{}
+		scn := bufio.NewScanner(strings.NewReader(s))
+		for scn.Scan() {
+			buf.WriteString("> ")
+			buf.WriteString(scn.Text())
+			buf.WriteString("\n")
+		}
+		return buf.String()
+	},
+}
+var tmpl = template.Must(template.New("output").
+	Funcs(fm).
+	Parse(`Q: {{ .Query }}
+
+{{.Output}}
+{{ with .References }}{{- range $i, $ref := . }}
+[{{$i}}]: [{{$ref.Title}}]({{$ref.URL}})
+{{ quote $ref.Snippet }}
+{{- end }}
+{{- end }}
+`))
+
+func main() {
+
+	q := Query{
+		Cache:     *cache,
+		WebSearch: *websearch,
+	}
+	buf := &bytes.Buffer{}
+
+	flag.Parse()
+	if *stdin {
+		if _, e := io.Copy(buf, os.Stdin); e != nil {
+			panic(e)
+		}
+		q.Query = buf.String()
+		buf.Reset()
+	} else {
+		q.Query = strings.Join(os.Args, " ")
+	}
+
+	if e := json.NewEncoder(buf).Encode(&q); e != nil {
+		panic(e)
+	}
+
+	req, e := http.NewRequest("POST",
+		"https://kagi.com/api/v0/fastgpt",
+		buf)
+	if e != nil {
+		panic(e)
+	}
+	req.Header.Add("Authorization", "Bot "+*token)
+	req.Header.Add("Content-Type", "application/json")
+
+	r, e := http.DefaultClient.Do(req)
+	if e != nil {
+		panic(e)
+	}
+
+	if r.StatusCode != 200 {
+		log.Printf("Non-200 status code: %d\nBody: \n", r.StatusCode)
+		io.Copy(os.Stdout, r.Body)
+		os.Exit(r.StatusCode)
+	}
+
+	buf.Reset()
+	res := &Result{}
+	if e := json.NewDecoder(io.TeeReader(r.Body, buf)).Decode(res); e != nil {
+		log.Printf("Couldn't decode result: %v\nBody:\n", e)
+		io.Copy(os.Stdout, buf)
+	}
+	res.Data.Query = q.Query
+
+	if *raw {
+		io.Copy(os.Stdout, buf)
+	} else {
+		tmpl.Execute(os.Stdout, res.Data)
+	}
+}