Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 41 additions & 0 deletions internal/ai/gemini/helpers.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
package gemini

import (
"strings"
)

// CleanLabels cleans and validates labels, keeping only the allowed ones.
// It accepts a list of labels to clean and a list of available labels from the repository.
// If availableLabels is empty, it falls back to a default list of common labels.
func CleanLabels(labels []string, availableLabels []string) []string {
allowedLabels := make(map[string]bool)

if len(availableLabels) > 0 {
for _, l := range availableLabels {
allowedLabels[strings.ToLower(l)] = true
}
} else {
// Fallback to default list if no repo labels provided
defaultLabels := []string{
"feature", "fix", "refactor", "docs", "test", "infra",
"enhancement", "bug", "good first issue", "help wanted",
"chore", "performance", "security", "tech-debt", "breaking-change",
}
for _, l := range defaultLabels {
allowedLabels[l] = true
}
}

cleaned := make([]string, 0)
seen := make(map[string]bool)

for _, label := range labels {
trimmed := strings.TrimSpace(strings.ToLower(label))
if trimmed != "" && allowedLabels[trimmed] && !seen[trimmed] {
cleaned = append(cleaned, trimmed)
seen[trimmed] = true
}
}

return cleaned
}
95 changes: 33 additions & 62 deletions internal/ai/gemini/issue_content_generator.go
Original file line number Diff line number Diff line change
Expand Up @@ -72,8 +72,33 @@ func NewGeminiIssueContentGenerator(ctx context.Context, cfg *config.Config, onC
return service, nil
}

func getIssueSchema() *genai.Schema {
return &genai.Schema{
Type: genai.TypeObject,
Required: []string{"title", "description", "labels"},
Properties: map[string]*genai.Schema{
"title": {
Type: genai.TypeString,
Description: "The title of the issue",
},
"description": {
Type: genai.TypeString,
Description: "The body of the issue in markdown format",
},
"labels": {
Type: genai.TypeArray,
Items: &genai.Schema{
Type: genai.TypeString,
},
Description: "List of labels (e.g. bug, feature, refactor, good first issue)",
},
},
}
}

func (s *GeminiIssueContentGenerator) defaultGenerate(ctx context.Context, mName string, p string) (interface{}, *models.TokenUsage, error) {
genConfig := GetGenerateConfig(mName, "", nil)
schema := getIssueSchema()
genConfig := GetGenerateConfig(mName, "application/json", schema)
log := logger.FromContext(ctx)

resp, err := s.Client.Models.GenerateContent(ctx, mName, genai.Text(p), genConfig)
Expand Down Expand Up @@ -167,6 +192,7 @@ func (s *GeminiIssueContentGenerator) GenerateIssueContent(ctx context.Context,
return nil, domainErrors.NewAppError(domainErrors.TypeAI, "error parsing AI response", err)
}

result.Labels = CleanLabels(result.Labels, request.AvailableLabels)
result.Usage = usage

log.Info("issue content generated successfully via gemini",
Expand All @@ -179,7 +205,7 @@ func (s *GeminiIssueContentGenerator) GenerateIssueContent(ctx context.Context,
// buildIssuePrompt builds the prompt to generate issue content.
func (s *GeminiIssueContentGenerator) buildIssuePrompt(request models.IssueGenerationRequest) string {
if request.Description != "" && request.Diff == "" && request.Hint == "" &&
request.Template == nil && len(request.ChangedFiles) == 0 {
request.Template == nil && len(request.ChangedFiles) == 0 && len(request.AvailableLabels) == 0 {
return request.Description
}

Expand Down Expand Up @@ -227,37 +253,8 @@ func (s *GeminiIssueContentGenerator) buildIssuePrompt(request models.IssueGener
return ""
}

if request.Template != nil {
rendered += `

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

🚨 FINAL REMINDER - CRITICAL OUTPUT REQUIREMENT 🚨

YOU MUST OUTPUT **ONLY** VALID JSON.

The template structure above should be used to FILL the "description" field with markdown content.

BUT your actual response MUST be a JSON object like this:
{
"title": "string here",
"description": "markdown content following the template structure",
"labels": ["array", "of", "strings"]
}

❌ DO NOT output prose like "Here is a high-quality GitHub issue..."
❌ DO NOT output markdown text directly
❌ DO NOT output explanations

✅ ONLY output the JSON object
✅ Use the template to structure the markdown in the "description" field
✅ Return valid parseable JSON

BEGIN YOUR JSON OUTPUT NOW:`

logger.Debug(context.Background(), "full prompt with template and final JSON reminder",
"prompt_length", len(rendered),
"prompt", rendered)
if len(request.AvailableLabels) > 0 {
rendered += fmt.Sprintf("\n\nAvailable Labels (Select ONLY from this list):\n%s", strings.Join(request.AvailableLabels, ", "))
}

return rendered
Expand All @@ -270,6 +267,8 @@ func (s *GeminiIssueContentGenerator) parseIssueResponse(content string) (*model
return nil, domainErrors.NewAppError(domainErrors.TypeAI, "empty response from AI", nil)
}

content = strings.TrimSpace(content)

if len(content) > 0 {
preview := content
if len(content) > 200 {
Expand Down Expand Up @@ -307,40 +306,12 @@ func (s *GeminiIssueContentGenerator) parseIssueResponse(content string) (*model
result := &models.IssueGenerationResult{
Title: strings.TrimSpace(jsonResult.Title),
Description: strings.TrimSpace(jsonResult.Description),
Labels: s.cleanLabels(jsonResult.Labels),
Labels: jsonResult.Labels,
}

if result.Title == "" {
result.Title = "Generated Issue"
}
if result.Description == "" {
result.Description = content
}

return result, nil
}

// cleanLabels cleans and validates labels, keeping only the allowed ones.
func (s *GeminiIssueContentGenerator) cleanLabels(labels []string) []string {
allowedLabels := map[string]bool{
"feature": true,
"fix": true,
"refactor": true,
"docs": true,
"test": true,
"infra": true,
}

cleaned := make([]string, 0)
seen := make(map[string]bool)

for _, label := range labels {
trimmed := strings.TrimSpace(strings.ToLower(label))
if trimmed != "" && allowedLabels[trimmed] && !seen[trimmed] {
cleaned = append(cleaned, trimmed)
seen[trimmed] = true
}
}

return cleaned
}
96 changes: 36 additions & 60 deletions internal/ai/gemini/issue_content_generator_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,15 @@ func TestBuildIssuePrompt(t *testing.T) {
},
contains: []string{"Code Changes (git diff)", "user description", "special hint"},
},
{
name: "with available labels",
request: models.IssueGenerationRequest{
Description: "user description",
Language: "en",
AvailableLabels: []string{"bug", "enhancement"},
},
contains: []string{"Available Labels", "bug, enhancement"},
},
}

for _, tt := range tests {
Expand Down Expand Up @@ -106,26 +115,6 @@ func TestBuildIssuePrompt_WithTemplate(t *testing.T) {

// Should contain the template
assert.Contains(t, prompt, "Bug Report")

// Should contain the final JSON reminder
assert.Contains(t, prompt, "🚨 FINAL REMINDER - CRITICAL OUTPUT REQUIREMENT 🚨")
assert.Contains(t, prompt, "YOU MUST OUTPUT **ONLY** VALID JSON")
assert.Contains(t, prompt, "BEGIN YOUR JSON OUTPUT NOW:")

// Should contain instructions about using template in description field
assert.Contains(t, prompt, "The template structure above should be used to FILL the \"description\" field")

// Should contain prohibitions
assert.Contains(t, prompt, "❌ DO NOT output prose like \"Here is a high-quality GitHub issue...\"")
assert.Contains(t, prompt, "❌ DO NOT output markdown text directly")

// Verify the reminder is at the end
lastIndex := len(prompt) - 500
if lastIndex < 0 {
lastIndex = 0
}
finalSection := prompt[lastIndex:]
assert.Contains(t, finalSection, "BEGIN YOUR JSON OUTPUT NOW:")
})

t.Run("does NOT add final reminder when no template", func(t *testing.T) {
Expand All @@ -137,9 +126,8 @@ func TestBuildIssuePrompt_WithTemplate(t *testing.T) {

prompt := gen.buildIssuePrompt(request)

// Should NOT contain the final JSON reminder
assert.NotContains(t, prompt, "🚨 FINAL REMINDER - CRITICAL OUTPUT REQUIREMENT 🚨")
assert.NotContains(t, prompt, "BEGIN YOUR JSON OUTPUT NOW:")
// Verification is just that prompt exists and is relevant
assert.Contains(t, prompt, "Code Changes")
})

t.Run("includes template in Spanish", func(t *testing.T) {
Expand All @@ -159,10 +147,6 @@ func TestBuildIssuePrompt_WithTemplate(t *testing.T) {

// Should contain the template
assert.Contains(t, prompt, "Reporte de Bug")

// Should still contain the final JSON reminder (in English for consistency)
assert.Contains(t, prompt, "🚨 FINAL REMINDER - CRITICAL OUTPUT REQUIREMENT 🚨")
assert.Contains(t, prompt, "BEGIN YOUR JSON OUTPUT NOW:")
})

t.Run("handles template with all fields", func(t *testing.T) {
Expand All @@ -188,28 +172,11 @@ func TestBuildIssuePrompt_WithTemplate(t *testing.T) {
// Should contain changed files
assert.Contains(t, prompt, "main.go")
assert.Contains(t, prompt, "test.go")

// Should contain the final reminder
assert.Contains(t, prompt, "🚨 FINAL REMINDER - CRITICAL OUTPUT REQUIREMENT 🚨")
assert.Contains(t, prompt, "BEGIN YOUR JSON OUTPUT NOW:")
})

t.Run("reminder contains complete JSON structure example", func(t *testing.T) {
template := &models.IssueTemplate{
Name: "Test Template",
}

request := models.IssueGenerationRequest{
Template: template,
Language: "en",
}

prompt := gen.buildIssuePrompt(request)

// Should show the expected JSON structure
assert.Contains(t, prompt, `"title": "string here"`)
assert.Contains(t, prompt, `"description": "markdown content following the template structure"`)
assert.Contains(t, prompt, `"labels": ["array", "of", "strings"]`)
// This test is now obsolete as structure is enforced by Schema, not prompt text.
// We can remove it or just check nothing.
})
}

Expand Down Expand Up @@ -268,33 +235,42 @@ func TestParseIssueResponse(t *testing.T) {
}

func TestCleanLabels(t *testing.T) {
gen := &GeminiIssueContentGenerator{}

tests := []struct {
name string
input []string
expected []string
name string
input []string
availableLabels []string
expected []string
}{
{
name: "only allowed labels",
input: []string{"fix", "feature", "bug", "invalid"},
expected: []string{"fix", "feature"},
name: "default whitelist - allowed",
input: []string{"fix", "feature", "bug", "invalid"},
availableLabels: nil,
expected: []string{"fix", "feature", "bug"},
},
{
name: "default whitelist - mixed case",
input: []string{" Fix ", "FEATURE", "test"},
availableLabels: nil,
expected: []string{"fix", "feature", "test"},
},
{
name: "mixed case and spaces",
input: []string{" Fix ", "FEATURE", "test"},
expected: []string{"fix", "feature", "test"},
name: "strict available labels",
input: []string{"custom-1", "custom-2", "fix"},
availableLabels: []string{"custom-1", "custom-2"},
expected: []string{"custom-1", "custom-2"},
},
{
name: "duplicates",
input: []string{"fix", "fix", "FIX"},
expected: []string{"fix"},
name: "strict available labels - excludes non-existent",
input: []string{"custom-1", "random"},
availableLabels: []string{"custom-1"},
expected: []string{"custom-1"},
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := gen.cleanLabels(tt.input)
result := CleanLabels(tt.input, tt.availableLabels)
assert.ElementsMatch(t, tt.expected, result)
})
}
Expand Down
15 changes: 10 additions & 5 deletions internal/ai/gemini/pull_requests_summarizer_service.go
Original file line number Diff line number Diff line change
Expand Up @@ -126,13 +126,14 @@ func (gps *GeminiPRSummarizer) defaultGenerate(ctx context.Context, mName string
return resp, usage, nil
}

func (gps *GeminiPRSummarizer) GeneratePRSummary(ctx context.Context, prContent string) (models.PRSummary, error) {
func (gps *GeminiPRSummarizer) GeneratePRSummary(ctx context.Context, prContent string, availableLabels []string) (models.PRSummary, error) {
log := logger.FromContext(ctx)

log.Info("generating PR summary via gemini",
"content_length", len(prContent))
"content_length", len(prContent),
"available_labels_count", len(availableLabels))

prompt := gps.generatePRPrompt(prContent)
prompt := gps.generatePRPrompt(prContent, availableLabels)

log.Debug("calling gemini API for PR summary",
"prompt_length", len(prompt))
Expand Down Expand Up @@ -199,12 +200,12 @@ func (gps *GeminiPRSummarizer) GeneratePRSummary(ctx context.Context, prContent
return models.PRSummary{
Title: jsonSummary.Title,
Body: jsonSummary.Body,
Labels: jsonSummary.Labels,
Labels: CleanLabels(jsonSummary.Labels, availableLabels),
Usage: usage,
}, nil
}

func (gps *GeminiPRSummarizer) generatePRPrompt(prContent string) string {
func (gps *GeminiPRSummarizer) generatePRPrompt(prContent string, availableLabels []string) string {
templateStr := ai.GetPRPromptTemplate(gps.config.Language)
data := ai.PromptData{
PRContent: prContent,
Expand All @@ -215,5 +216,9 @@ func (gps *GeminiPRSummarizer) generatePRPrompt(prContent string) string {
return ""
}

if len(availableLabels) > 0 {
rendered += fmt.Sprintf("\n\nAvailable Labels (Select ONLY from this list):\n%s", strings.Join(availableLabels, ", "))
}

return rendered
}
Loading
Loading