diff --git a/README.md b/README.md index 4885c3f..242d862 100644 --- a/README.md +++ b/README.md @@ -19,7 +19,7 @@

If OmniToken is useful in your Go project, a star helps others discover it.

-OmniToken is built for Go services that need fast local token accounting for prompt sizing, context-window planning, tokenizer experiments, and cache-boundary analysis without CGO, Rust, or Python runtime dependencies in the root module. +OmniToken is built for Go services that need fast local token accounting for prompt sizing, context-window planning, tokenizer experiments, and cacheflow analysis without CGO, Rust, or Python runtime dependencies in the root module. The root module supports Go 1.23+. Some optional comparison tooling in this repository uses dependencies that require newer Go versions. @@ -30,7 +30,7 @@ The root module supports Go 1.23+. Some optional comparison tooling in this repo - Local `Encode`, `EncodeOrdinary`, `CountTokens`, and `Decode` APIs. - Zero-allocation `CountTokens` hot path for supported OpenAI BPE workloads. - Custom WordPiece and SentencePiece-style vocabularies. -- Prompt-cache alignment planner for token block-boundary analysis. +- `cacheflow` package for prompt-cache boundary and trace analysis. - Optional adapter modules for Gemini, Llama 3, Mistral, Hugging Face `tokenizer.json`, OSS SentencePiece models, and Anthropic message token counting. ## Benchmarks @@ -115,21 +115,23 @@ Use `SpecialTokenID` or `SpecialTokens` on `*omnitoken.Engine` when constructing | Hugging Face WordPiece adapter | Optional module | | Anthropic message counter | Optional module | -## Cache Alignment +## Cacheflow ```go +import "github.com/ron2111/omnitoken/cacheflow" + engine, err := omnitoken.ForModel("gpt-4o") if err != nil { panic(err) } -report := omnitoken.NewCacheAligner(engine).AlignPromptToProfile( +report := cacheflow.NewAligner(engine).AlignPromptToProfile( systemPrompt, - omnitoken.CacheProfileOpenAI, + cacheflow.ProfileOpenAI, ) ``` -Cache alignment is informational: OmniToken does not edit prompts automatically. See [cache alignment](./docs/cache.md). +Cacheflow is informational: OmniToken does not edit prompts automatically or claim provider billing parity. See [cacheflow](./cacheflow/README.md). ## Custom Models @@ -148,7 +150,7 @@ err = omnitoken.RegisterModelPrefix("my-model-", "my_wordpiece") - [Architecture](./docs/architecture.md) - [Benchmarks and correctness](./docs/benchmarks.md) -- [Cache alignment](./docs/cache.md) +- [Cacheflow](./cacheflow/README.md) - [CLI](./docs/cli.md) - [Adapters](./adapters/README.md) diff --git a/cache.go b/cache.go deleted file mode 100644 index 14d4e98..0000000 --- a/cache.go +++ /dev/null @@ -1,118 +0,0 @@ -package omnitoken - -import "fmt" - -// CacheProfile describes token-boundary rules for cache planning. -type CacheProfile struct { - Name string - BlockSize int - MinimumTokens int -} - -// Common cache-planning profiles. Provider cache behavior can change; use these -// as local planning helpers, not billing guarantees. -var ( - CacheProfileGeneric = CacheProfile{Name: "generic", BlockSize: 1024} - CacheProfileOpenAI = CacheProfile{Name: "openai", BlockSize: 128, MinimumTokens: 1024} -) - -// CacheReport describes how close a prompt is to a cache block boundary. -type CacheReport struct { - CurrentTokens int - BlockSize int - MinimumTokens int - PreviousBlockSize int - NextBlockSize int - Remainder int - PaddingNeeded int - TokensUntilMinimum int - IsAligned bool - IsEligible bool - StrategyHint string -} - -// CacheAligner evaluates prompt lengths against provider cache block sizes. -type CacheAligner struct { - engine ModelEngine -} - -// NewCacheAligner creates a prompt cache alignment helper for an engine. -func NewCacheAligner(engine ModelEngine) *CacheAligner { - return &CacheAligner{engine: engine} -} - -// AlignPrompt evaluates prompt lengths to hit exact pricing or cache tier boundaries. -func (c *CacheAligner) AlignPrompt(text string, providerBlockSize int) CacheReport { - return c.AlignPromptToProfile(text, CacheProfile{Name: "custom", BlockSize: providerBlockSize}) -} - -// AlignPromptToProfile evaluates prompt length against a cache-planning profile. -func (c *CacheAligner) AlignPromptToProfile(text string, profile CacheProfile) CacheReport { - if c == nil || c.engine == nil { - return CacheReport{StrategyHint: "Invalid cache alignment configuration"} - } - - tokens := c.engine.CountTokens(text) - blockSize := profile.BlockSize - if blockSize <= 0 { - return CacheReport{CurrentTokens: tokens, StrategyHint: "Invalid cache alignment configuration"} - } - minimum := profile.MinimumTokens - if minimum < 0 { - minimum = 0 - } - - remainder := tokens % blockSize - previous := tokens - remainder - padding := 0 - if remainder == 0 { - previous = tokens - } else { - padding = blockSize - remainder - } - next := tokens + padding - eligible := tokens >= minimum - untilMinimum := 0 - if !eligible { - untilMinimum = minimum - tokens - if next < minimum { - next = roundUp(minimum, blockSize) - padding = next - tokens - } - } - - return CacheReport{ - CurrentTokens: tokens, - BlockSize: blockSize, - MinimumTokens: minimum, - PreviousBlockSize: previous, - NextBlockSize: next, - Remainder: remainder, - PaddingNeeded: padding, - TokensUntilMinimum: untilMinimum, - IsAligned: remainder == 0, - IsEligible: eligible, - StrategyHint: cacheStrategyHint(tokens, minimum, padding, remainder), - } -} - -func roundUp(value int, blockSize int) int { - if value <= 0 { - return 0 - } - remainder := value % blockSize - if remainder == 0 { - return value - } - return value + blockSize - remainder -} - -func cacheStrategyHint(tokens int, minimum int, padding int, remainder int) string { - if minimum > 0 && tokens < minimum { - return fmt.Sprintf("Prompt is %d tokens below the configured cache minimum", minimum-tokens) - } - if padding == 0 && remainder == 0 { - return "Prompt is aligned to the configured cache block boundary" - } - return fmt.Sprintf("Prompt is %d tokens from the next configured cache block boundary", padding) -} diff --git a/cacheflow/README.md b/cacheflow/README.md new file mode 100644 index 0000000..0747de9 --- /dev/null +++ b/cacheflow/README.md @@ -0,0 +1,81 @@ +# Cacheflow + +`cacheflow` is OmniToken's dependency-free prompt-cache analysis package. + +It is built for Go teams that want to understand whether rendered prompts have stable token prefixes before sending them to OpenAI, Anthropic, Gemini, or another provider. It does not claim billing parity and it does not call provider APIs. + +## What It Does + +- Counts prompt tokens with any `omnitoken.ModelEngine`. +- Calculates cache-boundary alignment for one prompt. +- Reads JSONL prompt traces. +- Finds common token prefixes across repeated prompts. +- Estimates reusable prefix tokens under a local cache profile. +- Emits best-effort cache-breaker hints for timestamps, UUIDs, request IDs, and dynamic metadata. + +## What It Does Not Do + +- It does not edit prompts automatically. +- It does not guarantee provider billing behavior. +- It does not require network calls or credentials. +- It does not add dependencies to the root module. + +## Align One Prompt + +```go +engine, err := omnitoken.ForModel("gpt-4o") +if err != nil { + panic(err) +} + +report := cacheflow.NewAligner(engine).AlignPromptToProfile(systemPrompt, cacheflow.ProfileOpenAI) +fmt.Println(report.CurrentTokens, report.PaddingNeeded) +``` + +## Simulate A Trace + +```go +items := []cacheflow.TraceItem{ + {ID: "1", Prompt: stablePrefix + "user question one"}, + {ID: "2", Prompt: stablePrefix + "user question two"}, +} + +report := cacheflow.Simulate(engine, items, cacheflow.SimulationOptions{ + Profile: cacheflow.ProfileOpenAI, + DetectBreakers: true, +}) +fmt.Println(report.ReusablePrefixTokens) +``` + +## JSONL Format + +Raw rendered prompts: + +```json +{"id":"1","model":"gpt-4o","prompt":"..."} +{"id":"2","model":"gpt-4o","prompt":"..."} +``` + +Structured prompt parts: + +```json +{"id":"1","model":"gpt-4o","parts":[{"name":"system","stable":true,"text":"..."},{"name":"user","stable":false,"text":"..."}]} +``` + +For structured parts, `cacheflow` concatenates `parts[].text` in order and uses `stable` only for diagnostics. + +## CLI + +```powershell +omni cache -model gpt-4o -profile openai "hello world" +omni cache-sim -model gpt-4o -profile openai -input prompts.jsonl -breakers +``` + +## Profiles + +```go +cacheflow.ProfileGeneric +cacheflow.ProfileOpenAI +``` + +Profiles are local planning helpers. Provider behavior can change, and final usage metadata remains authoritative. diff --git a/cacheflow/align.go b/cacheflow/align.go new file mode 100644 index 0000000..91caead --- /dev/null +++ b/cacheflow/align.go @@ -0,0 +1,126 @@ +// Package cacheflow provides dependency-free prompt-cache planning and trace analysis. +package cacheflow + +import ( + "fmt" + + "github.com/ron2111/omnitoken" +) + +// Profile describes local token-boundary rules for cache planning. +type Profile struct { + Name string `json:"name"` + BlockSize int `json:"block_size"` + MinimumTokens int `json:"minimum_tokens"` +} + +// Common cache-planning profiles. Provider behavior can change; use these as +// local planning helpers, not billing guarantees. +var ( + ProfileGeneric = Profile{Name: "generic", BlockSize: 1024} + ProfileOpenAI = Profile{Name: "openai", BlockSize: 128, MinimumTokens: 1024} +) + +// Alignment describes how close a prompt or prefix is to a cache block boundary. +type Alignment struct { + CurrentTokens int `json:"current_tokens"` + BlockSize int `json:"block_size"` + MinimumTokens int `json:"minimum_tokens"` + PreviousBlockSize int `json:"previous_block_size"` + NextBlockSize int `json:"next_block_size"` + Remainder int `json:"remainder"` + PaddingNeeded int `json:"padding_needed"` + TokensUntilMinimum int `json:"tokens_until_minimum"` + IsAligned bool `json:"is_aligned"` + IsEligible bool `json:"is_eligible"` + StrategyHint string `json:"strategy_hint"` +} + +// Aligner evaluates prompt lengths against provider cache block sizes. +type Aligner struct { + engine omnitoken.ModelEngine +} + +// NewAligner creates a prompt-cache alignment helper for an engine. +func NewAligner(engine omnitoken.ModelEngine) *Aligner { + return &Aligner{engine: engine} +} + +// AlignPrompt evaluates prompt length against a custom cache block size. +func (a *Aligner) AlignPrompt(text string, providerBlockSize int) Alignment { + return a.AlignPromptToProfile(text, Profile{Name: "custom", BlockSize: providerBlockSize}) +} + +// AlignPromptToProfile evaluates prompt length against a cache-planning profile. +func (a *Aligner) AlignPromptToProfile(text string, profile Profile) Alignment { + if a == nil || a.engine == nil { + return Alignment{StrategyHint: "Invalid cache alignment configuration"} + } + return AlignTokenCount(a.engine.CountTokens(text), profile) +} + +// AlignTokenCount evaluates an already-computed token count against a profile. +func AlignTokenCount(tokens int, profile Profile) Alignment { + blockSize := profile.BlockSize + if blockSize <= 0 { + return Alignment{CurrentTokens: tokens, StrategyHint: "Invalid cache alignment configuration"} + } + minimum := profile.MinimumTokens + if minimum < 0 { + minimum = 0 + } + + remainder := tokens % blockSize + previous := tokens - remainder + padding := 0 + if remainder == 0 { + previous = tokens + } else { + padding = blockSize - remainder + } + next := tokens + padding + eligible := tokens >= minimum + untilMinimum := 0 + if !eligible { + untilMinimum = minimum - tokens + if next < minimum { + next = roundUp(minimum, blockSize) + padding = next - tokens + } + } + + return Alignment{ + CurrentTokens: tokens, + BlockSize: blockSize, + MinimumTokens: minimum, + PreviousBlockSize: previous, + NextBlockSize: next, + Remainder: remainder, + PaddingNeeded: padding, + TokensUntilMinimum: untilMinimum, + IsAligned: remainder == 0, + IsEligible: eligible, + StrategyHint: strategyHint(tokens, minimum, padding, remainder), + } +} + +func roundUp(value int, blockSize int) int { + if value <= 0 { + return 0 + } + remainder := value % blockSize + if remainder == 0 { + return value + } + return value + blockSize - remainder +} + +func strategyHint(tokens int, minimum int, padding int, remainder int) string { + if minimum > 0 && tokens < minimum { + return fmt.Sprintf("Prompt is %d tokens below the configured cache minimum", minimum-tokens) + } + if padding == 0 && remainder == 0 { + return "Prompt is aligned to the configured cache block boundary" + } + return fmt.Sprintf("Prompt is %d tokens from the next configured cache block boundary", padding) +} diff --git a/cache_test.go b/cacheflow/align_test.go similarity index 70% rename from cache_test.go rename to cacheflow/align_test.go index d051967..e6ac1aa 100644 --- a/cache_test.go +++ b/cacheflow/align_test.go @@ -1,4 +1,4 @@ -package omnitoken +package cacheflow import "testing" @@ -8,18 +8,18 @@ func (e fixedCountEngine) EncodeOrdinary(string) []int { return nil } func (e fixedCountEngine) Decode([]int) string { return "" } func (e fixedCountEngine) CountTokens(string) int { return int(e) } -func TestCacheAlignerAlignPrompt(t *testing.T) { +func TestAlignPrompt(t *testing.T) { tests := []struct { name string count int block int - want CacheReport + want Alignment }{ { name: "already aligned", count: 1024, block: 128, - want: CacheReport{ + want: Alignment{ CurrentTokens: 1024, BlockSize: 128, PreviousBlockSize: 1024, @@ -32,7 +32,7 @@ func TestCacheAlignerAlignPrompt(t *testing.T) { name: "near boundary", count: 990, block: 128, - want: CacheReport{ + want: Alignment{ CurrentTokens: 990, BlockSize: 128, PreviousBlockSize: 896, @@ -46,20 +46,20 @@ func TestCacheAlignerAlignPrompt(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - got := NewCacheAligner(fixedCountEngine(tt.count)).AlignPrompt("ignored", tt.block) - assertCacheReport(t, got, tt.want) + got := NewAligner(fixedCountEngine(tt.count)).AlignPrompt("ignored", tt.block) + assertAlignment(t, got, tt.want) }) } } -func TestCacheAlignerProfileMinimum(t *testing.T) { - report := NewCacheAligner(fixedCountEngine(900)).AlignPromptToProfile("ignored", CacheProfile{ +func TestAlignPromptProfileMinimum(t *testing.T) { + report := NewAligner(fixedCountEngine(900)).AlignPromptToProfile("ignored", Profile{ Name: "test", BlockSize: 128, MinimumTokens: 1024, }) - want := CacheReport{ + want := Alignment{ CurrentTokens: 900, BlockSize: 128, MinimumTokens: 1024, @@ -70,20 +70,20 @@ func TestCacheAlignerProfileMinimum(t *testing.T) { TokensUntilMinimum: 124, IsEligible: false, } - assertCacheReport(t, report, want) + assertAlignment(t, report, want) } -func TestCacheAlignerInvalidConfig(t *testing.T) { - if got := NewCacheAligner(nil).AlignPrompt("x", 128); got.StrategyHint == "" { +func TestAlignInvalidConfig(t *testing.T) { + if got := NewAligner(nil).AlignPrompt("x", 128); got.StrategyHint == "" { t.Fatal("nil engine returned empty strategy hint") } - if got := NewCacheAligner(fixedCountEngine(10)).AlignPrompt("x", 0); got.StrategyHint == "" { + if got := NewAligner(fixedCountEngine(10)).AlignPrompt("x", 0); got.StrategyHint == "" { t.Fatal("invalid block returned empty strategy hint") } } -func TestCacheProfileOpenAI(t *testing.T) { - report := NewCacheAligner(fixedCountEngine(1025)).AlignPromptToProfile("ignored", CacheProfileOpenAI) +func TestProfileOpenAI(t *testing.T) { + report := NewAligner(fixedCountEngine(1025)).AlignPromptToProfile("ignored", ProfileOpenAI) if report.BlockSize != 128 { t.Fatalf("OpenAI block size = %d, want 128", report.BlockSize) } @@ -95,7 +95,7 @@ func TestCacheProfileOpenAI(t *testing.T) { } } -func assertCacheReport(t *testing.T, got CacheReport, want CacheReport) { +func assertAlignment(t *testing.T, got Alignment, want Alignment) { t.Helper() if got.CurrentTokens != want.CurrentTokens || got.BlockSize != want.BlockSize || @@ -107,7 +107,7 @@ func assertCacheReport(t *testing.T, got CacheReport, want CacheReport) { got.TokensUntilMinimum != want.TokensUntilMinimum || got.IsAligned != want.IsAligned || got.IsEligible != want.IsEligible { - t.Fatalf("CacheReport = %+v, want %+v", got, want) + t.Fatalf("Alignment = %+v, want %+v", got, want) } if got.StrategyHint == "" { t.Fatal("empty StrategyHint") diff --git a/cacheflow/breakers.go b/cacheflow/breakers.go new file mode 100644 index 0000000..1528f28 --- /dev/null +++ b/cacheflow/breakers.go @@ -0,0 +1,86 @@ +package cacheflow + +import "strings" + +// DetectBreakers returns best-effort warnings for dynamic data commonly placed +// before stable prompt content. These are heuristics, not provider guarantees. +func DetectBreakers(items []TraceItem) []CacheBreaker { + var out []CacheBreaker + for _, item := range items { + if len(item.Parts) > 0 { + for _, part := range item.Parts { + if part.Stable { + continue + } + out = append(out, detectTextBreakers(item.ID, part.Name, part.Text)...) + } + continue + } + out = append(out, detectTextBreakers(item.ID, "", item.Prompt)...) + } + return out +} + +func detectTextBreakers(itemID string, part string, text string) []CacheBreaker { + checks := []struct { + kind string + ok bool + msg string + }{ + {"timestamp", hasLikelyTimestamp(text), "Likely timestamp detected; keep dynamic timestamps after stable cacheable prompt sections."}, + {"uuid", hasLikelyUUID(text), "Likely UUID/request ID detected; keep per-request identifiers after stable cacheable prompt sections."}, + {"random_json", hasLikelyDynamicJSON(text), "Likely dynamic JSON metadata detected; keep changing metadata after stable instructions and context."}, + } + out := make([]CacheBreaker, 0, len(checks)) + for _, check := range checks { + if check.ok { + out = append(out, CacheBreaker{ItemID: itemID, Part: part, Kind: check.kind, Message: check.msg}) + } + } + return out +} + +func hasLikelyTimestamp(text string) bool { + for i := 0; i+19 <= len(text); i++ { + if isDigit(text[i]) && isDigit(text[i+1]) && isDigit(text[i+2]) && isDigit(text[i+3]) && + text[i+4] == '-' && isDigit(text[i+5]) && isDigit(text[i+6]) && text[i+7] == '-' && + isDigit(text[i+8]) && isDigit(text[i+9]) && (text[i+10] == 'T' || text[i+10] == ' ') && + isDigit(text[i+11]) && isDigit(text[i+12]) && text[i+13] == ':' && + isDigit(text[i+14]) && isDigit(text[i+15]) && text[i+16] == ':' && + isDigit(text[i+17]) && isDigit(text[i+18]) { + return true + } + } + return false +} + +func hasLikelyUUID(text string) bool { + for i := 0; i+36 <= len(text); i++ { + if isHexRun(text[i:i+8]) && text[i+8] == '-' && isHexRun(text[i+9:i+13]) && text[i+13] == '-' && + isHexRun(text[i+14:i+18]) && text[i+18] == '-' && isHexRun(text[i+19:i+23]) && text[i+23] == '-' && + isHexRun(text[i+24:i+36]) { + return true + } + } + return false +} + +func hasLikelyDynamicJSON(text string) bool { + lower := strings.ToLower(text) + return strings.Contains(lower, "\"request_id\"") || strings.Contains(lower, "\"timestamp\"") || strings.Contains(lower, "\"trace_id\"") || strings.Contains(lower, "\"session_id\"") +} + +func isDigit(b byte) bool { return '0' <= b && b <= '9' } + +func isHexRun(s string) bool { + if s == "" { + return false + } + for i := 0; i < len(s); i++ { + b := s[i] + if !isDigit(b) && !(('a' <= b && b <= 'f') || ('A' <= b && b <= 'F')) { + return false + } + } + return true +} diff --git a/cache_example_test.go b/cacheflow/example_test.go similarity index 56% rename from cache_example_test.go rename to cacheflow/example_test.go index 46d8e1e..3d1c75b 100644 --- a/cache_example_test.go +++ b/cacheflow/example_test.go @@ -1,18 +1,19 @@ -package omnitoken_test +package cacheflow_test import ( "fmt" "github.com/ron2111/omnitoken" + "github.com/ron2111/omnitoken/cacheflow" ) -func ExampleCacheAligner_AlignPromptToProfile() { +func ExampleAligner_AlignPromptToProfile() { engine, err := omnitoken.ForModel("gpt-4o") if err != nil { panic(err) } - report := omnitoken.NewCacheAligner(engine).AlignPromptToProfile("hello world", omnitoken.CacheProfileOpenAI) + report := cacheflow.NewAligner(engine).AlignPromptToProfile("hello world", cacheflow.ProfileOpenAI) fmt.Println(report.CurrentTokens) fmt.Println(report.PaddingNeeded) diff --git a/cacheflow/simulate.go b/cacheflow/simulate.go new file mode 100644 index 0000000..d3c87d1 --- /dev/null +++ b/cacheflow/simulate.go @@ -0,0 +1,193 @@ +package cacheflow + +import ( + "bufio" + "encoding/json" + "fmt" + "io" + "strings" + + "github.com/ron2111/omnitoken" +) + +// TraceItem is one rendered prompt or structured prompt in a workload trace. +type TraceItem struct { + ID string `json:"id,omitempty"` + Model string `json:"model,omitempty"` + Prompt string `json:"prompt,omitempty"` + Parts []PromptPart `json:"parts,omitempty"` +} + +// PromptPart describes a named prompt section. Stable is user-provided metadata +// used for diagnostics; simulation always analyzes the rendered token prefix. +type PromptPart struct { + Name string `json:"name,omitempty"` + Stable bool `json:"stable,omitempty"` + Text string `json:"text,omitempty"` +} + +// SimulationOptions controls prompt-cache trace analysis. +type SimulationOptions struct { + Profile Profile `json:"profile"` + DetectBreakers bool `json:"detect_breakers"` +} + +// SimulationReport summarizes token-level cache planning for a trace. +type SimulationReport struct { + Profile Profile `json:"profile"` + Items []ItemReport `json:"items"` + TotalTokens int `json:"total_tokens"` + CommonPrefixTokens int `json:"common_prefix_tokens"` + ReusablePrefixTokens int `json:"reusable_prefix_tokens"` + ReusablePrefixPercent float64 `json:"reusable_prefix_percent"` + CacheBreakerHints []CacheBreaker `json:"cache_breaker_hints,omitempty"` + Summary string `json:"summary"` +} + +// ItemReport describes one trace item in a simulation report. +type ItemReport struct { + ID string `json:"id,omitempty"` + Model string `json:"model,omitempty"` + PromptTokens int `json:"prompt_tokens"` + CommonPrefixTokens int `json:"common_prefix_tokens"` + DynamicSuffixTokens int `json:"dynamic_suffix_tokens"` + Alignment Alignment `json:"alignment"` +} + +// CacheBreaker is a best-effort warning about dynamic data that can reduce cache hits. +type CacheBreaker struct { + ItemID string `json:"item_id,omitempty"` + Part string `json:"part,omitempty"` + Kind string `json:"kind"` + Message string `json:"message"` +} + +// Simulate analyzes a batch of rendered prompts for stable token prefixes. +func Simulate(engine omnitoken.ModelEngine, items []TraceItem, opts SimulationOptions) SimulationReport { + profile := opts.Profile + if profile.BlockSize == 0 { + profile = ProfileOpenAI + } + if engine == nil || len(items) == 0 { + return SimulationReport{Profile: profile, Summary: "No trace items to analyze"} + } + + tokensByItem := make([][]int, len(items)) + reports := make([]ItemReport, len(items)) + totalTokens := 0 + for i, item := range items { + prompt := item.RenderedPrompt() + tokens := engine.EncodeOrdinary(prompt) + tokensByItem[i] = tokens + totalTokens += len(tokens) + reports[i] = ItemReport{ + ID: item.ID, + Model: item.Model, + PromptTokens: len(tokens), + Alignment: AlignTokenCount(len(tokens), profile), + } + } + + commonPrefix := commonTokenPrefix(tokensByItem) + reusablePrefix := reusablePrefix(commonPrefix, profile) + for i := range reports { + reports[i].CommonPrefixTokens = commonPrefix + reports[i].DynamicSuffixTokens = reports[i].PromptTokens - commonPrefix + if reports[i].DynamicSuffixTokens < 0 { + reports[i].DynamicSuffixTokens = 0 + } + } + + percent := 0.0 + if totalTokens > 0 && len(items) > 1 { + percent = float64(reusablePrefix*(len(items)-1)) / float64(totalTokens) * 100 + } + report := SimulationReport{ + Profile: profile, + Items: reports, + TotalTokens: totalTokens, + CommonPrefixTokens: commonPrefix, + ReusablePrefixTokens: reusablePrefix, + ReusablePrefixPercent: percent, + Summary: simulationSummary(commonPrefix, reusablePrefix, profile), + } + if opts.DetectBreakers { + report.CacheBreakerHints = DetectBreakers(items) + } + return report +} + +// RenderedPrompt returns Prompt when present, otherwise concatenates Parts. +func (i TraceItem) RenderedPrompt() string { + if i.Prompt != "" || len(i.Parts) == 0 { + return i.Prompt + } + var b strings.Builder + for _, part := range i.Parts { + b.WriteString(part.Text) + } + return b.String() +} + +// ReadJSONL reads cacheflow trace items from newline-delimited JSON. +func ReadJSONL(r io.Reader) ([]TraceItem, error) { + scanner := bufio.NewScanner(r) + scanner.Buffer(make([]byte, 0, 64*1024), 16*1024*1024) + var items []TraceItem + line := 0 + for scanner.Scan() { + line++ + raw := strings.TrimSpace(scanner.Text()) + raw = strings.TrimPrefix(raw, "\ufeff") + if raw == "" { + continue + } + var item TraceItem + if err := json.Unmarshal([]byte(raw), &item); err != nil { + return nil, fmt.Errorf("cacheflow JSONL line %d: %w", line, err) + } + items = append(items, item) + } + if err := scanner.Err(); err != nil { + return nil, err + } + return items, nil +} + +func commonTokenPrefix(tokensByItem [][]int) int { + if len(tokensByItem) == 0 { + return 0 + } + limit := len(tokensByItem[0]) + for _, tokens := range tokensByItem[1:] { + if len(tokens) < limit { + limit = len(tokens) + } + } + for i := 0; i < limit; i++ { + want := tokensByItem[0][i] + for _, tokens := range tokensByItem[1:] { + if tokens[i] != want { + return i + } + } + } + return limit +} + +func reusablePrefix(tokens int, profile Profile) int { + if profile.BlockSize <= 0 || tokens < profile.MinimumTokens { + return 0 + } + return tokens - tokens%profile.BlockSize +} + +func simulationSummary(commonPrefix int, reusablePrefix int, profile Profile) string { + if reusablePrefix > 0 { + return fmt.Sprintf("%d common prefix tokens are reusable at %s cache boundaries", reusablePrefix, profile.Name) + } + if commonPrefix > 0 { + return fmt.Sprintf("%d common prefix tokens found, below the configured reusable cache boundary", commonPrefix) + } + return "No common token prefix found across trace items" +} diff --git a/cacheflow/simulate_test.go b/cacheflow/simulate_test.go new file mode 100644 index 0000000..c6aff4a --- /dev/null +++ b/cacheflow/simulate_test.go @@ -0,0 +1,74 @@ +package cacheflow + +import ( + "strings" + "testing" +) + +type wordEngine struct{} + +func (wordEngine) EncodeOrdinary(text string) []int { + fields := strings.Fields(text) + ids := make([]int, len(fields)) + for i, field := range fields { + ids[i] = wordID(field) + } + return ids +} + +func (wordEngine) Decode([]int) string { return "" } +func (wordEngine) CountTokens(text string) int { + return len(strings.Fields(text)) +} + +func TestSimulateCommonPrefix(t *testing.T) { + items := []TraceItem{ + {ID: "a", Prompt: "stable prefix policy user one"}, + {ID: "b", Prompt: "stable prefix policy user two"}, + } + report := Simulate(wordEngine{}, items, SimulationOptions{Profile: Profile{Name: "test", BlockSize: 2, MinimumTokens: 2}}) + if report.CommonPrefixTokens != 4 { + t.Fatalf("CommonPrefixTokens = %d, want 4", report.CommonPrefixTokens) + } + if report.ReusablePrefixTokens != 4 { + t.Fatalf("ReusablePrefixTokens = %d, want 4", report.ReusablePrefixTokens) + } + if len(report.Items) != 2 || report.Items[0].DynamicSuffixTokens != 1 { + t.Fatalf("Items = %+v", report.Items) + } +} + +func TestReadJSONLAndBreakers(t *testing.T) { + input := strings.NewReader(`{"id":"1","parts":[{"name":"meta","stable":false,"text":"timestamp 2026-07-11T10:00:00Z request 123e4567-e89b-12d3-a456-426614174000"},{"name":"system","stable":true,"text":"stable"}]} +{"id":"2","prompt":"stable prompt"} +`) + items, err := ReadJSONL(input) + if err != nil { + t.Fatal(err) + } + if len(items) != 2 { + t.Fatalf("items len = %d", len(items)) + } + hints := DetectBreakers(items) + if len(hints) < 2 { + t.Fatalf("hints = %+v, want timestamp and uuid", hints) + } +} + +func TestStructuredPartsRender(t *testing.T) { + item := TraceItem{Parts: []PromptPart{{Text: "hello "}, {Text: "world"}}} + if got := item.RenderedPrompt(); got != "hello world" { + t.Fatalf("RenderedPrompt = %q", got) + } +} + +func wordID(s string) int { + h := 0 + for i := 0; i < len(s); i++ { + h = h*31 + int(s[i]) + } + if h < 0 { + return -h + } + return h +} diff --git a/cmd/omni/main.go b/cmd/omni/main.go index 59634ca..f49cb55 100644 --- a/cmd/omni/main.go +++ b/cmd/omni/main.go @@ -14,6 +14,7 @@ import ( "time" "github.com/ron2111/omnitoken" + "github.com/ron2111/omnitoken/cacheflow" ) func main() { @@ -32,6 +33,8 @@ func main() { err = runDecode(os.Args[2:]) case "cache": err = runCache(os.Args[2:]) + case "cache-sim": + err = runCacheSim(os.Args[2:]) case "bench": err = runBench(os.Args[2:]) case "encodings": @@ -58,6 +61,7 @@ Usage: omni encode [-model gpt-4o|-encoding o200k_base] [-file path] [text] omni decode [-model gpt-4o|-encoding o200k_base] [token ids] omni cache [-model gpt-4o|-encoding o200k_base] [-profile openai|generic] [-file path] [text] + omni cache-sim [-model gpt-4o|-encoding o200k_base] -input prompts.jsonl [-profile openai|generic] [-breakers] omni bench -input path -timings dir -name name [-model cl100k_base] [-iters 100] [-warmup 10] omni encodings [-json] omni models [-json] [-prefixes] @@ -67,6 +71,7 @@ Examples: omni encode -encoding o200k_base "hello world" omni decode -encoding o200k_base "24912 2375" omni cache -model gpt-4o -profile openai system-prompt.txt + omni cache-sim -model gpt-4o -profile openai -input prompts.jsonl -breakers omni encodings omni models -prefixes `) @@ -184,7 +189,7 @@ func runCache(args []string) error { if minimumTokens >= 0 { profile.MinimumTokens = minimumTokens } - report := omnitoken.NewCacheAligner(engine).AlignPromptToProfile(text, profile) + report := cacheflow.NewAligner(engine).AlignPromptToProfile(text, profile) encoded, err := json.MarshalIndent(report, "", " ") if err != nil { return err @@ -193,6 +198,65 @@ func runCache(args []string) error { return nil } +func runCacheSim(args []string) error { + fs := flag.NewFlagSet("cache-sim", flag.ExitOnError) + var flags engineFlags + profileName := "openai" + input := "" + jsonOutput := true + detectBreakers := false + blockSize := 0 + minimumTokens := -1 + addEngineFlags(fs, &flags) + fs.StringVar(&input, "input", input, "JSONL prompt trace input") + fs.StringVar(&profileName, "profile", profileName, "cache profile: openai, generic, custom") + fs.BoolVar(&jsonOutput, "json", jsonOutput, "emit JSON") + fs.BoolVar(&detectBreakers, "breakers", detectBreakers, "detect likely cache breakers") + fs.IntVar(&blockSize, "block", blockSize, "custom block size") + fs.IntVar(&minimumTokens, "min", minimumTokens, "custom minimum tokens") + if err := fs.Parse(args); err != nil { + return err + } + if input == "" { + return fmt.Errorf("input file is required") + } + engine, err := engineFromFlags(flags) + if err != nil { + return err + } + file, err := os.Open(input) + if err != nil { + return err + } + defer file.Close() + items, err := cacheflow.ReadJSONL(file) + if err != nil { + return err + } + profile := cacheProfile(profileName) + if blockSize > 0 { + profile.BlockSize = blockSize + } + if minimumTokens >= 0 { + profile.MinimumTokens = minimumTokens + } + report := cacheflow.Simulate(engine, items, cacheflow.SimulationOptions{Profile: profile, DetectBreakers: detectBreakers}) + if jsonOutput { + encoded, err := json.MarshalIndent(report, "", " ") + if err != nil { + return err + } + fmt.Println(string(encoded)) + return nil + } + fmt.Println(report.Summary) + fmt.Printf("items=%d total_tokens=%d common_prefix=%d reusable_prefix=%d\n", len(report.Items), report.TotalTokens, report.CommonPrefixTokens, report.ReusablePrefixTokens) + for _, hint := range report.CacheBreakerHints { + fmt.Printf("warning %s %s: %s\n", hint.ItemID, hint.Kind, hint.Message) + } + return nil +} + func runBench(args []string) error { fs := flag.NewFlagSet("bench", flag.ExitOnError) var name string @@ -372,13 +436,13 @@ func resolveEncoding(model string) string { } } -func cacheProfile(name string) omnitoken.CacheProfile { +func cacheProfile(name string) cacheflow.Profile { switch name { case "generic": - return omnitoken.CacheProfileGeneric + return cacheflow.ProfileGeneric case "custom": - return omnitoken.CacheProfile{Name: "custom", BlockSize: 1024} + return cacheflow.Profile{Name: "custom", BlockSize: 1024} default: - return omnitoken.CacheProfileOpenAI + return cacheflow.ProfileOpenAI } } diff --git a/docs/architecture.md b/docs/architecture.md index d0f2288..286c171 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -12,7 +12,7 @@ flowchart TD B --> I[Optional Adapter Modules] I --> J[Gemini / Llama 3 / Mistral / Hugging Face / Anthropic] C --> G[Encode / EncodeOrdinary / Decode / CountTokens] - G --> H[CacheAligner] + G --> H[cacheflow] ``` ## Registry diff --git a/docs/cache.md b/docs/cache.md deleted file mode 100644 index 7da6a36..0000000 --- a/docs/cache.md +++ /dev/null @@ -1,60 +0,0 @@ -# Cache Alignment - -OmniToken includes a small cache-planning helper for prompt-cache block boundaries. - -It does not edit prompts. It counts tokens, calculates boundary distance, and returns a report that users can apply in their own prompt-building code. - -## Usage - -```go -engine, err := omnitoken.ForModel("gpt-4o") -if err != nil { - panic(err) -} - -aligner := omnitoken.NewCacheAligner(engine) -report := aligner.AlignPromptToProfile(systemPrompt, omnitoken.CacheProfileOpenAI) -``` - -## Report - -| Field | Meaning | -| --- | --- | -| `CurrentTokens` | Token count for the input text. | -| `BlockSize` | Cache block size used for the calculation. | -| `MinimumTokens` | Minimum token threshold in the selected profile. | -| `PreviousBlockSize` | Previous block boundary. | -| `NextBlockSize` | Next block boundary. | -| `Remainder` | Tokens past the previous boundary. | -| `PaddingNeeded` | Additional tokens needed to reach the next boundary. | -| `TokensUntilMinimum` | Tokens needed before the prompt meets the profile minimum. | -| `IsAligned` | Whether the prompt is already on a boundary. | -| `IsEligible` | Whether the prompt meets the profile minimum. | -| `StrategyHint` | Human-readable planning hint. | - -## Profiles - -```go -omnitoken.CacheProfileGeneric -omnitoken.CacheProfileOpenAI -``` - -Profiles are local planning helpers, not billing guarantees. Provider cache behavior can change, and final usage should still be checked against provider usage metadata. - -## Custom Blocks - -For a custom block size: - -```go -report := aligner.AlignPrompt(prompt, 1024) -``` - -For a custom profile: - -```go -report := aligner.AlignPromptToProfile(prompt, omnitoken.CacheProfile{ - Name: "custom", - BlockSize: 512, - MinimumTokens: 2048, -}) -``` diff --git a/docs/cli.md b/docs/cli.md index 28fb071..31b9e01 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -39,6 +39,14 @@ omni decode -encoding o200k_base "24912 2375" omni cache -model gpt-4o -profile openai "hello world" ``` +Analyze a JSONL trace for stable token prefixes and cache-breaker hints: + +```powershell +omni cache-sim -model gpt-4o -profile openai -input prompts.jsonl -breakers +``` + +See [`cacheflow`](../cacheflow/README.md) for the JSONL format. + ## Registry Discovery List encodings registered in the current CLI build: diff --git a/docs/project-structure.md b/docs/project-structure.md index 22e2974..9e61b7d 100644 --- a/docs/project-structure.md +++ b/docs/project-structure.md @@ -9,7 +9,6 @@ Root files implement `github.com/ron2111/omnitoken`: - `engine.go`, `bpe.go`, `bpe_runtime.go`, `bpe_ranks.go`, `scanner.go`: OpenAI-compatible BPE runtime. - `registry.go`: model and encoding registry. - `options.go`, `special_tokens.go`: special-token-aware APIs. -- `cache.go`: cache-boundary planning. - `wordpiece.go`, `sentencepiece.go`: lightweight custom tokenizer engines. - `doc.go`, `example_test.go`: package documentation and examples. @@ -23,6 +22,10 @@ Embedded OpenAI-compatible vocab files live under `internal/openai/data`, with t Adapters are separate modules under `adapters/` so root users do not pull provider-specific dependencies. +## Cacheflow + +Prompt-cache boundary planning and trace simulation live under `cacheflow/`. It is a subpackage so the root tokenizer API stays focused while cache simulation can grow independently without adding dependencies. + ## CLI, Docs, Tools, Benchmarks - `cmd/omni`: CLI. diff --git a/registry_test.go b/registry_test.go index 0a69c63..77d921f 100644 --- a/registry_test.go +++ b/registry_test.go @@ -9,6 +9,12 @@ import ( var registryTestCounter uint64 +type fixedCountEngine int + +func (e fixedCountEngine) EncodeOrdinary(string) []int { return nil } +func (e fixedCountEngine) Decode([]int) string { return "" } +func (e fixedCountEngine) CountTokens(string) int { return int(e) } + func TestForEncodingBuildsCustomEncodingOnce(t *testing.T) { encoding := fmt.Sprintf("test_singleflight_%d", atomic.AddUint64(®istryTestCounter, 1)) started := make(chan struct{})