Skip to content

Commit 5c167de

Browse files
committed
refactor(tokenbudget): extract effectiveTokensPerWord
Code review on PR #732 flagged that the "resolve TokensPerWord, falling back to defaultTokensPerWord when unset" snippet was duplicated three times in this file (definitelyUnderBudget, tokenCount, modeLabel) — the new definitelyUnderBudget added a third copy instead of sharing the resolution tokenCount already had. Extract effectiveTokensPerWord and call it from all three; same behavior, one definition of the default instead of three.
1 parent a69750e commit 5c167de

1 file changed

Lines changed: 15 additions & 17 deletions

File tree

internal/rules/tokenbudget/rule.go

Lines changed: 15 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -140,19 +140,25 @@ func (r *Rule) activeBudget(path string) int {
140140
// never produces a false "under budget" when the real count would
141141
// exceed it.
142142
func (r *Rule) definitelyUnderBudget(source []byte, budget int) bool {
143-
n := float64(len(source))
144143
switch normalizeMode(r.Mode) {
145144
case "tokenizer":
146-
return n <= float64(budget)
145+
return len(source) <= budget
147146
default:
148-
tpw := r.TokensPerWord
149-
if tpw <= 0 {
150-
tpw = defaultTokensPerWord
151-
}
152-
return math.Round(n*tpw) <= float64(budget)
147+
return math.Round(float64(len(source))*r.effectiveTokensPerWord()) <= float64(budget)
153148
}
154149
}
155150

151+
// effectiveTokensPerWord returns r.TokensPerWord, falling back to
152+
// defaultTokensPerWord when unset — the same resolution definitelyUnderBudget,
153+
// tokenCount, and modeLabel all need for heuristic mode.
154+
func (r *Rule) effectiveTokensPerWord() float64 {
155+
tpw := r.TokensPerWord
156+
if tpw <= 0 {
157+
tpw = defaultTokensPerWord
158+
}
159+
return tpw
160+
}
161+
156162
// tokenCount estimates the token count of source without copying it.
157163
// The mode label is built separately by modeLabel, only when a
158164
// diagnostic is actually emitted — formatting it for every under-budget
@@ -164,12 +170,8 @@ func (r *Rule) tokenCount(source []byte) int {
164170
enc := normalizeEncoding(r.Encoding)
165171
return tokenizerCount(source, tok, enc)
166172
default:
167-
tpw := r.TokensPerWord
168-
if tpw <= 0 {
169-
tpw = defaultTokensPerWord
170-
}
171173
words := mdtext.CountWordsBytes(source)
172-
count := int(math.Round(float64(words) * tpw))
174+
count := int(math.Round(float64(words) * r.effectiveTokensPerWord()))
173175
if count < 0 {
174176
count = 0
175177
}
@@ -183,11 +185,7 @@ func (r *Rule) modeLabel() string {
183185
case "tokenizer":
184186
return "tokenizer:" + normalizeTokenizer(r.Tokenizer) + "/" + normalizeEncoding(r.Encoding)
185187
default:
186-
tpw := r.TokensPerWord
187-
if tpw <= 0 {
188-
tpw = defaultTokensPerWord
189-
}
190-
return "heuristic:tokens-per-word=" + strconv.FormatFloat(tpw, 'f', 2, 64)
188+
return "heuristic:tokens-per-word=" + strconv.FormatFloat(r.effectiveTokensPerWord(), 'f', 2, 64)
191189
}
192190
}
193191

0 commit comments

Comments
 (0)