velox-bot/internal/tldr/transcript.go

package tldr

import (
	"fmt"
	"net/url"
	"path"
	"regexp"
	"strings"
	"time"
	"unicode"

	"github.com/bwmarrin/discordgo"
)

const (
	maxMessages      = 800   // fetched per run, across all channels
	maxPerMsgChars   = 400   // longer messages are cut
	maxMergedChars   = 1000  // consecutive messages from one author are merged up to this
	chunkBudgetChars = 18000 // ~5-6k tokens, leaves room for prompt+reply inside an 8k context
	// Sized from measured CPU prefill speed (~30-60 tok/s on a 6-thread
	// container): 3 chunks is a few minutes, comfortably inside the time
	// Discord gives an interaction. Anything beyond is picked up next run.
	maxChunks = 3
)

type ChannelBatch struct {
	ID       string
	Name     string
	Messages []*discordgo.Message // oldest first
}

type Transcript struct {
	Chunks []string
	// LastProcessed maps channel ID to the newest message ID that was either
	// included or deliberately skipped (bot/system/filler), i.e. how far the
	// user's cursor may safely advance.
	LastProcessed map[string]string
	Messages      int // messages that made it into the transcript
	Channels      int
	Truncated     bool
}

var (
	urlRe = regexp.MustCompile(`<?https?://[^\s<>]+>?`)
	// tagRe matches the placeholders condense writes for non-text content.
	tagRe         = regexp.MustCompile(`\[(?:link|image|gif|video|audio|file|attachment)[^\]]*\]`)
	customEmojiRe = regexp.MustCompile(`<a?:\w+:\d+>`)
	channelRe     = regexp.MustCompile(`<#\d+>`)
	roleRe        = regexp.MustCompile(`<@&\d+>`)

	// Reactions that carry no topic. Short answers like "sim" or "não" are
	// kept on purpose: they can be votes or decisions.
	filler = map[string]bool{
		"ok": true, "okay": true, "lol": true, "lmao": true, "kk": true, "kkk": true, "kkkk": true,
		"haha": true, "hahaha": true, "hehe": true, "nice": true, "boa": true, "pois": true,
		"yep": true, "thx": true, "thanks": true, "obg": true, "obrigado": true, "obrigada": true,
		"gg": true, "hm": true, "hmm": true, "ah": true, "oh": true, "wow": true, "top": true,
		"fixe": true, "mm": true, "ya": true,
	}
)

func displayName(u *discordgo.User) string {
	if u.GlobalName != "" {
		return u.GlobalName
	}
	return u.Username
}

func alnumCount(s string) int {
	n := 0
	for _, r := range s {
		if unicode.IsLetter(r) || unicode.IsDigit(r) {
			n++
		}
	}
	return n
}

const (
	maxFileNameChars = 40
	maxPreviewChars  = 80
	maxAttachTags    = 3
	maxPreviews      = 2
)

var gifHosts = map[string]bool{"tenor.com": true, "media.tenor.com": true, "c.tenor.com": true, "giphy.com": true, "media.giphy.com": true}

// mediaTag names a file by its type: the model cannot open it, but "[image]"
// or "[gif]" still says what was posted.
func mediaTag(ext string) string {
	switch ext {
	case ".gif":
		return "[gif]"
	case ".png", ".jpg", ".jpeg", ".webp", ".heic":
		return "[image]"
	case ".mp4", ".mov", ".webm", ".mkv":
		return "[video]"
	case ".mp3", ".ogg", ".wav", ".m4a", ".flac":
		return "[audio]"
	}
	return ""
}

// linkTag turns a URL into a short placeholder that keeps the site, since a
// URL is many tokens but "a YouTube link" is useful context.
func linkTag(raw string) string {
	trimmed := strings.Trim(raw, "<>")
	u, err := url.Parse(strings.TrimRight(trimmed, ".,;:!?)]\"'"))
	if err != nil || u.Hostname() == "" {
		return "[link]"
	}
	host := strings.TrimPrefix(strings.ToLower(u.Hostname()), "www.")
	if gifHosts[host] {
		return "[gif]"
	}
	if t := mediaTag(strings.ToLower(path.Ext(u.Path))); t != "" {
		return t
	}
	return "[link: " + host + "]"
}

func clip(s string, n int) string {
	if r := []rune(s); len(r) > n {
		return string(r[:n]) + "…"
	}
	return s
}

// attachmentTag describes one attachment and reports whether it says anything
// about the conversation (a named document does; a photo does not).
func attachmentTag(a *discordgo.MessageAttachment) (tag string, informative bool) {
	ct := strings.ToLower(a.ContentType)
	switch {
	case ct == "image/gif":
		return "[gif]", false
	case strings.HasPrefix(ct, "image/"):
		return "[image]", false
	case strings.HasPrefix(ct, "video/"):
		return "[video]", false
	case strings.HasPrefix(ct, "audio/"):
		return "[audio]", false
	}
	if t := mediaTag(strings.ToLower(path.Ext(a.Filename))); t != "" {
		return t, false
	}
	if a.Filename != "" {
		return "[file: " + clip(a.Filename, maxFileNameChars) + "]", true
	}
	return "[attachment]", false
}

// condense reduces a message to the text worth sending to the model, or
// reports false when it has none (bots, system messages, reactions, emoji or
// link-only lines). Links collapse to "[link: site]", attachments to
// "[image]", "[file: name]" and so on, since a URL is many tokens and rarely
// the topic. A link preview's title is kept because it names what was shared.
// Replies carry "(reply to Name)" so interleaved topics stay attributable.
func condense(m *discordgo.Message) (author, text string, ok bool) {
	if m.Author == nil || m.Author.Bot {
		return "", "", false
	}
	if m.Type != discordgo.MessageTypeDefault && m.Type != discordgo.MessageTypeReply {
		return "", "", false
	}

	content := m.Content
	for _, u := range m.Mentions {
		name := "@" + displayName(u)
		content = strings.ReplaceAll(content, "<@"+u.ID+">", name)
		content = strings.ReplaceAll(content, "<@!"+u.ID+">", name)
	}
	content = customEmojiRe.ReplaceAllString(content, "")
	content = channelRe.ReplaceAllString(content, "#channel")
	content = roleRe.ReplaceAllString(content, "@role")
	content = urlRe.ReplaceAllStringFunc(content, linkTag)
	content = strings.Join(strings.Fields(content), " ")

	bare := strings.TrimSpace(tagRe.ReplaceAllString(content, ""))
	hasText := alnumCount(bare) >= 2 && !filler[strings.ToLower(strings.Trim(bare, ".!?,;: "))]

	var extras []string
	informative := false
	for _, a := range m.Attachments {
		if len(extras) >= maxAttachTags {
			break
		}
		tag, info := attachmentTag(a)
		extras = append(extras, tag)
		informative = informative || info
	}
	previews := 0
	for _, e := range m.Embeds {
		if previews >= maxPreviews {
			break
		}
		if e == nil || strings.TrimSpace(e.Title) == "" || e.Type == discordgo.EmbedTypeImage || e.Type == discordgo.EmbedTypeGifv {
			continue
		}
		extras = append(extras, fmt.Sprintf("(preview: %q)", clip(strings.Join(strings.Fields(e.Title), " "), maxPreviewChars)))
		informative = true
		previews++
	}

	if !hasText && !informative {
		return "", "", false
	}

	content = clip(content, maxPerMsgChars)
	if len(extras) > 0 {
		content = strings.TrimSpace(content + " " + strings.Join(extras, " "))
	}
	if m.Type == discordgo.MessageTypeReply && m.ReferencedMessage != nil && m.ReferencedMessage.Author != nil &&
		m.ReferencedMessage.Author.ID != m.Author.ID {
		content = "(reply to " + displayName(m.ReferencedMessage.Author) + ") " + content
	}
	return displayName(m.Author), content, true
}

type entry struct {
	author string
	text   string
	count  int
	// gap, when set, is a pause marker written on its own line before this one.
	gap string
}

func (e entry) line() string {
	l := e.author + ": " + e.text + "\n"
	if e.gap != "" {
		l = e.gap + "\n" + l
	}
	return l
}

// gapMin is how long the chat has to go quiet before a marker is written.
const gapMin = 3 * time.Hour

// gapMarker returns the pause marker to write before a message, or "" when
// none is needed. The first message of a run gets the date, so "Sunday" in the
// text can be placed; after that only a pause of gapMin or more gets one,
// because the talk before and after a long pause is rarely the same
// conversation, which is the cheapest cue a small model has for keeping
// topics apart. Messages without a timestamp never get markers.
func gapMarker(prev, cur time.Time) string {
	if cur.IsZero() {
		return ""
	}
	if prev.IsZero() {
		return "--- " + cur.Format("Mon 2 Jan") + " ---"
	}
	if cur.Sub(prev) < gapMin {
		return ""
	}
	if y1, m1, d1 := prev.In(cur.Location()).Date(); y1 != cur.Year() || m1 != cur.Month() || d1 != cur.Day() {
		return "--- " + cur.Format("Mon 2 Jan") + " ---"
	}
	return "--- later that day ---"
}

// BuildTranscript flattens per-channel batches into prompt-sized chunks. With
// a single channel (the /tldr case) no channel header is written: it would
// only spend tokens and invite the model to tag every bullet with the channel.
// Consecutive messages from one author are merged into a single line, which
// saves the repeated "Author:" tokens. Chunks never split a line, and each
// one re-states its channel header so it reads correctly on its own. Once the
// overall budget is spent it stops, and LastProcessed only covers what made
// it in, so the rest shows up next run.
func BuildTranscript(batches []ChannelBatch) Transcript {
	t := Transcript{LastProcessed: make(map[string]string)}
	multi := len(batches) > 1
	totalBudget := maxChunks * chunkBudgetChars
	total := 0

	var cur strings.Builder
	headerIn := "" // channel whose header the current chunk already has
	flushChunk := func() {
		if cur.Len() > 0 {
			t.Chunks = append(t.Chunks, cur.String())
			cur.Reset()
		}
		headerIn = ""
	}

	for _, b := range batches {
		header := fmt.Sprintf("## #%s\n", b.Name)
		included := 0
		var pending *entry

		// write adds a finished line to the current chunk, reporting false
		// when the overall budget has no room for it.
		write := func(e entry) bool {
			line := e.line()
			need := len(line)
			if multi && headerIn != b.ID {
				need += len(header)
			}
			if total+need > totalBudget {
				t.Truncated = true
				return false
			}
			if cur.Len() > 0 && cur.Len()+need > chunkBudgetChars {
				flushChunk()
				need = len(line)
				if multi {
					need += len(header)
				}
			}
			if multi && headerIn != b.ID {
				cur.WriteString(header)
				headerIn = b.ID
			}
			cur.WriteString(line)
			total += need
			included++
			t.Messages += e.count
			return true
		}

		prevID := ""
		var lastKept time.Time // timestamp of the previous message that made it in
		for _, m := range b.Messages {
			author, text, ok := condense(m)
			if !ok {
				if pending == nil {
					t.LastProcessed[b.ID] = m.ID
				}
				prevID = m.ID
				continue
			}

			gap := gapMarker(lastKept, m.Timestamp)
			if !m.Timestamp.IsZero() {
				lastKept = m.Timestamp
			}

			if pending != nil && gap == "" && pending.author == author && len(pending.text)+len(text)+3 <= maxMergedChars {
				pending.text += " / " + text
				pending.count++
				prevID = m.ID
				continue
			}

			if pending != nil {
				if !write(*pending) {
					pending = nil
					break
				}
				// Everything up to the message before this one is now covered.
				t.LastProcessed[b.ID] = prevID
			}
			pending = &entry{author: author, text: text, count: 1, gap: gap}
			prevID = m.ID
		}

		if pending != nil {
			if write(*pending) {
				t.LastProcessed[b.ID] = prevID
			}
		}

		if included > 0 {
			t.Channels++
		}
		if t.Truncated {
			break
		}
	}
	flushChunk()
	return t
}