From 5877aa2c6f734c9ed4707c232c67df1a7e8e17d5 Mon Sep 17 00:00:00 2001 From: cpendery Date: Wed, 13 Sep 2023 16:19:45 -0400 Subject: [PATCH] feat: add generic ai generator & fix spacing issues Signed-off-by: cpendery --- autocomplete/completion.go | 2 + autocomplete/generators/ai.go | 175 +++++++++++++++++++++++++++ autocomplete/generators/generator.go | 43 +++++++ autocomplete/generators/git/git.go | 25 ++++ autocomplete/generators/templates.go | 21 ++++ autocomplete/model/model.go | 9 +- autocomplete/suggestions.go | 24 ++-- 7 files changed, 280 insertions(+), 19 deletions(-) create mode 100644 autocomplete/generators/ai.go create mode 100644 autocomplete/generators/generator.go create mode 100644 autocomplete/generators/git/git.go diff --git a/autocomplete/completion.go b/autocomplete/completion.go index 979a579..8003741 100644 --- a/autocomplete/completion.go +++ b/autocomplete/completion.go @@ -118,6 +118,7 @@ func getSubcommandDrivenRecommendation(spec model.Subcommand, persistentOptions } if len(spec.Args) != 0 { activeArg := spec.Args[0] + getGeneratorDrivenRecommendations(activeArg.Generator, &suggestions) getSuggestionDrivenRecommendations(activeArg.Suggestions, &suggestions) getTemplateDrivenRecommendations(activeArg.Templates, &suggestions) } @@ -146,6 +147,7 @@ func getArgDrivenRecommendation(args []model.Arg, spec model.Subcommand, persist activeArg := args[0] allOptions := append(spec.Options, persistentOptions...) + getGeneratorDrivenRecommendations(activeArg.Generator, &suggestions) getSuggestionDrivenRecommendations(activeArg.Suggestions, &suggestions) getTemplateDrivenRecommendations(activeArg.Templates, &suggestions) diff --git a/autocomplete/generators/ai.go b/autocomplete/generators/ai.go new file mode 100644 index 0000000..1b178da --- /dev/null +++ b/autocomplete/generators/ai.go @@ -0,0 +1,175 @@ +package generators + +import ( + "bytes" + "encoding/json" + "io" + "log/slog" + "math" + "net/http" + "os" + "os/exec" + "strings" + + "github.com/cpendery/clac/autocomplete/model" +) + +const ( + enableAiEnvVar = "CLAC_ENABLE_AUTOCOMPLETE_AI" + apiKeyEnvVar = "CLAC_AI_TOKEN" + apiEndpoint = "https://api.openai.com/v1/chat/completions" + maxTokens = 4097 + tokensToRuneRatio = 4 + marginRatio = 0.8 + maxRunes = maxTokens * tokensToRuneRatio * marginRatio +) + +var ( + suggestionCache = make(map[string][]model.TermSuggestion) +) + +type PromptFunction func(executeShellCommand func(string) string) string +type MessageFunction func(executeShellCommand func(string) string) string + +type apiResponse struct { + Choices []choice `json:"choices"` +} +type choice struct { + Message message `json:"message"` + Info map[string]int `json:"info"` +} +type message struct { + Content *string `json:"content,omitempty"` +} + +type apiRequest struct { + Model string `json:"model"` + Messages []apiMessage `json:"messages"` +} + +type apiMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +func enabled() bool { + enabled := os.Getenv(enableAiEnvVar) + token := os.Getenv(apiKeyEnvVar) + s1 := strings.ToLower(enabled) == "true" + s2 := len(token) != 0 + return s1 && s2 +} + +func executeShellCommand(script string) string { + args := strings.Split(script, " ") + var cmd *exec.Cmd = nil + if len(args) > 1 { + cmd = exec.Command(args[0], args[1:]...) + } else { + cmd = exec.Command(args[0]) + } + output, err := cmd.Output() + if err != nil { + slog.Error("failed to run script in generator", slog.String("script", script), slog.String("error", err.Error()), slog.String("output", string(output))) + return "" + } + return string(output) +} + +func request(name, prompt, message, splitOn string) []model.TermSuggestion { + cacheKey := name + "|" + prompt + "|" + message + if suggestions, exists := suggestionCache[cacheKey]; exists { + return suggestions + } + suggestions := []model.TermSuggestion{} + jsonBytes, err := json.Marshal(apiRequest{ + Model: "gpt-3.5-turbo", + Messages: []apiMessage{ + {Role: "system", Content: prompt}, + {Role: "user", Content: message}, + }, + }) + if err != nil { + slog.Error("unable to marshal ai request", slog.String("error", err.Error())) + return suggestions + } + + apiKey := os.Getenv(apiKeyEnvVar) + request, _ := http.NewRequest("POST", apiEndpoint, bytes.NewBuffer(jsonBytes)) + request.Header.Set("Content-Type", "application/json") + request.Header.Set("Authorization", "Bearer "+apiKey) + client := http.Client{} + response, err := client.Do(request) + if err != nil { + slog.Error("failed to make ai request", slog.String("error", err.Error())) + return suggestions + } + if response.StatusCode != 200 { + contentBytes, _ := io.ReadAll(response.Body) + slog.Error("failed to make ai request", slog.String("status", response.Status), slog.String("content", string(contentBytes))) + return suggestions + } + + responseBody, err := io.ReadAll(response.Body) + if err != nil { + slog.Error("failed to read ai response", slog.String("error", err.Error())) + return suggestions + } + + var responseData apiResponse + if err := json.Unmarshal(responseBody, &responseData); err != nil { + slog.Error("invalid json response", slog.String("error", err.Error())) + return suggestions + } + + aiSuggestions := []string{} + for _, choice := range responseData.Choices { + content := choice.Message.Content + if content == nil { + continue + } + if splitOn != "" { + for _, s := range strings.Split(*content, splitOn) { + if strings.TrimSpace(s) != "" { + aiSuggestions = append(aiSuggestions, s) + } + } + + } else { + aiSuggestions = append(aiSuggestions, *content) + } + } + for _, aiSuggestion := range aiSuggestions { + suggestions = append(suggestions, model.TermSuggestion{ + Name: `"` + aiSuggestion + `"`, + Description: "Generated by Clac AI\n\n" + aiSuggestion, + Type: model.TermSuggestionTypeAI, + }) + } + suggestionCache[cacheKey] = suggestions + return suggestions +} + +func AI(name string, prompt PromptFunction, message MessageFunction, splitOn string) *model.Generator { + return &model.Generator{ + Function: func() []model.TermSuggestion { + suggestions := []model.TermSuggestion{} + if !enabled() { + return suggestions + } + + promptContent := prompt(executeShellCommand) + messageContent := "" + if message != nil { + messageContent = message(executeShellCommand) + } + + budget := int(math.Floor(maxRunes)) - len(promptContent) + if len(messageContent) > budget { + messageContent = messageContent[:budget] + } + + return request(name, promptContent, messageContent, splitOn) + }, + } +} diff --git a/autocomplete/generators/generator.go b/autocomplete/generators/generator.go new file mode 100644 index 0000000..795ff6c --- /dev/null +++ b/autocomplete/generators/generator.go @@ -0,0 +1,43 @@ +package generators + +import ( + "log/slog" + "os/exec" + "strings" + + "github.com/cpendery/clac/autocomplete/model" +) + +func Run(g model.Generator) []model.TermSuggestion { + suggestions := []model.TermSuggestion{} + if g.Script != "" { + args := strings.Split(g.Script, " ") + var cmd *exec.Cmd = nil + if len(args) > 1 { + cmd = exec.Command(args[0], args[1:]...) + } else { + cmd = exec.Command(args[0]) + } + output, err := cmd.Output() + if err != nil { + slog.Error("failed to run script in generator", slog.String("script", g.Script), slog.String("error", err.Error()), slog.String("output", string(output))) + } else if g.PostProcess != nil { + suggestions = append(suggestions, g.PostProcess(string(output))...) + } else { + tokens := strings.Split(string(output), g.SplitOn) + for _, token := range tokens { + suggestions = append(suggestions, model.TermSuggestion{ + Name: token, + }) + } + } + } + + if g.Function != nil { + suggestions = append(suggestions, g.Function()...) + } + + suggestions = append(suggestions, RunTemplates(g.Template)...) + + return suggestions +} diff --git a/autocomplete/generators/git/git.go b/autocomplete/generators/git/git.go new file mode 100644 index 0000000..0c813f8 --- /dev/null +++ b/autocomplete/generators/git/git.go @@ -0,0 +1,25 @@ +package git + +import ( + "github.com/cpendery/clac/autocomplete/generators" + "github.com/cpendery/clac/autocomplete/model" +) + +func CommitMessageGenerator() *model.Generator { + return generators.AI( + "git commit -m", + func(executeShellCommand func(string) string) string { + gitLogShortMessages := executeShellCommand( + "git log --pretty=format:%s --abbrev-commit --max-count=20", + ) + + return "Generate a git commit message summary based on this git diff, the \"summary\" must be no more " + + "than 70-75 characters, and it must describe what the patch changes" + + "\n\nHere are some examples from the repo:\n" + gitLogShortMessages + }, + func(executeShellCommand func(string) string) string { + return executeShellCommand("git diff --staged") + }, + "\n", + ) +} diff --git a/autocomplete/generators/templates.go b/autocomplete/generators/templates.go index d7d6af2..fc3d932 100644 --- a/autocomplete/generators/templates.go +++ b/autocomplete/generators/templates.go @@ -3,6 +3,7 @@ package generators import ( "log/slog" "os" + "slices" "github.com/cpendery/clac/autocomplete/model" ) @@ -65,3 +66,23 @@ func History() []model.TermSuggestion { func Help() []model.TermSuggestion { return []model.TermSuggestion{} } + +func RunTemplates(templates []model.Template) []model.TermSuggestion { + suggestions := []model.TermSuggestion{} + for _, template := range templates { + switch template { + case model.TemplateFilepaths: + suggestions = append(suggestions, Filepaths()...) + case model.TemplateFolders: + if slices.Contains(templates, model.TemplateFilepaths) { + continue + } + suggestions = append(suggestions, Folders()...) + case model.TemplateHelp: + suggestions = append(suggestions, Help()...) + case model.TemplateHistory: + suggestions = append(suggestions, History()...) + } + } + return suggestions +} diff --git a/autocomplete/model/model.go b/autocomplete/model/model.go index 76bd7dd..3bcc66b 100644 --- a/autocomplete/model/model.go +++ b/autocomplete/model/model.go @@ -53,6 +53,7 @@ const ( TermSuggestionTypeArg TermSuggestionType = "arg" TermSuggestionTypeSubcommand TermSuggestionType = "subcommand" TermSuggestionTypeOption TermSuggestionType = "option" + TermSuggestionTypeAI TermSuggestionType = "ai" TermSuggestionTypeDefault TermSuggestionType = "" ) @@ -63,9 +64,10 @@ type ProcessedToken struct { type Generator struct { Script string - Function func() string - PostProcess func(string) []Suggestion + Function func() []TermSuggestion + PostProcess func(string) []TermSuggestion Template []Template + SplitOn string } type Template string @@ -91,8 +93,9 @@ var ( TermSuggestionTypeFolder: "📁", TermSuggestionTypeFile: "📄", TermSuggestionTypeSubcommand: "📦", - TermSuggestionTypeOption: "⚙️ ", + TermSuggestionTypeOption: "💲", TermSuggestionTypeArg: "💪", TermSuggestionTypeDefault: "💪", + TermSuggestionTypeAI: "🔮", } ) diff --git a/autocomplete/suggestions.go b/autocomplete/suggestions.go index 92d1000..93ddb9e 100644 --- a/autocomplete/suggestions.go +++ b/autocomplete/suggestions.go @@ -1,7 +1,6 @@ package autocomplete import ( - "slices" "sort" "strings" @@ -64,22 +63,15 @@ func getPrefixFilteredRecommendations(input string, suggestions *[]model.TermSug *suggestions = results } -func getTemplateDrivenRecommendations(templates []model.Template, suggestions *[]model.TermSuggestion) { - for _, template := range templates { - switch template { - case model.TemplateFilepaths: - *suggestions = append(*suggestions, generators.Filepaths()...) - case model.TemplateFolders: - if slices.Contains(templates, model.TemplateFilepaths) { - continue - } - *suggestions = append(*suggestions, generators.Folders()...) - case model.TemplateHelp: - *suggestions = append(*suggestions, generators.Help()...) - case model.TemplateHistory: - *suggestions = append(*suggestions, generators.History()...) - } +func getGeneratorDrivenRecommendations(g *model.Generator, suggestions *[]model.TermSuggestion) { + if g != nil { + *suggestions = append(*suggestions, generators.Run(*g)...) } + +} + +func getTemplateDrivenRecommendations(templates []model.Template, suggestions *[]model.TermSuggestion) { + *suggestions = append(*suggestions, generators.RunTemplates(templates)...) } func getSuggestionDrivenRecommendations(suggestionSet []model.Suggestion, suggestions *[]model.TermSuggestion) {