better token selectionHEADmain

1 files changed, 18 insertions, 1 deletions

diff --git a/main.go b/main.go
index 59f3648..8fec3fc 100644
--- a/main.go
+++ b/main.go
@@ -14,7 +14,7 @@ import (
 )
 
 var (
-	token     = flag.String("token", os.Getenv("FASTGPT_API_TOKEN"), "API-Token, default from FASTGPT_API_TOKEN environment")
+	token     = flag.String("token", "", "API-Token to use instead of $FASTGPT_API_TOKEN or ~/fastgpt_token")
 	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.")
@@ -81,6 +81,23 @@ func main() {
 	buf := &bytes.Buffer{}
 
 	flag.Parse()
+
+	if *token == "" {
+		*token = os.Getenv("FASTGPT_API_TOKEN")
+		if *token == "" {
+			h, _ := os.UserHomeDir()
+			b, e := os.ReadFile(h + "/.fastgpt_token")
+			if e != nil {
+				panic(e)
+			}
+			*token = string(b)
+		}
+	}
+
+	if *token == "" {
+		panic("token needed")
+	}
+
 	if *stdin {
 		if _, e := io.Copy(buf, os.Stdin); e != nil {
 			panic(e)