From 43d9e4cd4d1fcd7274fc154fa99445945b99ebbc Mon Sep 17 00:00:00 2001
From: Ammar Bandukwala <ammar@ammar.io>
Date: Mon, 2 Sep 2024 12:30:21 -0500
Subject: [PATCH] Add feature to save OpenAI API key locally

- Extend `README.md` with instructions to save the API key.
- Modify `main.go` to include a `--save-key` flag.
- Implement functions in `savekey.go` for saving/loading the key.
---
 README.md               |  8 +++++++
 cmd/aicommit/main.go    | 30 +++++++++++++++++++++++++
 cmd/aicommit/savekey.go | 50 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 88 insertions(+)
 create mode 100644 cmd/aicommit/savekey.go

diff --git a/README.md b/README.md
index 8f3d0e6..42683be 100644
--- a/README.md
+++ b/README.md
@@ -61,6 +61,14 @@ aicommit -c "improved HTTP performance by 50%"
 aicommit -c "bad code but need for urgent customer fix"
 ```
 
+When tired of setting environment variables, you can save your key to disk:
+
+```bash
+export OPENAI_API_KEY="..."
+aicommit --save-key
+# The environment variable will override the saved key.
+```
+
 ## Style Guide
 
 `aicommit` will read the `COMMITS.md` file in the root of the repository to
diff --git a/cmd/aicommit/main.go b/cmd/aicommit/main.go
index 27fc6d1..cfd0590 100644
--- a/cmd/aicommit/main.go
+++ b/cmd/aicommit/main.go
@@ -204,14 +204,38 @@ func main() {
 	var (
 		opts      runOptions
 		openAIKey string
+		doSaveKey bool
 	)
 	cmd := &serpent.Command{
 		Use:   "aicommit [ref]",
 		Short: "aicommit is a tool for generating commit messages",
 		Handler: func(inv *serpent.Invocation) error {
+			savedKey, err := loadKey()
+			if err != nil && !os.IsNotExist(err) {
+				return err
+			}
+			if savedKey != "" && openAIKey == "" {
+				openAIKey = savedKey
+			}
 			if openAIKey == "" {
 				return errors.New("$OPENAI_API_KEY is not set")
 			}
+
+			if doSaveKey {
+				err := saveKey(openAIKey)
+				if err != nil {
+					return err
+				}
+
+				kp, err := keyPath()
+				if err != nil {
+					return err
+				}
+
+				fmt.Fprintf(inv.Stdout, "Saved OpenAI API key to %s\n", kp)
+				return nil
+			}
+
 			client := openai.NewClient(openAIKey)
 			opts.client = client
 			if len(inv.Args) > 0 {
@@ -226,6 +250,12 @@ func main() {
 				Env:         "OPENAI_API_KEY",
 				Value:       serpent.StringOf(&openAIKey),
 			},
+			{
+				Name:        "save-key",
+				Description: "Save the OpenAI API key to persistent local configuration and exit.",
+				Flag:        "save-key",
+				Value:       serpent.BoolOf(&doSaveKey),
+			},
 			{
 				Name:          "dry-run",
 				Flag:          "dry",
diff --git a/cmd/aicommit/savekey.go b/cmd/aicommit/savekey.go
new file mode 100644
index 0000000..5c89510
--- /dev/null
+++ b/cmd/aicommit/savekey.go
@@ -0,0 +1,50 @@
+package main
+
+import (
+	"errors"
+	"os"
+	"path/filepath"
+)
+
+func configDir() (string, error) {
+	cdir, err := os.UserConfigDir()
+	if err != nil {
+		return "", err
+	}
+	err = os.MkdirAll(filepath.Join(cdir, "aicommit"), 0o700)
+	if err != nil {
+		return "", err
+	}
+	return filepath.Join(cdir, "aicommit"), nil
+}
+
+func keyPath() (string, error) {
+	cdir, err := configDir()
+	if err != nil {
+		return "", err
+	}
+	return filepath.Join(cdir, "openai.key"), nil
+}
+
+func saveKey(key string) error {
+	if key == "" {
+		return errors.New("key is empty")
+	}
+	kp, err := keyPath()
+	if err != nil {
+		return err
+	}
+	return os.WriteFile(kp, []byte(key), 0o600)
+}
+
+func loadKey() (string, error) {
+	kp, err := keyPath()
+	if err != nil {
+		return "", err
+	}
+	b, err := os.ReadFile(kp)
+	if err != nil {
+		return "", err
+	}
+	return string(b), nil
+}