import { useCallback, useEffect, useRef, useState } from 'react'; import { useTranslation } from 'react-i18next'; import { ApiError, deletePricing, fetchPricing, fetchUsedModels, updatePricing } from '@/lib/api'; import { useNotificationStore } from '@/stores'; import { loadModelPrices, saveModelPrices, type ModelPrice } from '@/utils/usage'; export interface UsePricingDataOptions { onAuthRequired?: () => void; enabled?: boolean; } export interface UsePricingDataReturn { modelNames: string[]; modelPrices: Record; loading: boolean; error: string; lastRefreshedAt: Date | null; loadPricing: () => Promise; setModelPrices: (prices: Record) => Promise; } const pricingToModelPrice = (entry: { model: string; prompt_price_per_1m: number; completion_price_per_1m: number; cache_price_per_1m: number; }): ModelPrice => ({ prompt: entry.prompt_price_per_1m, completion: entry.completion_price_per_1m, cache: entry.cache_price_per_1m, }); export function usePricingData(options: UsePricingDataOptions = {}): UsePricingDataReturn { const { onAuthRequired, enabled = true } = options; const { t } = useTranslation(); const { showNotification } = useNotificationStore(); const [modelNames, setModelNames] = useState([]); const [modelPrices, setModelPricesState] = useState>(() => loadModelPrices()); const [loading, setLoading] = useState(false); const [error, setError] = useState(''); const [lastRefreshedAt, setLastRefreshedAt] = useState(null); const requestControllerRef = useRef(null); const loadPricing = useCallback(async () => { requestControllerRef.current?.abort(); const controller = new AbortController(); requestControllerRef.current = controller; setLoading(true); setError(''); try { const [pricingResponse, usedModelsResponse] = await Promise.all([ fetchPricing(controller.signal), fetchUsedModels(controller.signal), ]); if (requestControllerRef.current !== controller) { return; } const prices = Object.fromEntries( pricingResponse.pricing.map((entry) => [entry.model, pricingToModelPrice(entry)]) ); saveModelPrices(prices); setModelPricesState(prices); setModelNames(usedModelsResponse.models); setLastRefreshedAt(new Date()); } catch (error) { if (controller.signal.aborted) { return; } if (error instanceof ApiError && error.status === 401) { onAuthRequired?.(); return; } setModelPricesState(loadModelPrices()); setError(error instanceof Error ? error.message : 'Failed to load pricing'); } finally { if (requestControllerRef.current === controller) { setLoading(false); requestControllerRef.current = null; } } }, [onAuthRequired]); useEffect(() => { if (!enabled) { requestControllerRef.current?.abort(); requestControllerRef.current = null; setLoading(false); return; } void loadPricing(); return () => { requestControllerRef.current?.abort(); requestControllerRef.current = null; }; }, [enabled, loadPricing]); const setModelPrices = useCallback(async (prices: Record) => { const previousPrices = modelPrices; setModelPricesState(prices); saveModelPrices(prices); try { const previousModels = new Set(Object.keys(previousPrices)); const nextModels = new Set(Object.keys(prices)); await Promise.all([ ...Object.entries(prices).map(([model, pricing]) => updatePricing(model, { prompt_price_per_1m: pricing.prompt, completion_price_per_1m: pricing.completion, cache_price_per_1m: pricing.cache, }) ), ...Array.from(previousModels) .filter((model) => !nextModels.has(model)) .map((model) => deletePricing(model)), ]); setLastRefreshedAt(new Date()); } catch (error) { setModelPricesState(previousPrices); saveModelPrices(previousPrices); if (error instanceof ApiError && error.status === 401) { onAuthRequired?.(); return; } const message = error instanceof Error ? error.message : ''; showNotification( `${t('notification.upload_failed')}${message ? `: ${message}` : ''}`, 'error' ); } }, [modelPrices, onAuthRequired, showNotification, t]); return { modelNames, modelPrices, loading, error, lastRefreshedAt, loadPricing, setModelPrices, }; }