Allow setting of AI config using ai command (#272)

* Allow setting of AI config using ai command

* We're going through changes

* Add support for multiple OpenAI tokens

* Run gofumft and linter

* Revert "Add support for multiple OpenAI tokens"

This reverts commit 2f33aef.

---------

Co-authored-by: Jarryd Tilbrook <jarryd@buildkite.com>

fixtures/config/user.basic.yaml

@@ -1,4 +1,5 @@
+openai_token: jumanji
 selected_org: buildkite-org
 organizations:
   buildkite-test:
-    api_token: test-token-abcd
+    api_token: test-token-abcd

internal/config/config.go

@@ -101,6 +101,17 @@ func (conf *Config) SelectOrganization(org string) error {
 	return conf.localConfig.WriteConfig()
 }
 
+// SetOpenAIToken sets the user's OpenAI token in the user config file
+func (conf *Config) SetOpenAIToken(token string) error {
+	conf.userConfig.Set("openai_token", token)
+	return conf.userConfig.WriteConfig()
+}
+
+// GetOpenAIToken reads the user's OpenAI token from the user config file
+func (conf *Config) GetOpenAIToken() string {
+	return conf.userConfig.GetString("openai_token")
+}
+
 // APIToken gets the API token configured for the currently selected organization
 func (conf *Config) APIToken() string {
 	slug := conf.OrganizationSlug()

pkg/cmd/ai/add/add.go

@@ -0,0 +1,41 @@
+package add
+
+import (
+	"github.com/AlecAivazis/survey/v2"
+	"github.com/buildkite/cli/v3/pkg/cmd/factory"
+	"github.com/spf13/cobra"
+)
+
+func NewCmdAIAdd(f *factory.Factory) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "add",
+		Args:  cobra.NoArgs,
+		Short: "Add an OpenAI token.",
+		Long:  "Add an OpenAI token for use with CLI AI tooling",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			return ConfigureRun(f)
+		},
+	}
+
+	return cmd
+}
+
+func ConfigureRun(f *factory.Factory) error {
+	qs := []*survey.Question{
+		{
+			Name:     "token",
+			Prompt:   &survey.Password{Message: "Paste your OpenAI token:"},
+			Validate: survey.Required,
+		},
+	}
+	answers := struct{ Token string }{}
+
+	err := survey.Ask(qs, &answers)
+	if err != nil {
+		return err
+	}
+
+	err = f.Config.SetOpenAIToken(answers.Token)
+
+	return err
+}

pkg/cmd/ai/ai.go

@@ -0,0 +1,23 @@
+package ai
+
+import (
+	"github.com/MakeNowJust/heredoc"
+	"github.com/buildkite/cli/v3/pkg/cmd/factory"
+	"github.com/spf13/cobra"
+)
+
+func NewCmdAI(f *factory.Factory) *cobra.Command {
+	cmd := cobra.Command{
+		Use:   "ai <command>",
+		Short: "Manage AI integration.",
+		Long:  "Work with Buildkite AI.",
+		Example: heredoc.Doc(`
+			# To configure your AI token
+			$ bk ai configure
+		`),
+	}
+
+	cmd.AddCommand(NewCmdAIConfigure(f))
+
+	return &cmd
+}

pkg/cmd/ai/configure.go

@@ -0,0 +1,34 @@
+package ai
+
+import (
+	"errors"
+
+	addCmd "github.com/buildkite/cli/v3/pkg/cmd/ai/add"
+	"github.com/buildkite/cli/v3/pkg/cmd/factory"
+	"github.com/spf13/cobra"
+)
+
+func NewCmdAIConfigure(f *factory.Factory) *cobra.Command {
+	var force bool
+
+	cmd := &cobra.Command{
+		Use:     "configure",
+		Aliases: []string{"config"},
+		Args:    cobra.NoArgs,
+		Short:   "Configure OpenAI token.",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			// if the token already exists and --force is not used
+			if !force && f.Config.GetOpenAIToken() != "" {
+				return errors.New("OpenAI token already configured. You must use --force.")
+			}
+
+			return addCmd.ConfigureRun(f)
+		},
+	}
+
+	cmd.Flags().BoolVar(&force, "force", false, "Force setting a new token")
+
+	cmd.AddCommand(addCmd.NewCmdAIAdd(f))
+
+	return cmd
+}

pkg/cmd/root/root.go

@@ -3,6 +3,7 @@ package root
 import (
 	"github.com/MakeNowJust/heredoc"
 	agentCmd "github.com/buildkite/cli/v3/pkg/cmd/agent"
+	aiCmd "github.com/buildkite/cli/v3/pkg/cmd/ai"
 	buildCmd "github.com/buildkite/cli/v3/pkg/cmd/build"
 	configureCmd "github.com/buildkite/cli/v3/pkg/cmd/configure"
 	"github.com/buildkite/cli/v3/pkg/cmd/factory"
@@ -30,6 +31,7 @@ func NewCmdRoot(f *factory.Factory) (*cobra.Command, error) {
 	cmd.AddCommand(useCmd.NewCmdUse(f))
 	cmd.AddCommand(initCmd.NewCmdInit(f))
 	cmd.AddCommand(agentCmd.NewCmdAgent(f))
+	cmd.AddCommand(aiCmd.NewCmdAI(f))
 	cmd.AddCommand(versionCmd.NewCmdVersion(f))
 	cmd.AddCommand(buildCmd.NewCmdBuild(f))