context

package
v0.20.1 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Apr 11, 2026 License: MIT Imports: 16 Imported by: 0

Documentation

Overview

Package context provides the budget fitting algorithm for the Context Window Optimizer.

The budget fitter is responsible for selecting which source code blocks to include in the final context output, given a token budget constraint. It uses a greedy algorithm with priority-based truncation to maximize the usefulness of the context.

Algorithm Overview

The fitting process works in three phases:

  1. Sort blocks by relevance score (highest first)
  2. Greedily include blocks until budget is exhausted
  3. Apply smart truncation for overflow situations

Priority Tiers

Blocks are handled differently based on their relevance score:

  • Score >= 100 (Ring 0 definitions): Always included, truncated if needed
  • Score >= 60 (Ring 1-2 dependencies): Signature-only if space is tight
  • Score < 60 (Ring 3+ context): Excluded when budget is constrained

Truncation Strategies

When a block doesn't fit completely, two strategies are available:

  • truncateBlock: Binary search for maximum lines that fit, adds "... omitted" marker
  • truncateToSignature: Extracts just the function/type signature with body placeholder

Package context provides context window optimization for AI agents.

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func CountFileLines 

func CountFileLines(path string) (int, error)

CountFileLines returns the total line count of a file.

func EstimateFromChars 

func EstimateFromChars(chars int) int

EstimateFromChars provides a rough token estimate from character count. Useful when you want to avoid the overhead of actual tokenization. Rule of thumb: ~4 characters per token for English/code.

func Format 

func Format(result *FitResult, format OutputFormat) string

Format formats a FitResult for output.

func FormatManifest 

func FormatManifest(result *FitResult) string

FormatManifest produces a manifest-only output (no source content).

func Gather 

func Gather(ctx context.Context, storage graph.Storage, opts Options) (string, error)

Gather collects context for a task and formats it for output.

func GatherToWriter 

func GatherToWriter(ctx context.Context, storage graph.Storage, opts Options, w io.Writer) error

GatherToWriter collects context and writes directly to a writer.

Types

type ExcludedInfo 

type ExcludedInfo struct {
	File   string  // Source file path
	Tokens int     // Token count of the excluded block
	Score  float64 // Relevance score (lower scores are excluded first)
	Reason string  // Why this block was relevant (for diagnostics)
}

ExcludedInfo records details about a block that was excluded entirely. Low-priority blocks are excluded when the budget is tight.

type FitResult 

type FitResult struct {
	Task        string          // Original task description from user input
	Blocks      []SourceBlock   // Included blocks, sorted by file path for readability
	Summary     Summary         // Aggregated statistics about the fit
	Truncated   []TruncatedInfo // Blocks that were shortened to fit
	Excluded    []ExcludedInfo  // Blocks that didn't fit at all
	UsedTokens  int             // Actual tokens consumed by included blocks
	TotalBudget int             // Original token budget (includes header reserve)
}

FitResult holds the complete result of the budget fitting process. It contains both the selected blocks and metadata about what was truncated or excluded, useful for diagnostics and manifest output.

func FitBudget 

func FitBudget(blocks []SourceBlock, maxTokens int, task string, counter *TokenCounter) *FitResult

FitBudget selects source blocks that fit within the given token budget.

The algorithm processes blocks in descending score order (greedy approach), applying different strategies based on priority:

  • High priority (score >= 100): Truncate to fit if needed
  • Medium priority (score >= 60): Include signature only if tight on space
  • Low priority (score < 60): Exclude when budget is constrained

The function reserves 200 tokens for the output header/summary, so the actual content budget is maxTokens - 200.

Parameters:

  • blocks: Source blocks to consider, typically from ReadSources
  • maxTokens: Total token budget (including header reserve)
  • task: Original task description for the result
  • counter: Token counter for measuring block sizes

Returns a FitResult containing the selected blocks and metadata about what was truncated or excluded.

type JSONBlock 

type JSONBlock struct {
	File      string   `json:"file"`
	StartLine int      `json:"start_line"`
	EndLine   int      `json:"end_line"`
	Content   string   `json:"content"`
	Tokens    int      `json:"tokens"`
	Score     float64  `json:"score"`
	Reason    string   `json:"reason,omitempty"`
	Symbols   []string `json:"symbols,omitempty"`
}

JSONBlock is a source block in JSON format.

type JSONExclude 

type JSONExclude struct {
	File   string  `json:"file"`
	Tokens int     `json:"tokens"`
	Score  float64 `json:"score"`
	Reason string  `json:"reason,omitempty"`
}

JSONExclude is an excluded block in JSON format.

type JSONOutput 

type JSONOutput struct {
	Task     string        `json:"task"`
	Budget   int           `json:"budget"`
	Used     int           `json:"used"`
	Summary  JSONSummary   `json:"summary"`
	Blocks   []JSONBlock   `json:"blocks"`
	Excluded []JSONExclude `json:"excluded,omitempty"`
}

JSONOutput is the JSON output structure.

type JSONSummary 

type JSONSummary struct {
	FileCount     int      `json:"file_count"`
	SymbolCount   int      `json:"symbol_count"`
	Symbols       []string `json:"symbols"`
	Truncated     int      `json:"truncated"`
	ExcludedCount int      `json:"excluded"`
}

JSONSummary is the JSON summary structure.

type Options 

type Options struct {
	Task         string       // Task description (required)
	MaxTokens    int          // Token budget (default: 12000)
	MaxRing      Ring         // Maximum ring to expand (default: RingSiblings)
	IncludeTests bool         // Include test files (default: true)
	ScopeNodeID  string       // Optional: limit to descendants of this node
	Output       OutputFormat // Output format (default: FormatText)
	ManifestOnly bool         // Output manifest only, no source content
}

Options configures the context gathering.

func DefaultOptions 

func DefaultOptions() Options

DefaultOptions returns the default options.

type OutputFormat 

type OutputFormat int

OutputFormat represents the output format. 

const (
	FormatText OutputFormat = iota
	FormatJSON
)

type ParsedTask 

type ParsedTask struct {
	Raw      string     // Original task description
	Symbols  []string   // Extracted PascalCase/camelCase identifiers (e.g., "Storage", "NewNode")
	Keywords []string   // Meaningful lowercase words after stop-word removal
	Intent   TaskIntent // Classified intent
}

ParsedTask holds the parsed components of a task description.

func ParseTask 

func ParseTask(description string) *ParsedTask

ParseTask parses a task description and extracts symbols, keywords, and intent.

type Position 

type Position struct {
	File    string
	Line    int
	EndLine int
}

Position holds extracted position info from node data.

type ReadSourcesOptions 

type ReadSourcesOptions struct {
	ContextLines   int // Lines of context to add above each range (default: 2)
	MergeThreshold int // Merge ranges within this many lines (default: 3)
}

ReadSourcesOptions configures source reading behavior.

func DefaultReadSourcesOptions 

func DefaultReadSourcesOptions() ReadSourcesOptions

DefaultReadSourcesOptions returns the default options.

type RelevanceItem 

type RelevanceItem struct {
	Node      *graph.Node
	Score     float64
	Ring      Ring
	File      string // Resolved file path
	StartLine int
	EndLine   int
	Reason    string // Human-readable reason, e.g., "definition of Storage"
}

RelevanceItem represents a node with relevance scoring.

func Walk 

func Walk(ctx context.Context, storage graph.Storage, task *ParsedTask, opts WalkOptions) ([]RelevanceItem, error)

Walk expands outward from task symbols to find relevant nodes.

type Ring 

type Ring int

Ring represents the distance from the seed symbols in the relevance graph. 

const (
	RingDefinition Ring = 0 // The symbol itself (definitions)
	RingDirect     Ring = 1 // Direct children, implementations, containing package
	RingCallers    Ring = 2 // References/callers, test files
	RingSiblings   Ring = 3 // Same-package siblings, related types
	RingTransitive Ring = 4 // Callers of callers (one hop further)
)

func (Ring) String 

func (r Ring) String() string

String returns the string representation of the ring.

type SourceBlock 

type SourceBlock struct {
	File      string          // File path
	StartLine int             // Start line (1-indexed, inclusive)
	EndLine   int             // End line (1-indexed, inclusive)
	Content   string          // The actual source code
	Tokens    int             // Token count for this block
	Items     []RelevanceItem // Symbols contributing to this block
	MaxScore  float64         // Highest relevance score among items
	Reason    string          // Aggregated reason for this block
}

SourceBlock represents a contiguous block of source code.

func ReadSources 

func ReadSources(items []RelevanceItem, counter *TokenCounter, opts ReadSourcesOptions) ([]SourceBlock, error)

ReadSources reads source code for relevance items, merging overlapping ranges.

type Summary 

type Summary struct {
	FileCount     int      // Number of unique files with included blocks
	SymbolCount   int      // Total symbols found across all blocks
	Symbols       []string // Top symbol names (capped at 10 for display)
	TruncateCount int      // Number of blocks that were truncated
	ExcludeCount  int      // Number of blocks excluded entirely
}

Summary provides aggregated statistics about the context fit. This is displayed in the output header to give users a quick overview of what's included.

type TaskIntent 

type TaskIntent int

TaskIntent represents the inferred intent of a task. 

const (
	IntentGeneral    TaskIntent = iota // Default/unknown
	IntentModify                       // "add", "change", "update", "refactor"
	IntentFix                          // "fix", "debug", "resolve", "repair"
	IntentUnderstand                   // "explain", "how does", "what is", "understand"
	IntentAdd                          // "implement", "create", "new", "write"
	IntentRemove                       // "remove", "delete", "drop"
)

func (TaskIntent) String 

func (i TaskIntent) String() string

String returns the string representation of the intent.

type TokenCounter 

type TokenCounter struct {
	// contains filtered or unexported fields
}

TokenCounter provides token counting using cl100k_base encoding. This is an approximation for Claude tokenization (~5-15% variance).

func NewTokenCounter 

func NewTokenCounter() (*TokenCounter, error)

NewTokenCounter creates a new token counter. Uses the offline BPE loader to avoid network calls at runtime.

func (*TokenCounter) Count 

func (tc *TokenCounter) Count(text string) int

Count returns the number of tokens in the given text.

func (*TokenCounter) CountLines 

func (tc *TokenCounter) CountLines(lines []string) int

CountLines returns the total token count for a slice of lines.

type TruncatedInfo 

type TruncatedInfo struct {
	File       string // Source file path
	FullTokens int    // Original token count before truncation
	KeptTokens int    // Token count after truncation
	Reason     string // Why truncation was applied (e.g., "definition priority", "signature only")
}

TruncatedInfo records details about a block that was shortened. This helps users understand what context might be missing.

type WalkOptions 

type WalkOptions struct {
	MaxRing      Ring   // Maximum ring to expand to (default: RingSiblings)
	IncludeTests bool   // Include test files (default: true)
	ScopeNodeID  string // Optional: limit to descendants of this node
}

WalkOptions configures the graph walker.

func DefaultWalkOptions 

func DefaultWalkOptions() WalkOptions

DefaultWalkOptions returns the default walk options.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.