feat(autofix): update gemini sdk and add anthropic claude

* upgrade gemini sdk to google.golang.org/genai v1.25.0
* support newer gemini models
* add anthropic claude

Author: MatteoCalabro-TomTom

autofix/ai.go

@@ -4,97 +4,53 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"strings"
 	"time"
 
-	"github.com/google/generative-ai-go/genai"
-	"google.golang.org/api/option"
-
 	"github.com/securego/gosec/v2/issue"
 )
 
 const (
-	GeminiModel = "gemini-1.5-flash"
-	AIPrompt    = `Provide a brief explanation and a solution to fix this security issue
+	AIProviderFlagHelp = `AI API provider to generate auto fixes to issues. Valid options are:
+- gemini-2.5-pro, gemini-2.5-flash, gemini-2.5-flash-lite, gemini-2.0-flash, gemini-2.0-flash-lite (gemini, default);
+- claude-sonnet-4-0 (claude, default), claude-opus-4-0, claude-opus-4-1, claude-sonnet-3-7`
+
+	AIPrompt = `Provide a brief explanation and a solution to fix this security issue
   in Go programming language: %q.
   Answer in markdown format and keep the response limited to 200 words.`
-	GeminiProvider = "gemini"
 
 	timeout = 30 * time.Second
 )
 
-// GenAIClient defines the interface for the GenAI client.
 type GenAIClient interface {
-	// Close clean up and close the client.
-	Close() error
-	// GenerativeModel build the generative mode.
-	GenerativeModel(name string) GenAIGenerativeModel
+	GenerateSolution(ctx context.Context, prompt string) (string, error)
 }
 
-// GenAIGenerativeModel defines the interface for the Generative Model.
-type GenAIGenerativeModel interface {
-	// GenerateContent generates an response for given prompt.
-	GenerateContent(ctx context.Context, prompt string) (string, error)
-}
-
-// genAIClientWrapper wraps the genai.Client to implement GenAIClient.
-type genAIClientWrapper struct {
-	client *genai.Client
-}
-
-// Close closes the gen AI client.
-func (w *genAIClientWrapper) Close() error {
-	return w.client.Close()
-}
-
-// GenerativeModel builds the generative Model.
-func (w *genAIClientWrapper) GenerativeModel(name string) GenAIGenerativeModel {
-	return &genAIGenerativeModelWrapper{model: w.client.GenerativeModel(name)}
-}
-
-// genAIGenerativeModelWrapper wraps the genai.GenerativeModel to implement GenAIGenerativeModel
-type genAIGenerativeModelWrapper struct {
-	// model is the underlying generative model
-	model *genai.GenerativeModel
-}
-
-// GenerateContent generates a response for the given prompt using gemini API.
-func (w *genAIGenerativeModelWrapper) GenerateContent(ctx context.Context, prompt string) (string, error) {
-	resp, err := w.model.GenerateContent(ctx, genai.Text(prompt))
-	if err != nil {
-		return "", fmt.Errorf("generating autofix: %w", err)
-	}
-	if len(resp.Candidates) == 0 {
-		return "", errors.New("no autofix returned by gemini")
-	}
-
-	if len(resp.Candidates[0].Content.Parts) == 0 {
-		return "", errors.New("nothing found in the first autofix returned by gemini")
-	}
-
-	// Return the first candidate
-	return fmt.Sprintf("%+v", resp.Candidates[0].Content.Parts[0]), nil
-}
+// GenerateSolution generates a solution for the given issues using the specified AI provider
+func GenerateSolution(model, aiAPIKey string, issues []*issue.Issue) (err error) {
+	var client GenAIClient
 
-// NewGenAIClient creates a new gemini API client.
-func NewGenAIClient(ctx context.Context, aiAPIKey, endpoint string) (GenAIClient, error) {
-	clientOptions := []option.ClientOption{option.WithAPIKey(aiAPIKey)}
-	if endpoint != "" {
-		clientOptions = append(clientOptions, option.WithEndpoint(endpoint))
+	switch {
+	case strings.HasPrefix(model, "claude"):
+		client, err = NewClaudeClient(model, aiAPIKey)
+	case strings.HasPrefix(model, "gemini"):
+		client, err = NewGeminiClient(model, aiAPIKey)
 	}
 
-	client, err := genai.NewClient(ctx, clientOptions...)
-	if err != nil {
-		return nil, fmt.Errorf("calling gemini API: %w", err)
+	switch {
+	case err != nil:
+		return fmt.Errorf("initializing AI client: %w", err)
+	case client == nil:
+		return fmt.Errorf("unsupported AI backend: %s", model)
 	}
 
-	return &genAIClientWrapper{client: client}, nil
+	return generateSolution(client, issues)
 }
 
-func generateSolutionByGemini(client GenAIClient, issues []*issue.Issue) error {
+func generateSolution(client GenAIClient, issues []*issue.Issue) error {
 	ctx, cancel := context.WithTimeout(context.Background(), timeout)
 	defer cancel()
 
-	model := client.GenerativeModel(GeminiModel)
 	cachedAutofix := make(map[string]string)
 	for _, issue := range issues {
 		if val, ok := cachedAutofix[issue.What]; ok {
@@ -103,7 +59,7 @@ func generateSolutionByGemini(client GenAIClient, issues []*issue.Issue) error {
 		}
 
 		prompt := fmt.Sprintf(AIPrompt, issue.What)
-		resp, err := model.GenerateContent(ctx, prompt)
+		resp, err := client.GenerateSolution(ctx, prompt)
 		if err != nil {
 			return fmt.Errorf("generating autofix with gemini: %w", err)
 		}
@@ -117,26 +73,3 @@ func generateSolutionByGemini(client GenAIClient, issues []*issue.Issue) error {
 	}
 	return nil
 }
-
-// GenerateSolution generates a solution for the given issues using the specified AI provider
-func GenerateSolution(aiAPIProvider, aiAPIKey, endpoint string, issues []*issue.Issue) error {
-	ctx, cancel := context.WithTimeout(context.Background(), timeout)
-	defer cancel()
-
-	var client GenAIClient
-
-	switch aiAPIProvider {
-	case GeminiProvider:
-		var err error
-		client, err = NewGenAIClient(ctx, aiAPIKey, endpoint)
-		if err != nil {
-			return fmt.Errorf("generating autofix: %w", err)
-		}
-	default:
-		return errors.New("ai provider not supported")
-	}
-
-	defer client.Close()
-
-	return generateSolutionByGemini(client, issues)
-}

autofix/ai_test.go

@@ -17,22 +17,7 @@ type MockGenAIClient struct {
 	mock.Mock
 }
 
-func (m *MockGenAIClient) Close() error {
-	args := m.Called()
-	return args.Error(0)
-}
-
-func (m *MockGenAIClient) GenerativeModel(name string) GenAIGenerativeModel {
-	args := m.Called(name)
-	return args.Get(0).(GenAIGenerativeModel)
-}
-
-// MockGenAIGenerativeModel is a mock of the GenAIGenerativeModel interface
-type MockGenAIGenerativeModel struct {
-	mock.Mock
-}
-
-func (m *MockGenAIGenerativeModel) GenerateContent(ctx context.Context, prompt string) (string, error) {
+func (m *MockGenAIClient) GenerateSolution(ctx context.Context, prompt string) (string, error) {
 	args := m.Called(ctx, prompt)
 	return args.String(0), args.Error(1)
 }
@@ -44,17 +29,15 @@ func TestGenerateSolutionByGemini_Success(t *testing.T) {
 	}
 
 	mockClient := new(MockGenAIClient)
-	mockModel := new(MockGenAIGenerativeModel)
-	mockClient.On("GenerativeModel", GeminiModel).Return(mockModel).Once()
-	mockModel.On("GenerateContent", mock.Anything, mock.Anything).Return("Autofix for issue 1", nil).Once()
+	mockClient.On("GenerateSolution", mock.Anything, mock.Anything).Return("Autofix for issue 1", nil).Once()
 
 	// Act
-	err := generateSolutionByGemini(mockClient, issues)
+	err := generateSolution(mockClient, issues)
 
 	// Assert
 	require.NoError(t, err)
 	assert.Equal(t, []*issue.Issue{{What: "Example issue 1", Autofix: "Autofix for issue 1"}}, issues)
-	mock.AssertExpectationsForObjects(t, mockClient, mockModel)
+	mock.AssertExpectationsForObjects(t, mockClient)
 }
 
 func TestGenerateSolutionByGemini_NoCandidates(t *testing.T) {
@@ -64,16 +47,14 @@ func TestGenerateSolutionByGemini_NoCandidates(t *testing.T) {
 	}
 
 	mockClient := new(MockGenAIClient)
-	mockModel := new(MockGenAIGenerativeModel)
-	mockClient.On("GenerativeModel", GeminiModel).Return(mockModel).Once()
-	mockModel.On("GenerateContent", mock.Anything, mock.Anything).Return("", nil).Once()
+	mockClient.On("GenerateSolution", mock.Anything, mock.Anything).Return("", nil).Once()
 
 	// Act
-	err := generateSolutionByGemini(mockClient, issues)
+	err := generateSolution(mockClient, issues)
 
 	// Assert
 	require.EqualError(t, err, "no autofix returned by gemini")
-	mock.AssertExpectationsForObjects(t, mockClient, mockModel)
+	mock.AssertExpectationsForObjects(t, mockClient)
 }
 
 func TestGenerateSolutionByGemini_APIError(t *testing.T) {
@@ -83,16 +64,14 @@ func TestGenerateSolutionByGemini_APIError(t *testing.T) {
 	}
 
 	mockClient := new(MockGenAIClient)
-	mockModel := new(MockGenAIGenerativeModel)
-	mockClient.On("GenerativeModel", GeminiModel).Return(mockModel).Once()
-	mockModel.On("GenerateContent", mock.Anything, mock.Anything).Return("", errors.New("API error")).Once()
+	mockClient.On("GenerateSolution", mock.Anything, mock.Anything).Return("", errors.New("API error")).Once()
 
 	// Act
-	err := generateSolutionByGemini(mockClient, issues)
+	err := generateSolution(mockClient, issues)
 
 	// Assert
 	require.EqualError(t, err, "generating autofix with gemini: API error")
-	mock.AssertExpectationsForObjects(t, mockClient, mockModel)
+	mock.AssertExpectationsForObjects(t, mockClient)
 }
 
 func TestGenerateSolution_UnsupportedProvider(t *testing.T) {
@@ -102,8 +81,8 @@ func TestGenerateSolution_UnsupportedProvider(t *testing.T) {
 	}
 
 	// Act
-	err := GenerateSolution("unsupported-provider", "test-api-key", "", issues)
+	err := GenerateSolution("unsupported-provider", "test-api-key", issues)
 
 	// Assert
-	require.EqualError(t, err, "ai provider not supported")
+	require.EqualError(t, err, "unsupported AI backend: unsupported-provider")
 }

autofix/claude.go

@@ -0,0 +1,74 @@
+package autofix
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/anthropics/anthropic-sdk-go"
+	"github.com/anthropics/anthropic-sdk-go/option"
+)
+
+const (
+	ModelClaudeOpus4_0   = anthropic.ModelClaudeOpus4_0
+	ModelClaudeOpus4_1   = anthropic.ModelClaudeOpus4_1_20250805
+	ModelClaudeSonnet4_0 = anthropic.ModelClaudeSonnet4_0
+)
+
+var _ GenAIClient = (*claudeWrapper)(nil)
+
+type claudeWrapper struct {
+	client anthropic.Client
+	model  anthropic.Model
+}
+
+func NewClaudeClient(model, apiKey string) (GenAIClient, error) {
+	var options []option.RequestOption
+
+	if apiKey != "" {
+		options = append(options, option.WithAPIKey(apiKey))
+	}
+
+	anthropicModel := parseAnthropicModel(model)
+
+	return &claudeWrapper{
+		client: anthropic.NewClient(options...),
+		model:  anthropicModel,
+	}, nil
+}
+
+func (c *claudeWrapper) GenerateSolution(ctx context.Context, prompt string) (string, error) {
+	resp, err := c.client.Messages.New(ctx, anthropic.MessageNewParams{
+		Model:     anthropic.Model(c.model),
+		MaxTokens: 1024,
+		Messages: []anthropic.MessageParam{
+			anthropic.NewUserMessage(anthropic.NewTextBlock(prompt)),
+		},
+	})
+	if err != nil {
+		return "", fmt.Errorf("generating autofix: %w", err)
+	}
+
+	if resp == nil || len(resp.Content) == 0 {
+		return "", errors.New("no autofix returned by claude")
+	}
+
+	if len(resp.Content[0].Text) == 0 {
+		return "", errors.New("nothing found in the first autofix returned by claude")
+	}
+
+	return resp.Content[0].Text, nil
+}
+
+func parseAnthropicModel(model string) anthropic.Model {
+	switch model {
+	case "claude-sonnet-3-7":
+		return anthropic.ModelClaude3_7SonnetLatest
+	case "claude-opus", "claude-opus-4-0":
+		return anthropic.ModelClaudeOpus4_0
+	case "claude-opus-4-1":
+		return anthropic.ModelClaudeOpus4_1_20250805
+	}
+
+	return anthropic.ModelClaudeSonnet4_0
+}