refactor(prompt): fit builders through template context

Author: mariusvniekerk

internal/prompt/prompt.go

@@ -216,7 +216,7 @@ func (b *Builder) BuildDirty(repoPath, diff string, repoID int64, contextCount i
 		}
 	}
 
-	body, err := fitDirtyPrompt(bodyLimit, view)
+	body, err := fitDirtyPromptContext(bodyLimit, templateContextFromDirtyView(view))
 	if err != nil {
 		return "", err
 	}
@@ -565,13 +565,13 @@ func (b *Builder) buildSinglePrompt(repoPath, sha string, repoID int64, contextC
 		diffView.Body = inlineDiff
 	}
 
-	body, err := fitSinglePrompt(
+	body, err := fitSinglePromptContext(
 		bodyLimit,
-		singlePromptView{
+		templateContextFromSingleView(singlePromptView{
 			Optional: optional,
 			Current:  currentView,
 			Diff:     diffView,
-		},
+		}),
 	)
 	if err != nil {
 		return "", err
@@ -683,13 +683,13 @@ func (b *Builder) buildRangePrompt(repoPath, rangeRef string, repoID int64, cont
 		diffView.Body = inlineDiff
 	}
 
-	body, err := fitRangePrompt(
+	body, err := fitRangePromptContext(
 		bodyLimit,
-		rangePromptView{
+		templateContextFromRangeView(rangePromptView{
 			Optional: optional,
 			Current:  currentView,
 			Diff:     diffView,
-		},
+		}),
 	)
 	if err != nil {
 		return "", err

internal/prompt/prompt_body_templates.go

@@ -276,33 +276,49 @@ func templateContextFromSystemView(view systemPromptView) TemplateContext {
 	return TemplateContext{System: &SystemTemplateContext{NoSkillsInstruction: view.NoSkillsInstruction, CurrentDate: view.CurrentDate}}
 }
 
+func renderSinglePromptContext(ctx TemplateContext) (string, error) {
+	return executePromptTemplate("assembled_single.md.gotmpl", ctx)
+}
+
+func renderRangePromptContext(ctx TemplateContext) (string, error) {
+	return executePromptTemplate("assembled_range.md.gotmpl", ctx)
+}
+
+func renderDirtyPromptContext(ctx TemplateContext) (string, error) {
+	return executePromptTemplate("assembled_dirty.md.gotmpl", ctx)
+}
+
+func renderAddressPromptContext(ctx TemplateContext) (string, error) {
+	return executePromptTemplate("assembled_address.md.gotmpl", ctx)
+}
+
 func renderSinglePrompt(view singlePromptView) (string, error) {
-	return executePromptTemplate("assembled_single.md.gotmpl", templateContextFromSingleView(view))
+	return renderSinglePromptContext(templateContextFromSingleView(view))
 }
 
 func renderRangePrompt(view rangePromptView) (string, error) {
-	return executePromptTemplate("assembled_range.md.gotmpl", templateContextFromRangeView(view))
+	return renderRangePromptContext(templateContextFromRangeView(view))
 }
 
 func renderDirtyPrompt(view dirtyPromptView) (string, error) {
-	return executePromptTemplate("assembled_dirty.md.gotmpl", templateContextFromDirtyView(view))
+	return renderDirtyPromptContext(templateContextFromDirtyView(view))
 }
 
 func renderAddressPrompt(view addressPromptView) (string, error) {
-	return executePromptTemplate("assembled_address.md.gotmpl", templateContextFromAddressView(view))
+	return renderAddressPromptContext(templateContextFromAddressView(view))
 }
 
-func fitSinglePrompt(limit int, view singlePromptView) (string, error) {
-	body, err := renderSinglePrompt(view)
+func fitSinglePromptContext(limit int, ctx TemplateContext) (string, error) {
+	body, err := renderSinglePromptContext(ctx)
 	if err != nil {
 		return "", err
 	}
 	if len(body) <= limit {
 		return body, nil
 	}
 
-	for trimOptionalSections(&view.Optional) {
-		body, err = renderSinglePrompt(view)
+	for ctx.Review != nil && ctx.Review.Optional.TrimNext() {
+		body, err = renderSinglePromptContext(ctx)
 		if err != nil {
 			return "", err
 		}
@@ -311,9 +327,8 @@ func fitSinglePrompt(limit int, view singlePromptView) (string, error) {
 		}
 	}
 
-	if view.Current.Message != "" {
-		view.Current.Message = ""
-		body, err = renderSinglePrompt(view)
+	if ctx.Review != nil && ctx.Review.Subject.TrimSingleMessage() {
+		body, err = renderSinglePromptContext(ctx)
 		if err != nil {
 			return "", err
 		}
@@ -322,9 +337,8 @@ func fitSinglePrompt(limit int, view singlePromptView) (string, error) {
 		}
 	}
 
-	if view.Current.Author != "" {
-		view.Current.Author = ""
-		body, err = renderSinglePrompt(view)
+	if ctx.Review != nil && ctx.Review.Subject.TrimSingleAuthor() {
+		body, err = renderSinglePromptContext(ctx)
 		if err != nil {
 			return "", err
 		}
@@ -333,10 +347,10 @@ func fitSinglePrompt(limit int, view singlePromptView) (string, error) {
 		}
 	}
 
-	for len(body) > limit && view.Current.Subject != "" {
+	for ctx.Review != nil && ctx.Review.Subject.Single != nil && len(body) > limit && ctx.Review.Subject.Single.Subject != "" {
 		overflow := len(body) - limit
-		view.Current.Subject = truncateUTF8(view.Current.Subject, max(0, len(view.Current.Subject)-overflow))
-		body, err = renderSinglePrompt(view)
+		ctx.Review.Subject.TrimSingleSubjectTo(max(0, len(ctx.Review.Subject.Single.Subject)-overflow))
+		body, err = renderSinglePromptContext(ctx)
 		if err != nil {
 			return "", err
 		}
@@ -345,76 +359,95 @@ func fitSinglePrompt(limit int, view singlePromptView) (string, error) {
 	return hardCapPrompt(body, limit), nil
 }
 
-func fitRangePrompt(limit int, view rangePromptView) (string, error) {
-	_, body, err := trimRangePromptView(limit, view)
+func fitSinglePrompt(limit int, view singlePromptView) (string, error) {
+	return fitSinglePromptContext(limit, templateContextFromSingleView(view))
+}
+
+func fitRangePromptContext(limit int, ctx TemplateContext) (string, error) {
+	_, body, err := trimRangePromptContext(limit, ctx)
 	if err != nil {
 		return "", err
 	}
 	return hardCapPrompt(body, limit), nil
 }
 
-func cloneCommitRangeSectionView(view commitRangeSectionView) commitRangeSectionView {
-	cloned := view
-	if len(view.Entries) == 0 {
-		return cloned
-	}
-	cloned.Entries = append([]commitRangeEntryView(nil), view.Entries...)
-	return cloned
+func fitRangePrompt(limit int, view rangePromptView) (string, error) {
+	return fitRangePromptContext(limit, templateContextFromRangeView(view))
 }
 
-func trimRangePromptView(limit int, view rangePromptView) (rangePromptView, string, error) {
-	view.Current = cloneCommitRangeSectionView(view.Current)
-	body, err := renderRangePrompt(view)
+func trimRangePromptContext(limit int, ctx TemplateContext) (TemplateContext, string, error) {
+	ctx = ctx.Clone()
+	body, err := renderRangePromptContext(ctx)
 	if err != nil {
-		return rangePromptView{}, "", err
+		return TemplateContext{}, "", err
 	}
 	if len(body) <= limit {
-		return view, body, nil
+		return ctx, body, nil
 	}
 
-	for trimOptionalSections(&view.Optional) {
-		body, err = renderRangePrompt(view)
+	for ctx.Review != nil && ctx.Review.Optional.TrimNext() {
+		body, err = renderRangePromptContext(ctx)
 		if err != nil {
-			return rangePromptView{}, "", err
+			return TemplateContext{}, "", err
 		}
 		if len(body) <= limit {
-			return view, body, nil
+			return ctx, body, nil
 		}
 	}
 
-	for i := len(view.Current.Entries) - 1; i >= 0 && len(body) > limit; i-- {
-		if view.Current.Entries[i].Subject == "" {
-			continue
-		}
-		view.Current.Entries[i].Subject = ""
-		body, err = renderRangePrompt(view)
+	for ctx.Review != nil && len(body) > limit && ctx.Review.Subject.BlankNextRangeSubject() {
+		body, err = renderRangePromptContext(ctx)
 		if err != nil {
-			return rangePromptView{}, "", err
+			return TemplateContext{}, "", err
 		}
 	}
 
-	for len(view.Current.Entries) > 0 && len(body) > limit {
-		view.Current.Entries = view.Current.Entries[:len(view.Current.Entries)-1]
-		body, err = renderRangePrompt(view)
+	for ctx.Review != nil && len(body) > limit && ctx.Review.Subject.DropLastRangeEntry() {
+		body, err = renderRangePromptContext(ctx)
 		if err != nil {
-			return rangePromptView{}, "", err
+			return TemplateContext{}, "", err
 		}
 	}
 
-	return view, body, nil
+	return ctx, body, nil
 }
 
-func fitDirtyPrompt(limit int, view dirtyPromptView) (string, error) {
-	body, err := renderDirtyPrompt(view)
+func trimRangePromptView(limit int, view rangePromptView) (rangePromptView, string, error) {
+	trimmed, body, err := trimRangePromptContext(limit, templateContextFromRangeView(view))
+	if err != nil {
+		return rangePromptView{}, "", err
+	}
+	rangeCtx := trimmed.Review.Subject.Range
+	if trimmed.Review == nil || rangeCtx == nil {
+		return rangePromptView{}, body, nil
+	}
+	entries := make([]commitRangeEntryView, 0, len(rangeCtx.Entries))
+	for _, entry := range rangeCtx.Entries {
+		entries = append(entries, commitRangeEntryView(entry))
+	}
+	return rangePromptView{
+		Optional: optionalSectionsView{
+			ProjectGuidelines: buildProjectGuidelinesSectionView(trimmed.Review.Optional.ProjectGuidelinesBody()),
+			AdditionalContext: trimmed.Review.Optional.AdditionalContext,
+			PreviousReviews:   previousReviewViewsFromTemplateContext(trimmed.Review.Optional.PreviousReviews),
+			PreviousAttempts:  reviewAttemptViewsFromTemplateContext(trimmed.Review.Optional.PreviousAttempts),
+		},
+		Current: commitRangeSectionView{Count: rangeCtx.Count, Entries: entries},
+		Diff:    diffSectionView{Heading: trimmed.Review.Diff.Heading, Body: trimmed.Review.Diff.Body, Fallback: trimmed.Review.Fallback.Rendered()},
+	}, body, nil
+}
+
+func fitDirtyPromptContext(limit int, ctx TemplateContext) (string, error) {
+	body, err := renderDirtyPromptContext(ctx)
 	if err != nil {
 		return "", err
 	}
 	if len(body) <= limit {
 		return body, nil
 	}
 
-	for trimOptionalSections(&view.Optional) {
-		body, err = renderDirtyPrompt(view)
+	for ctx.Review != nil && ctx.Review.Optional.TrimNext() {
+		body, err = renderDirtyPromptContext(ctx)
 		if err != nil {
 			return "", err
 		}
@@ -426,19 +459,22 @@ func fitDirtyPrompt(limit int, view dirtyPromptView) (string, error) {
 	return hardCapPrompt(body, limit), nil
 }
 
+func fitDirtyPrompt(limit int, view dirtyPromptView) (string, error) {
+	return fitDirtyPromptContext(limit, templateContextFromDirtyView(view))
+}
+
 func trimOptionalSections(view *optionalSectionsView) bool {
-	switch {
-	case len(view.PreviousAttempts) > 0:
-		view.PreviousAttempts = nil
-	case len(view.PreviousReviews) > 0:
-		view.PreviousReviews = nil
-	case view.AdditionalContext != "":
-		view.AdditionalContext = ""
-	case view.ProjectGuidelines != nil:
-		view.ProjectGuidelines = nil
-	default:
+	if view == nil {
 		return false
 	}
+	ctx := reviewOptionalContextFromView(*view)
+	if !ctx.TrimNext() {
+		return false
+	}
+	view.ProjectGuidelines = buildProjectGuidelinesSectionView(ctx.ProjectGuidelinesBody())
+	view.AdditionalContext = ctx.AdditionalContext
+	view.PreviousReviews = previousReviewViewsFromTemplateContext(ctx.PreviousReviews)
+	view.PreviousAttempts = reviewAttemptViewsFromTemplateContext(ctx.PreviousAttempts)
 	return true
 }
 
@@ -524,6 +560,21 @@ func renderDirtyTruncatedDiffFallback(body string) (string, error) {
 	return executePromptTemplate("dirty_truncated_diff_fallback", dirtyTruncatedDiffFallbackView{Body: body})
 }
 
+func previousReviewViewsFromTemplateContext(contexts []PreviousReviewTemplateContext) []previousReviewView {
+	views := make([]previousReviewView, 0, len(contexts))
+	for _, ctx := range contexts {
+		view := previousReviewView{Commit: ctx.Commit, Available: ctx.Available, Output: ctx.Output}
+		if len(ctx.Comments) > 0 {
+			view.Comments = make([]reviewCommentView, 0, len(ctx.Comments))
+			for _, comment := range ctx.Comments {
+				view.Comments = append(view.Comments, reviewCommentView(comment))
+			}
+		}
+		views = append(views, view)
+	}
+	return views
+}
+
 func previousReviewViews(contexts []HistoricalReviewContext) []previousReviewView {
 	views := make([]previousReviewView, 0, len(contexts))
 	for _, ctx := range contexts {
@@ -568,6 +619,21 @@ func renderPreviousAttemptsFromReviews(reviews []storage.Review) (string, error)
 	return renderOptionalSectionsFromView(optionalSectionsView{PreviousAttempts: reviewAttemptViews(reviews)})
 }
 
+func reviewAttemptViewsFromTemplateContext(attempts []ReviewAttemptTemplateContext) []reviewAttemptView {
+	views := make([]reviewAttemptView, 0, len(attempts))
+	for _, attempt := range attempts {
+		view := reviewAttemptView{Label: attempt.Label, Agent: attempt.Agent, When: attempt.When, Output: attempt.Output}
+		if len(attempt.Comments) > 0 {
+			view.Comments = make([]reviewCommentView, 0, len(attempt.Comments))
+			for _, comment := range attempt.Comments {
+				view.Comments = append(view.Comments, reviewCommentView(comment))
+			}
+		}
+		views = append(views, view)
+	}
+	return views
+}
+
 func previousAttemptViewsFromContexts(attempts []reviewAttemptContext) []reviewAttemptView {
 	views := make([]reviewAttemptView, 0, len(attempts))
 	for i, attempt := range attempts {

internal/prompt/template_context.go

@@ -88,6 +88,13 @@ func (o ReviewOptionalContext) IsEmpty() bool {
 		len(o.PreviousAttempts) == 0
 }
 
+func (o ReviewOptionalContext) ProjectGuidelinesBody() string {
+	if o.ProjectGuidelines == nil {
+		return ""
+	}
+	return o.ProjectGuidelines.Body
+}
+
 func (o *ReviewOptionalContext) TrimNext() bool {
 	if o == nil {
 		return false