feat(#90): Add per-agent statistics collection and reporting (#97)

Author: volodya-lombrozo

README.md

@@ -131,6 +131,18 @@ If multiple sources are provided, the following priority order is applied (highe
 3. `TOKEN` environment variable (deprecated)
 4. `.env` file (`DEEPSEEK_TOKEN` > `TOKEN`)
 
+## Statistics
+
+To gather interaction statistics, you can use the following command:
+
+```sh
+refrax refactor . --ai=deepseek --stats --stats-format=csv --stats-output=stats.csv
+```
+
+This command generates a `stats.csv` file containing the interaction statistics.
+The `--stats-output` and `--stats-format` parameters are optional.
+If you omit them, `refrax` will output the statistics directly to the console.
+
 ## License
 
 Licensed under the [MIT](LICENSE.txt) License.

internal/client/refrax_client.go

@@ -53,31 +53,42 @@ func (c *RefraxClient) Refactor(proj project.Project) (project.Project, error) {
 		return proj, fmt.Errorf("no java classes found in the project %s, add java files to the appropriate directory", proj)
 	}
 	log.Debug("found %d classes in the project: %v", len(classes), classes)
-	counter := &stats.Stats{}
-	ai, err := mind(c.params, counter)
+
+	criticStats := &stats.Stats{Name: "critic"}
+	criticBrain, err := mind(c.params, criticStats)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create AI instance: %w", err)
 	}
 	criticPort, err := util.FreePort()
 	if err != nil {
 		return nil, fmt.Errorf("failed to find free port for critic: %w", err)
 	}
-	ctc := critic.NewCritic(ai, criticPort)
-	ctc.Handler(countStats(counter))
+	ctc := critic.NewCritic(criticBrain, criticPort)
+	ctc.Handler(countStats(criticStats))
 
+	fixerStats := &stats.Stats{Name: "fixer"}
+	fixerBrain, err := mind(c.params, fixerStats)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create AI instance: %w", err)
+	}
 	fixerPort, err := util.FreePort()
 	if err != nil {
 		return nil, fmt.Errorf("failed to find free port for fixer: %w", err)
 	}
-	fxr := fixer.NewFixer(ai, fixerPort)
-	fxr.Handler(countStats(counter))
+	fxr := fixer.NewFixer(fixerBrain, fixerPort)
+	fxr.Handler(countStats(fixerStats))
 
+	facilitatorStats := &stats.Stats{Name: "facilitator"}
+	facilitatorBrain, err := mind(c.params, facilitatorStats)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create AI instance: %w", err)
+	}
 	facilitatorPort, err := util.FreePort()
 	if err != nil {
 		return nil, fmt.Errorf("failed to find free port for facilitator: %w", err)
 	}
-	fclttor := facilitator.NewFacilitator(ai, ctc, fxr, facilitatorPort)
-	fclttor.Handler(countStats(counter))
+	fclttor := facilitator.NewFacilitator(facilitatorBrain, ctc, fxr, facilitatorPort)
+	fclttor.Handler(countStats(facilitatorStats))
 
 	go func() {
 		faerr := fclttor.ListenAndServe()
@@ -128,7 +139,7 @@ func (c *RefraxClient) Refactor(proj project.Project) (project.Project, error) {
 		}
 	}
 	log.Info("refactoring is finished")
-	err = printStats(c.params, counter)
+	err = printStats(c.params, criticStats, fixerStats, facilitatorStats)
 	if err != nil {
 		return nil, fmt.Errorf("failed to print statistics: %w", err)
 	}
@@ -186,7 +197,7 @@ func initLogger(params *Params) {
 	}
 }
 
-func printStats(p Params, s *stats.Stats) error {
+func printStats(p Params, s ...*stats.Stats) error {
 	if p.Stats {
 		var swriter stats.Writer
 		if p.Format == "csv" {
@@ -200,7 +211,15 @@ func printStats(p Params, s *stats.Stats) error {
 			log.Info("using stdout format for statistics output")
 			swriter = stats.NewStdWriter(log.Default())
 		}
-		return swriter.Print(s)
+		var res []*stats.Stats
+		total := &stats.Stats{}
+		for _, st := range s {
+			res = append(res, st)
+			total = total.Add(st)
+		}
+		total.Name = "total"
+		res = append(res, total)
+		return swriter.Print(res...)
 	}
 	return nil
 }

internal/stats/csv_writer.go

@@ -21,21 +21,42 @@ func NewCSVWriter(path string) Writer {
 }
 
 // Print writes the statistics to a CSV file.
-func (c *csvWriter) Print(stats *Stats) error {
+func (c *csvWriter) Print(stats ...*Stats) error {
 	file, err := os.Create(c.path)
 	if err != nil {
 		return fmt.Errorf("failed to create file: %v", err)
 	}
 	defer func() { _ = file.Close() }()
 	w := csv.NewWriter(file)
 	defer w.Flush()
-	if err = w.Write([]string{"Metric", "Value"}); err != nil {
+	header := make([]string, 0)
+	header = append(header, "metric")
+	for _, s := range stats {
+		header = append(header, s.Name)
+	}
+	if err = w.Write(header); err != nil {
 		return fmt.Errorf("failed to write header: %v", err)
 	}
-	for _, v := range stats.Entries() {
-		err = w.Write([]string{v.Title, v.Value})
+	values := make(map[string][]string, 0)
+	order := make([]string, 0)
+	for _, s := range stats {
+		for _, v := range s.Entries() {
+			if _, ok := values[v.Title]; !ok {
+				line := make([]string, 0)
+				values[v.Title] = line
+				order = append(order, v.Title)
+			}
+			values[v.Title] = append(values[v.Title], v.Value)
+		}
+	}
+	for _, k := range order {
+		v := values[k]
+		line := make([]string, 0)
+		line = append(line, k)
+		line = append(line, v...)
+		err = w.Write(line)
 		if err != nil {
-			return fmt.Errorf("failed to write entry %s: %v", v.Title, err)
+			return fmt.Errorf("failed to write entry %s: %v", k, err)
 		}
 	}
 	abs, err := filepath.Abs(c.path)

internal/stats/csv_writer_test.go

@@ -19,7 +19,7 @@ func TestCSVWriter_Print_Success(t *testing.T) {
 	dir := t.TempDir()
 	p := filepath.Join(dir, "output.csv")
 	w := NewCSVWriter(p)
-	var stats Stats
+	stats := Stats{Name: "test-stats"}
 	stats.LLMReq(1*time.Second, 0, 0, 0, 0)
 	stats.LLMReq(2*time.Second, 0, 0, 0, 0)
 	stats.LLMReq(3*time.Second, 0, 0, 0, 0)
@@ -33,8 +33,32 @@ func TestCSVWriter_Print_Success(t *testing.T) {
 	lines, err := csv.NewReader(file).ReadAll()
 	require.NoError(t, err)
 	require.Len(t, lines, 27)
-	assert.Equal(t, []string{"Metric", "Value"}, lines[0])
+	assert.Equal(t, []string{"metric", "test-stats"}, lines[0])
 	assert.Equal(t, []string{"Total LLM messages asked", "3"}, lines[1])
 	assert.Equal(t, []string{"Total LLM request duration", "6s"}, lines[2])
 	assert.Equal(t, []string{"Total LLM tokens", "0"}, lines[3])
 }
+
+func TestCSVWriter_PrintSeveral_Success(t *testing.T) {
+	dir := t.TempDir()
+	p := filepath.Join(dir, "output.csv")
+	w := NewCSVWriter(p)
+	first := Stats{Name: "first-stats"}
+	first.LLMReq(3*time.Second, 0, 0, 0, 0)
+	second := Stats{Name: "second-stats"}
+	second.LLMReq(3*time.Second, 1, 1, 1, 1)
+
+	err := w.Print(&first, &second)
+
+	require.NoError(t, err)
+	file, err := os.Open(filepath.Clean(p))
+	require.NoError(t, err)
+	defer func() { _ = file.Close() }()
+	lines, err := csv.NewReader(file).ReadAll()
+	require.NoError(t, err)
+	require.Len(t, lines, 27)
+	assert.Equal(t, []string{"metric", "first-stats", "second-stats"}, lines[0])
+	assert.Equal(t, []string{"Total LLM messages asked", "1", "1"}, lines[1])
+	assert.Equal(t, []string{"Total LLM request duration", "3s", "3s"}, lines[2])
+	assert.Equal(t, []string{"Total LLM tokens", "0", "2"}, lines[3])
+}

internal/stats/stats.go

@@ -11,6 +11,9 @@ import (
 
 // Stats is a struct that contains interaction statistics between components.
 type Stats struct {
+	// name is a string that represents the name of the statistics.
+	Name string
+
 	// mu is a mutex to protect concurrent write access to stats.
 	mu sync.Mutex
 
@@ -45,9 +48,6 @@ type Stats struct {
 	a2arespbytes int
 }
 
-// StatEntry is a placeholder type. Add fields or remove if unnecessary.
-type StatEntry struct{}
-
 // AverageA2ARespBytes calculates the average number of A2A response bytes.
 func (s *Stats) AverageA2ARespBytes() float64 {
 	s.mu.Lock()
@@ -298,6 +298,41 @@ func (s *Stats) A2AReq(duration time.Duration, reqt, respt, reqb, respb int) {
 	s.a2arespbytes += respb
 }
 
+// Add combines the current Stats instance with another and returns a new Stats instance.
+// It does not mutate either of the original Stats instances.
+func (s *Stats) Add(other *Stats) *Stats {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	other.mu.Lock()
+	defer other.mu.Unlock()
+	combined := &Stats{
+		Name:          fmt.Sprintf("%s + %s", s.Name, other.Name),
+		llmreq:        append([]time.Duration{}, s.llmreq...),
+		llmreqtokens:  s.llmreqtokens,
+		llmresptokens: s.llmresptokens,
+		llmreqbytes:   s.llmreqbytes,
+		llmrespbytes:  s.llmrespbytes,
+		a2areqs:       append([]time.Duration{}, s.a2areqs...),
+		a2areqtokens:  s.a2areqtokens,
+		a2aresptokens: s.a2aresptokens,
+		a2areqbytes:   s.a2areqbytes,
+		a2arespbytes:  s.a2arespbytes,
+	}
+	// Append other durations
+	combined.llmreq = append(combined.llmreq, other.llmreq...)
+	combined.a2areqs = append(combined.a2areqs, other.a2areqs...)
+	// Sum other values
+	combined.llmreqtokens += other.llmreqtokens
+	combined.llmresptokens += other.llmresptokens
+	combined.llmreqbytes += other.llmreqbytes
+	combined.llmrespbytes += other.llmrespbytes
+	combined.a2areqtokens += other.a2areqtokens
+	combined.a2aresptokens += other.a2aresptokens
+	combined.a2areqbytes += other.a2areqbytes
+	combined.a2arespbytes += other.a2arespbytes
+	return combined
+}
+
 // Tokens counts the number of tokens in a given text using the tiktoken library.
 func Tokens(text string) (int, error) {
 	tiktoken.SetBpeLoader(tiktoken_loader.NewOfflineLoader())

internal/stats/stats_test.go

@@ -2,6 +2,7 @@ package stats
 
 import (
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -12,3 +13,130 @@ func TestTokens_Counts(t *testing.T) {
 	require.NoError(t, err, "Expected to create tokens without error")
 	assert.Equal(t, 3, tokens)
 }
+
+func TestStats_Add_Success(t *testing.T) {
+	s1 := &Stats{
+		Name:          "Stats1",
+		llmreq:        []time.Duration{time.Millisecond, 2 * time.Millisecond},
+		llmreqtokens:  10,
+		llmresptokens: 15,
+		llmreqbytes:   100,
+		llmrespbytes:  200,
+		a2areqs:       []time.Duration{3 * time.Millisecond},
+		a2areqtokens:  5,
+		a2aresptokens: 8,
+		a2areqbytes:   50,
+		a2arespbytes:  80,
+	}
+	s2 := &Stats{
+		Name:          "Stats2",
+		llmreq:        []time.Duration{500 * time.Microsecond},
+		llmreqtokens:  20,
+		llmresptokens: 25,
+		llmreqbytes:   300,
+		llmrespbytes:  400,
+		a2areqs:       []time.Duration{4 * time.Millisecond},
+		a2areqtokens:  15,
+		a2aresptokens: 18,
+		a2areqbytes:   60,
+		a2arespbytes:  90,
+	}
+	combined := s1.Add(s2)
+	require.NotNil(t, combined)
+	assert.Equal(t, "Stats1 + Stats2", combined.Name)
+	assert.Equal(t, []time.Duration{time.Millisecond, 2 * time.Millisecond, 500 * time.Microsecond}, combined.llmreq)
+	assert.Equal(t, []time.Duration{3 * time.Millisecond, 4 * time.Millisecond}, combined.a2areqs)
+	assert.Equal(t, 30, combined.llmreqtokens)
+	assert.Equal(t, 40, combined.llmresptokens)
+	assert.Equal(t, 400, combined.llmreqbytes)
+	assert.Equal(t, 600, combined.llmrespbytes)
+	assert.Equal(t, 20, combined.a2areqtokens)
+	assert.Equal(t, 26, combined.a2aresptokens)
+	assert.Equal(t, 110, combined.a2areqbytes)
+	assert.Equal(t, 170, combined.a2arespbytes)
+}
+
+func TestStats_Add_EmptyOther(t *testing.T) {
+	s := &Stats{
+		Name:          "Stats1",
+		llmreq:        []time.Duration{time.Millisecond},
+		llmreqtokens:  10,
+		llmresptokens: 15,
+		llmreqbytes:   100,
+		llmrespbytes:  200,
+		a2areqs:       []time.Duration{2 * time.Millisecond},
+		a2areqtokens:  5,
+		a2aresptokens: 8,
+		a2areqbytes:   50,
+		a2arespbytes:  80,
+	}
+	empty := &Stats{}
+
+	combined := s.Add(empty)
+
+	require.NotNil(t, combined)
+	assert.Equal(t, "Stats1 + ", combined.Name)
+	assert.Equal(t, []time.Duration{time.Millisecond}, combined.llmreq)
+	assert.Equal(t, []time.Duration{2 * time.Millisecond}, combined.a2areqs)
+	assert.Equal(t, 10, combined.llmreqtokens)
+	assert.Equal(t, 15, combined.llmresptokens)
+	assert.Equal(t, 100, combined.llmreqbytes)
+	assert.Equal(t, 200, combined.llmrespbytes)
+	assert.Equal(t, 5, combined.a2areqtokens)
+	assert.Equal(t, 8, combined.a2aresptokens)
+	assert.Equal(t, 50, combined.a2areqbytes)
+	assert.Equal(t, 80, combined.a2arespbytes)
+}
+
+func TestStats_Add_EmptySelf(t *testing.T) {
+	empty := &Stats{}
+	s := &Stats{
+		Name:          "Stats2",
+		llmreq:        []time.Duration{500 * time.Microsecond},
+		llmreqtokens:  20,
+		llmresptokens: 25,
+		llmreqbytes:   300,
+		llmrespbytes:  400,
+		a2areqs:       []time.Duration{4 * time.Millisecond},
+		a2areqtokens:  15,
+		a2aresptokens: 18,
+		a2areqbytes:   60,
+		a2arespbytes:  90,
+	}
+
+	combined := empty.Add(s)
+
+	require.NotNil(t, combined)
+	assert.Equal(t, " + Stats2", combined.Name)
+	assert.Equal(t, []time.Duration{500 * time.Microsecond}, combined.llmreq)
+	assert.Equal(t, []time.Duration{4 * time.Millisecond}, combined.a2areqs)
+	assert.Equal(t, 20, combined.llmreqtokens)
+	assert.Equal(t, 25, combined.llmresptokens)
+	assert.Equal(t, 300, combined.llmreqbytes)
+	assert.Equal(t, 400, combined.llmrespbytes)
+	assert.Equal(t, 15, combined.a2areqtokens)
+	assert.Equal(t, 18, combined.a2aresptokens)
+	assert.Equal(t, 60, combined.a2areqbytes)
+	assert.Equal(t, 90, combined.a2arespbytes)
+}
+
+func TestStats_Add_BothEmpty(t *testing.T) {
+	empty1 := &Stats{}
+	empty2 := &Stats{}
+
+	combined := empty1.Add(empty2)
+
+	combined.Name = "total"
+	require.NotNil(t, combined)
+	assert.Equal(t, "total", combined.Name)
+	assert.Empty(t, combined.llmreq)
+	assert.Empty(t, combined.a2areqs)
+	assert.Equal(t, 0, combined.llmreqtokens)
+	assert.Equal(t, 0, combined.llmresptokens)
+	assert.Equal(t, 0, combined.llmreqbytes)
+	assert.Equal(t, 0, combined.llmrespbytes)
+	assert.Equal(t, 0, combined.a2areqtokens)
+	assert.Equal(t, 0, combined.a2aresptokens)
+	assert.Equal(t, 0, combined.a2areqbytes)
+	assert.Equal(t, 0, combined.a2arespbytes)
+}