Skip to content

feat: implement accurate token-based counting with progressive warnings #137

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
AnthonyRonning merged 2 commits into from
Jul 8, 2025
Merged

feat: implement accurate token-based counting with progressive warnings #137

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
68325c1
feat: implement accurate token-based counting with progressive warnings
AnthonyRonning Jul 8, 2025
8a107e5
refactor: extract image token estimate as a constant
AnthonyRonning Jul 8, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions frontend/bun.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions frontend/package.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
"@tauri-apps/plugin-os": "^2.2.1",
"class-variance-authority": "^0.7.0",
"clsx": "^2.1.1",
"gpt-tokenizer": "^3.0.1",
Comment thread
AnthonyRonning marked this conversation as resolved.
"lucide-react": "^0.436.0",
"openai": "^4.56.1",
"react": "^18.3.1",
Expand Down
166 changes: 112 additions & 54 deletions frontend/src/components/ChatBox.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,18 +7,18 @@ import {
DropdownMenuItem,
DropdownMenuTrigger
} from "@/components/ui/dropdown-menu";
import { useEffect, useRef, useState } from "react";
import { useEffect, useRef, useState, useMemo } from "react";
import { useLocalState } from "@/state/useLocalState";
import { cn, useIsMobile } from "@/utils/utils";
import { useQuery } from "@tanstack/react-query";
import { getBillingService } from "@/billing/billingService";
import { BillingStatus } from "@/billing/billingApi";
import { Route as ChatRoute } from "@/routes/_auth.chat.$chatId";
import { ChatMessage } from "@/state/LocalStateContext";
import { useNavigate, useRouter } from "@tanstack/react-router";
import { ModelSelector, MODEL_CONFIG } from "@/components/ModelSelector";
import { ModelSelector, MODEL_CONFIG, getModelTokenLimit } from "@/components/ModelSelector";
import { useOpenSecret } from "@opensecret/react";
import type { DocumentResponse } from "@opensecret/react";
import { encode } from "gpt-tokenizer";

interface ParsedDocument {
document: {
Expand All @@ -35,30 +35,18 @@ interface ParsedDocument {
timings: Record<string, unknown>;
}

// Rough token estimation function
// Accurate token counting using gpt-tokenizer
function estimateTokenCount(text: string): number {
// A very rough estimation: ~4 characters per token on average
return Math.ceil(text.length / 4);
// Use gpt-tokenizer for accurate token counting
return encode(text).length;
}
Comment on lines +38 to 42

@coderabbitai coderabbitai Bot Jul 8, 2025

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Improve function naming and add error handling for token counting.

The function name estimateTokenCount is misleading since it now provides accurate counting using gpt-tokenizer.

-// Accurate token counting using gpt-tokenizer
-function estimateTokenCount(text: string): number {
-  // Use gpt-tokenizer for accurate token counting
-  return encode(text).length;
-}
+// Accurate token counting using gpt-tokenizer
+function countTokens(text: string): number {
+  try {
+    return encode(text).length;
+  } catch (error) {
+    console.error('Token counting failed:', error);
+    // Fallback to rough estimate
+    return Math.ceil(text.length / 4);
+  }
+}
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
// Accurate token counting using gpt-tokenizer
function estimateTokenCount(text: string): number {
// A very rough estimation: ~4 characters per token on average
return Math.ceil(text.length / 4);
// Use gpt-tokenizer for accurate token counting
return encode(text).length;
}
// Accurate token counting using gpt-tokenizer
function countTokens(text: string): number {
try {
return encode(text).length;
} catch (error) {
console.error('Token counting failed:', error);
// Fallback to rough estimate
return Math.ceil(text.length / 4);
}
}
🤖 Prompt for AI Agents 
In frontend/src/components/ChatBox.tsx around lines 38 to 42, rename the
function from estimateTokenCount to a name that reflects accurate token
counting, such as getTokenCount. Additionally, add error handling inside the
function to catch and handle any exceptions that might occur during token
encoding, ensuring the function does not fail silently or crash the app.


function TokenWarning({
messages,
currentInput,
chatId,
className,
billingStatus,
onCompress,
isCompressing = false
}: {
messages: ChatMessage[];
currentInput: string;
chatId?: string;
className?: string;
billingStatus?: BillingStatus;
onCompress?: () => void;
isCompressing?: boolean;
}) {
const totalTokens =
// Estimated token count for images (varies by model and image size)
const IMAGE_TOKEN_ESTIMATE = 85;
Comment thread
AnthonyRonning marked this conversation as resolved.

// Calculate total tokens for messages and current input
function calculateTotalTokens(messages: ChatMessage[], currentInput: string): number {
return (
messages.reduce((acc, msg) => {
if (typeof msg.content === "string") {
return acc + estimateTokenCount(msg.content);
Expand All @@ -71,26 +59,52 @@ function TokenWarning({
return sum + estimateTokenCount(part.text);
}
// Rough estimate for images
return sum + 85;
return sum + IMAGE_TOKEN_ESTIMATE;
}, 0)
);
}
}, 0) + (currentInput ? estimateTokenCount(currentInput) : 0);
}, 0) + (currentInput ? estimateTokenCount(currentInput) : 0)
);
}

const navigate = useNavigate();
// Custom hook for debouncing values
function useDebounce<T>(value: T, delay: number): T {
const [debouncedValue, setDebouncedValue] = useState(value);

useEffect(() => {
const handler = setTimeout(() => {
setDebouncedValue(value);
}, delay);

// Check if user is on starter plan
const isStarter = billingStatus?.product_name?.toLowerCase().includes("starter") || false;
return () => {
clearTimeout(handler);
};
}, [value, delay]);

return debouncedValue;
}

// Token thresholds for different plan types
const STARTER_WARNING_THRESHOLD = 4000;
const PRO_WARNING_THRESHOLD = 10000;
function TokenWarning({
chatId,
className,
onCompress,
isCompressing = false,
tokenPercentage
}: {
chatId?: string;
className?: string;
onCompress?: () => void;
isCompressing?: boolean;
tokenPercentage: number;
}) {
const navigate = useNavigate();

// Different thresholds for starter vs pro users
const warningThreshold = isStarter ? STARTER_WARNING_THRESHOLD : PRO_WARNING_THRESHOLD;
// Only show warning if above 50%
if (tokenPercentage < 50) return null;

// Only show warning if above the threshold
if (totalTokens < warningThreshold) return null;
// Determine the severity and behavior based on percentage
const isAt95Percent = tokenPercentage >= 95;
const isAt99Percent = tokenPercentage >= 99;

const handleNewChat = async (e: React.MouseEvent) => {
e.preventDefault();
Expand All @@ -103,7 +117,17 @@ function TokenWarning({
}
};

// Determine button text based on compression state
// Get appropriate message and styling based on threshold
const getMessage = () => {
if (isAt99Percent) {
return "This chat is too long to continue.";
} else if (isAt95Percent) {
return "Chat is at capacity. Compress to continue.";
} else {
return "This chat is getting long. Compress it to save tokens.";
}
};

const getButtonText = () => {
if (isCompressing) {
return { desktop: "Compressing...", mobile: "Compressing..." };
Expand All @@ -116,26 +140,46 @@ function TokenWarning({

const buttonText = getButtonText();

// Determine background color based on severity
const bgClass = isAt99Percent
? "bg-destructive/20 border border-destructive/30"
: isAt95Percent
? "bg-warning/20 border border-warning/30"
: "bg-muted/50";

return (
<div
className={cn(
"flex items-center justify-between px-3 py-1.5 mb-1",
"bg-muted/50 backdrop-blur-sm rounded-t-lg",
"text-xs text-muted-foreground/90",
"backdrop-blur-sm rounded-t-lg",
"text-xs",
bgClass,
isAt99Percent
? "text-destructive"
: isAt95Percent
? "text-warning-foreground"
: "text-muted-foreground/90",
className
)}
>
<div className="flex items-center gap-2 min-w-0">
<span className="text-[11px] font-semibold text-foreground/70 shrink-0">Tip:</span>
<span className="min-w-0">This chat is getting long. Compress it to save tokens.</span>
<span className="text-[11px] font-semibold shrink-0">
{isAt99Percent ? "Error:" : isAt95Percent ? "Warning:" : "Tip:"}
</span>
<span className="min-w-0">{getMessage()}</span>
</div>
{chatId && (
{chatId && !isAt99Percent && (
<button
onClick={!isCompressing ? onCompress || handleNewChat : undefined}
disabled={isCompressing}
className={cn(
"font-medium text-primary transition-colors whitespace-nowrap shrink-0 ml-4",
isCompressing ? "opacity-70 cursor-default" : "hover:text-primary/80 hover:underline"
"font-medium transition-colors whitespace-nowrap shrink-0 ml-4",
isCompressing ? "opacity-70 cursor-default" : "hover:underline",
isAt99Percent
? "text-destructive"
: isAt95Percent
? "text-warning-foreground hover:text-warning-foreground/80"
: "text-primary hover:text-primary/80"
)}
>
<span className="hidden md:inline">{buttonText.desktop}</span>
Expand Down Expand Up @@ -556,6 +600,18 @@ export default function Component({
}
}, [systemPromptValue]);

// Debounce input for token calculations to avoid lag while typing
const debouncedInputValue = useDebounce(inputValue, 300);

// Calculate token usage percentage
const totalTokens = useMemo(
() => calculateTotalTokens(messages, debouncedInputValue),
[messages, debouncedInputValue]
);
const tokenLimit = getModelTokenLimit(model);
const tokenPercentage = (totalTokens / tokenLimit) * 100;
const isAt99Percent = tokenPercentage >= 99;
Comment thread
AnthonyRonning marked this conversation as resolved.

// Update current input ref when input value changes
useEffect(() => {
currentInputRef.current = inputValue;
Expand Down Expand Up @@ -607,7 +663,8 @@ export default function Component({
(!freshBillingStatus.can_chat ||
(freshBillingStatus.chats_remaining !== null &&
freshBillingStatus.chats_remaining <= 0))) ||
isStreaming;
isStreaming ||
isAt99Percent;

// Disable the input box only when the user is out of chats or when streaming
const isInputDisabled =
Expand Down Expand Up @@ -646,6 +703,9 @@ export default function Component({
// No longer need token calculation or plan type check since we removed the hard limit
// Just keeping the TokenWarning component which handles its own calculations
const placeholderText = (() => {
if (isAt99Percent) {
return "Chat is too long to continue.";
}
if (billingStatus === null || freshBillingStatus === undefined)
return "Type your message here...";
if (freshBillingStatus.can_chat === false) {
Expand All @@ -656,15 +716,6 @@ export default function Component({

return (
<div className="flex flex-col w-full">
<TokenWarning
messages={messages}
currentInput={inputValue}
chatId={chatId}
billingStatus={freshBillingStatus}
onCompress={onCompress}
isCompressing={isSummarizing}
/>

{/* Simple System Prompt Section - just a gear button and input when expanded */}
{canEditSystemPrompt && (
<div className="mb-2">
Expand Down Expand Up @@ -704,6 +755,13 @@ export default function Component({
</div>
)}

<TokenWarning
chatId={chatId}
onCompress={onCompress}
isCompressing={isSummarizing}
tokenPercentage={tokenPercentage}
/>

<form
className={cn(
"p-2 rounded-lg border border-primary bg-background/80 backdrop-blur-lg focus-within:ring-1 focus-within:ring-ring",
Expand Down
23 changes: 18 additions & 5 deletions frontend/src/components/ModelSelector.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,38 +12,51 @@ import { useEffect, useRef } from "react";
import { useNavigate } from "@tanstack/react-router";
import type { Model } from "openai/resources/models.js";

// Model configuration for display names and badges
// Model configuration for display names, badges, and token limits
type ModelCfg = {
displayName: string;
badge?: string;
disabled?: boolean;
requiresPro?: boolean;
requiresStarter?: boolean;
supportsVision?: boolean;
tokenLimit: number;
};

export const MODEL_CONFIG: Record<string, ModelCfg> = {
"ibnzterrell/Meta-Llama-3.3-70B-Instruct-AWQ-INT4": {
displayName: "Llama 3.3 70B"
displayName: "Llama 3.3 70B",
tokenLimit: 70000
},
"google/gemma-3-27b-it": {
displayName: "Gemma 3 27B",
badge: "Starter",
requiresStarter: true
requiresStarter: true,
tokenLimit: 70000
},
"leon-se/gemma-3-27b-it-fp8-dynamic": {
displayName: "Gemma 3 27B",
badge: "Starter",
requiresStarter: true,
supportsVision: true
supportsVision: true,
tokenLimit: 70000
},
"deepseek-r1-70b": {
displayName: "DeepSeek R1 70B",
badge: "Pro",
requiresPro: true
requiresPro: true,
tokenLimit: 64000
}
};

// Default token limit for unknown models
export const DEFAULT_TOKEN_LIMIT = 64000;

// Get token limit for a specific model
export function getModelTokenLimit(modelId: string): number {
return MODEL_CONFIG[modelId]?.tokenLimit || DEFAULT_TOKEN_LIMIT;
}

import { ChatMessage } from "@/state/LocalStateContextDef";
Comment thread
AnthonyRonning marked this conversation as resolved.

export function ModelSelector({
Expand Down
Loading