config file allows AI providers to be switched

cmd/main.go

@@ -63,11 +63,20 @@ func main() {
 
 	// Set system prompt
 	system_prompt := cfg.Ai.Commands[*command].System_prompt
+	logger.Println("Prompt:", system_prompt, user_prompt)
 
 	// Get response from OpenAI
-	logger.Println("Prompt:", system_prompt, user_prompt)
-	ai := ai.NewOpenAIClient(*oai_key, cfg.Ai.Model)
-	comment, _ := ai.GetResponse(system_prompt + user_prompt)
+	var aic ai.Ai
+	if cfg.Ai.Provider == "openai" {
+		aic = ai.NewOpenAIClient(*oai_key, cfg.Ai.OpenAI.Model)
+		logger.Println("Using OpenAI API")
+		logger.Println("OpenAI model:", cfg.Ai.OpenAI.Model)
+	} else {
+		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)
+	}
+	comment, _ := aic.GetResponse(system_prompt + user_prompt)
 	logger.Println("Response:", comment)
 
 	// Post a comment on the Issue

internal/ai/ai.go

@@ -0,0 +1,5 @@
+package ai
+
+type Ai interface {
+	GetResponse(prompt string) (string, error)
+}

internal/config/config.yaml

@@ -8,7 +8,14 @@ github:
   repo: "alert-menta"
 
 ai:
-  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"`
+  provider: "vertexai" # "openai" or "vertexai"
+  openai:
+    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"`
+
+  vertexai:
+    project: "<YOUR_PROJECT_ID>"
+    location: "us-central1"
+    model: "gemini-1.5-flash-001"
 
   commands:
     - describe:

internal/utils/utils.go

@@ -29,8 +29,10 @@ type SystemDebug struct {
 }
 
 type Ai struct {
-	Model    string             `yaml:"model"`
 	Commands map[string]Command `yaml:"commands"`
+	Provider string             `yaml:"provider"`
+	OpenAI   OpenAI             `yaml:"openai"`
+	VertexAI VertexAI           `yaml:"vertexai"`
 }
 
 type Command struct {
@@ -44,6 +46,16 @@ type Github struct {
 	Repo  string `yaml:"repo"`
 }
 
+type OpenAI struct {
+	Model string `yaml:"model"`
+}
+
+type VertexAI struct {
+	Model   string `yaml:"model"`
+	Project string `yaml:"project"`
+	Region  string `yaml:"region"`
+}
+
 func NewConfig(filename string) (*Config, error) {
 	// Initialize a logger
 	logger := log.New(