Skip to content

Commit

Permalink
update main.go github.go
Browse files Browse the repository at this point in the history
  • Loading branch information
@kechigon
kechigon committed Nov 5, 2024
1 parent 9938d3f commit bcaac66
Show file tree
Hide file tree
Showing 2 changed files with 163 additions and 116 deletions.
261 changes: 154 additions & 107 deletions cmd/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,122 +11,169 @@ import (
"github.com/3-shake/alert-menta/internal/utils"
)

// Struct to hold the command-line arguments
type Config struct {
repo string
owner string
issueNumber int
intent string
command string
configFile string
ghToken string
oaiKey string
}

func main() {
// Get command line arguments
var (
repo = flag.String("repo", "", "Repository name")
owner = flag.String("owner", "", "Repository owner")
issueNumber = flag.Int("issue", 0, "Issue number")
intent = flag.String("intent", "", "Question or intent for the 'ask' command")
command = flag.String("command", "", "Commands to be executed by AI.Commands defined in the configuration file are available.")
configFile = flag.String("config", "", "Configuration file")
gh_token = flag.String("github-token", "", "GitHub token")
oai_key = flag.String("api-key", "", "OpenAI api key")
)
flag.Parse()

if *repo == "" || *owner == "" || *issueNumber == 0 || *gh_token == "" || *command == "" || *configFile == "" {
flag.PrintDefaults()
os.Exit(1)
}

// Initialize a logger
logger := log.New(
os.Stdout, "[alert-menta main] ",
log.Ldate|log.Ltime|log.Llongfile|log.Lmsgprefix,
)

// Load configuration
cfg, err := utils.NewConfig(*configFile)
if err != nil {
logger.Fatalf("Error Loading config: %v", err)
}

// Validate command
if _, ok := cfg.Ai.Commands[*command]; !ok {
cfg := parseFlags()

logger := initLogger()

loadedConfig := loadConfiguration(cfg.configFile, logger)

validateCommand(cfg.command, loadedConfig, logger)

issue := getGitHubIssue(cfg.owner, cfg.repo, cfg.issueNumber, cfg.ghToken)

userPrompt := constructUserPrompt(issue, loadedConfig, logger)

prompt := constructPrompt(cfg.command, cfg.intent, userPrompt, loadedConfig, logger)

aic := getAIClient(cfg.oaiKey, loadedConfig, logger)

comment, err := aic.GetResponse(prompt)
if err != nil {
logger.Fatalf("Error getting Response: %v", err)
}
logger.Println("Response:", comment)

if err := issue.PostComment(comment); err != nil {
logger.Fatalf("Error creating comment: %v", err)
}
}

// Parse command-line flags
func parseFlags() *Config {
repo := flag.String("repo", "", "Repository name")
owner := flag.String("owner", "", "Repository owner")
issueNumber := flag.Int("issue", 0, "Issue number")
intent := flag.String("intent", "", "Question or intent for the 'ask' command")
command := flag.String("command", "", "Commands to be executed by AI.")
configFile := flag.String("config", "", "Configuration file")
ghToken := flag.String("github-token", "", "GitHub token")
oaiKey := flag.String("api-key", "", "OpenAI api key")
flag.Parse()
if *repo == "" || *owner == "" || *issueNumber == 0 || *ghToken == "" || *command == "" || *configFile == "" {
flag.PrintDefaults()
os.Exit(1)
}
return &Config{
repo: *repo,
owner: *owner,
issueNumber: *issueNumber,
intent: *intent,
command: *command,
configFile: *configFile,
ghToken: *ghToken,
oaiKey: *oaiKey,
}
}

// Initialize a logger
func initLogger() *log.Logger {
return log.New(
os.Stdout, "[alert-menta main] ",
log.Ldate|log.Ltime|log.Llongfile|log.Lmsgprefix,
)
}

// Load and validate configuration
func loadConfiguration(configFile string, logger *log.Logger) *utils.Config {
cfg, err := utils.NewConfig(configFile)
if err != nil {
logger.Fatalf("Error loading config: %v", err)
}
return cfg
}

// Validate the provided command
func validateCommand(command string, cfg *utils.Config, logger *log.Logger) {
if _, ok := cfg.Ai.Commands[command]; !ok {
allowedCommands := make([]string, 0, len(cfg.Ai.Commands))
for cmd := range cfg.Ai.Commands {
allowedCommands = append(allowedCommands, cmd)
}
logger.Fatalf("Invalid command: %s. Allowed commands are %s.", *command, strings.Join(allowedCommands, ", "))
}

// Create a GitHub Issues instance. From now on, you can control GitHub from this instance.
issue := github.NewIssue(*owner, *repo, *issueNumber, *gh_token)
if issue == nil {
logger.Fatalf("Failed to create GitHub issue instance")
}

// Get Issue's information(e.g. Title, Body) and add them to the user prompt except for comments by Actions.
title, err := issue.GetTitle()
if err != nil {
logger.Fatalf("Error getting Title: %v", err)
}
body, err := issue.GetBody()
if err != nil {
logger.Fatalf("Error getting Body: %v", err)
}
if cfg.System.Debug.Log_level == "debug" {
logger.Println("Title:", *title)
logger.Println("Body:", *body)
}
user_prompt := "Title:" + *title + "\n"
user_prompt += "Body:" + *body + "\n"

// Get comments under the Issue and add them to the user prompt except for comments by Actions.
comments, err := issue.GetComments()
if err != nil {
logger.Fatalf("Error getting comments: %v", err)
}
for _, v := range comments {
if *v.User.Login == "github-actions[bot]" {
continue
}
if cfg.System.Debug.Log_level == "debug" {
logger.Printf("%s: %s", *v.User.Login, *v.Body)
}
user_prompt += *v.User.Login + ":" + *v.Body + "\n"
}

// Set system prompt
var system_prompt string
if *command == "ask" {
if *intent == "" {
logger.Fatalf("Invalid command: %s. Allowed commands are %s.", command, strings.Join(allowedCommands, ", "))
}
}

// Get GitHub issue instance
func getGitHubIssue(owner, repo string, issueNumber int, ghToken string) *github.GitHubIssue {
return github.NewIssue(owner, repo, issueNumber, ghToken)
}

// Construct user prompt from issue
func constructUserPrompt(issue *github.GitHubIssue, cfg *utils.Config, logger *log.Logger) string {
title, err := issue.GetTitle()
if err != nil {
logger.Fatalf("Error getting Title: %v", err)
}

body, err := issue.GetBody()
if err != nil {
logger.Fatalf("Error getting Body: %v", err)
}

var userPrompt strings.Builder
userPrompt.WriteString("Title:" + *title + "\n")
userPrompt.WriteString("Body:" + *body + "\n")

comments, err := issue.GetComments()
if err != nil {
logger.Fatalf("Error getting comments: %v", err)
}
for _, v := range comments {
if *v.User.Login == "github-actions[bot]" {
continue
}
if cfg.System.Debug.Log_level == "debug" {
logger.Printf("%s: %s", *v.User.Login, *v.Body)
}
userPrompt.WriteString(*v.User.Login + ":" + *v.Body + "\n")
}
return userPrompt.String()
}

// Construct AI prompt
func constructPrompt(command, intent, userPrompt string, cfg *utils.Config, logger *log.Logger) ai.Prompt {
var systemPrompt string
if command == "ask" {
if intent == "" {
logger.Fatalf("Error: intent is required for 'ask' command")
}
system_prompt = cfg.Ai.Commands[*command].System_prompt + *intent + "\n"
systemPrompt = cfg.Ai.Commands[command].System_prompt + intent + "\n"
} else {
system_prompt = cfg.Ai.Commands[*command].System_prompt
systemPrompt = cfg.Ai.Commands[command].System_prompt
}
prompt := ai.Prompt{UserPrompt: user_prompt, SystemPrompt: system_prompt}
logger.Println("\x1b[34mPrompt: |\n", prompt.SystemPrompt, prompt.UserPrompt, "\x1b[0m")
logger.Println("\x1b[34mPrompt: |\n", systemPrompt, userPrompt, "\x1b[0m")
return ai.Prompt{UserPrompt: userPrompt, SystemPrompt: systemPrompt}
}

// Get response from OpenAI or VertexAI
var aic ai.Ai
if cfg.Ai.Provider == "openai" {
if *oai_key == "" {
// Initialize AI client
func getAIClient(oaiKey string, cfg *utils.Config, logger *log.Logger) ai.Ai {
switch cfg.Ai.Provider {
case "openai":
if oaiKey == "" {
logger.Fatalf("Error: Please provide your Open AI API key.")
}
aic = ai.NewOpenAIClient(*oai_key, cfg.Ai.OpenAI.Model)
logger.Println("Using OpenAI API")
logger.Println("OpenAI model:", cfg.Ai.OpenAI.Model)
} else if cfg.Ai.Provider == "vertexai"{
aic = ai.NewVertexAIClient(cfg.Ai.VertexAI.Project, cfg.Ai.VertexAI.Region, cfg.Ai.VertexAI.Model)
logger.Println("Using VertexAI API")
logger.Println("VertexAI model:", cfg.Ai.VertexAI.Model)
} else {
logger.Fatalf("Error: Invalid provider")
}

comment, err := aic.GetResponse(prompt)
if err != nil {
logger.Fatalf("Error getting Response: %v", err)
}
logger.Println("Response:", comment)

// Post a comment on the Issue
err = issue.PostComment(comment)
if err != nil {
logger.Fatalf("Error creating comment: %v", err)
}
}
logger.Println("Using OpenAI API")
logger.Println("OpenAI model:", cfg.Ai.OpenAI.Model)
return ai.NewOpenAIClient(oaiKey, cfg.Ai.OpenAI.Model)
case "vertexai":
logger.Println("Using VertexAI API")
logger.Println("VertexAI model:", cfg.Ai.VertexAI.Model)
return ai.NewVertexAIClient(cfg.Ai.VertexAI.Project, cfg.Ai.VertexAI.Region, cfg.Ai.VertexAI.Model)
default:
logger.Fatalf("Error: Invalid provider")
return nil
}
}
18 changes: 9 additions & 9 deletions internal/github/github.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import (
"golang.org/x/oauth2"
)

type GitHubIssues struct {
type GitHubIssue struct {
owner string
repo string
issueNumber int
Expand All @@ -20,7 +20,7 @@ type GitHubIssues struct {
logger *log.Logger
}

func (gh *GitHubIssues) GetIssue() (*github.Issue, error) {
func (gh *GitHubIssue) GetIssue() (*github.Issue, error) {
// Only the first call retrieves information from GitHub, all other calls use cache
if gh.cache == nil {
issue, _, err := gh.client.Issues.Get(gh.ctx, gh.owner, gh.repo, gh.issueNumber)
Expand All @@ -32,23 +32,23 @@ func (gh *GitHubIssues) GetIssue() (*github.Issue, error) {
return gh.cache, nil
}

func (gh *GitHubIssues) GetBody() (*string, error) {
func (gh *GitHubIssue) GetBody() (*string, error) {
issue, err := gh.GetIssue()
if err != nil {
return nil, err
}
return issue.Body, err
}

func (gh *GitHubIssues) GetTitle() (*string, error) {
func (gh *GitHubIssue) GetTitle() (*string, error) {
issue, err := gh.GetIssue()
if err != nil {
return nil, err
}
return issue.Title, err
}

func (gh *GitHubIssues) GetComments() ([]*github.IssueComment, error) {
func (gh *GitHubIssue) GetComments() ([]*github.IssueComment, error) {
// Options
opt := &github.IssueListCommentsOptions{Direction: "asc", Sort: "created"}
opt.Page = 1
Expand All @@ -58,7 +58,7 @@ func (gh *GitHubIssues) GetComments() ([]*github.IssueComment, error) {
return comments, err
}

func (gh *GitHubIssues) PostComment(commentBody string) error {
func (gh *GitHubIssue) PostComment(commentBody string) error {
comment := &github.IssueComment{Body: github.String(commentBody)}
_, _, err := gh.client.Issues.CreateComment(gh.ctx, gh.owner, gh.repo, gh.issueNumber, comment)
if err != nil {
Expand All @@ -68,7 +68,7 @@ func (gh *GitHubIssues) PostComment(commentBody string) error {
return err
}

func NewIssue(owner string, repo string, issueNumber int, token string) *GitHubIssues {
func NewIssue(owner string, repo string, issueNumber int, token string) *GitHubIssue {
// Create GitHub client with OAuth2 token
ctx := context.Background()
ts := oauth2.StaticTokenSource(
Expand All @@ -83,7 +83,7 @@ func NewIssue(owner string, repo string, issueNumber int, token string) *GitHubI
log.Ldate|log.Ltime|log.Llongfile|log.Lmsgprefix,
)

// Create a new GitHubIssues instance
issue := &GitHubIssues{owner: owner, repo: repo, issueNumber: issueNumber, token: token, client: client, ctx: ctx, logger: logger}
// Create a new GitHubIssue instance
issue := &GitHubIssue{owner: owner, repo: repo, issueNumber: issueNumber, token: token, client: client, ctx: ctx, logger: logger}
return issue
}

0 comments on commit bcaac66

Please sign in to comment.