Add error handling

README.md

@@ -18,8 +18,14 @@ $ go install github.com/noboruma/prompt-ai@latest
 ```
 
 ## Usage
-Find your key [here](https://platform.openai.com/account/api-keys)
+Find your key `<API_KEY>` [here](https://platform.openai.com/account/api-keys)
+Then run:
 
 ```
-$ API_KEY=blah prompt-ai
+$ OPENAI_KEY=<API_KEY> prompt-ai
+```
+or
+```
+$ export OPENAI_KEY=<API_KEY>
+$ prompt-ai
 ```

ais/chatgpt.go

@@ -3,6 +3,7 @@ package ais
 import (
 	"bytes"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"log"
@@ -18,9 +19,9 @@ var (
 var apiKey string
 
 func init() {
-	apiKey = os.Getenv("API_KEY")
+	apiKey = os.Getenv("OPENAI_KEY")
 	if apiKey == "" {
-		log.Fatal("API_KEY not set")
+		log.Fatal("OPENAI_KEY not set")
 	}
 }
 
@@ -86,6 +87,10 @@ func SendPrompt(prompt string, max_tokens int) ([]string, error) {
 	for i := range r.Choices {
 		res = append(res, r.Choices[i].Message.Content)
 	}
+
+	if len(res) == 0 {
+		return res, errors.New("Communication failed")
+	}
 	return res, nil
 }
 
@@ -125,7 +130,6 @@ func ListEngines() ([]Engine, error) {
 	if err != nil {
 		return res, err
 	}
-	fmt.Println(string(body))
 
 	listresp := ListEnginesResponse{}
 	err = json.Unmarshal(body, &listresp)
@@ -159,7 +163,6 @@ func GetOpenAIQuotaUsage() (QuotaUsage, error) {
 		return QuotaUsage{}, fmt.Errorf("error making request: %v", err)
 	}
 	defer resp.Body.Close()
-	fmt.Printf("%v\n", resp)
 
 	var usage struct {
 		UsageUsd float64 `json:"current_usage_usd"`

main.go

@@ -27,26 +27,29 @@ func prepareChatViewSection() *tview.TextView {
 
 func prepareInputSection(app *tview.Application, chatView *tview.TextView, copy_clipboards *clipboards.CopyClipboards) *tview.Flex {
 
-	errorText := tview.NewTextView()
-	errorText.SetTextColor(tcell.ColorRed)
-	statusText := tview.NewTextView().
+	errorText := tview.NewTextView().
 		SetScrollable(false).
 		SetToggleHighlights(false).
-		SetDynamicColors(true).
-		SetRegions(true).
 		SetWordWrap(true)
+	errorText.SetTextColor(tcell.ColorRed)
 	errorLog := func(msg string) {
 		app.QueueUpdateDraw(func() {
 			errorText.SetText(msg)
 		})
 	}
+	statusText := tview.NewTextView().
+		SetScrollable(false).
+		SetToggleHighlights(false).
+		SetDynamicColors(true).
+		SetRegions(true).
+		SetWordWrap(true)
 
 	updateQuota := func() {
 		usage, err := ais.GetOpenAIQuotaUsage()
 		if err != nil {
 			errorLog(err.Error())
 		}
-		fmt.Fprintf(statusText, "[yellow]Usage: $%v", usage.Used)
+		statusText.SetText(fmt.Sprintf("[yellow]Usage: $%v", usage.Used))
 	}
 	updateQuota()
 
@@ -64,10 +67,10 @@ func prepareInputSection(app *tview.Application, chatView *tview.TextView, copy_
 				defer func() { spinctl <- struct{}{} }()
 				responses, err := ais.SendPrompt(inputTextArea.GetText(), 100)
 				if err != nil {
-					chatView.SetText(err.Error())
-					chatView.SetBackgroundColor(tcell.ColorRed)
+					errorLog(err.Error())
 					return
 				}
+				errorLog("")
 
 				sb.WriteString("[yellow]>>")
 				sb.WriteString(inputTextArea.GetText())
@@ -92,6 +95,7 @@ func prepareInputSection(app *tview.Application, chatView *tview.TextView, copy_
 				sb.WriteString("\n")
 				chatView.SetText(sb.String())
 				inputTextArea.SetText("")
+				updateQuota()
 			}()
 		}
 	})
@@ -104,7 +108,8 @@ func prepareInputSection(app *tview.Application, chatView *tview.TextView, copy_
 		AddItem(tview.NewTextView().
 			SetDynamicColors(true).
 			SetText("[green]TAB[white]: switch panes"), 0, 1, false).
-		AddItem(statusText, 0, 1, true)
+		AddItem(statusText, 0, 1, false).
+		AddItem(errorText, 0, 1, false)
 
 	inputField.SetBorder(true).SetTitle("Input")
 	return inputField
@@ -138,7 +143,7 @@ func main() {
 	mainLayout := tview.NewFlex().
 		AddItem(tview.NewFlex().SetDirection(tview.FlexRow).
 			AddItem(chatView, 0, 2, false).
-			AddItem(inputField, 6, 1, true), 0, 1, true)
+			AddItem(inputField, 7, 1, true), 0, 1, true)
 
 	if err := app.SetRoot(mainLayout, true).EnableMouse(true).Run(); err != nil {
 		panic(err)