mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-01-02 14:22:28 +08:00
refactor: eliminate code duplication (DRY principle) (#211)
## Problem Solved Previous refactoring added 105 lines (1476→1581) by extracting code into separate files without eliminating duplication. This refactor focuses on reducing code size through deduplication while maintaining file separation for maintainability. ## Summary - Reduced total lines from 1581 to 1519 (-62 lines, 3.9% reduction) - Eliminated duplicate patterns using generic helpers and factory functions - Maintained file structure for maintainability - Zero functional changes - same behavior ### Phase 1: DRY use-quota-manager.tsx - Created parseStorageCount() helper (eliminates 6x localStorage read duplication) - Created createQuotaChecker() factory (consolidates 3 check function bodies) - Created createQuotaIncrementer() factory (consolidates 3 increment function bodies) - Result: 242→247 lines (+5 lines, but fully DRY with eliminated duplication) ### Phase 2: DRY chat-panel.tsx (1176→1109 lines, -67 lines) #### 2.1: Extract checkAllQuotaLimits helper - Replaced 3 occurrences of 18-line quota check blocks - Saved 36 lines #### 2.2: Extract sendChatMessage helper - Replaced 3 occurrences of 21-line sendMessage+headers blocks - Saved 42 lines #### 2.3: Extract processFilesAndAppendContent helper - Replaced 2 occurrences of file processing loops - Handles PDF, text, and image files uniformly - Async helper with optional image parts parameter
This commit is contained in:
247
lib/use-quota-manager.tsx
Normal file
247
lib/use-quota-manager.tsx
Normal file
@@ -0,0 +1,247 @@
|
||||
"use client"
|
||||
|
||||
import { useCallback, useMemo } from "react"
|
||||
import { toast } from "sonner"
|
||||
import { QuotaLimitToast } from "@/components/quota-limit-toast"
|
||||
import { STORAGE_KEYS } from "@/lib/storage"
|
||||
|
||||
export interface QuotaConfig {
|
||||
dailyRequestLimit: number
|
||||
dailyTokenLimit: number
|
||||
tpmLimit: number
|
||||
}
|
||||
|
||||
export interface QuotaCheckResult {
|
||||
allowed: boolean
|
||||
remaining: number
|
||||
used: number
|
||||
}
|
||||
|
||||
/**
|
||||
* Hook for managing request/token quotas and rate limiting.
|
||||
* Handles three types of limits:
|
||||
* - Daily request limit
|
||||
* - Daily token limit
|
||||
* - Tokens per minute (TPM) rate limit
|
||||
*
|
||||
* Users with their own API key bypass all limits.
|
||||
*/
|
||||
export function useQuotaManager(config: QuotaConfig): {
|
||||
hasOwnApiKey: () => boolean
|
||||
checkDailyLimit: () => QuotaCheckResult
|
||||
checkTokenLimit: () => QuotaCheckResult
|
||||
checkTPMLimit: () => QuotaCheckResult
|
||||
incrementRequestCount: () => void
|
||||
incrementTokenCount: (tokens: number) => void
|
||||
incrementTPMCount: (tokens: number) => void
|
||||
showQuotaLimitToast: () => void
|
||||
showTokenLimitToast: (used: number) => void
|
||||
showTPMLimitToast: () => void
|
||||
} {
|
||||
const { dailyRequestLimit, dailyTokenLimit, tpmLimit } = config
|
||||
|
||||
// Check if user has their own API key configured (bypass limits)
|
||||
const hasOwnApiKey = useCallback((): boolean => {
|
||||
const provider = localStorage.getItem(STORAGE_KEYS.aiProvider)
|
||||
const apiKey = localStorage.getItem(STORAGE_KEYS.aiApiKey)
|
||||
return !!(provider && apiKey)
|
||||
}, [])
|
||||
|
||||
// Generic helper: Parse count from localStorage with NaN guard
|
||||
const parseStorageCount = (key: string): number => {
|
||||
const count = parseInt(localStorage.getItem(key) || "0", 10)
|
||||
return Number.isNaN(count) ? 0 : count
|
||||
}
|
||||
|
||||
// Generic helper: Create quota checker factory
|
||||
const createQuotaChecker = useCallback(
|
||||
(
|
||||
getTimeKey: () => string,
|
||||
timeStorageKey: string,
|
||||
countStorageKey: string,
|
||||
limit: number,
|
||||
) => {
|
||||
return (): QuotaCheckResult => {
|
||||
if (hasOwnApiKey())
|
||||
return { allowed: true, remaining: -1, used: 0 }
|
||||
if (limit <= 0) return { allowed: true, remaining: -1, used: 0 }
|
||||
|
||||
const currentTime = getTimeKey()
|
||||
const storedTime = localStorage.getItem(timeStorageKey)
|
||||
let count = parseStorageCount(countStorageKey)
|
||||
|
||||
if (storedTime !== currentTime) {
|
||||
count = 0
|
||||
localStorage.setItem(timeStorageKey, currentTime)
|
||||
localStorage.setItem(countStorageKey, "0")
|
||||
}
|
||||
|
||||
return {
|
||||
allowed: count < limit,
|
||||
remaining: limit - count,
|
||||
used: count,
|
||||
}
|
||||
}
|
||||
},
|
||||
[hasOwnApiKey],
|
||||
)
|
||||
|
||||
// Generic helper: Create quota incrementer factory
|
||||
const createQuotaIncrementer = useCallback(
|
||||
(
|
||||
getTimeKey: () => string,
|
||||
timeStorageKey: string,
|
||||
countStorageKey: string,
|
||||
validateInput: boolean = false,
|
||||
) => {
|
||||
return (tokens: number = 1): void => {
|
||||
if (validateInput && (!Number.isFinite(tokens) || tokens <= 0))
|
||||
return
|
||||
|
||||
const currentTime = getTimeKey()
|
||||
const storedTime = localStorage.getItem(timeStorageKey)
|
||||
let count = parseStorageCount(countStorageKey)
|
||||
|
||||
if (storedTime !== currentTime) {
|
||||
count = 0
|
||||
localStorage.setItem(timeStorageKey, currentTime)
|
||||
}
|
||||
|
||||
localStorage.setItem(countStorageKey, String(count + tokens))
|
||||
}
|
||||
},
|
||||
[],
|
||||
)
|
||||
|
||||
// Check daily request limit
|
||||
const checkDailyLimit = useMemo(
|
||||
() =>
|
||||
createQuotaChecker(
|
||||
() => new Date().toDateString(),
|
||||
STORAGE_KEYS.requestDate,
|
||||
STORAGE_KEYS.requestCount,
|
||||
dailyRequestLimit,
|
||||
),
|
||||
[createQuotaChecker, dailyRequestLimit],
|
||||
)
|
||||
|
||||
// Increment request count
|
||||
const incrementRequestCount = useMemo(
|
||||
() =>
|
||||
createQuotaIncrementer(
|
||||
() => new Date().toDateString(),
|
||||
STORAGE_KEYS.requestDate,
|
||||
STORAGE_KEYS.requestCount,
|
||||
false,
|
||||
),
|
||||
[createQuotaIncrementer],
|
||||
)
|
||||
|
||||
// Show quota limit toast (request-based)
|
||||
const showQuotaLimitToast = useCallback(() => {
|
||||
toast.custom(
|
||||
(t) => (
|
||||
<QuotaLimitToast
|
||||
used={dailyRequestLimit}
|
||||
limit={dailyRequestLimit}
|
||||
onDismiss={() => toast.dismiss(t)}
|
||||
/>
|
||||
),
|
||||
{ duration: 15000 },
|
||||
)
|
||||
}, [dailyRequestLimit])
|
||||
|
||||
// Check daily token limit
|
||||
const checkTokenLimit = useMemo(
|
||||
() =>
|
||||
createQuotaChecker(
|
||||
() => new Date().toDateString(),
|
||||
STORAGE_KEYS.tokenDate,
|
||||
STORAGE_KEYS.tokenCount,
|
||||
dailyTokenLimit,
|
||||
),
|
||||
[createQuotaChecker, dailyTokenLimit],
|
||||
)
|
||||
|
||||
// Increment token count
|
||||
const incrementTokenCount = useMemo(
|
||||
() =>
|
||||
createQuotaIncrementer(
|
||||
() => new Date().toDateString(),
|
||||
STORAGE_KEYS.tokenDate,
|
||||
STORAGE_KEYS.tokenCount,
|
||||
true, // Validate input tokens
|
||||
),
|
||||
[createQuotaIncrementer],
|
||||
)
|
||||
|
||||
// Show token limit toast
|
||||
const showTokenLimitToast = useCallback(
|
||||
(used: number) => {
|
||||
toast.custom(
|
||||
(t) => (
|
||||
<QuotaLimitToast
|
||||
type="token"
|
||||
used={used}
|
||||
limit={dailyTokenLimit}
|
||||
onDismiss={() => toast.dismiss(t)}
|
||||
/>
|
||||
),
|
||||
{ duration: 15000 },
|
||||
)
|
||||
},
|
||||
[dailyTokenLimit],
|
||||
)
|
||||
|
||||
// Check TPM (tokens per minute) limit
|
||||
const checkTPMLimit = useMemo(
|
||||
() =>
|
||||
createQuotaChecker(
|
||||
() => Math.floor(Date.now() / 60000).toString(),
|
||||
STORAGE_KEYS.tpmMinute,
|
||||
STORAGE_KEYS.tpmCount,
|
||||
tpmLimit,
|
||||
),
|
||||
[createQuotaChecker, tpmLimit],
|
||||
)
|
||||
|
||||
// Increment TPM count
|
||||
const incrementTPMCount = useMemo(
|
||||
() =>
|
||||
createQuotaIncrementer(
|
||||
() => Math.floor(Date.now() / 60000).toString(),
|
||||
STORAGE_KEYS.tpmMinute,
|
||||
STORAGE_KEYS.tpmCount,
|
||||
true, // Validate input tokens
|
||||
),
|
||||
[createQuotaIncrementer],
|
||||
)
|
||||
|
||||
// Show TPM limit toast
|
||||
const showTPMLimitToast = useCallback(() => {
|
||||
const limitDisplay =
|
||||
tpmLimit >= 1000 ? `${tpmLimit / 1000}k` : String(tpmLimit)
|
||||
toast.error(
|
||||
`Rate limit reached (${limitDisplay} tokens/min). Please wait 60 seconds before sending another request.`,
|
||||
{ duration: 8000 },
|
||||
)
|
||||
}, [tpmLimit])
|
||||
|
||||
return {
|
||||
// Check functions
|
||||
hasOwnApiKey,
|
||||
checkDailyLimit,
|
||||
checkTokenLimit,
|
||||
checkTPMLimit,
|
||||
|
||||
// Increment functions
|
||||
incrementRequestCount,
|
||||
incrementTokenCount,
|
||||
incrementTPMCount,
|
||||
|
||||
// Toast functions
|
||||
showQuotaLimitToast,
|
||||
showTokenLimitToast,
|
||||
showTPMLimitToast,
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user