Merge pull request #71 from thomas-vilte/dev

feat(ai): Leverage Repository Labels for Enhanced AI Generation

Author: thomas-vilte

internal/ai/gemini/helpers.go

@@ -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
+}

internal/ai/gemini/issue_content_generator.go

@@ -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)
@@ -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",
@@ -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
 	}
 
@@ -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
@@ -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 {
@@ -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
-}

internal/ai/gemini/issue_content_generator_test.go

@@ -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 {
@@ -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) {
@@ -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) {
@@ -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) {
@@ -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.
 	})
 }
 
@@ -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)
 		})
 	}

internal/ai/gemini/pull_requests_summarizer_service.go

@@ -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))
@@ -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,
@@ -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
 }