tok

func DefaultChunkOptions ¶

func DefaultChunkOptions() ChunkOptions

DefaultChunkOptions returns sensible defaults for code chunking.

func ChunkCode ¶

func ChunkCode(source string, opts ChunkOptions) []CodeChunk

ChunkCode splits source code into semantically meaningful chunks based on language-aware boundary detection (function/class/method definitions). If a custom chunker is registered for the file extension in opts.Language, it is used instead of the default pipeline.

func ChunkCodePath ¶

func ChunkCodePath(path, source string, opts ChunkOptions) []CodeChunk

ChunkCodePath is like ChunkCode but accepts a file path for registry lookup.

func ParseCompactionResponse ¶ added in v0.1.1

func ParseCompactionResponse(response string) (*CompactionSchema, error)

ParseCompactionResponse parses an LLM's JSON response into a CompactionSchema.

func NewCompressionAdvisor ¶ added in v0.1.1

func NewCompressionAdvisor() *CompressionAdvisor

NewCompressionAdvisor creates a new CompressionAdvisor instance.

func NewCompressor ¶

func NewCompressor(opts ...Option) *Compressor

NewCompressor creates a reusable compressor.

package main

import (
	"fmt"

	"github.com/GrayCodeAI/tok"
)

func main() {
	// Create a reusable compressor for multiple calls
	c := tok.NewCompressor(
		tok.WithBudget(1000),
	)

	text1 := "First document to compress..."
	text2 := "Second document to compress..."

	_, stats1 := c.Compress(text1)
	_, stats2 := c.Compress(text2)

	fmt.Printf("Doc 1 saved: %d tokens\n", stats1.TokensSaved)
	fmt.Printf("Doc 2 saved: %d tokens\n", stats2.TokensSaved)
}

Output:

func CompressBlock ¶ added in v0.1.1

func CompressBlock(block ContentBlock, targetTokens int) ContentBlock

CompressBlock applies compression to a block proportional to how much we need to save. targetTokens is the desired token count after compression.

func NewContextOptimizer ¶ added in v0.1.1

func NewContextOptimizer(budget int) *ContextOptimizer

NewContextOptimizer creates an optimizer with the given token budget.

func NewCustomFilter ¶ added in v0.1.1

func NewCustomFilter(rules []FilterRule) *CustomFilter

NewCustomFilter compiles rules into an applicable filter. Disabled rules and rules with invalid patterns are dropped (the latter with a logged warning). Rules are stored in ascending priority order regardless of input ordering.

func GetEncodingForModel ¶ added in v0.1.1

func GetEncodingForModel(model string) Encoding

GetEncodingForModel returns the BPE encoding for the given model name. If the model is not recognized, it falls back to Cl100kBase.

func LoadFilterRules ¶ added in v0.1.1

func LoadFilterRules(path string) ([]FilterRule, error)

LoadFilterRules parses a TOML rules file at path and returns the rules in ascending priority order.

A rule with `enabled` omitted defaults to enabled. Rules with an empty pattern are dropped with a logged warning. Regex validity is NOT checked here — it is checked (once) when building a CustomFilter, so that an invalid pattern is skipped rather than failing the whole load. A malformed TOML file returns a non-nil error.

func NewHeuristicPerplexityScorer ¶ added in v0.1.1

func NewHeuristicPerplexityScorer() *HeuristicPerplexityScorer

NewHeuristicPerplexityScorer returns the default heuristic scorer with tuned weights. It is the zero-dependency default used by WithPerplexityGuided when no scorer is supplied.

func ExtractAllJSON ¶ added in v0.1.1

func ExtractAllJSON(text string) []JSONExtract

ExtractAllJSON returns every top-level balanced JSON object in text, in source order. Useful for parsing LLM output that contains multiple JSON snippets.

func ExtractJSON ¶ added in v0.1.1

func ExtractJSON(text string) JSONExtract

ExtractJSON returns the first complete JSON object found in text, or JSONExtract{Found: false} if no balanced object exists.

Handles nested objects, arrays, escaped quotes, and double-quoted strings. Skips preceding and trailing prose. Apostrophes in English prose (e.g. "Here's the JSON:") are NOT treated as string delimiters (only standard JSON double-quotes are).

Example:

out := tok.ExtractJSON(`Sure! Here's the data: {"name": "alice"}`)
// out.JSON = `{"name": "alice"}`, out.Found = true

func ExtractJSONArray ¶ added in v0.1.1

func ExtractJSONArray(text string) JSONExtract

ExtractJSONArray is the array counterpart of ExtractJSON. It looks for a balanced `[...]` instead of `{...}`.

func GetModelPricing ¶ added in v0.1.1

func GetModelPricing(model string) (ModelPricing, bool)

GetModelPricing returns the pricing for a model. Returns the pricing and true if found, or zero pricing and false if unknown.

func BalancedOptimize ¶ added in v0.1.1

func BalancedOptimize(blocks []ContentBlock, budget int) *OptimizationResult

BalancedOptimize ensures each category gets minimum representation, then fills by priority.

func GreedyOptimize ¶ added in v0.1.1

func GreedyOptimize(blocks []ContentBlock, budget int) *OptimizationResult

GreedyOptimize takes highest priority blocks until budget is full.

func PriorityOptimize ¶ added in v0.1.1

func PriorityOptimize(blocks []ContentBlock, budget int) *OptimizationResult

PriorityOptimize uses strict priority ordering, compressing before dropping.

func WithBudget ¶

func WithBudget(tokens int) Option

WithBudget sets a hard token limit for the output.

func WithCodeAware ¶ added in v0.1.1

func WithCodeAware(language string) Option

WithCodeAware marks the input as source code and activates symbol preservation. Lines that declare semantically significant symbols (function/type/export declarations and their signatures) are identified via the configured SymbolProvider and guarded so that downstream compression stages can never strip them: any protected line removed by the pipeline is re-injected into the output.

The lang argument is a language name as understood by the chunker (e.g. "go", "python", "typescript", "rust", "java"). An empty string triggers best-effort auto-detection.

The default SymbolProvider is regex-backed (RegexSymbolProvider, using the chunker's symbol patterns). Pass a custom provider — e.g. an LSP-backed one — with WithSymbolProvider.

func WithCustomFilters ¶ added in v0.1.1

func WithCustomFilters(rules []FilterRule) Option

WithCustomFilters registers user-defined regex find/replace rules (see FilterRule / LoadFilterRules) to run as a post-pipeline stage. The rules are compiled once into a CustomFilter and applied in ascending priority order to the compressed output; invalid patterns are skipped with a logged warning.

Passing nil or an empty slice is a no-op.

func WithLLMLinguaProse ¶ added in v0.1.1

func WithLLMLinguaProse() Option

WithLLMLinguaProse enables the opt-in LLMLingua prose compression layer. It drops low-information sentences from natural-language prose using a local perplexity proxy (no model calls); code and short inputs are passed through unchanged. It is disabled unless explicitly requested.

func WithMode ¶

func WithMode(m Mode) Option

WithMode sets the compression mode.

func WithModel ¶ added in v0.1.1

func WithModel(model string) Option

WithModel sets the target LLM model and auto-selects the appropriate BPE encoding. This configures both the model name (for cost calculations) and the tokenizer encoding.

Supported model prefixes:

• OpenAI: gpt-4, gpt-4o, gpt-4.1, gpt-3.5-turbo, o1, o3, o4
• Anthropic: claude-3, claude-3.5, claude-4
• Google: gemini-1.5, gemini-2.0, gemini-2.5

Unknown models fall back to cl100k_base encoding.

func WithModelEncoding ¶ added in v0.1.1

func WithModelEncoding(encoding Encoding) Option

WithModelEncoding sets a specific BPE encoding for tokenization. Use this when you need direct control over the encoding, bypassing model detection.

func WithPerplexityGuided ¶ added in v0.1.1

func WithPerplexityGuided(scorer PerplexityScorer, dropRatio float64) Option

WithPerplexityGuided enables an opt-in compression layer that drops the lowest-importance tokens (per scorer) up to dropRatio of the token count, while preserving sentence boundaries and protected symbols.

dropRatio is clamped to [0, 0.9]; values <= 0 disable the layer. If scorer is nil, the default HeuristicPerplexityScorer is used.

This runs as a final stage after the standard pipeline (it operates on the pipeline output), so it composes with tiers, modes, and budgets.

func WithQuery ¶

func WithQuery(intent string) Option

WithQuery provides intent context for goal-driven filtering.

func WithSymbolProvider ¶ added in v0.1.1

func WithSymbolProvider(p SymbolProvider) Option

WithSymbolProvider overrides the SymbolProvider used by WithCodeAware. It lets an embedder (e.g. hawk) inject a real LSP-backed provider in place of the default regex implementation. Passing a nil provider is a no-op.

func WithTier ¶

func WithTier(t Tier) Option

WithTier selects a pipeline profile.

func BuiltinProfile ¶ added in v0.1.1

func BuiltinProfile(name string) (*Profile, bool)

BuiltinProfile returns a copy of the named built-in profile and whether it exists. Returning a copy keeps the registry immutable to callers.

func LoadProfile ¶ added in v0.1.1

func LoadProfile(path string) (*Profile, error)

LoadProfile reads a single profile from a TOML file. If the profile omits a name, the file's base name (without extension) is used.

func PromptCompress ¶ added in v0.1.1

func PromptCompress(text string, intensity Intensity) (string, PromptStats)

PromptCompress applies the prompt-compression algorithm to text.

Behavior:

• Empty input returns ("", zero stats).
• Sensitive segments (security/destructive keywords) are preserved verbatim.
• CJK text is passed through unchanged.
• Multi-word drop-list entries are matched as whole phrases.
• Dictionary substitutions preserve the case of the first character.

This is independent of tok.Compress() and does not consume options/presets.

func NewRestorationTracker ¶ added in v0.1.1

func NewRestorationTracker() *RestorationTracker

NewRestorationTracker creates a tracker with default settings.

func DefaultSecretDetector ¶ added in v0.1.1

func DefaultSecretDetector() *SecretDetector

DefaultSecretDetector returns the singleton SecretDetector instance with all built-in patterns. It is safe for concurrent use.

func NewSecretDetector ¶ added in v0.1.1

func NewSecretDetector() *SecretDetector

NewSecretDetector creates a SecretDetector with all built-in patterns compiled and ready.

func Compress ¶

func Compress(text string, opts ...Option) (string, Stats)

Compress applies tok's compression pipeline to the input text.

package main

import (
	"fmt"

	"github.com/GrayCodeAI/tok"
)

func main() {
	text := `This is a very long piece of text that contains redundant information.
The purpose of this text is to demonstrate how tok can compress it.
We repeat the same ideas multiple times to inflate the token count.
The key insight is that tok uses entropy-based filtering to remove low-information tokens.
This allows us to reduce the token count while preserving the core meaning.`

	compressed, stats := tok.Compress(text)

	fmt.Printf("Original tokens: %d\n", stats.OriginalTokens)
	fmt.Printf("Final tokens: %d\n", stats.FinalTokens)
	fmt.Printf("Tokens saved: %d\n", stats.TokensSaved)
	fmt.Printf("Reduction: %.1f%%\n", stats.ReductionPercent)
	fmt.Printf("Compressed text length: %d\n", len(compressed))
}

Output:

package main

import (
	"fmt"

	"github.com/GrayCodeAI/tok"
)

func main() {
	text := `Long document with many details that we want to fit within a token budget.
We can use the budget option to limit the output to a specific number of tokens.
This is useful when working with context windows that have strict limits.`

	compressed, stats := tok.Compress(text, tok.WithBudget(50))

	fmt.Printf("Final tokens: %d (budget: 50)\n", stats.FinalTokens)
	fmt.Printf("Compressed text: %s\n", compressed[:50]+"...")
}

Output:

package main

import (
	"fmt"

	"github.com/GrayCodeAI/tok"
)

func main() {
	text := `The application uses a REST API with JSON payloads.
Database migrations are handled by the migration package.
The frontend is built with React and TypeScript.
Authentication uses JWT tokens with refresh token rotation.
The CI/CD pipeline runs on GitHub Actions.`

	// Query-aware compression preserves content relevant to the query
	_, stats := tok.Compress(text, tok.WithQuery("authentication"))

	fmt.Printf("Tokens saved: %d\n", stats.TokensSaved)
}

Output:

func CompressPerplexityGuided ¶ added in v0.1.1

func CompressPerplexityGuided(text string, scorer PerplexityScorer, dropRatio float64) (string, Stats)

CompressPerplexityGuided applies ONLY perplexity-guided token dropping to text (no other pipeline layers). It is a self-contained entry point for callers who want the LLMLingua-style drop in isolation.

If scorer is nil, the default HeuristicPerplexityScorer is used. dropRatio is clamped to [0, 0.9].

func NewStreamCompressor ¶

func NewStreamCompressor(threshold int, opts ...Option) *StreamCompressor

NewStreamCompressor creates a background compressor that keeps compressed output ready at all times. Threshold is the token count that triggers background re-compression. If threshold <= 0, it defaults to 2000 tokens.

func NewTracker ¶ added in v0.1.1

func NewTracker(ctx context.Context) (*Tracker, error)

NewTracker creates a tracker at the default path (~/.tok/tracker.db).

func NewTrackerAt ¶ added in v0.1.1

func NewTrackerAt(ctx context.Context, path string) (*Tracker, error)

NewTrackerAt creates a tracker at a custom path.

func NewUsageTracker ¶ added in v0.1.1

func NewUsageTracker() *UsageTracker

NewUsageTracker creates a UsageTracker with sensible defaults.