Skip to content

Commit 007fbfb

Browse files
kechigonkechigon
authored
Merge pull request #6 from 3-shake/feature/insert_image_to_prompt
Feature/insert image to prompt
2 parents 445e51a + 7e6c770 commit 007fbfb

File tree

8 files changed

+376
-261
lines changed

8 files changed

+376
-261
lines changed

‎.alert-menta.user.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ system:
55
ai:
66
provider: "openai" # "openai" or "vertexai"
77
openai:
8-
model: "gpt-3.5-turbo" # Check the list of available models by `curl https://api.openai.com/v1/models -H "Authorization: Bearer $OPENAI_API_KEY"`
8+
model: "gpt-4o-mini-2024-07-18" # Check the list of available models by `curl https://api.openai.com/v1/models -H "Authorization: Bearer $OPENAI_API_KEY"`
99

1010
vertexai:
1111
project: "<YOUR_PROJECT_ID>"

‎cmd/main.go‎

Lines changed: 151 additions & 135 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,11 @@ package main
22

33
import (
44
"flag"
5+
"fmt"
56
"log"
67
"os"
8+
"regexp"
79
"strings"
8-
"fmt"
910

1011
"github.com/3-shake/alert-menta/internal/ai"
1112
"github.com/3-shake/alert-menta/internal/github"
@@ -14,154 +15,169 @@ import (
1415

1516
// Struct to hold the command-line arguments
1617
type Config struct {
17-
repo string
18-
owner string
19-
issueNumber int
20-
intent string
21-
command string
22-
configFile string
23-
ghToken string
24-
oaiKey string
18+
repo string
19+
owner string
20+
issueNumber int
21+
intent string
22+
command string
23+
configFile string
24+
ghToken string
25+
oaiKey string
2526
}
2627

2728
func main() {
28-
cfg := &Config{}
29-
flag.StringVar(&cfg.repo, "repo", "", "Repository name")
30-
flag.StringVar(&cfg.owner, "owner", "", "Repository owner")
31-
flag.IntVar(&cfg.issueNumber, "issue", 0, "Issue number")
32-
flag.StringVar(&cfg.intent, "intent", "", "Question or intent for the 'ask' command")
33-
flag.StringVar(&cfg.command, "command", "", "Commands to be executed by AI. Commands defined in the configuration file are available.")
34-
flag.StringVar(&cfg.configFile, "config", "", "Configuration file")
35-
flag.StringVar(&cfg.ghToken, "github-token", "", "GitHub token")
36-
flag.StringVar(&cfg.oaiKey, "api-key", "", "OpenAI api key")
37-
flag.Parse()
38-
39-
if cfg.repo == "" || cfg.owner == "" || cfg.issueNumber == 0 || cfg.ghToken == "" || cfg.command == "" || cfg.configFile == "" {
40-
flag.PrintDefaults()
41-
os.Exit(1)
42-
}
43-
44-
logger := log.New(
45-
os.Stdout, "[alert-menta main] ",
46-
log.Ldate|log.Ltime|log.Llongfile|log.Lmsgprefix,
47-
)
48-
49-
loadedcfg, err := utils.NewConfig(cfg.configFile)
50-
if err != nil {
51-
logger.Fatalf("Error loading config: %v", err)
52-
}
53-
54-
err = validateCommand(cfg.command, loadedcfg)
55-
if err != nil {
56-
logger.Fatalf("Error validating command: %v", err)
57-
}
58-
59-
issue := github.NewIssue(cfg.owner, cfg.repo, cfg.issueNumber, cfg.ghToken)
60-
61-
userPrompt, err := constructUserPrompt(issue, loadedcfg, logger)
62-
if err != nil {
63-
logger.Fatalf("Erro constructing userPrompt: %v", err)
64-
}
65-
66-
prompt, err := constructPrompt(cfg.command, cfg.intent, userPrompt, loadedcfg, logger)
67-
if err != nil {
68-
logger.Fatalf("Error constructing prompt: %v", err)
69-
}
70-
71-
aic, err := getAIClient(cfg.oaiKey, loadedcfg, logger)
72-
if err != nil {
73-
logger.Fatalf("Error geting AI client: %v", err)
74-
}
75-
76-
comment, err := aic.GetResponse(prompt)
77-
if err != nil {
78-
logger.Fatalf("Error getting Response: %v", err)
79-
}
80-
logger.Println("Response:", comment)
81-
82-
if err := issue.PostComment(comment); err != nil {
83-
logger.Fatalf("Error creating comment: %v", err)
84-
}
29+
cfg := &Config{}
30+
flag.StringVar(&cfg.repo, "repo", "", "Repository name")
31+
flag.StringVar(&cfg.owner, "owner", "", "Repository owner")
32+
flag.IntVar(&cfg.issueNumber, "issue", 0, "Issue number")
33+
flag.StringVar(&cfg.intent, "intent", "", "Question or intent for the 'ask' command")
34+
flag.StringVar(&cfg.command, "command", "", "Commands to be executed by AI. Commands defined in the configuration file are available.")
35+
flag.StringVar(&cfg.configFile, "config", "", "Configuration file")
36+
flag.StringVar(&cfg.ghToken, "github-token", "", "GitHub token")
37+
flag.StringVar(&cfg.oaiKey, "api-key", "", "OpenAI api key")
38+
flag.Parse()
39+
40+
if cfg.repo == "" || cfg.owner == "" || cfg.issueNumber == 0 || cfg.ghToken == "" || cfg.command == "" || cfg.configFile == "" {
41+
flag.PrintDefaults()
42+
os.Exit(1)
43+
}
44+
45+
logger := log.New(
46+
os.Stdout, "[alert-menta main] ",
47+
log.Ldate|log.Ltime|log.Llongfile|log.Lmsgprefix,
48+
)
49+
50+
loadedcfg, err := utils.NewConfig(cfg.configFile)
51+
if err != nil {
52+
logger.Fatalf("Error loading config: %v", err)
53+
}
54+
55+
err = validateCommand(cfg.command, loadedcfg)
56+
if err != nil {
57+
logger.Fatalf("Error validating command: %v", err)
58+
}
59+
60+
issue := github.NewIssue(cfg.owner, cfg.repo, cfg.issueNumber, cfg.ghToken)
61+
62+
userPrompt, imgs, err := constructUserPrompt(cfg.ghToken, issue, loadedcfg, logger)
63+
if err != nil {
64+
logger.Fatalf("Erro constructing userPrompt: %v", err)
65+
}
66+
67+
prompt, err := constructPrompt(cfg.command, cfg.intent, userPrompt, imgs, loadedcfg, logger)
68+
if err != nil {
69+
logger.Fatalf("Error constructing prompt: %v", err)
70+
}
71+
72+
aic, err := getAIClient(cfg.oaiKey, loadedcfg, logger)
73+
if err != nil {
74+
logger.Fatalf("Error geting AI client: %v", err)
75+
}
76+
77+
comment, err := aic.GetResponse(prompt)
78+
if err != nil {
79+
logger.Fatalf("Error getting Response: %v", err)
80+
}
81+
logger.Println("Response:", comment)
82+
83+
if err := issue.PostComment(comment); err != nil {
84+
logger.Fatalf("Error creating comment: %v", err)
85+
}
8586
}
8687

8788
// Validate the provided command
8889
func validateCommand(command string, cfg *utils.Config) error {
89-
if _, ok := cfg.Ai.Commands[command]; !ok {
90-
allowedCommands := make([]string, 0, len(cfg.Ai.Commands))
91-
for cmd := range cfg.Ai.Commands {
92-
allowedCommands = append(allowedCommands, cmd)
93-
}
94-
return fmt.Errorf("Invalid command: %s. Allowed commands are %s", command, strings.Join(allowedCommands, ", "))
95-
}
96-
return nil
90+
if _, ok := cfg.Ai.Commands[command]; !ok {
91+
allowedCommands := make([]string, 0, len(cfg.Ai.Commands))
92+
for cmd := range cfg.Ai.Commands {
93+
allowedCommands = append(allowedCommands, cmd)
94+
}
95+
return fmt.Errorf("Invalid command: %s. Allowed commands are %s", command, strings.Join(allowedCommands, ", "))
96+
}
97+
return nil
9798
}
9899

99100
// Construct user prompt from issue
100-
func constructUserPrompt(issue *github.GitHubIssue, cfg *utils.Config, logger *log.Logger) (string, error) {
101-
title, err := issue.GetTitle()
102-
if err != nil {
103-
return "", fmt.Errorf("Error getting Title: %w", err)
104-
}
105-
106-
body, err := issue.GetBody()
107-
if err != nil {
108-
return "", fmt.Errorf("Error getting Body: %w", err)
109-
}
110-
111-
var userPrompt strings.Builder
112-
userPrompt.WriteString("Title:" + *title + "\n")
113-
userPrompt.WriteString("Body:" + *body + "\n")
114-
115-
comments, err := issue.GetComments()
116-
if err != nil {
117-
return "", fmt.Errorf("Error getting comments: %w", err)
118-
}
119-
for _, v := range comments {
120-
if *v.User.Login == "github-actions[bot]" {
121-
continue
122-
}
123-
if cfg.System.Debug.Log_level == "debug" {
124-
logger.Printf("%s: %s", *v.User.Login, *v.Body)
125-
}
126-
userPrompt.WriteString(*v.User.Login + ":" + *v.Body + "\n")
127-
}
128-
return userPrompt.String(), nil
101+
func constructUserPrompt(ghToken string, issue *github.GitHubIssue, cfg *utils.Config, logger *log.Logger) (string, []ai.Image, error) {
102+
title, err := issue.GetTitle()
103+
if err != nil {
104+
return "", nil, fmt.Errorf("Error getting Title: %w", err)
105+
}
106+
107+
body, err := issue.GetBody()
108+
if err != nil {
109+
return "", nil, fmt.Errorf("Error getting Body: %w", err)
110+
}
111+
112+
var userPrompt strings.Builder
113+
userPrompt.WriteString("Title:" + *title + "\n")
114+
userPrompt.WriteString("Body:" + *body + "\n")
115+
116+
comments, err := issue.GetComments()
117+
if err != nil {
118+
return "", nil, fmt.Errorf("Error getting comments: %w", err)
119+
}
120+
121+
var images []ai.Image
122+
imageRegex := regexp.MustCompile(`!\[(.*?)\]\((.*?)\)`)
123+
124+
for _, v := range comments {
125+
if *v.User.Login == "github-actions[bot]" {
126+
continue
127+
}
128+
if cfg.System.Debug.Log_level == "debug" {
129+
logger.Printf("%s: %s", *v.User.Login, *v.Body)
130+
}
131+
userPrompt.WriteString(*v.User.Login + ":" + *v.Body + "\n")
132+
133+
matches := imageRegex.FindAllStringSubmatch(*v.Body, -1)
134+
for _, match := range matches {
135+
logger.Println("Image URL:", match[2]) // Log the URL of the image
136+
imgData, ext, err := utils.DownloadImage(match[2], ghToken)
137+
if err != nil {
138+
return "", nil, fmt.Errorf("Error downloading image: %w", err)
139+
}
140+
141+
images = append(images, ai.Image{Data: imgData, Extension: ext})
142+
}
143+
}
144+
return userPrompt.String(), images, nil
129145
}
130146

131147
// Construct AI prompt
132-
func constructPrompt(command, intent, userPrompt string, cfg *utils.Config, logger *log.Logger) (*ai.Prompt, error){
133-
var systemPrompt string
134-
if command == "ask" {
135-
if intent == "" {
136-
return nil, fmt.Errorf("Error: intent is required for 'ask' command")
137-
}
138-
systemPrompt = cfg.Ai.Commands[command].System_prompt + intent + "\n"
139-
} else {
140-
systemPrompt = cfg.Ai.Commands[command].System_prompt
141-
}
142-
logger.Println("\x1b[34mPrompt: |\n", systemPrompt, userPrompt, "\x1b[0m")
143-
return &ai.Prompt{UserPrompt: userPrompt, SystemPrompt: systemPrompt}, nil
148+
func constructPrompt(command, intent, userPrompt string, imgs []ai.Image, cfg *utils.Config, logger *log.Logger) (*ai.Prompt, error) {
149+
var systemPrompt string
150+
if command == "ask" {
151+
if intent == "" {
152+
return nil, fmt.Errorf("Error: intent is required for 'ask' command")
153+
}
154+
systemPrompt = cfg.Ai.Commands[command].System_prompt + intent + "\n"
155+
} else {
156+
systemPrompt = cfg.Ai.Commands[command].System_prompt
157+
}
158+
logger.Println("\x1b[34mPrompt: |\n", systemPrompt, userPrompt, "\x1b[0m")
159+
return &ai.Prompt{UserPrompt: userPrompt, SystemPrompt: systemPrompt, Images: imgs}, nil
144160
}
145161

146162
// Initialize AI client
147163
func getAIClient(oaiKey string, cfg *utils.Config, logger *log.Logger) (ai.Ai, error) {
148-
switch cfg.Ai.Provider {
149-
case "openai":
150-
if oaiKey == "" {
151-
return nil, fmt.Errorf("Error: Please provide your Open AI API key")
152-
}
153-
logger.Println("Using OpenAI API")
154-
logger.Println("OpenAI model:", cfg.Ai.OpenAI.Model)
155-
return ai.NewOpenAIClient(oaiKey, cfg.Ai.OpenAI.Model), nil
156-
case "vertexai":
157-
logger.Println("Using VertexAI API")
158-
logger.Println("VertexAI model:", cfg.Ai.VertexAI.Model)
159-
aic, err := ai.NewVertexAIClient(cfg.Ai.VertexAI.Project, cfg.Ai.VertexAI.Region, cfg.Ai.VertexAI.Model)
160-
if err != nil {
161-
return nil, fmt.Errorf("Error: new Vertex AI client: %w", err)
162-
}
163-
return aic, nil
164-
default:
165-
return nil, fmt.Errorf("Error: Invalid provider")
166-
}
167-
}
164+
switch cfg.Ai.Provider {
165+
case "openai":
166+
if oaiKey == "" {
167+
return nil, fmt.Errorf("Error: Please provide your Open AI API key")
168+
}
169+
logger.Println("Using OpenAI API")
170+
logger.Println("OpenAI model:", cfg.Ai.OpenAI.Model)
171+
return ai.NewOpenAIClient(oaiKey, cfg.Ai.OpenAI.Model), nil
172+
case "vertexai":
173+
logger.Println("Using VertexAI API")
174+
logger.Println("VertexAI model:", cfg.Ai.VertexAI.Model)
175+
aic, err := ai.NewVertexAIClient(cfg.Ai.VertexAI.Project, cfg.Ai.VertexAI.Region, cfg.Ai.VertexAI.Model)
176+
if err != nil {
177+
return nil, fmt.Errorf("Error: new Vertex AI client: %w", err)
178+
}
179+
return aic, nil
180+
default:
181+
return nil, fmt.Errorf("Error: Invalid provider")
182+
}
183+
}

0 commit comments

Comments
 (0)