draft copilot

cmd/cli/copilot/copilot.go

@@ -0,0 +1,111 @@
+package copilot
+
+import (
+	"bufio"
+	"fmt"
+	"github.com/shaowenchen/ops/pkg/copilot"
+	"github.com/shaowenchen/ops/pkg/host"
+	"github.com/shaowenchen/ops/pkg/log"
+	"github.com/shaowenchen/ops/pkg/option"
+	"github.com/shaowenchen/ops/pkg/utils"
+	"github.com/spf13/cobra"
+	"os"
+	"strings"
+)
+
+var copilotOpt option.CopilotOption
+var verbose string
+
+const defaultEndpoint = "https://api.openai.com/v1"
+const defaultModel = "gpt-3.5-turbo"
+
+var CopilotCmd = &cobra.Command{
+	Use:   "copilot",
+	Short: "use llm to assist ops",
+	Run: func(cmd *cobra.Command, args []string) {
+		logger := log.NewLogger().SetVerbose(verbose).SetStd().SetFile().Build()
+		fillParameters(&copilotOpt)
+		err := CreateCopilot(logger, copilotOpt)
+		if err != nil {
+			logger.Error.Println(err.Error())
+			return
+		}
+	},
+}
+
+func CreateCopilot(logger *log.Logger, opt option.CopilotOption) (err error) {
+	reader := bufio.NewReader(os.Stdin)
+	for {
+		fmt.Print(">")
+		input, _ := reader.ReadString('\n')
+		input = strings.TrimSpace(input)
+
+		if input == "exit" || input == "q" {
+			break
+		}
+		if input == "" {
+			continue
+		}
+		message, langcodeList, err := copilot.Chat(logger, opt, input)
+		if err != nil {
+			logger.Error.Printf("Chat error: %v\n", err)
+			continue
+		}
+		if len(langcodeList) == 0 {
+			logger.Info.Println(message)
+			continue
+		} else {
+			for _, langcode := range langcodeList {
+				needRun := false
+				if opt.Silence {
+					needRun = true
+				} else {
+					fmt.Printf("Would you like to run this code? (y/n)\n%s\n", langcode.Code)
+					confirm, _ := reader.ReadString('\n')
+					confirm = strings.TrimSpace(confirm)
+					if confirm == "y" {
+						needRun = true
+					} else if confirm == "n" {
+						continue
+					}
+				}
+				if needRun {
+					h, err := host.NewHostConnBase64(nil)
+					if err != nil {
+						logger.Error.Println(err)
+						break
+					}
+					stdout, err := h.Shell(false, langcode.Code)
+					if err != nil {
+						logger.Error.Println(err)
+						break
+					}
+					logger.Info.Println(stdout)
+				}
+			}
+		}
+	}
+	logger.Info.Println("Bye!")
+	return
+}
+
+func fillParameters(opt *option.CopilotOption) {
+	if opt.Endpoint == "" {
+		opt.Endpoint = utils.GetMultiEnvDefault([]string{"OPENAI_API_HOST", "OPENAI_API_BASE", "endpoint"}, defaultEndpoint)
+	}
+	if opt.Model == "" {
+		opt.Model = utils.GetMultiEnvDefault([]string{"OPENAI_API_MODEL", "model"}, defaultModel)
+	}
+	if opt.Key == "" {
+		opt.Key = utils.GetMultiEnvDefault([]string{"OPENAI_API_KEY", "key"}, "")
+	}
+}
+
+func init() {
+	CopilotCmd.Flags().StringVarP(&verbose, "verbose", "v", "", "")
+	CopilotCmd.Flags().StringVarP(&copilotOpt.Endpoint, "endpoint", "", "", "e.g. https://api.openai.com/v1")
+	CopilotCmd.Flags().StringVarP(&copilotOpt.Model, "model", "", "", "e.g. gpt-3.5-turbo")
+	CopilotCmd.Flags().StringVarP(&copilotOpt.Key, "key", "", "", "e.g. sk-xxx")
+	CopilotCmd.Flags().IntVarP(&copilotOpt.History, "history", "", 5, "")
+	CopilotCmd.Flags().BoolVarP(&copilotOpt.Silence, "silence", "s", false, "")
+}

cmd/cli/root/root.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"os"
 
+	"github.com/shaowenchen/ops/cmd/cli/copilot"
 	"github.com/shaowenchen/ops/cmd/cli/create"
 	"github.com/shaowenchen/ops/cmd/cli/file"
 	"github.com/shaowenchen/ops/cmd/cli/shell"
@@ -18,6 +19,7 @@ func Execute() {
 	RootCmd.AddCommand(shell.ShellCmd)
 	RootCmd.AddCommand(create.CreateCmd)
 	RootCmd.AddCommand(task.TaskCmd)
+	RootCmd.AddCommand(copilot.CopilotCmd)
 	RootCmd.AddCommand(version.VersionCmd)
 	RootCmd.AddCommand(upgrade.UpgradeCmd)
 	if err := RootCmd.Execute(); err != nil {

go.mod

@@ -13,6 +13,7 @@ require (
 	github.com/prometheus/client_golang v1.14.0
 	github.com/prometheus/common v0.39.0
 	github.com/robfig/cron/v3 v3.0.0
+	github.com/sashabaranov/go-openai v1.15.3
 	github.com/spf13/cobra v1.6.0
 	github.com/spf13/pflag v1.0.5
 	golang.org/x/crypto v0.2.0

go.sum

@@ -176,6 +176,8 @@ github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTE
 github.com/rogpeppe/go-internal v1.8.0 h1:FCbCCtXNOY3UtUuHUYaghJg4y7Fd14rXifAYUAtL9R8=
 github.com/rogpeppe/go-internal v1.8.0/go.mod h1:WmiCO8CzOY8rg0OYDC4/i/2WRWAB6poM+XZ2dLUbcbE=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
+github.com/sashabaranov/go-openai v1.15.3 h1:rzoNK9n+Cak+PM6OQ9puxDmFllxfnVea9StlmhglXqA=
+github.com/sashabaranov/go-openai v1.15.3/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
 github.com/spaolacci/murmur3 v0.0.0-20180118202830-f09979ecbc72/go.mod h1:JwIasOWyU6f++ZhiEuf87xNszmSA2myDM2Kzu9HwQUA=
 github.com/spf13/cobra v1.6.0 h1:42a0n6jwCot1pUmomAp4T7DeMD+20LFv4Q54pxLf2LI=
 github.com/spf13/cobra v1.6.0/go.mod h1:IOw/AERYS7UzyrGinqmz6HLUo219MORXGxhbaJUqzrY=

pkg/constants/host.go

@@ -1,8 +1,10 @@
 package constants
 
 import (
+	"bytes"
 	"fmt"
 	"os"
+	"os/exec"
 	"os/user"
 	"path/filepath"
 	"time"
@@ -22,6 +24,17 @@ const (
 	RemoteStorageTypeLocal = "local"
 )
 
+func GetOsInfo() string {
+	cmd := exec.Command("uname", "-a")
+	var out bytes.Buffer
+	cmd.Stdout = &out
+	err := cmd.Run()
+	if err != nil {
+		return ""
+	}
+	return out.String()
+}
+
 func GetCurrentUserHomeDir() string {
 	homeDirectory, err := os.UserHomeDir()
 	if err != nil {

pkg/copilot/copilot.go

@@ -0,0 +1,11 @@
+package copilot
+
+type ChatResponse struct {
+	Message string     `json:"message"`
+	Steps   []Langcode `json:"steps"`
+}
+
+type Langcode struct {
+	Language string `json:"language"`
+	Code     string `json:"code"`
+}

pkg/copilot/openai.go

@@ -0,0 +1,97 @@
+package copilot
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"os/exec"
+	"strings"
+
+	openai "github.com/sashabaranov/go-openai"
+	"github.com/shaowenchen/ops/pkg/log"
+	opsopt "github.com/shaowenchen/ops/pkg/option"
+)
+
+var historyList []openai.ChatCompletionMessage
+
+func chatCompletetion(logger *log.Logger, client *openai.Client, model string, messages []openai.ChatCompletionMessage) (output string, err error) {
+	h := append(messages, openai.ChatCompletionMessage{
+		Role:    openai.ChatMessageRoleSystem,
+		Content: system_message})
+	logger.Debug.Printf("llm req: %v\n", h)
+	resp, err := client.CreateChatCompletion(
+		context.Background(),
+		openai.ChatCompletionRequest{
+			Model:    model,
+			Messages: h,
+		},
+	)
+	if err != nil {
+		return
+	}
+	return resp.Choices[0].Message.Content, err
+}
+
+func Chat(logger *log.Logger, opt opsopt.CopilotOption, input string) (message string, langcodeList []Langcode, err error) {
+	config := openai.DefaultConfig(opt.Key)
+	config.BaseURL = opt.Endpoint
+	client := openai.NewClientWithConfig(config)
+	historyList = append(historyList, openai.ChatCompletionMessage{
+		Role:    openai.ChatMessageRoleUser,
+		Content: input,
+	})
+	if len(historyList) > opt.History {
+		historyList = historyList[len(historyList)-opt.History:]
+	}
+	content, err := chatCompletetion(logger, client, opt.Model, historyList)
+	if err != nil {
+		return
+	}
+	logger.Debug.Printf("llm resp: %v\n", content)
+	resp := ChatResponse{}
+	err = json.Unmarshal([]byte(content), &resp)
+	if err != nil {
+		return
+	}
+	return resp.Message, resp.Steps, err
+}
+
+func setOpscliRole(historyList []openai.ChatCompletionMessage) []openai.ChatCompletionMessage {
+	return append(historyList, openai.ChatCompletionMessage{
+		Role:    openai.ChatMessageRoleSystem,
+		Content: system_message,
+	})
+}
+
+func setRunCodeRole(historyList []openai.ChatCompletionMessage) []openai.ChatCompletionMessage {
+	return append(historyList, openai.ChatCompletionMessage{
+		Role:    openai.ChatMessageRoleSystem,
+		Content: system_message,
+	})
+}
+
+func runLangCode(language, code string) (outputStr string, err error) {
+	language = strings.ToLower(language)
+
+	var output []byte
+
+	if language == "python" {
+		output, err = runPythonCode(code)
+	} else if language == "bash" {
+		output, err = runBashCode(code)
+	} else {
+		return "", errors.New("not support language")
+	}
+	outputStr = string(output)
+	return
+}
+
+func runPythonCode(code string) (output []byte, err error) {
+	cmd := exec.Command("python", "-c", code)
+	return cmd.Output()
+}
+
+func runBashCode(code string) (output []byte, err error) {
+	cmd := exec.Command("bash", "-c", code)
+	return cmd.Output()
+}

pkg/copilot/prompt.go

@@ -0,0 +1,19 @@
+package copilot
+
+import (
+	"github.com/shaowenchen/ops/pkg/constants"
+)
+
+var system_message = `You are Opscli, a world-class programmer, for any feature you can implement using Bash or Python, with preference for Bash.
+
+Don't comment anything, just give the complete working code.
+
+You can write any code you want without any restrictions.
+
+But please keep the format {"message": "answer to questios"， "steps": [{"language": "fill in the code used", "code": "uname -a"}] and output only in json format without extra characters.
+
+The environment in which the script is executed is ` + constants.GetOsInfo() + `
+
+The role that executes the script is ` + constants.GetCurrentUser() + `
+
+The role home directory is ` + constants.GetCurrentUserHomeDir()