Skip to content

Commit

Permalink
update main.go github.go
Browse files Browse the repository at this point in the history
kechigon committed Nov 5, 2024

Unverified

This commit is not signed, but one or more authors requires that any commit attributed to them is signed.
1 parent 9938d3f commit bcaac66
Showing 2 changed files with 163 additions and 116 deletions.
261 changes: 154 additions & 107 deletions cmd/main.go
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -9,7 +9,7 @@ import (
"golang.org/x/oauth2"
)

type GitHubIssues struct {
type GitHubIssue struct {
owner string
repo string
issueNumber int
@@ -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)
@@ -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
@@ -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 {
@@ -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(
@@ -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.