perf: use strings.Builder for string concatenation in loops (#1819)

Carlos Alexandro Becker created 1 month ago

Change summary

internal/agent/agent.go          | 25 ++++++++++++------
internal/agent/agent_test.go     | 36 +++++++++++++++++++++++++++
internal/message/content.go      | 16 +++++++----
internal/message/content_test.go | 45 ++++++++++++++++++++++++++++++++++
4 files changed, 107 insertions(+), 15 deletions(-)

Detailed changes

internal/agent/agent.go

@@ -561,15 +561,7 @@ func (a *sessionAgent) Summarize(ctx context.Context, sessionID string, opts fan
 		return err
 	}
 
-	summaryPromptText := "Provide a detailed summary of our conversation above."
-	if len(currentSession.Todos) > 0 {
-		summaryPromptText += "\n\n## Current Todo List\n\n"
-		for _, t := range currentSession.Todos {
-			summaryPromptText += fmt.Sprintf("- [%s] %s\n", t.Status, t.Content)
-		}
-		summaryPromptText += "\nInclude these tasks and their statuses in your summary. "
-		summaryPromptText += "Instruct the resuming assistant to use the `todos` tool to continue tracking progress on these tasks."
-	}
+	summaryPromptText := buildSummaryPrompt(currentSession.Todos)
 
 	resp, err := agent.Stream(genCtx, fantasy.AgentStreamCall{
 		Prompt:          summaryPromptText,
@@ -1101,3 +1093,18 @@ func (a *sessionAgent) workaroundProviderMediaLimitations(messages []fantasy.Mes
 
 	return convertedMessages
 }
+
+// buildSummaryPrompt constructs the prompt text for session summarization.
+func buildSummaryPrompt(todos []session.Todo) string {
+	var sb strings.Builder
+	sb.WriteString("Provide a detailed summary of our conversation above.")
+	if len(todos) > 0 {
+		sb.WriteString("\n\n## Current Todo List\n\n")
+		for _, t := range todos {
+			fmt.Fprintf(&sb, "- [%s] %s\n", t.Status, t.Content)
+		}
+		sb.WriteString("\nInclude these tasks and their statuses in your summary. ")
+		sb.WriteString("Instruct the resuming assistant to use the `todos` tool to continue tracking progress on these tasks.")
+	}
+	return sb.String()
+}

internal/agent/agent_test.go

@@ -1,6 +1,7 @@
 package agent
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
 	"runtime"
@@ -11,6 +12,7 @@ import (
 	"charm.land/x/vcr"
 	"github.com/charmbracelet/crush/internal/agent/tools"
 	"github.com/charmbracelet/crush/internal/message"
+	"github.com/charmbracelet/crush/internal/session"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -619,3 +621,37 @@ func TestCoderAgent(t *testing.T) {
 		})
 	}
 }
+
+func makeTestTodos(n int) []session.Todo {
+	todos := make([]session.Todo, n)
+	for i := range n {
+		todos[i] = session.Todo{
+			Status:  session.TodoStatusPending,
+			Content: fmt.Sprintf("Task %d: Implement feature with some description that makes it realistic", i),
+		}
+	}
+	return todos
+}
+
+func BenchmarkBuildSummaryPrompt(b *testing.B) {
+	cases := []struct {
+		name     string
+		numTodos int
+	}{
+		{"0todos", 0},
+		{"5todos", 5},
+		{"10todos", 10},
+		{"50todos", 50},
+	}
+
+	for _, tc := range cases {
+		todos := makeTestTodos(tc.numTodos)
+
+		b.Run(tc.name, func(b *testing.B) {
+			b.ReportAllocs()
+			for range b.N {
+				_ = buildSummaryPrompt(todos)
+			}
+		})
+	}
+}

internal/message/content.go

@@ -437,23 +437,27 @@ func (m *Message) AddBinary(mimeType string, data []byte) {
 }
 
 func PromptWithTextAttachments(prompt string, attachments []Attachment) string {
+	var sb strings.Builder
+	sb.WriteString(prompt)
 	addedAttachments := false
 	for _, content := range attachments {
 		if !content.IsText() {
 			continue
 		}
 		if !addedAttachments {
-			prompt += "\n<system_info>The files below have been attached by the user, consider them in your response</system_info>\n"
+			sb.WriteString("\n<system_info>The files below have been attached by the user, consider them in your response</system_info>\n")
 			addedAttachments = true
 		}
-		tag := `<file>\n`
 		if content.FilePath != "" {
-			tag = fmt.Sprintf("<file path='%s'>\n", content.FilePath)
+			fmt.Fprintf(&sb, "<file path='%s'>\n", content.FilePath)
+		} else {
+			sb.WriteString("<file>\n")
 		}
-		prompt += tag
-		prompt += "\n" + string(content.Content) + "\n</file>\n"
+		sb.WriteString("\n")
+		sb.Write(content.Content)
+		sb.WriteString("\n</file>\n")
 	}
-	return prompt
+	return sb.String()
 }
 
 func (m *Message) ToAIMessage() []fantasy.Message {

internal/message/content_test.go

@@ -0,0 +1,45 @@
+package message
+
+import (
+	"fmt"
+	"strings"
+	"testing"
+)
+
+func makeTestAttachments(n int, contentSize int) []Attachment {
+	attachments := make([]Attachment, n)
+	content := []byte(strings.Repeat("x", contentSize))
+	for i := range n {
+		attachments[i] = Attachment{
+			FilePath: fmt.Sprintf("/path/to/file%d.txt", i),
+			MimeType: "text/plain",
+			Content:  content,
+		}
+	}
+	return attachments
+}
+
+func BenchmarkPromptWithTextAttachments(b *testing.B) {
+	cases := []struct {
+		name        string
+		numFiles    int
+		contentSize int
+	}{
+		{"1file_100bytes", 1, 100},
+		{"5files_1KB", 5, 1024},
+		{"10files_10KB", 10, 10 * 1024},
+		{"20files_50KB", 20, 50 * 1024},
+	}
+
+	for _, tc := range cases {
+		attachments := makeTestAttachments(tc.numFiles, tc.contentSize)
+		prompt := "Process these files"
+
+		b.Run(tc.name, func(b *testing.B) {
+			b.ReportAllocs()
+			for range b.N {
+				_ = PromptWithTextAttachments(prompt, attachments)
+			}
+		})
+	}
+}