velox-bot/internal/tldr/bench_test.go

//go:build bench

// This is the local /tldr benchmark. It is excluded from normal test runs by
// the "bench" build tag; see tldrbench/README.md for how to run it. It reads a
// chat dump from disk, runs the real summarizing pipeline against an Ollama
// server (normally reached through an SSH tunnel) and writes a dashboard to
// disk. The chat is never uploaded or committed.
package tldr

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"testing"
	"time"

	"velox-bot/internal/llm"
	"velox-bot/internal/tldr/tldrbench"
)

func benchEnv(key, def string) string {
	if v := strings.TrimSpace(os.Getenv(key)); v != "" {
		return v
	}
	return def
}

// benchPath resolves a relative path against the module root (velox-bot/),
// not the package directory `go test` runs in, so the paths you type are the
// ones you would use from a shell in velox-bot/.
func benchPath(t *testing.T, p string) string {
	t.Helper()
	if p == "" || filepath.IsAbs(p) {
		return p
	}
	dir, err := os.Getwd()
	if err != nil {
		t.Fatal(err)
	}
	for {
		if _, err := os.Stat(filepath.Join(dir, "go.mod")); err == nil {
			return filepath.Join(dir, p)
		}
		parent := filepath.Dir(dir)
		if parent == dir {
			t.Fatal("go.mod not found")
		}
		dir = parent
	}
}

func benchInt(t *testing.T, key string, def int) int {
	t.Helper()
	v := benchEnv(key, "")
	if v == "" {
		return def
	}
	n, err := strconv.Atoi(v)
	if err != nil || n < 1 {
		t.Fatalf("%s must be a positive number, got %q", key, v)
	}
	return n
}

// benchClient stands in for the real client in one run: it swaps in the
// prompt being tested and puts a unique line first so Ollama's prompt cache
// cannot make a repeat run look faster than a real one.
type benchClient struct {
	c      *llm.Client
	system string // replaces systemPrompt when set
	suffix string // appended to every user message when set
}

func (b benchClient) prep(system, user string) (string, string) {
	if b.system != "" && system == systemPrompt {
		system = b.system
	}
	if b.suffix != "" {
		user += "\n" + b.suffix
	}
	return system, fmt.Sprintf("nonce %d\n%s", time.Now().UnixNano(), user)
}

func (b benchClient) Chat(ctx context.Context, system, user string) (string, error) {
	s, u := b.prep(system, user)
	return b.c.Chat(ctx, s, u)
}

func (b benchClient) ChatStream(ctx context.Context, system, user string, f func(string)) (string, error) {
	s, u := b.prep(system, user)
	return b.c.ChatStream(ctx, s, u, f)
}

func TestBench(t *testing.T) {
	chatPath := benchPath(t, benchEnv("TLDR_BENCH_CHAT", ""))
	models := splitList(benchEnv("TLDR_BENCH_MODELS", ""))
	if chatPath == "" || len(models) == 0 {
		t.Fatal("set TLDR_BENCH_CHAT (a chat dump) and TLDR_BENCH_MODELS (comma separated Ollama model names); see internal/tldr/tldrbench/README.md")
	}
	url := benchEnv("TLDR_BENCH_URL", "http://localhost:11500")
	runs := benchInt(t, "TLDR_BENCH_RUNS", 2)
	threads := benchInt(t, "TLDR_BENCH_THREADS", 8)
	outDir := benchPath(t, benchEnv("TLDR_BENCH_OUT", "bench-out"))
	omitText := benchEnv("TLDR_BENCH_OMIT_TEXT", "") == "1"

	msgs, err := tldrbench.LoadChat(chatPath)
	if err != nil {
		t.Fatal(err)
	}
	var facts []tldrbench.Fact
	if p := benchPath(t, benchEnv("TLDR_BENCH_FACTS", "")); p != "" {
		if facts, err = tldrbench.LoadFacts(p); err != nil {
			t.Fatal(err)
		}
	}

	// Prompt variants: the bot's own prompt, plus every *.txt in TLDR_BENCH_PROMPTS
	// used as a replacement system prompt.
	type variant struct{ name, system, suffix string }
	variants := []variant{{name: "current"}}
	if benchEnv("TLDR_BENCH_SKIP_CURRENT", "") == "1" {
		variants = nil // only the prompts in TLDR_BENCH_PROMPTS
	}
	if dir := benchPath(t, benchEnv("TLDR_BENCH_PROMPTS", "")); dir != "" {
		files, _ := filepath.Glob(filepath.Join(dir, "*.txt"))
		sort.Strings(files)
		for _, f := range files {
			raw, err := os.ReadFile(f)
			if err != nil {
				t.Fatal(err)
			}
			// An optional "---USER---" line splits the file into the system
			// prompt and a reminder appended to every user message.
			system, suffix, _ := strings.Cut(string(raw), "\n---USER---\n")
			variants = append(variants, variant{strings.TrimSuffix(filepath.Base(f), ".txt"), strings.TrimSpace(system), strings.TrimSpace(suffix)})
		}
	}

	if err := os.MkdirAll(outDir, 0o700); err != nil {
		t.Fatal(err)
	}
	// Handy starting point for writing variants.
	_ = os.WriteFile(filepath.Join(outDir, "current-prompt.txt"), []byte(systemPrompt), 0o600)

	if len(variants) == 0 {
		t.Fatal("TLDR_BENCH_SKIP_CURRENT=1 needs prompt files in TLDR_BENCH_PROMPTS")
	}

	tr := BuildTranscript([]ChannelBatch{{ID: "1", Name: "chat", Messages: msgs}})
	if len(tr.Chunks) == 0 {
		t.Fatal("nothing left to summarize after trimming the chat")
	}
	inChars := 0
	for _, c := range tr.Chunks {
		inChars += len(c)
	}
	source := strings.Join(tr.Chunks, "\n")

	rep := &tldrbench.Report{
		When:    time.Now(),
		Threads: threads,
		Chat: tldrbench.ChatInfo{
			Source: filepath.Base(chatPath), Messages: len(msgs), Kept: tr.Messages,
			Chunks: len(tr.Chunks), InputChars: inChars, Truncated: tr.Truncated,
		},
		Models: models,
	}
	for _, f := range facts {
		rep.Facts = append(rep.Facts, f.Name)
	}
	for _, v := range variants {
		rep.Variants = append(rep.Variants, v.name)
	}

	stamp := rep.When.Format("20060102-150405")
	htmlPath := filepath.Join(outDir, "tldr-bench-"+stamp+".html")
	jsonPath := filepath.Join(outDir, "tldr-bench-"+stamp+".json")
	save := func() {
		raw := *rep
		if omitText {
			raw.Runs = append([]tldrbench.Run(nil), rep.Runs...)
			for i := range raw.Runs {
				raw.Runs[i].Output = ""
			}
		}
		if err := raw.WriteJSON(jsonPath); err != nil {
			t.Errorf("writing %s: %v", jsonPath, err)
		}
		f, err := os.OpenFile(htmlPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
		if err != nil {
			t.Errorf("writing %s: %v", htmlPath, err)
			return
		}
		defer f.Close()
		if err := raw.WriteHTML(f); err != nil {
			t.Errorf("rendering %s: %v", htmlPath, err)
		}
	}
	// Whatever finished is saved even if the run is interrupted or a test fails.
	defer func() {
		save()
		fmt.Printf("\nDashboard: %s\nResults:   %s\n", htmlPath, jsonPath)
	}()

	fmt.Printf("chat: %d messages, %d kept, %d chars, %d chunk(s); %d model(s) x %d prompt(s) x %d run(s)\n",
		len(msgs), tr.Messages, inChars, len(tr.Chunks), len(models), len(variants), runs)

	for _, model := range models {
		client := llm.New(url, model, threads)
		for _, v := range variants {
			for n := 1; n <= runs; n++ {
				run := tldrbench.Run{Model: model, Variant: v.name, N: n}
				ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
				start := time.Now()
				var first time.Duration
				out, err := Summarize(ctx, benchClient{c: client, system: v.system, suffix: v.suffix}, tr.Chunks, nil, func(string) {
					if first == 0 {
						first = time.Since(start)
					}
				})
				cancel()
				run.Seconds = time.Since(start).Seconds()
				run.FirstText = first.Seconds()
				if err != nil {
					run.Error = err.Error()
				} else {
					run.Output = out
					run.Score = tldrbench.Evaluate(out, source, facts)
				}
				rep.Runs = append(rep.Runs, run)

				switch {
				case err != nil:
					fmt.Printf("  %-34s %-14s run %d  ERROR %v\n", model, v.name, n, err)
				case len(facts) > 0:
					fmt.Printf("  %-34s %-14s run %d  %5.1f s  facts %d/%d\n", model, v.name, n, run.Seconds, run.FactsHit, run.FactsTotal)
				default:
					fmt.Printf("  %-34s %-14s run %d  %5.1f s  %d bullets\n", model, v.name, n, run.Seconds, run.Bullets)
				}
				save() // keep the dashboard current while a long run is going
			}
		}
	}
}

func splitList(s string) []string {
	var out []string
	for _, p := range strings.Split(s, ",") {
		if p = strings.TrimSpace(p); p != "" {
			out = append(out, p)
		}
	}
	return out
}

// TestFactsSkeleton writes a starting point for the facts file: the chat lines
// that carry figures, times, days, money or links, each with suggested
// patterns. Prune and merge them by hand; see tldrbench/README.md.
func TestFactsSkeleton(t *testing.T) {
	chatPath := benchPath(t, benchEnv("TLDR_BENCH_CHAT", ""))
	if chatPath == "" {
		t.Fatal("set TLDR_BENCH_CHAT to the chat dump")
	}
	outDir := benchPath(t, benchEnv("TLDR_BENCH_OUT", "bench-out"))
	max := benchInt(t, "TLDR_BENCH_CANDIDATES", 40)

	msgs, err := tldrbench.LoadChat(chatPath)
	if err != nil {
		t.Fatal(err)
	}
	// The same trimmed lines the model will see, so patterns match its input.
	tr := BuildTranscript([]ChannelBatch{{ID: "1", Name: "chat", Messages: msgs}})
	var lines []string
	for _, c := range tr.Chunks {
		lines = append(lines, strings.Split(c, "\n")...)
	}
	facts := tldrbench.SuggestFacts(lines, max)
	if len(facts) == 0 {
		t.Fatal("no lines with figures, times, days or links found; write the facts by hand")
	}

	if err := os.MkdirAll(outDir, 0o700); err != nil {
		t.Fatal(err)
	}
	path := filepath.Join(outDir, "facts-skeleton.json")
	raw, err := json.MarshalIndent(facts, "", "  ")
	if err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(path, raw, 0o600); err != nil {
		t.Fatal(err)
	}

	fmt.Printf("%d candidate lines from %d messages:\n\n", len(facts), len(msgs))
	for i, f := range facts {
		fmt.Printf("%2d. %s\n      patterns: %s\n", i+1, f.Note, strings.Join(f.All, "  +  "))
	}
	fmt.Printf("\nSkeleton: %s\nCopy it to a new file, delete the weak candidates, merge lines about the same topic,\nrename each fact, and pass it as TLDR_BENCH_FACTS.\n", path)
}