feat(dragoman.go): enhance document markers for clarity in translation prompts

refactor(dragoman.go): streamline response processing with newline addition
feat(internal/chunks): introduce chunking utility for text segmentation
feat(internal/chunks/chunks_test.go): add tests for text chunking functionality
refactor(internal/cli/cli.go): implement text chunking in CLI with new option
feat(internal/cli/cli.go): add `SplitChunks` option for customizable text segmentation

feat(internal/markdown): remove chunks.go to streamline markdown processing

The removal of `chunks.go` from the internal markdown package suggests a strategic decision to streamline markdown processing capabilities, possibly due to redundancy or the adoption of a more efficient approach in handling markdown segmentation. This change could indicate a shift towards utilizing external libraries or a simplification of the codebase to improve maintainability and performance.

Author: bounoable

dragoman.go

@@ -130,9 +130,9 @@ func (t *Translator) Translate(ctx context.Context, document string, opts ...Tra
 
 	prompt := heredoc.Docf(`
 		Translate the following document %sto %s:
-		---
+		---<DOC_BEGIN>---
 		%s
-		---
+		---<DOC_END>---
 
 		%s
 
@@ -149,7 +149,10 @@ func (t *Translator) Translate(ctx context.Context, document string, opts ...Tra
 		return "", err
 	}
 
-	return trimDividers(response), nil
+	response = trimDividers(response)
+	response = addNewline(response)
+
+	return response, nil
 }
 
 func trimDividers(text string) string {
@@ -160,13 +163,25 @@ func trimDividers(text string) string {
 		return text
 	}
 
-	if out[0] == "---" {
+	if out[0] == "---<DOC_BEGIN>---" {
 		out = out[1:]
 	}
 
-	if len(out) > 0 && out[len(out)-1] == "---" {
+	if len(out) > 0 && out[len(out)-1] == "---<DOC_END>---" {
 		out = out[:len(out)-1]
 	}
 
 	return strings.TrimSpace(strings.Join(out, "\n"))
 }
+
+func addNewline(text string) string {
+	if text == "" {
+		return text
+	}
+
+	if !strings.HasSuffix(text, "\n") {
+		return text + "\n"
+	}
+
+	return text
+}

internal/chunks/chunks.go

@@ -0,0 +1,48 @@
+package chunks
+
+import (
+	"strings"
+)
+
+// Chunks splits a string into segments based on line prefixes specified in a
+// slice. If no prefixes are provided, it returns the entire string as a single
+// segment. Each segment is trimmed of leading and trailing whitespace.
+func Chunks(source string, splitPrefixes []string) []string {
+	if len(splitPrefixes) == 0 {
+		return []string{source}
+	}
+
+	lines := strings.Split(source, "\n")
+
+	var chunks []string
+	var currentChunk []string
+
+	appendChunk := func() {
+		if len(currentChunk) == 0 {
+			return
+		}
+
+		chunks = append(chunks, strings.TrimSpace(strings.Join(currentChunk, "\n")))
+		currentChunk = currentChunk[:0]
+	}
+
+	for _, line := range lines {
+		if len(currentChunk) == 0 {
+			currentChunk = append(currentChunk, line)
+			continue
+		}
+
+		for _, prefix := range splitPrefixes {
+			if strings.HasPrefix(line, prefix) {
+				appendChunk()
+				break
+			}
+		}
+
+		currentChunk = append(currentChunk, line)
+	}
+
+	appendChunk()
+
+	return chunks
+}

internal/markdown/chunks_test.go internal/chunks/chunks_test.go

@@ -1,12 +1,12 @@
-package markdown_test
+package chunks_test
 
 import (
 	"strings"
 	"testing"
 
 	"github.com/MakeNowJust/heredoc/v2"
 	"github.com/google/go-cmp/cmp"
-	"github.com/modernice/dragoman/internal/markdown"
+	"github.com/modernice/dragoman/internal/chunks"
 )
 
 func TestChunks(t *testing.T) {
@@ -37,22 +37,22 @@ func TestChunks(t *testing.T) {
 	`))
 
 	tests := []struct {
-		name        string
-		splitLevels []int
-		expected    []string
+		name          string
+		splitPrefixes []string
+		expected      []string
 	}{
 		{
 			name:     "no levels",
 			expected: []string{source},
 		},
 		{
-			name:        "heading #1",
-			splitLevels: []int{1},
-			expected:    []string{source},
+			name:          "heading #1",
+			splitPrefixes: []string{"# "},
+			expected:      []string{source},
 		},
 		{
-			name:        "heading #2",
-			splitLevels: []int{2},
+			name:          "heading #2",
+			splitPrefixes: []string{"## "},
 			expected: []string{
 				takeLines(source, 3),
 				skipAndTakeLines(source, 4, 3),
@@ -62,16 +62,16 @@ func TestChunks(t *testing.T) {
 			},
 		},
 		{
-			name:        "heading #3",
-			splitLevels: []int{3},
+			name:          "heading #3",
+			splitPrefixes: []string{"### "},
 			expected: []string{
 				takeLines(source, 11),
 				skipAndTakeLines(source, 12, 11),
 			},
 		},
 		{
-			name:        "heading #1 and #2",
-			splitLevels: []int{1, 2},
+			name:          "heading #1 and #2",
+			splitPrefixes: []string{"# ", "## "},
 			expected: []string{
 				takeLines(source, 3),
 				skipAndTakeLines(source, 4, 3),
@@ -81,8 +81,8 @@ func TestChunks(t *testing.T) {
 			},
 		},
 		{
-			name:        "heading #2 and #3",
-			splitLevels: []int{2, 3},
+			name:          "heading #2 and #3",
+			splitPrefixes: []string{"## ", "### "},
 			expected: []string{
 				takeLines(source, 3),
 				skipAndTakeLines(source, 4, 3),
@@ -96,7 +96,7 @@ func TestChunks(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			chunks := markdown.Chunks(source, tt.splitLevels)
+			chunks := chunks.Chunks(source, tt.splitPrefixes)
 
 			if len(tt.expected) != len(chunks) {
 				t.Fatalf("unexpected number of chunks. want %d; got %d", len(tt.expected), len(chunks))

internal/cli/cli.go

@@ -10,23 +10,26 @@ import (
 	"io/fs"
 	"os"
 	"os/signal"
+	"strings"
 	"syscall"
 	"time"
 
 	"github.com/alecthomas/kong"
 	"github.com/modernice/dragoman"
+	"github.com/modernice/dragoman/internal/chunks"
 	"github.com/modernice/dragoman/openai"
 )
 
-var options struct {
-	SourcePath string   `arg:"source" name:"source" optional:"" help:"Source file" type:"path" env:"DRAGOMAN_SOURCE"`
-	SourceLang string   `name:"from" short:"f" help:"Source language" env:"DRAGOMAN_SOURCE_LANG" default:"auto"`
-	TargetLang string   `name:"to" short:"t" help:"Target language" env:"DRAGOMAN_TARGET_LANG" default:"English"`
-	Preserve   []string `short:"p" help:"Preserve the specified terms/words" env:"DRAGOMAN_PRESERVE"`
-	Rules      []string `name:"rule" short:"r" help:"Additional rules for the prompt" env:"DRAGOMAN_RULES"`
-	Out        string   `short:"o" help:"Output file" type:"path" env:"DRAGOMAN_OUT"`
-	Update     bool     `short:"u" help:"Only translate missing fields in output file (requires JSON files)" env:"DRAGOMAN_UPDATE"`
-	Dry        bool     `help:"Write the result to stdout" env:"DRAGOMAN_DRY_RUN"`
+type cliOptions struct {
+	SourcePath  string   `arg:"source" name:"source" optional:"" help:"Source file" type:"path" env:"DRAGOMAN_SOURCE"`
+	SourceLang  string   `name:"from" short:"f" help:"Source language" env:"DRAGOMAN_SOURCE_LANG" default:"auto"`
+	TargetLang  string   `name:"to" short:"t" help:"Target language" env:"DRAGOMAN_TARGET_LANG" default:"English"`
+	Preserve    []string `short:"p" help:"Preserve the specified terms/words" env:"DRAGOMAN_PRESERVE"`
+	Rules       []string `name:"rule" short:"r" help:"Additional rules for the prompt" env:"DRAGOMAN_RULES"`
+	Out         string   `short:"o" help:"Output file" type:"path" env:"DRAGOMAN_OUT"`
+	Update      bool     `short:"u" help:"Only translate missing fields in output file (requires JSON files)" env:"DRAGOMAN_UPDATE"`
+	Dry         bool     `help:"Write the result to stdout" env:"DRAGOMAN_DRY_RUN"`
+	SplitChunks []string `name:"split-chunks" help:"Chunk source file at lines that start with one of the provided prefixes" env:"DRAGOMAN_SPLIT_CHUNKS"`
 
 	OpenAIKey            string  `name:"openai-key" help:"OpenAI API key" env:"OPENAI_KEY"`
 	OpenAIModel          string  `name:"openai-model" help:"OpenAI model" env:"OPENAI_MODEL" default:"gpt-3.5-turbo"`
@@ -39,6 +42,8 @@ var options struct {
 	Stream  bool          `short:"s" help:"Stream output to stdout"`
 }
 
+var options cliOptions
+
 // App coordinates the translation of structured text using AI language models.
 // It sets up a command-line interface with various options to specify source
 // and target languages, preserve certain terms, apply translation rules, and
@@ -158,15 +163,23 @@ func (app *App) Run() {
 		options.SourceLang = ""
 	}
 
-	result, err := translator.Translate(
-		ctx,
-		string(source),
-		dragoman.Source(options.SourceLang),
-		dragoman.Target(options.TargetLang),
-		dragoman.Preserve(options.Preserve...),
-		dragoman.Rules(options.Rules...),
-	)
-	app.kong.FatalIfErrorf(err)
+	chunks := getChunks(string(source), options)
+
+	var results []string
+	for _, chunk := range chunks {
+		chunkResult, err := translator.Translate(
+			ctx,
+			chunk,
+			dragoman.Source(options.SourceLang),
+			dragoman.Target(options.TargetLang),
+			dragoman.Preserve(options.Preserve...),
+			dragoman.Rules(options.Rules...),
+		)
+		app.kong.FatalIfErrorf(err)
+		results = append(results, chunkResult)
+	}
+
+	result := strings.Join(results, "\n\n")
 
 	if options.Dry {
 		fmt.Fprintf(os.Stdout, "%s\n", result)
@@ -262,3 +275,15 @@ func jsonMarshal(v any) ([]byte, error) {
 	err := enc.Encode(v)
 	return buf.Bytes(), err
 }
+
+func getChunks(source string, opts cliOptions) []string {
+	if len(opts.SplitChunks) == 0 {
+		return []string{string(source)}
+	}
+
+	if opts.Verbose {
+		fmt.Fprintf(os.Stderr, "Splitting source file at lines with prefixes: %v\n", opts.SplitChunks)
+	}
+
+	return chunks.Chunks(string(source), opts.SplitChunks)
+}