rony-chat-bot/bench/bench.go

// Package bench holds one-time validation benchmarks used to drive library
// decisions. Run with:
//
//	go test -bench=. ./bench/                  (modernc only, pure Go)
//	CGO_ENABLED=1 go test -bench=. ./bench/    (mattn + modernc, requires gcc)
package bench

import (
	"context"
	"database/sql"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

const schema = `
CREATE VIRTUAL TABLE IF NOT EXISTS portfolio_chunks USING fts5(
    id UNINDEXED,
    project_id UNINDEXED,
    source_file UNINDEXED,
    chunk_index UNINDEXED,
    content,
    tokenize = 'unicode61 remove_diacritics 2'
);
`

var queries = []string{
	"rony-llm-agent",
	"AI agent harness",
	"portfolio projects",
	"machine learning",
	"CLI tool for development",
	"vector database",
	"agent loop",
	"streaming response",
}

func loadChunks(tb testing.TB, dataPath string) []chunkRow {
	tb.Helper()
	files, err := filepath.Glob(filepath.Join(dataPath, "*.md"))
	if err != nil {
		tb.Fatal(err)
	}
	var rows []chunkRow
	for _, f := range files {
		b, err := os.ReadFile(f)
		if err != nil {
			tb.Fatal(err)
		}
		project := strings.TrimSuffix(filepath.Base(f), ".md")
		for i, c := range splitIntoChunks(string(b), 500, 50) {
			rows = append(rows, chunkRow{
				ID:        fmt.Sprintf("%s-chunk-%d", project, i),
				ProjectID: project,
				Source:    f,
				Index:     i,
				Content:   c,
			})
		}
	}
	return rows
}

type chunkRow struct {
	ID, ProjectID, Source string
	Index                 int
	Content               string
}

func splitIntoChunks(text string, size, overlap int) []string {
	if size <= 0 {
		return []string{text}
	}
	if overlap < 0 || overlap >= size {
		overlap = size / 10
	}
	var chunks []string
	for i := 0; i < len(text); i += size - overlap {
		end := i + size
		if end > len(text) {
			end = len(text)
		}
		if i >= end {
			break
		}
		chunks = append(chunks, text[i:end])
		if end == len(text) {
			break
		}
	}
	return chunks
}

func insertAll(ctx context.Context, db *sql.DB, rows []chunkRow) error {
	if _, err := db.ExecContext(ctx, `DELETE FROM portfolio_chunks`); err != nil {
		return err
	}
	tx, err := db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()
	stmt, err := tx.PrepareContext(ctx,
		`INSERT INTO portfolio_chunks (id, project_id, source_file, chunk_index, content) VALUES (?,?,?,?,?)`)
	if err != nil {
		return err
	}
	defer stmt.Close()
	for _, r := range rows {
		if _, err := stmt.ExecContext(ctx, r.ID, r.ProjectID, r.Source, r.Index, r.Content); err != nil {
			return err
		}
	}
	return tx.Commit()
}

func runQuery(ctx context.Context, db *sql.DB, q string) (int, time.Duration, error) {
	start := time.Now()
	rows, err := db.QueryContext(ctx, fmt.Sprintf(`
		SELECT project_id, source_file, content
		FROM portfolio_chunks
		WHERE portfolio_chunks MATCH '%s'
		ORDER BY bm25(portfolio_chunks)
		LIMIT 5
	`, sanitizeFTS5(q)))
	if err != nil {
		return 0, 0, err
	}
	defer rows.Close()
	n := 0
	for rows.Next() {
		n++
	}
	return n, time.Since(start), rows.Err()
}

// sanitizeFTS5 is the same simple wrapper used by the production code path
// (see docs/architecture.md §4.4). It keeps the benchmark comparable.
func sanitizeFTS5(q string) string {
	tokens := strings.FieldsFunc(strings.ToLower(q), func(r rune) bool {
		return !(r == '-' || r == '_' || (r >= '0' && r <= '9') ||
			(r >= 'a' && r <= 'z') || r > 0x7F)
	})
	if len(tokens) == 0 {
		return `""`
	}
	for i, t := range tokens {
		tokens[i] = `"` + t + `"*`
	}
	return strings.Join(tokens, " ")
}