github.com/xyproto/orbiton/v2@v2.65.12-0.20240516144430-e10a419274ec/apikey.go (about)

     1  package main
     2  
     3  import (
     4  	"errors"
     5  	"os"
     6  	"path/filepath"
     7  	"strings"
     8  
     9  	"github.com/xyproto/env/v2"
    10  )
    11  
    12  // KeyHolder holds an API key and a cache filename
    13  type KeyHolder struct {
    14  	Key      string
    15  	Filename string
    16  }
    17  
    18  var openAIKeyHolder = NewKeyHolder()
    19  
    20  // NewKeyHolder creates a new struct for storing the OpenAI API Key + a key cache filename
    21  // Can return nil if the key ends up being empty!
    22  // Use NewKeyHolderWithKey instead to allow empty keys and never return nil.
    23  func NewKeyHolder() *KeyHolder {
    24  	key := env.StrAlt("OPENAI_API_KEY", "OPENAI_KEY", env.Str("CHATGPT_API_KEY"))
    25  	kh := NewKeyHolderWithKey(key)
    26  	if kh.Key == "" {
    27  		if !kh.ReadAPIKey() {
    28  			return nil // !
    29  		}
    30  	}
    31  	return kh
    32  }
    33  
    34  // NewKeyHolderWithKey creates a new struct for storing the OpenAI API Key + a key cache filename,
    35  // and takes an initial key string. Will always return a struct, never nil.
    36  func NewKeyHolderWithKey(key string) *KeyHolder {
    37  	var kh KeyHolder
    38  	kh.Filename = filepath.Join(userCacheDir, "o", "openai_key.txt") // just for caching the key, if it's entered via the menu
    39  	kh.Key = key
    40  	return &kh
    41  }
    42  
    43  // ReadAPIKey tries to read the Open AI API Key from file.
    44  // An empty string is returned if the file could not be read.
    45  // Return true if a key is exists, or false if the key is empty.
    46  func (kh *KeyHolder) ReadAPIKey() bool {
    47  	if kh.Filename == "" {
    48  		return kh.Key != ""
    49  	}
    50  	data, err := os.ReadFile(kh.Filename)
    51  	if err != nil {
    52  		return kh.Key != ""
    53  	}
    54  	kh.Key = strings.TrimSpace(string(data))
    55  	return kh.Key != ""
    56  }
    57  
    58  // WriteAPIKey writes the given OpenAI API key to file
    59  func (kh *KeyHolder) WriteAPIKey() error {
    60  	if noWriteToCache {
    61  		return nil
    62  	}
    63  	if kh.Key == "" {
    64  		return errors.New("no API Key to write")
    65  	}
    66  	if kh.Filename == "" {
    67  		return errors.New("no API filename to write to")
    68  	}
    69  	return os.WriteFile(kh.Filename, []byte(kh.Key+"\n"), 0o600)
    70  }