From 4164da453976a4eeac3c5b52a77b6d044dd5265c Mon Sep 17 00:00:00 2001 From: Youngjin Jo Date: Tue, 12 Nov 2024 18:06:55 +0900 Subject: [PATCH] feat: add ai feature to cfctl Signed-off-by: Youngjin Jo --- cmd/ai.go | 185 ++++++++++++++++++++++++++++++++++++++++++++++++++ cmd/config.go | 86 +++++++++++------------ go.mod | 1 + go.sum | 2 + 4 files changed, 232 insertions(+), 42 deletions(-) create mode 100644 cmd/ai.go diff --git a/cmd/ai.go b/cmd/ai.go new file mode 100644 index 0000000..0a55934 --- /dev/null +++ b/cmd/ai.go @@ -0,0 +1,185 @@ +package cmd + +import ( + "bufio" + "context" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "strings" + + "github.com/pterm/pterm" + openai "github.com/sashabaranov/go-openai" + "github.com/spf13/cobra" +) + +var ( + apiToken string + configPath = filepath.Join(os.Getenv("HOME"), ".spaceone", "config") +) + +// aiCmd represents the ai command +var aiCmd = &cobra.Command{ + Use: "ai", + Short: "Run AI-powered tasks", + Long: `Run various AI-powered tasks, including general text processing or natural language +queries using OpenAI's API.`, + Run: func(cmd *cobra.Command, args []string) { + inputText, _ := cmd.Flags().GetString("input") + isNaturalLanguage, _ := cmd.Flags().GetBool("natural") + + if inputText == "" { + cmd.Help() + return + } + + result, err := runAIWithOpenAI(inputText, isNaturalLanguage) + if err != nil { + fmt.Printf("Error: %v\n", err) + return + } + fmt.Println("AI Result:", result) + }, +} + +// aiConfigCmd represents the ai config command +var aiConfigCmd = &cobra.Command{ + Use: "config", + Short: "Set up OpenAI Secret Key", + Run: func(cmd *cobra.Command, args []string) { + pterm.Info.Println("Setting up OpenAI API key for AI commands...") + + // Prompt user for OpenAI API key + apiKey, err := pterm.DefaultInteractiveTextInput.WithMultiLine(false).Show("Enter your OpenAI API Key") + if err != nil { + pterm.Error.Println("Failed to read API key:", err) + return + } + + // Write to config file + if err := saveAPIKeyToConfig(apiKey); err != nil { + pterm.Error.Println("Error saving API key:", err) + return + } + + pterm.Success.Println("OpenAI API key saved successfully to", configPath) + }, +} + +// runAIWithOpenAI processes input with OpenAI's streaming API +func runAIWithOpenAI(input string, natural bool) (string, error) { + // Load the API key from config if it's not set in the environment + apiToken = os.Getenv("OPENAI_API_TOKEN") + if apiToken == "" { + apiToken, _ = readAPIKeyFromConfig() + } + if apiToken == "" { + return "", errors.New("OpenAI API key is not set. Run `cfctl ai config` to configure it.") + } + + client := openai.NewClient(apiToken) + ctx := context.Background() + + // Set up the request based on the mode (natural language or standard) + model := openai.GPT3Babbage002 + if natural { + model = openai.GPT3Babbage002 + } + + req := openai.CompletionRequest{ + Model: model, + MaxTokens: 5, + Prompt: input, + Stream: true, + } + + stream, err := client.CreateCompletionStream(ctx, req) + if err != nil { + return "", fmt.Errorf("completion stream error: %v", err) + } + defer stream.Close() + + // Capture the streamed response + var responseText string + for { + response, err := stream.Recv() + if errors.Is(err, io.EOF) { + fmt.Println("Stream finished") + break + } + if err != nil { + return "", fmt.Errorf("stream error: %v", err) + } + responseText += response.Choices[0].Text + fmt.Printf("Stream response: %s", response.Choices[0].Text) + } + return responseText, nil +} + +// saveAPIKeyToConfig saves or updates the OpenAI API key in the config file +func saveAPIKeyToConfig(apiKey string) error { + // Read the existing config file content + content := "" + if _, err := os.Stat(configPath); err == nil { + data, err := os.ReadFile(configPath) + if err != nil { + return fmt.Errorf("failed to read config file: %v", err) + } + content = string(data) + } + + // Check if OPENAI_SECRET_KEY is already present + if strings.Contains(content, "OPENAI_SECRET_KEY=") { + // Update the existing key + content = strings.ReplaceAll(content, + "OPENAI_SECRET_KEY="+getAPIKeyFromContent(content), + "OPENAI_SECRET_KEY="+apiKey) + } else { + // Append the key if not present + content += "\nOPENAI_SECRET_KEY=" + apiKey + } + + // Write the updated content back to the config file + return os.WriteFile(configPath, []byte(content), 0600) +} + +// getAPIKeyFromContent extracts the existing API key from content if available +func getAPIKeyFromContent(content string) string { + scanner := bufio.NewScanner(strings.NewReader(content)) + for scanner.Scan() { + line := scanner.Text() + if strings.HasPrefix(line, "OPENAI_SECRET_KEY=") { + return line[17:] + } + } + return "" +} + +// readAPIKeyFromConfig reads the OpenAI API key from the config file +func readAPIKeyFromConfig() (string, error) { + file, err := os.Open(configPath) + if err != nil { + return "", err + } + defer file.Close() + + scanner := bufio.NewScanner(file) + for scanner.Scan() { + line := scanner.Text() + if len(line) > 0 && line[0:17] == "OPENAI_SECRET_KEY=" { + return line[17:], nil + } + } + return "", errors.New("API key not found in config file") +} + +func init() { + rootCmd.AddCommand(aiCmd) + aiCmd.Flags().String("input", "", "Input text for the AI to process") + aiCmd.Flags().BoolP("natural", "n", false, "Enable natural language mode for the AI") + + // Add config command as a subcommand to aiCmd + aiCmd.AddCommand(aiConfigCmd) +} diff --git a/cmd/config.go b/cmd/config.go index ce65b79..fc5c74f 100644 --- a/cmd/config.go +++ b/cmd/config.go @@ -272,9 +272,20 @@ func updateGlobalConfig() { log.Fatalf("Unable to list environments: %v", err) } - // Open ~/.spaceone/config for writing (will overwrite existing contents) configPath := filepath.Join(getConfigDir(), "config") - file, err := os.Create(configPath) + + // Read existing config file content to avoid overwriting + var content string + if _, err := os.Stat(configPath); err == nil { + data, err := os.ReadFile(configPath) + if err != nil { + log.Fatalf("Failed to read config file: %v", err) + } + content = string(data) + } + + // Open the config file for writing + file, err := os.OpenFile(configPath, os.O_WRONLY|os.O_CREATE, 0600) if err != nil { log.Fatalf("Failed to open config file: %v", err) } @@ -283,70 +294,61 @@ func updateGlobalConfig() { writer := bufio.NewWriter(file) defer writer.Flush() - // Write each environment that currently exists in ~/.spaceone/environments + // Add each environment without duplicates for _, entry := range entries { if !entry.IsDir() && filepath.Ext(entry.Name()) == ".yml" { name := strings.TrimSuffix(entry.Name(), filepath.Ext(entry.Name())) - _, err := writer.WriteString(fmt.Sprintf("[%s]\ncfctl environments -s %s\n\n", name, name)) - if err != nil { - log.Fatalf("Failed to write to config file: %v", err) + envEntry := fmt.Sprintf("[%s]\ncfctl environments -s %s\n\n", name, name) + if !strings.Contains(content, fmt.Sprintf("[%s]", name)) { + _, err := writer.WriteString(envEntry) + if err != nil { + log.Fatalf("Failed to write to config file: %v", err) + } } } } + pterm.Success.Println("Global config updated with existing environments.") } -// updateGlobalConfigWithEnvironment adds the new environment command to the global config file +// updateGlobalConfigWithEnvironment adds or updates the environment entry in the global config file func updateGlobalConfigWithEnvironment(environment string) { configPath := filepath.Join(getConfigDir(), "config") - var file *os.File - var err error - - // Check if the config file already exists - if _, err := os.Stat(configPath); os.IsNotExist(err) { - // Create a new config file if it does not exist - file, err = os.Create(configPath) - if err != nil { - log.Fatalf("Failed to create config file: %v", err) - } - pterm.Info.Printf("Created global config file with environment '%s'.\n", environment) - } else { - // Read the existing config file to check for duplicates - content, err := os.ReadFile(configPath) + // Read the existing config content if it exists + var content string + if _, err := os.Stat(configPath); err == nil { + data, err := os.ReadFile(configPath) if err != nil { log.Fatalf("Failed to read config file: %v", err) } - if strings.Contains(string(content), fmt.Sprintf("[%s]", environment)) { - pterm.Info.Printf("Environment '%s' already exists in the config file.\n", environment) - return - } + content = string(data) + } - // Open the existing config file for appending - file, err = os.OpenFile(configPath, os.O_APPEND|os.O_WRONLY, 0600) - if err != nil { - log.Fatalf("Failed to open config file: %v", err) - } + // Check if the environment already exists + envEntry := fmt.Sprintf("[%s]\ncfctl environments -s %s\n", environment, environment) + if strings.Contains(content, fmt.Sprintf("[%s]", environment)) { + pterm.Info.Printf("Environment '%s' already exists in the config file.\n", environment) + return + } - // Ensure the last line ends with a newline - if len(content) > 0 && content[len(content)-1] != '\n' { - _, err = file.WriteString("\n") - if err != nil { - log.Fatalf("Failed to write newline to config file: %v", err) - } - } + // Append the new environment entry + file, err := os.OpenFile(configPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0600) + if err != nil { + log.Fatalf("Failed to open config file: %v", err) } defer file.Close() - writer := bufio.NewWriter(file) - defer writer.Flush() + // Ensure a newline at the end of existing content + if len(content) > 0 && content[len(content)-1] != '\n' { + _, _ = file.WriteString("\n") + } - // Append the new environment to the config file - _, err = writer.WriteString(fmt.Sprintf("[%s]\ncfctl environments -s %s\n\n", environment, environment)) + _, err = file.WriteString(envEntry + "\n") if err != nil { log.Fatalf("Failed to write to config file: %v", err) } - //pterm.Success.Printf("Added environment '%s' to global config file.\n", environment) + pterm.Success.Printf("Added environment '%s' to global config file.\n", environment) } func init() { diff --git a/go.mod b/go.mod index 65f0fd5..3d693ec 100644 --- a/go.mod +++ b/go.mod @@ -6,6 +6,7 @@ require ( github.com/atotto/clipboard v0.1.4 github.com/jhump/protoreflect v1.17.0 github.com/pterm/pterm v0.12.79 + github.com/sashabaranov/go-openai v1.35.6 github.com/spf13/cobra v1.8.1 github.com/spf13/viper v1.19.0 google.golang.org/grpc v1.62.1 diff --git a/go.sum b/go.sum index 2ba4436..4dc569f 100644 --- a/go.sum +++ b/go.sum @@ -90,6 +90,8 @@ github.com/sagikazarmark/locafero v0.4.0 h1:HApY1R9zGo4DBgr7dqsTH/JJxLTTsOt7u6ke github.com/sagikazarmark/locafero v0.4.0/go.mod h1:Pe1W6UlPYUk/+wc/6KFhbORCfqzgYEpgQ3O5fPuL3H4= github.com/sagikazarmark/slog-shim v0.1.0 h1:diDBnUNK9N/354PgrxMywXnAwEr1QZcOr6gto+ugjYE= github.com/sagikazarmark/slog-shim v0.1.0/go.mod h1:SrcSrq8aKtyuqEI1uvTDTK1arOWRIczQRv+GVI1AkeQ= +github.com/sashabaranov/go-openai v1.35.6 h1:oi0rwCvyxMxgFALDGnyqFTyCJm6n72OnEG3sybIFR0g= +github.com/sashabaranov/go-openai v1.35.6/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg= github.com/sergi/go-diff v1.2.0 h1:XU+rvMAioB0UC3q1MFrIQy4Vo5/4VsRDQQXHsEya6xQ= github.com/sergi/go-diff v1.2.0/go.mod h1:STckp+ISIX8hZLjrqAeVduY0gWCT9IjLuqbuNXdaHfM= github.com/sourcegraph/conc v0.3.0 h1:OQTbbt6P72L20UqAkXXuLOj79LfEanQ+YQFNpLA9ySo=