Vectrum
Sign In
88 lines · 3.2 KB · Go
package tools

import (
	"context"
	"fmt"
	"strings"
)

// ContextReport is what the agent knows about its own budget. It is a plain
// struct so this package stays free of any dependency on the agent: the loop
// hands in a closure, and the formatting lives here with the other tools.
type ContextReport struct {
	// UsedTokens is the size of the last request — the whole conversation as
	// the model saw it, which is what the window is actually spent on.
	UsedTokens int
	// WindowTokens is the model's context window.
	WindowTokens int
	// CompactAt is the fraction at which the conversation gets folded into a
	// summary, so the model can see how much room it has before that happens.
	CompactAt float64
	// Turns is how many model turns this session has taken.
	Turns int
	// CostUSD is the session's spend so far, and CostKnown reports whether
	// this model's pricing is known at all. A made-up number would be worse
	// than none.
	CostUSD   float64
	CostKnown bool
}

// NewContextStatus builds the context_status tool. The model cannot otherwise
// see what it is spending: it knows neither how full the window is nor how
// many turns it has taken, and both change how it should work — a task that
// has room can be done thoroughly, one that is close to the fold should be
// finished and reported before the summary loses the detail.
func NewContextStatus(report func() ContextReport) Tool {
	if report == nil {
		return nil
	}
	return New("context_status", strings.TrimSpace(`
Report how full the context window is, how many turns this session has taken
and what it has cost. Check it before starting a long stretch of work, and when
deciding whether to finish and report rather than keep going.`),
		func(_ context.Context, _ struct{}) (string, error) {
			return formatContext(report()), nil
		})
}

func formatContext(r ContextReport) string {
	var b strings.Builder
	if r.WindowTokens > 0 {
		pct := float64(r.UsedTokens) * 100 / float64(r.WindowTokens)
		fmt.Fprintf(&b, "Context: %s of %s tokens (%.0f%% full), %s left.",
			thousands(r.UsedTokens), thousands(r.WindowTokens), pct,
			thousands(max(r.WindowTokens-r.UsedTokens, 0)))
		if r.CompactAt > 0 {
			at := int(r.CompactAt * float64(r.WindowTokens))
			if r.UsedTokens >= at {
				b.WriteString("\nThe next turn may fold the older half of this conversation into a summary.")
			} else {
				fmt.Fprintf(&b, "\nThat fold happens at %s tokens, %s away.",
					thousands(at), thousands(at-r.UsedTokens))
			}
		}
	} else {
		fmt.Fprintf(&b, "Context: %s tokens used; the window size is not known.", thousands(r.UsedTokens))
	}

	fmt.Fprintf(&b, "\nTurns so far: %d.", r.Turns)
	if r.CostKnown {
		fmt.Fprintf(&b, " Cost so far: $%.2f.", r.CostUSD)
	}
	// Said plainly because it is the thing the number is for: the whole
	// conversation is re-sent every turn, so turn count is the lever.
	b.WriteString("\nEvery turn re-sends the whole conversation, so fewer, fuller turns cost less than many small ones.")
	return b.String()
}

// thousands groups an integer for reading: 970000 becomes 970,000.
func thousands(n int) string {
	if n < 0 {
		return "0"
	}
	s := fmt.Sprintf("%d", n)
	for i := len(s) - 3; i > 0; i -= 3 {
		s = s[:i] + "," + s[i:]
	}
	return s
}