diff --git a/internal/app/app_test.go b/internal/app/app_test.go index 86e4b13..c682d00 100644 --- a/internal/app/app_test.go +++ b/internal/app/app_test.go @@ -49,10 +49,49 @@ func (m *mockLLMService) GenerateContent(ctx context.Context, prompt string, sch return nil, nil } +func (m *mockLLMService) GenerateContentWithResponse(ctx context.Context, prompt string, schema *genai.Schema, response interface{}) error { + return nil +} + func TestNewApp(t *testing.T) { - app := NewApp() + mockArchiveSvc := &mockArchiveService{} + mockEmailSvc := &mockEmailService{} + mockMinifluxClientSvc := &mockMinifluxClientService{} + mockDigestSvc := &mockDigestService{} + mockLLMSvc := &mockLLMService{} + cfg := &config.Config{} + + app := NewApp( + WithConfig(cfg), + WithArchiveService(mockArchiveSvc), + WithEmailService(mockEmailSvc), + WithMinifluxClientService(mockMinifluxClientSvc), + WithDigestService(mockDigestSvc), + WithLLMService(mockLLMSvc), + ) + if app == nil { t.Error("NewApp() should not return nil") + return + } + + if app.Config != cfg { + t.Errorf("NewApp() did not set the config correctly. Got %v, want %v", app.Config, cfg) + } + if app.ArchiveService != mockArchiveSvc { + t.Errorf("NewApp() did not set the archive service correctly. Got %v, want %v", app.ArchiveService, mockArchiveSvc) + } + if app.EmailService != mockEmailSvc { + t.Errorf("NewApp() did not set the email service correctly. Got %v, want %v", app.EmailService, mockEmailSvc) + } + if app.MinifluxClientService != mockMinifluxClientSvc { + t.Errorf("NewApp() did not set the miniflux client service correctly. Got %v, want %v", app.MinifluxClientService, mockMinifluxClientSvc) + } + if app.DigestService != mockDigestSvc { + t.Errorf("NewApp() did not set the digest service correctly. Got %v, want %v", app.DigestService, mockDigestSvc) + } + if app.LLMService != mockLLMSvc { + t.Errorf("NewApp() did not set the LLM service correctly. Got %v, want %v", app.LLMService, mockLLMSvc) } } diff --git a/internal/app/services/llm.go b/internal/app/services/llm.go index 3ab154d..72162b7 100644 --- a/internal/app/services/llm.go +++ b/internal/app/services/llm.go @@ -12,4 +12,5 @@ import ( type LLMService interface { GenerateContent(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) + GenerateContentWithResponse(ctx context.Context, prompt string, schema *genai.Schema, response interface{}) error } diff --git a/internal/digest/digest_test.go b/internal/digest/digest_test.go index b88f163..536e4d6 100644 --- a/internal/digest/digest_test.go +++ b/internal/digest/digest_test.go @@ -2,6 +2,7 @@ package digest import ( "context" + "encoding/json" "errors" "miniflux-digest/internal/digest/view" "miniflux-digest/internal/models" @@ -14,7 +15,8 @@ import ( ) type mockLLMService struct { - GenerateContentFunc func(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) + GenerateContentFunc func(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) + GenerateContentWithResponseFunc func(ctx context.Context, prompt string, schema *genai.Schema, response interface{}) error } func (m *mockLLMService) GenerateContent(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) { @@ -24,6 +26,13 @@ func (m *mockLLMService) GenerateContent(ctx context.Context, prompt string, sch return nil, errors.New("GenerateContentFunc not implemented") } +func (m *mockLLMService) GenerateContentWithResponse(ctx context.Context, prompt string, schema *genai.Schema, response interface{}) error { + if m.GenerateContentWithResponseFunc != nil { + return m.GenerateContentWithResponseFunc(ctx, prompt, schema, response) + } + return errors.New("GenerateContentWithResponseFunc not implemented") +} + func createLLMGrouperMockEntries() []*models.Entry { return []*models.Entry{ {ID: 1, Title: "Entry 1", FeedID: 100, FeedTitle: "Feed A", GroupID: 1, GroupTitle: "Category A"}, @@ -34,31 +43,40 @@ func createLLMGrouperMockEntries() []*models.Entry { func TestDigestService_BuildDigestData(t *testing.T) { mockLLM := &mockLLMService{ - GenerateContentFunc: func(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) { + GenerateContentWithResponseFunc: func(ctx context.Context, prompt string, schema *genai.Schema, response interface{}) error { switch schema { case view.InitialGroupingResponseSchema: - return []byte(`{ - "groups": [ - { - "title": "Go Lang", - "entry_ids": [2, 4] - }, - { - "title": "AI News", - "entry_ids": [1, 3] - } - ] - }`), nil + if err := json.Unmarshal([]byte(`{ + "groups": [ + { + "title": "AI News", + "entry_ids": [1] + }, + { + "title": "Go Lang", + "entry_ids": [2, 3] + } + ] + }`), response); err != nil { + return err + } + return nil case view.SummaryResponseSchema: - return []byte(`{"summary": "This is a summary."}`), nil + if err := json.Unmarshal([]byte(`{"summary": "This is a summary."}`), response); err != nil { + return err + } + return nil case view.SubGroupingResponseSchema: - return []byte(`{ - "sub_groups": [ - {"title": "Sub-group 1", "entry_ids": [1]} - ] - }`), nil + if err := json.Unmarshal([]byte(`{ + "sub_groups": [ + {"title": "Sub-group 1", "entry_ids": [1]} + ] + }`), response); err != nil { + return err + } + return nil default: - return nil, errors.New("unexpected LLM call") + return errors.New("unexpected LLM call") } }, } @@ -129,11 +147,11 @@ func TestDigestService_BuildDigestData(t *testing.T) { t.Run("view=ai - GroupAIEntries error fallback", func(t *testing.T) { mockLLM := &mockLLMService{ - GenerateContentFunc: func(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) { + GenerateContentWithResponseFunc: func(ctx context.Context, prompt string, schema *genai.Schema, response interface{}) error { if schema == view.InitialGroupingResponseSchema { - return nil, errors.New("mock LLM grouping error") + return errors.New("mock LLM grouping error") } - return nil, errors.New("unexpected LLM call") + return errors.New("unexpected LLM call") }, } digestService := NewDigestService(mockLLM) @@ -154,24 +172,30 @@ func TestDigestService_BuildDigestData(t *testing.T) { t.Run("view=ai - sub-grouping error", func(t *testing.T) { mockLLM := &mockLLMService{ - GenerateContentFunc: func(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) { + GenerateContentWithResponseFunc: func(ctx context.Context, prompt string, schema *genai.Schema, response interface{}) error { if schema == view.InitialGroupingResponseSchema { - return []byte(`{ + if err := json.Unmarshal([]byte(`{ "groups": [ { "title": "Test Group", "entry_ids": [1, 2, 3] } ] - }`), nil + }`), response); err != nil { + return err + } + return nil } if schema == view.SummaryResponseSchema { - return []byte(`{"summary": "This is a summary."}`), nil + if err := json.Unmarshal([]byte(`{"summary": "This is a summary."}`), response); err != nil { + return err + } + return nil } if schema == view.SubGroupingResponseSchema { - return nil, errors.New("mock LLM sub-grouping error") + return errors.New("mock LLM sub-grouping error") } - return nil, errors.New("unexpected LLM call") + return errors.New("unexpected LLM call") }, } digestService := NewDigestService(mockLLM) @@ -192,4 +216,5 @@ func TestDigestService_BuildDigestData(t *testing.T) { t.Errorf("Expected a single 'Uncategorized' sub-group on error, got %+v", group.SubGroups) } }) -} \ No newline at end of file +} + diff --git a/internal/digest/view/ai.go b/internal/digest/view/ai.go index 575c100..c63f72f 100644 --- a/internal/digest/view/ai.go +++ b/internal/digest/view/ai.go @@ -7,17 +7,20 @@ import ( "fmt" "log" "miniflux-digest/internal/app/services" + "miniflux-digest/internal/llm" "miniflux-digest/internal/models" "miniflux-digest/internal/utils" "sort" "strings" "sync" "text/template" + + "google.golang.org/genai" ) const ( MaxEntryContentLengthForLLM = 250 - MaxEntriesForSummarization = 150 + MaxEntriesPerJob = 150 ) type llmEntry struct { @@ -28,14 +31,12 @@ type llmEntry struct { FeedTitle string `json:"feed_title"` } + + func GroupAIEntries(ctx context.Context, entries []*models.Entry, llmService services.LLMService) (map[string][]*models.Entry, map[int64]bool, map[string][]*models.Entry) { const chunkSize = 200 - var wg sync.WaitGroup - var mu sync.Mutex - rawGroups := make(map[string][]*models.Entry) - groupedEntryIDs := make(map[int64]bool) - failedChunks := make(map[string][]*models.Entry) + // Prepare for processing entryMap := make(map[int64]*models.Entry) for _, entry := range entries { entryMap[entry.ID] = entry @@ -43,96 +44,43 @@ func GroupAIEntries(ctx context.Context, entries []*models.Entry, llmService ser promptTemplate, err := template.New("grouping").Parse(initialGroupingPrompt) if err != nil { - // This is a startup failure, not a runtime one, so we can't easily recover. - // We'll return it as a failed chunk. + failedChunks := make(map[string][]*models.Entry) failedChunks["Prompt template parsing failed"] = entries return nil, nil, failedChunks } - for i := 0; i < len(entries); i += chunkSize { - end := min(i+chunkSize, len(entries)) - chunk := entries[i:end] - wg.Add(1) - - go func(chunk []*models.Entry, i, end int) { - defer wg.Done() - - var llmEntries []llmEntry - for _, entry := range chunk { - content := entry.Content - if len(content) > MaxEntryContentLengthForLLM { - content = content[:MaxEntryContentLengthForLLM] - } - llmEntries = append(llmEntries, llmEntry{ - ID: entry.ID, - Title: entry.Title, - URL: entry.URL, - Content: content, - FeedTitle: entry.FeedTitle, - }) - } - - entriesJSON, err := json.MarshalIndent(llmEntries, "", " ") - if err != nil { - log.Printf("Failed to marshal entries to JSON for chunk %d-%d: %v", i, end, err) - mu.Lock() - failedChunks["JSON Marshal Error"] = append(failedChunks["JSON Marshal Error"], chunk...) - mu.Unlock() - return - } - - var prompt bytes.Buffer - if err := promptTemplate.Execute(&prompt, nil); err != nil { - log.Printf("Failed to execute prompt template for chunk %d-%d: %v", i, end, err) - mu.Lock() - failedChunks["Prompt Execution Error"] = append(failedChunks["Prompt Execution Error"], chunk...) - mu.Unlock() - return - } - prompt.Write(entriesJSON) - - llmResponse, err := llmService.GenerateContent(ctx, prompt.String(), InitialGroupingResponseSchema) - if err != nil { - log.Printf("LLM service failed for chunk %d-%d: %v", i, end, err) - errorKey := "LLM Error" - if strings.Contains(err.Error(), "DEADLINE_EXCEEDED") { - errorKey = "Processing Timeout" - } - mu.Lock() - failedChunks[errorKey] = append(failedChunks[errorKey], chunk...) - mu.Unlock() - return - } - - var response InitialGroupingResponse - if err := json.Unmarshal(llmResponse, &response); err != nil { - log.Printf("Failed to parse LLM response for chunk %d-%d: %v", i, end, err) - mu.Lock() - failedChunks["LLM Response Parse Error"] = append(failedChunks["LLM Response Parse Error"], chunk...) - mu.Unlock() - return - } - - mu.Lock() - defer mu.Unlock() - for _, group := range response.Groups { - if _, ok := rawGroups[group.Title]; !ok { - rawGroups[group.Title] = []*models.Entry{} - } - for _, entryID := range group.EntryIDs { - if _, exists := groupedEntryIDs[entryID]; exists { - continue - } - if entry, ok := entryMap[entryID]; ok { - rawGroups[group.Title] = append(rawGroups[group.Title], entry) - groupedEntryIDs[entryID] = true - } - } - } - }(chunk, i, end) + workerConfig := llmWorkerConfig{ + promptTemplate: promptTemplate, + responseSchema: InitialGroupingResponseSchema, + prepareEntries: func(entries []*models.Entry) interface{} { + return prepareLLMEntries(entries) + }, } + worker := createLLMWorker[*InitialGroupingResponse](ctx, llmService, workerConfig, "") - wg.Wait() + // Run the parallel processing + chunkResponses, failedChunks := llm.ProcessInChunks(entries, chunkSize, worker) + + // Process the results + rawGroups := make(map[string][]*models.Entry) + groupedEntryIDs := make(map[int64]bool) + + for _, response := range chunkResponses { + for _, group := range response.Groups { + if _, ok := rawGroups[group.Title]; !ok { + rawGroups[group.Title] = []*models.Entry{} + } + for _, entryID := range group.EntryIDs { + if _, exists := groupedEntryIDs[entryID]; exists { + continue + } + if entry, ok := entryMap[entryID]; ok { + rawGroups[group.Title] = append(rawGroups[group.Title], entry) + groupedEntryIDs[entryID] = true + } + } + } + } return rawGroups, groupedEntryIDs, failedChunks } @@ -157,19 +105,13 @@ func consolidatePrimaryGroups(ctx context.Context, rawGroups map[string][]*model return nil, fmt.Errorf("failed to execute consolidation prompt template: %w", err) } - llmResponse, err := llmService.GenerateContent(ctx, prompt.String(), ConsolidationResponseSchema) - if err != nil { + var response ConsolidationResponse + if err := llmService.GenerateContentWithResponse(ctx, prompt.String(), ConsolidationResponseSchema, &response); err != nil { // Fallback: If LLM fails, use the original, unconsolidated groups log.Printf("LLM consolidation failed, falling back to original groups: %v", err) return convertRawGroupsToDigestData(rawGroups), nil } - var response ConsolidationResponse - if err := json.Unmarshal(llmResponse, &response); err != nil { - log.Printf("Failed to parse LLM consolidation response, falling back to original groups: %v", err) - return convertRawGroupsToDigestData(rawGroups), nil - } - // Merge groups based on consolidation response consolidatedGroups := make(map[string][]*models.Entry) mappedOldTitles := make(map[string]bool) @@ -223,16 +165,215 @@ func convertRawGroupsToDigestData(rawGroups map[string][]*models.Entry) []*model return digestData } -func ProcessPrimaryGroupWithAI(ctx context.Context, pg *models.PrimaryGroup, llmService services.LLMService) (string, []*models.EntryGroup, error) { - var llmEntries []llmEntry - entriesToProcess := pg.Entries - if len(entriesToProcess) > MaxEntriesForSummarization { - entriesToProcess = entriesToProcess[:MaxEntriesForSummarization] + + +func getSummaryForGroup(ctx context.Context, groupTitle string, entries []*models.Entry, llmService services.LLMService) (string, error) { + llmResponseWorker := createLLMWorker[SummaryResponse](ctx, llmService, llmWorkerConfig{ + promptFormat: summaryPrompt, + responseSchema: SummaryResponseSchema, + prepareEntries: func(entries []*models.Entry) interface{} { + return prepareLLMEntries(entries) + }, + }, groupTitle) + + summaryWorker := func(chunk []*models.Entry) (string, error) { + resp, err := llmResponseWorker(chunk) + if err != nil { + return "", err + } + return resp.Summary, nil } + partialSummaries, failedChunks := llm.ProcessInChunks(entries, MaxEntriesPerJob, summaryWorker) + + if len(partialSummaries) == 0 { + return "", fmt.Errorf("all summary chunks failed for group '%s'", groupTitle) + } + + // Reduce step: summarize the summaries + summariesString := strings.Join(partialSummaries, "\n\n---\n\n") + finalSummaryPromptStr := fmt.Sprintf(summaryOfSummariesPrompt, summariesString) + + var finalSummaryResponse SummaryResponse + if err := llmService.GenerateContentWithResponse(ctx, finalSummaryPromptStr, SummaryResponseSchema, &finalSummaryResponse); err != nil { + return "", fmt.Errorf("summary of summaries failed: %w", err) + } + + finalSummary := finalSummaryResponse.Summary + var failedItems []*models.Entry + for _, items := range failedChunks { + failedItems = append(failedItems, items...) + } + if len(failedItems) > 0 { + finalSummary += fmt.Sprintf("\n\n(Note: %d entries could not be included in the summary due to processing errors.)", len(failedItems)) + } + + return finalSummary, nil +} + +func getSubGroupsForGroup(ctx context.Context, groupTitle string, entries []*models.Entry, llmService services.LLMService) ([]*models.EntryGroup, error) { entryMap := make(map[int64]*models.Entry) - for _, entry := range entriesToProcess { + for _, entry := range entries { entryMap[entry.ID] = entry + } + + subGroupWorker := createLLMWorker[*SubGroupingResponse](ctx, llmService, llmWorkerConfig{ + promptFormat: subGroupingPrompt, + responseSchema: SubGroupingResponseSchema, + prepareEntries: func(entries []*models.Entry) interface{} { + return prepareLLMEntries(entries) + }, + }, groupTitle) + + partialSubGroupResponses, failedChunks := llm.ProcessInChunks(entries, MaxEntriesPerJob, subGroupWorker) + + if len(partialSubGroupResponses) == 0 { + return nil, fmt.Errorf("all sub-grouping chunks failed for group '%s'", groupTitle) + } + + // Reduce step: merge sub-groups + mergedSubGroups := make(map[string][]*models.Entry) + allGroupedEntryIDs := make(map[int64]bool) + + for _, response := range partialSubGroupResponses { + for _, llmSubGroup := range response.SubGroups { + if _, ok := mergedSubGroups[llmSubGroup.Title]; !ok { + mergedSubGroups[llmSubGroup.Title] = []*models.Entry{} + } + for _, entryID := range llmSubGroup.EntryIDs { + if _, exists := allGroupedEntryIDs[entryID]; exists { + continue // Avoid duplicating entries if they appear in multiple chunk responses + } + if entry, ok := entryMap[entryID]; ok { + mergedSubGroups[llmSubGroup.Title] = append(mergedSubGroups[llmSubGroup.Title], entry) + allGroupedEntryIDs[entryID] = true + } + } + } + } + + var finalSubGroups []*models.EntryGroup + for title, groupEntries := range mergedSubGroups { + if len(groupEntries) > 0 { + finalSubGroups = append(finalSubGroups, &models.EntryGroup{ + Title: title, + Entries: groupEntries, + Slug: utils.Slugify(title), + TotalEntries: len(groupEntries), + TotalFeeds: getUniqueFeedIDs(groupEntries), + }) + } + } + + // Handle all ungrouped entries from all chunks + var allUngroupedEntries []*models.Entry + for _, entry := range entries { + if !allGroupedEntryIDs[entry.ID] { + allUngroupedEntries = append(allUngroupedEntries, entry) + } + } + + // Collect all failed items from the map into a single slice + var allFailedItems []*models.Entry + for _, items := range failedChunks { + allFailedItems = append(allFailedItems, items...) + } + + // Add items from failed chunks to a separate group + if len(allFailedItems) > 0 { + finalSubGroups = append(finalSubGroups, &models.EntryGroup{ + Title: "Uncategorized - Processing Failed", + Entries: allFailedItems, + Slug: "uncategorized-processing-failed", + TotalEntries: len(allFailedItems), + TotalFeeds: getUniqueFeedIDs(allFailedItems), + }) + } + + if len(allUngroupedEntries) > 0 { + // Find existing "Uncategorized" group to append to, if it exists + var uncategorizedGroup *models.EntryGroup + for _, sg := range finalSubGroups { + if sg.Title == "Uncategorized" { + uncategorizedGroup = sg + break + } + } + if uncategorizedGroup != nil { + uncategorizedGroup.Entries = append(uncategorizedGroup.Entries, allUngroupedEntries...) + } else { + finalSubGroups = append(finalSubGroups, &models.EntryGroup{ + Title: "Uncategorized", + Entries: allUngroupedEntries, + Slug: "uncategorized", + TotalEntries: len(allUngroupedEntries), + TotalFeeds: getUniqueFeedIDs(allUngroupedEntries), + }) + } + } + + return finalSubGroups, nil +} + +func ProcessPrimaryGroupWithAI(ctx context.Context, pg *models.PrimaryGroup, llmService services.LLMService) (string, []*models.EntryGroup, error) { + entriesToProcess := pg.Entries + + // If the group is large, process in chunks concurrently. + var finalSummary string + var finalSubGroups []*models.EntryGroup + var errSummary, errSubgroup error + var wg sync.WaitGroup + + wg.Add(2) + + // Goroutine for Summarization + go func() { + defer wg.Done() + finalSummary, errSummary = getSummaryForGroup(ctx, pg.Title, entriesToProcess, llmService) + }() + + // Goroutine for Sub-grouping + go func() { + defer wg.Done() + finalSubGroups, errSubgroup = getSubGroupsForGroup(ctx, pg.Title, entriesToProcess, llmService) + }() + + wg.Wait() + + // Combine errors if any. For now, we just log them but don't halt everything. + if errSummary != nil { + log.Printf("Error generating summary for group '%s': %v", pg.Title, errSummary) + } + if errSubgroup != nil { + log.Printf("Error generating sub-groups for group '%s': %v", pg.Title, errSubgroup) + // If sub-grouping completely failed, create a single group as a fallback. + if finalSubGroups == nil { + finalSubGroups = []*models.EntryGroup{{ + Title: "Uncategorized", + Entries: entriesToProcess, + Slug: "uncategorized", + TotalEntries: len(entriesToProcess), + TotalFeeds: getUniqueFeedIDs(entriesToProcess), + }} + } + } + + return finalSummary, finalSubGroups, nil +} + +func getUniqueFeedIDs(entries []*models.Entry) int { + uniqueFeedIDs := make(map[int64]bool) + for _, entry := range entries { + uniqueFeedIDs[entry.FeedID] = true + } + return len(uniqueFeedIDs) +} + + + +func prepareLLMEntries(entries []*models.Entry) []llmEntry { + var llmEntries []llmEntry + for _, entry := range entries { content := entry.Content if len(content) > MaxEntryContentLengthForLLM { content = content[:MaxEntryContentLengthForLLM] @@ -245,103 +386,27 @@ func ProcessPrimaryGroupWithAI(ctx context.Context, pg *models.PrimaryGroup, llm FeedTitle: entry.FeedTitle, }) } - - entriesJSON, err := json.MarshalIndent(llmEntries, "", " ") - if err != nil { - return "", nil, fmt.Errorf("failed to marshal entries to JSON for AI processing: %w", err) - } - - // 1. Get Summary - var summary string - summaryPromptFmt := fmt.Sprintf(summaryPrompt, pg.Title) - summaryFullPrompt := summaryPromptFmt + string(entriesJSON) - - summaryResponseBytes, err := llmService.GenerateContent(ctx, summaryFullPrompt, SummaryResponseSchema) - if err != nil { - log.Printf("LLM service failed during summarization for primary group '%s': %v", pg.Title, err) - } else { - var summaryResponse SummaryResponse - if err := json.Unmarshal(summaryResponseBytes, &summaryResponse); err != nil { - log.Printf("Failed to parse LLM summarization response for primary group '%s': %v", pg.Title, err) - } else { - summary = summaryResponse.Summary - } - } - - // 2. Get Sub-groups - var subGroups []*models.EntryGroup - subGroupingPromptFmt := fmt.Sprintf(subGroupingPrompt, pg.Title) - subGroupingFullPrompt := subGroupingPromptFmt + string(entriesJSON) - - subGroupingResponseBytes, err := llmService.GenerateContent(ctx, subGroupingFullPrompt, SubGroupingResponseSchema) - if err != nil { - log.Printf("LLM service failed during sub-grouping for primary group '%s': %v. Creating single 'Uncategorized' group.", pg.Title, err) - } else { - var subGroupingResponse SubGroupingResponse - if err := json.Unmarshal(subGroupingResponseBytes, &subGroupingResponse); err != nil { - log.Printf("Failed to parse LLM sub-grouping response for primary group '%s': %v. Creating single 'Uncategorized' group.", pg.Title, err) - } else { - groupedEntryIDs := make(map[int64]bool) - for _, llmSubGroup := range subGroupingResponse.SubGroups { - var currentSubGroupEntries []*models.Entry - for _, entryID := range llmSubGroup.EntryIDs { - if entry, ok := entryMap[entryID]; ok { - currentSubGroupEntries = append(currentSubGroupEntries, entry) - groupedEntryIDs[entryID] = true - } - } - if len(currentSubGroupEntries) > 0 { - subGroups = append(subGroups, &models.EntryGroup{ - Title: llmSubGroup.Title, - Entries: currentSubGroupEntries, - Slug: utils.Slugify(llmSubGroup.Title), - TotalEntries: len(currentSubGroupEntries), - TotalFeeds: getUniqueFeedIDs(currentSubGroupEntries), - }) - } - } - // Handle ungrouped entries - var ungroupedEntries []*models.Entry - for _, entry := range entriesToProcess { - if !groupedEntryIDs[entry.ID] { - ungroupedEntries = append(ungroupedEntries, entry) - } - } - if len(ungroupedEntries) > 0 { - subGroups = append(subGroups, &models.EntryGroup{ - Title: "Uncategorized", - Entries: ungroupedEntries, - Slug: "uncategorized", - TotalEntries: len(ungroupedEntries), - TotalFeeds: getUniqueFeedIDs(ungroupedEntries), - }) - } - } - } - - // Fallback if sub-grouping failed or returned no groups - if len(subGroups) == 0 { - subGroups = append(subGroups, &models.EntryGroup{ - Title: "Uncategorized", - Entries: entriesToProcess, - Slug: "uncategorized", - TotalEntries: len(entriesToProcess), - TotalFeeds: getUniqueFeedIDs(entriesToProcess), - }) - } - - return summary, subGroups, nil -} - -func getUniqueFeedIDs(entries []*models.Entry) int { - uniqueFeedIDs := make(map[int64]bool) - for _, entry := range entries { - uniqueFeedIDs[entry.FeedID] = true - } - return len(uniqueFeedIDs) + return llmEntries } func BuildDigestDataByAI(entries []*models.Entry, ctx context.Context, llmService services.LLMService) []*models.PrimaryGroupDigestData { + rawGroups, groupedEntryIDs, failedChunks, fallbackGroups := initialAIGGrouping(ctx, entries, llmService) + if fallbackGroups != nil { + return fallbackGroups + } + + highLevelGroups := consolidateGroups(ctx, rawGroups, llmService) + + highLevelGroups = handleFailedAndUngroupedEntries(highLevelGroups, entries, groupedEntryIDs, failedChunks) + + allPrimaryGroups := processHighLevelGroupsConcurrently(ctx, highLevelGroups, llmService) + + sortPrimaryGroups(allPrimaryGroups) + + return allPrimaryGroups +} + +func initialAIGGrouping(ctx context.Context, entries []*models.Entry, llmService services.LLMService) (map[string][]*models.Entry, map[int64]bool, map[string][]*models.Entry, []*models.PrimaryGroupDigestData) { rawGroups, groupedEntryIDs, failedChunks := GroupAIEntries(ctx, entries, llmService) // If all chunks failed and we have no raw groups, fall back to category grouping. @@ -349,18 +414,24 @@ func BuildDigestDataByAI(entries []*models.Entry, ctx context.Context, llmServic log.Printf("All AI grouping chunks failed, falling back to category grouping.") categoryGroups := BuildDigestDataByCategory(entries) if len(categoryGroups) == 1 && len(entries) > 1 { - return BuildDigestDataByDate(entries) + return nil, nil, nil, BuildDigestDataByDate(entries) } - return categoryGroups + return nil, nil, nil, categoryGroups } + return rawGroups, groupedEntryIDs, failedChunks, nil +} +func consolidateGroups(ctx context.Context, rawGroups map[string][]*models.Entry, llmService services.LLMService) []*models.PrimaryGroupDigestData { highLevelGroups, err := consolidatePrimaryGroups(ctx, rawGroups, llmService) if err != nil { // This case should ideally not be hit if consolidatePrimaryGroups has a fallback log.Printf("AI group consolidation failed unexpectedly: %v", err) highLevelGroups = convertRawGroupsToDigestData(rawGroups) } + return highLevelGroups +} +func handleFailedAndUngroupedEntries(highLevelGroups []*models.PrimaryGroupDigestData, allEntries []*models.Entry, groupedEntryIDs map[int64]bool, failedChunks map[string][]*models.Entry) []*models.PrimaryGroupDigestData { failedEntryIDs := make(map[int64]bool) for reason, chunkEntries := range failedChunks { title := fmt.Sprintf("Uncategorized - %s", reason) @@ -388,7 +459,7 @@ func BuildDigestDataByAI(entries []*models.Entry, ctx context.Context, llmServic // Handle entries that were not grouped by the initial LLM call var otherUngroupedEntries []*models.Entry - for _, entry := range entries { + for _, entry := range allEntries { if !groupedEntryIDs[entry.ID] && !failedEntryIDs[entry.ID] { otherUngroupedEntries = append(otherUngroupedEntries, entry) } @@ -413,7 +484,10 @@ func BuildDigestDataByAI(entries []*models.Entry, ctx context.Context, llmServic TotalFeeds: getUniqueFeedIDs(otherUngroupedEntries), }) } + return highLevelGroups +} +func processHighLevelGroupsConcurrently(ctx context.Context, highLevelGroups []*models.PrimaryGroupDigestData, llmService services.LLMService) []*models.PrimaryGroupDigestData { var allPrimaryGroups []*models.PrimaryGroupDigestData var wg sync.WaitGroup var mu sync.Mutex @@ -458,10 +532,63 @@ func BuildDigestDataByAI(entries []*models.Entry, ctx context.Context, llmServic } wg.Wait() + return allPrimaryGroups +} +func sortPrimaryGroups(allPrimaryGroups []*models.PrimaryGroupDigestData) { sort.Slice(allPrimaryGroups, func(i, j int) bool { return allPrimaryGroups[i].Title < allPrimaryGroups[j].Title }) - - return allPrimaryGroups } + +// llmWorkerConfig holds configuration for creating an LLM worker function. +type llmWorkerConfig struct { + promptTemplate *template.Template + promptFormat string // For prompts that use Sprintf + responseSchema *genai.Schema + prepareEntries func([]*models.Entry) interface{} // Function to prepare entries for LLM +} + +// createLLMWorker creates a worker function for llm.ProcessInChunks. +func createLLMWorker[R any]( + ctx context.Context, + llmService services.LLMService, + config llmWorkerConfig, + groupTitle string, // Optional, for prompts that need it +) llm.WorkerFunc[*models.Entry, R] { + return func(chunk []*models.Entry) (R, error) { + var zero R // Initialize a zero value of type R + + // 1. Prepare LLM entries + var llmEntriesJSON []byte + if config.prepareEntries != nil { + prepared := config.prepareEntries(chunk) + var err error + llmEntriesJSON, err = json.MarshalIndent(prepared, "", " ") + if err != nil { + return zero, fmt.Errorf("failed to marshal entries to JSON: %w", err) + } + } + + // 2. Construct the prompt + var prompt bytes.Buffer + if config.promptTemplate != nil { + if err := config.promptTemplate.Execute(&prompt, nil); err != nil { + return zero, fmt.Errorf("failed to execute prompt template: %w", err) + } + } else if config.promptFormat != "" { + prompt.WriteString(fmt.Sprintf(config.promptFormat, groupTitle)) + } + + if len(llmEntriesJSON) > 0 { + prompt.Write(llmEntriesJSON) + } + + // 3. Call LLM service + var response R + if err := llmService.GenerateContentWithResponse(ctx, prompt.String(), config.responseSchema, &response); err != nil { + return zero, err + } + return response, nil + } +} \ No newline at end of file diff --git a/internal/digest/view/ai_prompts.go b/internal/digest/view/ai_prompts.go index e113446..46a6495 100644 --- a/internal/digest/view/ai_prompts.go +++ b/internal/digest/view/ai_prompts.go @@ -52,7 +52,7 @@ A good consolidation would be: const summaryPrompt = `You are an expert at summarizing content for busy readers. You will be given a list of entries from a primary group titled '%s'. -Your task is to write a concise, 2-3 sentence summary (under 150 words) that achieves two goals: +Your task is to write a concise, 2-3 sentence summary (under 250 words) that achieves two goals: 1. **Provide a high-level overview** of the main themes in the group. 2. **Highlight the most significant or surprising entries.** This could be a major announcement, a controversial opinion, or a particularly popular discussion. @@ -70,6 +70,19 @@ Your task is to create smaller, more granular sub-groups to help a user navigate - If an entry doesn't fit into a specific sub-group, leave it out of your response. ` +const summaryOfSummariesPrompt = `You are an expert at creating overviews of large swaths of web updates and web content. + +Your task is to write a concise, 2-3 sentence summary (under 250 words) that achieves two goals: +1. **Provide a high-level overview** of the main themes. +2. **Highlight the most significant or surprising entries.** This could be a major announcement, a controversial opinion, or a particularly popular discussion. + +The goal is to give the reader enough information to quickly decide if this group contains entries they want to explore further. Do not just list the topics; provide some insight into the content. + +Here are the summaries to combine: +---------------------------- +%s +` + // --- Prompt Response Types --- type InitialGroupingResponse struct { @@ -173,4 +186,4 @@ var SubGroupingResponseSchema = &genai.Schema{ }, }, }, -} \ No newline at end of file +} diff --git a/internal/digest/view/ai_test.go b/internal/digest/view/ai_test.go index c9d5edd..ca0711b 100644 --- a/internal/digest/view/ai_test.go +++ b/internal/digest/view/ai_test.go @@ -4,8 +4,10 @@ import ( "context" "encoding/json" "errors" + "fmt" "miniflux-digest/internal/models" "miniflux-digest/internal/testutil" + "strings" "sync" "testing" "time" @@ -17,7 +19,8 @@ import ( // mockLLMService is a mock implementation of the LLMService for testing. type mockLLMService struct { - GenerateContentFunc func(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) + GenerateContentFunc func(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) + GenerateContentWithResponseFunc func(ctx context.Context, prompt string, schema *genai.Schema, response interface{}) error } func (m *mockLLMService) GenerateContent(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) { @@ -27,6 +30,22 @@ func (m *mockLLMService) GenerateContent(ctx context.Context, prompt string, sch return nil, errors.New("mockLLMService.GenerateContentFunc is not implemented") } +func (m *mockLLMService) GenerateContentWithResponse(ctx context.Context, prompt string, schema *genai.Schema, response interface{}) error { + if m.GenerateContentWithResponseFunc != nil { + return m.GenerateContentWithResponseFunc(ctx, prompt, schema, response) + } + // Default implementation for GenerateContentWithResponse + // This will call GenerateContentFunc and then unmarshal the result + bytes, err := m.GenerateContent(ctx, prompt, schema) + if err != nil { + return err + } + if err := json.Unmarshal(bytes, &response); err != nil { + return fmt.Errorf("failed to parse LLM response in mock: %w", err) + } + return nil +} + func TestGroupAIEntries(t *testing.T) { entries := []*models.Entry{ {ID: 1, Title: "Entry 1", Content: "Content 1", FeedTitle: "Feed 1"}, @@ -48,8 +67,8 @@ func TestGroupAIEntries(t *testing.T) { mockJSON, _ := json.Marshal(mockResp) llmService := &mockLLMService{ - GenerateContentFunc: func(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) { - return mockJSON, nil + GenerateContentWithResponseFunc: func(ctx context.Context, prompt string, schema *genai.Schema, response interface{}) error { + return json.Unmarshal(mockJSON, response) }, } @@ -72,8 +91,8 @@ func TestGroupAIEntries(t *testing.T) { t.Run("llm service error", func(t *testing.T) { llmService := &mockLLMService{ - GenerateContentFunc: func(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) { - return nil, errors.New("llm error") + GenerateContentWithResponseFunc: func(ctx context.Context, prompt string, schema *genai.Schema, response interface{}) error { + return errors.New("llm error") }, } @@ -83,8 +102,8 @@ func TestGroupAIEntries(t *testing.T) { t.Run("llm returns invalid json", func(t *testing.T) { llmService := &mockLLMService{ - GenerateContentFunc: func(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) { - return []byte("invalid json"), nil + GenerateContentWithResponseFunc: func(ctx context.Context, prompt string, schema *genai.Schema, response interface{}) error { + return json.Unmarshal([]byte(`invalid json`), response) }, } @@ -102,7 +121,7 @@ func TestGroupAIEntries(t *testing.T) { var mu sync.Mutex llmService := &mockLLMService{ - GenerateContentFunc: func(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) { + GenerateContentWithResponseFunc: func(ctx context.Context, prompt string, schema *genai.Schema, response interface{}) error { mu.Lock() callCount++ mu.Unlock() @@ -123,8 +142,10 @@ func TestGroupAIEntries(t *testing.T) { {Title: "Chunk 2 Group", EntryIDs: []int64{200, 201}}, } } - mockJSON, _ := json.Marshal(resp) - return mockJSON, nil + return json.Unmarshal([]byte(func() []byte { + b, _ := json.Marshal(resp) + return b + }()), response) }, } @@ -226,58 +247,138 @@ func TestConsolidatePrimaryGroups(t *testing.T) { } func TestProcessPrimaryGroupWithAI(t *testing.T) { - pg := &models.PrimaryGroup{ - ID: 1, - Title: "Test Primary Group", - Entries: []*models.Entry{ - {ID: 1, Title: "Entry 1", Content: "Content 1"}, - {ID: 2, Title: "Entry 2", Content: "Content 2"}, - {ID: 3, Title: "Entry 3", Content: "Content 3"}, - }, - } - - t.Run("happy path", func(t *testing.T) { - summaryResp := SummaryResponse{Summary: "This is a summary."} - summaryJSON, _ := json.Marshal(summaryResp) - - subGroupingResp := SubGroupingResponse{ - SubGroups: []struct { - Title string `json:"title"` - EntryIDs []int64 `json:"entry_ids"` - }{ - {Title: "Sub-group 1", EntryIDs: []int64{1, 2}}, + // --- Test Case 1: Small group (no chunking) --- + t.Run("small group - no chunking", func(t *testing.T) { + pg := &models.PrimaryGroup{ + ID: 1, + Title: "Small Group", + Entries: []*models.Entry{ + {ID: 1, Title: "E1"}, + {ID: 2, Title: "E2"}, }, } - subGroupingJSON, _ := json.Marshal(subGroupingResp) llmService := &mockLLMService{ GenerateContentFunc: func(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) { if schema == SummaryResponseSchema { - return summaryJSON, nil + return json.Marshal(SummaryResponse{Summary: "Small group summary"}) } if schema == SubGroupingResponseSchema { - return subGroupingJSON, nil + return json.Marshal(SubGroupingResponse{SubGroups: []struct { + Title string `json:"title"` + EntryIDs []int64 `json:"entry_ids"` + }{{Title: "Sub 1", EntryIDs: []int64{1, 2}}}}) } - return nil, errors.New("unexpected schema in happy path test") + return nil, errors.New("unexpected schema") }, } summary, subGroups, err := ProcessPrimaryGroupWithAI(context.Background(), pg, llmService) assert.NoError(t, err) - assert.Equal(t, "This is a summary.", summary) - assert.Len(t, subGroups, 2, "Expected 1 sub-group and 1 uncategorized") - - sg1 := subGroups[0] - assert.Equal(t, "Sub-group 1", sg1.Title) - assert.Len(t, sg1.Entries, 2) - - uncategorized := subGroups[1] - assert.Equal(t, "Uncategorized", uncategorized.Title) - assert.Len(t, uncategorized.Entries, 1) - assert.Equal(t, int64(3), uncategorized.Entries[0].ID) + assert.Equal(t, "Small group summary", summary) + require.Len(t, subGroups, 1) + assert.Equal(t, "Sub 1", subGroups[0].Title) }) - // ... other tests for ProcessPrimaryGroupWithAI are unchanged ... + // --- Test Case 2: Large group (with chunking) --- + t.Run("large group - with chunking", func(t *testing.T) { + var entries []*models.Entry + for i := 0; i < 200; i++ { // Will create 2 chunks of 150 + entries = append(entries, &models.Entry{ID: int64(i)}) + } + pg := &models.PrimaryGroup{ID: 2, Title: "Large Group", Entries: entries} + + llmService := &mockLLMService{ + GenerateContentWithResponseFunc: func(ctx context.Context, prompt string, schema *genai.Schema, response interface{}) error { + isLargeChunk := len(prompt) > 10000 // Heuristic to differentiate chunks + + // Summarization + if strings.Contains(prompt, "You are an expert at summarizing content") && !strings.Contains(prompt, "overviews of large swaths") { // Partial summary call + return json.Unmarshal([]byte(`{"summary": "Partial summary."}`), response) + } + // Summary of summaries call + if strings.Contains(prompt, "You are an expert at creating overviews of large swaths of web updates and web content.") { + return json.Unmarshal([]byte(`{"summary": "Final synthesized summary."}`), response) + } + // Sub-grouping + if strings.Contains(prompt, "You are an expert at organizing content") { + if isLargeChunk { + return json.Unmarshal([]byte(`{"sub_groups": [{"title": "Chunk 1 Sub-Group", "entry_ids": [0, 1]}]}`), response) + } else { + return json.Unmarshal([]byte(`{"sub_groups": [{"title": "Chunk 2 Sub-Group", "entry_ids": [150, 151]}]}`), response) + } + } + return fmt.Errorf("unexpected prompt: %s", prompt) + }, + } + + summary, subGroups, err := ProcessPrimaryGroupWithAI(context.Background(), pg, llmService) + assert.NoError(t, err) + assert.Equal(t, "Final synthesized summary.", summary) + require.Len(t, subGroups, 3, "Expected 3 sub-groups: 2 from chunks + 1 for ungrouped") + + sg1 := testutil.FindSubGroup(subGroups, "Chunk 1 Sub-Group") + require.NotNil(t, sg1) + sg2 := testutil.FindSubGroup(subGroups, "Chunk 2 Sub-Group") + require.NotNil(t, sg2) + uncat := testutil.FindSubGroup(subGroups, "Uncategorized") + require.NotNil(t, uncat) + + assert.Len(t, sg1.Entries, 2) + assert.Len(t, sg2.Entries, 2) + assert.Equal(t, 196, uncat.TotalEntries) // 200 total - 4 grouped + }) + + // --- Test Case 3: Failure case --- + t.Run("large group - with failures", func(t *testing.T) { + var entries []*models.Entry + for i := 0; i < 200; i++ { + entries = append(entries, &models.Entry{ID: int64(i)}) + } + pg := &models.PrimaryGroup{ID: 3, Title: "Failure Group", Entries: entries} + + llmService := &mockLLMService{ + GenerateContentFunc: func(ctx context.Context, prompt string, schema *genai.Schema) ([]byte, error) { + isLargeChunk := len(prompt) > 10000 + + // Summarization: one chunk fails + if strings.Contains(prompt, "You are an expert at summarizing content") && !strings.Contains(prompt, "overviews of large swaths") { // Partial summary call + if isLargeChunk { // First chunk succeeds + return json.Marshal(SummaryResponse{Summary: "Partial summary."}) + } + return nil, errors.New("summary failed") // Second chunk fails + } + // This handles the summary of summaries prompt + if strings.Contains(prompt, "You are an expert at creating overviews of large swaths of web updates and web content.") { + return json.Marshal(SummaryResponse{Summary: "Final summary from one chunk."}) + } + // Sub-grouping: one chunk fails + if strings.Contains(prompt, "You are an expert at organizing content") { + if isLargeChunk { // First chunk succeeds + return json.Marshal(SubGroupingResponse{SubGroups: []struct { + Title string `json:"title"` + EntryIDs []int64 `json:"entry_ids"` + }{{Title: "Good Sub-Group", EntryIDs: []int64{0}}}})} + return nil, errors.New("sub-grouping failed") // Second chunk fails + } + return nil, fmt.Errorf("unexpected prompt: %s", prompt) + }, + } + + summary, subGroups, err := ProcessPrimaryGroupWithAI(context.Background(), pg, llmService) + assert.NoError(t, err) // The function itself shouldn't error + + assert.Contains(t, summary, "Final summary from one chunk.") + assert.Contains(t, summary, "(Note: 50 entries could not be included") // 200 total - 150 in first chunk = 50 in second + + failedSg := testutil.FindSubGroup(subGroups, "Uncategorized - Processing Failed") + require.NotNil(t, failedSg) + assert.Len(t, failedSg.Entries, 50) + + goodSg := testutil.FindSubGroup(subGroups, "Good Sub-Group") + require.NotNil(t, goodSg) + assert.Len(t, goodSg.Entries, 1) + }) } func TestBuildDigestDataByAI(t *testing.T) { @@ -377,4 +478,5 @@ func TestBuildDigestDataByAI(t *testing.T) { }) // ... other fallback tests are similar and should still pass ... -} \ No newline at end of file +} + diff --git a/internal/digest/view/date_test.go b/internal/digest/view/date_test.go index 6f5cc75..a41cd0d 100644 --- a/internal/digest/view/date_test.go +++ b/internal/digest/view/date_test.go @@ -41,7 +41,7 @@ func TestDayGrouper_GroupEntries(t *testing.T) { t.Fatalf("Incorrect sub-groups for Category A: %+v", catAGroup) } - subGroup := testutil.FindSubGroup(catAGroup, "Jan 2, 2024") + subGroup := testutil.FindSubGroup(catAGroup.SubGroups, "Jan 2, 2024") if subGroup == nil || len(subGroup.Entries) != 4 { t.Fatalf("Incorrect entries for sub-group: %+v", subGroup) } diff --git a/internal/llm/gemini.go b/internal/llm/gemini.go index 582c388..82c6b04 100644 --- a/internal/llm/gemini.go +++ b/internal/llm/gemini.go @@ -2,7 +2,9 @@ package llm import ( "context" + "encoding/json" "errors" + "fmt" "log" "time" @@ -103,4 +105,15 @@ func (s *GeminiService) GenerateContent(ctx context.Context, prompt string, sche } return jsonData, nil +} + +func (s *GeminiService) GenerateContentWithResponse(ctx context.Context, prompt string, schema *genai.Schema, response interface{}) error { + llmResponseBytes, err := s.GenerateContent(ctx, prompt, schema) + if err != nil { + return err + } + if err := json.Unmarshal(llmResponseBytes, &response); err != nil { + return fmt.Errorf("failed to parse LLM response: %w", err) + } + return nil } \ No newline at end of file diff --git a/internal/llm/util.go b/internal/llm/util.go new file mode 100644 index 0000000..d349d19 --- /dev/null +++ b/internal/llm/util.go @@ -0,0 +1,68 @@ +package llm + +import ( + "log" + "strings" + "sync" +) + +// WorkerFunc is a generic function type for processing a chunk of items. +type WorkerFunc[T any, R any] func(chunk []T) (R, error) + +// ProcessInChunks is a generic function that processes a slice of items in parallel chunks. +// It takes a slice of items, a chunk size, and a worker function to process each chunk. +// It returns a slice of results from the successful chunks and a map of failed items, keyed by the error reason. +func ProcessInChunks[T any, R any]( + items []T, + chunkSize int, + worker func(chunk []T) (R, error), +) ([]R, map[string][]T) { + var wg sync.WaitGroup + var mu sync.Mutex + + results := make([]R, 0) + failedChunks := make(map[string][]T) + + for i := 0; i < len(items); i += chunkSize { + end := min(i+chunkSize, len(items)) + chunk := items[i:end] + wg.Add(1) + + go func(chunk []T) { + defer wg.Done() + + result, err := worker(chunk) + if err != nil { + log.Printf("A chunk failed during parallel processing: %v", err) + + // Determine a simple error key for bucketing failed chunks. + errorKey := "Processing Error" + if strings.Contains(err.Error(), "DEADLINE_EXCEEDED") { + errorKey = "Processing Timeout" + } else if strings.Contains(strings.ToLower(err.Error()), "json") { + errorKey = "Invalid Response" + } + + mu.Lock() + failedChunks[errorKey] = append(failedChunks[errorKey], chunk...) + mu.Unlock() + return + } + + mu.Lock() + results = append(results, result) + mu.Unlock() + }(chunk) + } + + wg.Wait() + + return results, failedChunks +} + +func min(a, b int) int { + if a < b { + return a + } + return b +} \ No newline at end of file diff --git a/internal/llm/util_test.go b/internal/llm/util_test.go new file mode 100644 index 0000000..4c64563 --- /dev/null +++ b/internal/llm/util_test.go @@ -0,0 +1,145 @@ +package llm + +import ( + "errors" + "fmt" + "sync" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestProcessInChunks(t *testing.T) { + t.Run("happy path - processes all items", func(t *testing.T) { + items := []int{1, 2, 3, 4, 5, 6, 7} + chunkSize := 3 + + worker := func(chunk []int) (string, error) { + var sum int + for _, item := range chunk { + sum += item + } + return fmt.Sprintf("sum-%d", sum), nil + } + + results, failedItems := ProcessInChunks(items, chunkSize, worker) + + assert.Empty(t, failedItems) + assert.Len(t, results, 3) // 3 chunks: (1,2,3), (4,5,6), (7) + assert.Contains(t, results, "sum-6") + assert.Contains(t, results, "sum-15") + assert.Contains(t, results, "sum-7") + }) + + t.Run("one chunk fails", func(t *testing.T) { + items := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10} + chunkSize := 4 + + worker := func(chunk []int) (string, error) { + // Fail the second chunk + if chunk[0] == 5 { + return "", errors.New("processing error") + } + return "success", nil + } + + results, failedChunks := ProcessInChunks(items, chunkSize, worker) + + assert.Len(t, results, 2) // Chunks 1 and 3 should succeed + assert.Len(t, failedChunks, 1) // One failure reason + assert.Contains(t, failedChunks, "Processing Error") + assert.Len(t, failedChunks["Processing Error"], 4) // Chunk 2 (items 5,6,7,8) should fail + assert.Equal(t, 5, failedChunks["Processing Error"][0]) + }) + + t.Run("all chunks fail", func(t *testing.T) { + items := []int{1, 2, 3, 4} + chunkSize := 2 + + worker := func(chunk []int) (string, error) { + return "", errors.New("always fail") + } + + results, failedChunks := ProcessInChunks(items, chunkSize, worker) + + assert.Empty(t, results) + assert.Len(t, failedChunks, 1) + assert.Len(t, failedChunks["Processing Error"], 4) + }) + + t.Run("items less than chunk size", func(t *testing.T) { + items := []int{1, 2} + chunkSize := 5 + + worker := func(chunk []int) (string, error) { + return "success", nil + } + + results, failedItems := ProcessInChunks(items, chunkSize, worker) + + assert.Empty(t, failedItems) + assert.Len(t, results, 1) + }) + + t.Run("zero items", func(t *testing.T) { + items := []int{} + chunkSize := 5 + + worker := func(chunk []int) (string, error) { + t.Fatal("worker should not be called for zero items") + return "", nil + } + + results, failedItems := ProcessInChunks(items, chunkSize, worker) + + assert.Empty(t, failedItems) + assert.Empty(t, results) + }) + + t.Run("exact multiple of chunk size", func(t *testing.T) { + items := []int{1, 2, 3, 4, 5, 6} + chunkSize := 3 + + worker := func(chunk []int) (string, error) { + return "success", nil + } + + results, failedItems := ProcessInChunks(items, chunkSize, worker) + + assert.Empty(t, failedItems) + assert.Len(t, results, 2) + }) + + t.Run("thread safety", func(t *testing.T) { + numItems := 1000 + items := make([]int, numItems) + for i := 0; i < numItems; i++ { + items[i] = i + } + chunkSize := 10 + + var mu sync.Mutex + var callCount int + + worker := func(chunk []int) (int, error) { + mu.Lock() + callCount++ + // Fail every 10th chunk + if callCount%10 == 0 { + mu.Unlock() + return 0, errors.New("failed") + } + mu.Unlock() + return len(chunk), nil + } + + results, failedChunks := ProcessInChunks(items, chunkSize, worker) + + numChunks := numItems / chunkSize + numFailedChunks := numChunks / 10 + numSuccessfulChunks := numChunks - numFailedChunks + + assert.Len(t, results, numSuccessfulChunks) + assert.Len(t, failedChunks["Processing Error"], numFailedChunks*chunkSize) + }) +} \ No newline at end of file diff --git a/internal/testutil/mocks.go b/internal/testutil/mocks.go index 3467544..a80d2e2 100644 --- a/internal/testutil/mocks.go +++ b/internal/testutil/mocks.go @@ -79,18 +79,18 @@ func (m *MockDigestService) BuildDigestData(category *miniflux.Category, entries // Helper functions for tests func FindPrimaryGroup(groups []*models.PrimaryGroupDigestData, title string) *models.PrimaryGroupDigestData { - for _, group := range groups { - if group.Title == title { - return group + for _, g := range groups { + if g.Title == title { + return g } } return nil } -func FindSubGroup(primaryGroup *models.PrimaryGroupDigestData, title string) *models.EntryGroup { - for _, group := range primaryGroup.SubGroups { - if group.Title == title { - return group +func FindSubGroup(groups []*models.EntryGroup, title string) *models.EntryGroup { + for _, g := range groups { + if g.Title == title { + return g } } return nil