Integrate pricing with canonical model (#6130)

2025-12-18 11:44:36 -05:00
parent 473f269daa
commit 6a0b8c25d5
17 changed files with 263 additions and 1034 deletions
@@ -1,14 +1,9 @@
 import { useState, useEffect } from 'react';
 import { useModelAndProvider } from '../ModelAndProviderContext';
-import { useConfig } from '../ConfigContext';
 import { CoinIcon } from '../icons';
 import { Tooltip, TooltipContent, TooltipTrigger } from '../ui/Tooltip';
-import {
-  getCostForModel,
-  initializeCostDatabase,
-  updateAllModelCosts,
-  fetchAndCachePricing,
-} from '../../utils/costDatabase';
+import { fetchModelPricing } from '../../utils/pricing';
+import { PricingData } from '../../api';

 interface CostTrackerProps {
  inputTokens?: number;
@@ -24,18 +19,10 @@ interface CostTrackerProps {

 export function CostTracker({ inputTokens = 0, outputTokens = 0, sessionCosts }: CostTrackerProps) {
  const { currentModel, currentProvider } = useModelAndProvider();
-  const { getProviders } = useConfig();
-  const [costInfo, setCostInfo] = useState<{
-    input_token_cost?: number;
-    output_token_cost?: number;
-    currency?: string;
-  } | null>(null);
+  const [costInfo, setCostInfo] = useState<PricingData | null>(null);
  const [isLoading, setIsLoading] = useState(true);
  const [showPricing, setShowPricing] = useState(true);
  const [pricingFailed, setPricingFailed] = useState(false);
-  const [modelNotFound, setModelNotFound] = useState(false);
-  const [hasAttemptedFetch, setHasAttemptedFetch] = useState(false);
-  const [initialLoadComplete, setInitialLoadComplete] = useState(false);

  // Check if pricing is enabled
  useEffect(() => {
@@ -44,33 +31,11 @@ export function CostTracker({ inputTokens = 0, outputTokens = 0, sessionCosts }:
      setShowPricing(stored !== 'false');
    };

-    // Check on mount
    checkPricingSetting();
-
-    // Listen for storage changes
    window.addEventListener('storage', checkPricingSetting);
    return () => window.removeEventListener('storage', checkPricingSetting);
  }, []);

-  // Set initial load complete after a short delay
-  useEffect(() => {
-    const timer = setTimeout(() => {
-      setInitialLoadComplete(true);
-    }, 3000); // Give 3 seconds for initial load
-
-    return () => window.clearTimeout(timer);
-  }, []);
-
-  // Debug log props removed
-
-  // Initialize cost database on mount
-  useEffect(() => {
-    initializeCostDatabase();
-
-    // Update costs for all models in background
-    updateAllModelCosts().catch(() => {});
-  }, [getProviders]);
-
  useEffect(() => {
    const loadCostInfo = async () => {
      if (!currentModel || !currentProvider) {
@@ -78,49 +43,20 @@ export function CostTracker({ inputTokens = 0, outputTokens = 0, sessionCosts }:
        return;
      }

+      setIsLoading(true);
      try {
-        // First check sync cache
-        let costData = getCostForModel(currentProvider, currentModel);
-
+        const costData = await fetchModelPricing(currentProvider, currentModel);
        if (costData) {
-          // We have cached data
          setCostInfo(costData);
          setPricingFailed(false);
-          setModelNotFound(false);
-          setIsLoading(false);
-          setHasAttemptedFetch(true);
        } else {
-          // Need to fetch from backend
-          setIsLoading(true);
-          const result = await fetchAndCachePricing(currentProvider, currentModel);
-          setHasAttemptedFetch(true);
-
-          if (result && result.costInfo) {
-            setCostInfo(result.costInfo);
-            setPricingFailed(false);
-            setModelNotFound(false);
-          } else if (result && result.error === 'model_not_found') {
-            // Model not found in pricing database, but API call succeeded
-            setModelNotFound(true);
-            setPricingFailed(false);
-          } else {
-            // API call failed or other error
-            const freeProviders = ['ollama', 'local', 'localhost'];
-            if (!freeProviders.includes(currentProvider.toLowerCase())) {
-              setPricingFailed(true);
-              setModelNotFound(false);
-            }
-          }
-          setIsLoading(false);
+          setPricingFailed(true);
+          setCostInfo(null);
        }
      } catch {
-        setHasAttemptedFetch(true);
-        // Only set pricing failed if we're not dealing with a known free provider
-        const freeProviders = ['ollama', 'local', 'localhost'];
-        if (!freeProviders.includes(currentProvider.toLowerCase())) {
-          setPricingFailed(true);
-          setModelNotFound(false);
-        }
+        setPricingFailed(true);
+        setCostInfo(null);
+      } finally {
        setIsLoading(false);
      }
    };
@@ -221,10 +157,9 @@ export function CostTracker({ inputTokens = 0, outputTokens = 0, sessionCosts }:

    // Otherwise show as unavailable
    const getUnavailableTooltip = () => {
-      if (pricingFailed && hasAttemptedFetch && initialLoadComplete) {
-        return `Pricing data unavailable - OpenRouter connection failed. Click refresh in settings to retry.`;
+      if (pricingFailed) {
+        return `Pricing data unavailable for ${currentModel}`;
      }
-      // If we reach here, it must be modelNotFound (since we only get here after attempting fetch)
      return `Cost data not available for ${currentModel} (${inputTokens.toLocaleString()} input, ${outputTokens.toLocaleString()} output tokens)`;
    };

@@ -249,12 +184,8 @@ export function CostTracker({ inputTokens = 0, outputTokens = 0, sessionCosts }:
  // Build tooltip content
  const getTooltipContent = (): string => {
    // Handle error states first
-    if (pricingFailed && hasAttemptedFetch && initialLoadComplete) {
-      return `Pricing data unavailable - OpenRouter connection failed. Click refresh in settings to retry.`;
-    }
-
-    if (modelNotFound && hasAttemptedFetch && initialLoadComplete) {
-      return `Pricing not available for ${currentProvider}/${currentModel}. This model may not be supported by the pricing service.`;
+    if (pricingFailed) {
+      return `Pricing data unavailable for ${currentProvider}/${currentModel}`;
    }

    // Handle session costs
@@ -1,13 +1,12 @@
 import { useState, useEffect, useRef } from 'react';
 import { Switch } from '../../ui/switch';
 import { Button } from '../../ui/button';
-import { Settings, RefreshCw, ExternalLink } from 'lucide-react';
+import { Settings } from 'lucide-react';
 import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from '../../ui/dialog';
 import UpdateSection from './UpdateSection';
 import TunnelSection from '../tunnel/TunnelSection';

 import { COST_TRACKING_ENABLED, UPDATES_ENABLED } from '../../../updates';
-import { getApiUrl } from '../../../config';
 import { Card, CardContent, CardDescription, CardHeader, CardTitle } from '../../ui/card';
 import ThemeSelector from '../../GooseSidebar/ThemeSelector';
 import BlockLogoBlack from './icons/block-lockup_black.png';
@@ -26,9 +25,6 @@ export default function AppSettingsSection({ scrollToSection }: AppSettingsSecti
  const [isMacOS, setIsMacOS] = useState(false);
  const [isDockSwitchDisabled, setIsDockSwitchDisabled] = useState(false);
  const [showNotificationModal, setShowNotificationModal] = useState(false);
-  const [pricingStatus, setPricingStatus] = useState<'loading' | 'success' | 'error'>('loading');
-  const [lastFetchTime, setLastFetchTime] = useState<Date | null>(null);
-  const [isRefreshing, setIsRefreshing] = useState(false);
  const [showPricing, setShowPricing] = useState(true);
  const [isDarkMode, setIsDarkMode] = useState(false);
  const updateSectionRef = useRef<HTMLDivElement>(null);
@@ -66,71 +62,6 @@ export default function AppSettingsSection({ scrollToSection }: AppSettingsSecti
    setShowPricing(stored !== 'false');
  }, []);

-  // Check pricing status on mount
-  useEffect(() => {
-    checkPricingStatus();
-  }, []);
-
-  const checkPricingStatus = async () => {
-    try {
-      const apiUrl = getApiUrl('/config/pricing');
-      const secretKey = await window.electron.getSecretKey();
-
-      const headers: HeadersInit = { 'Content-Type': 'application/json' };
-      if (secretKey) {
-        headers['X-Secret-Key'] = secretKey;
-      }
-
-      const response = await fetch(apiUrl, {
-        method: 'POST',
-        headers,
-        body: JSON.stringify({ configured_only: true }),
-      });
-
-      if (response.ok) {
-        await response.json();
-        setPricingStatus('success');
-        setLastFetchTime(new Date());
-      } else {
-        setPricingStatus('error');
-      }
-    } catch {
-      setPricingStatus('error');
-    }
-  };
-
-  const handleRefreshPricing = async () => {
-    setIsRefreshing(true);
-    try {
-      const apiUrl = getApiUrl('/config/pricing');
-      const secretKey = await window.electron.getSecretKey();
-
-      const headers: HeadersInit = { 'Content-Type': 'application/json' };
-      if (secretKey) {
-        headers['X-Secret-Key'] = secretKey;
-      }
-
-      const response = await fetch(apiUrl, {
-        method: 'POST',
-        headers,
-        body: JSON.stringify({ configured_only: false }),
-      });
-
-      if (response.ok) {
-        setPricingStatus('success');
-        setLastFetchTime(new Date());
-        // Trigger a reload of the cost database
-        window.dispatchEvent(new CustomEvent('pricing-updated'));
-      } else {
-        setPricingStatus('error');
-      }
-    } catch {
-      setPricingStatus('error');
-    } finally {
-      setIsRefreshing(false);
-    }
-  };
-
  // Handle scrolling to update section
  useEffect(() => {
    if (scrollToSection === 'update' && updateSectionRef.current) {
@@ -326,69 +257,6 @@ export default function AppSettingsSection({ scrollToSection }: AppSettingsSecti
            </div>
          )}

-          {/* Pricing Status - only show if cost tracking is enabled */}
-          {COST_TRACKING_ENABLED && showPricing && (
-            <>
-              <div className="flex items-center justify-between text-xs mb-2 px-4">
-                <span className="text-textSubtle">Pricing Source:</span>
-                <a
-                  href="https://openrouter.ai/docs#models"
-                  target="_blank"
-                  rel="noopener noreferrer"
-                  className="text-blue-600 dark:text-blue-400 hover:underline flex items-center gap-1"
-                >
-                  OpenRouter Docs
-                  <ExternalLink size={10} />
-                </a>
-              </div>
-
-              <div className="flex items-center justify-between text-xs mb-2 px-4">
-                <span className="text-textSubtle">Status:</span>
-                <div className="flex items-center gap-2">
-                  <span
-                    className={`font-medium ${
-                      pricingStatus === 'success'
-                        ? 'text-green-600 dark:text-green-400'
-                        : pricingStatus === 'error'
-                          ? 'text-red-600 dark:text-red-400'
-                          : 'text-textSubtle'
-                    }`}
-                  >
-                    {pricingStatus === 'success'
-                      ? '✓ Connected'
-                      : pricingStatus === 'error'
-                        ? '✗ Failed'
-                        : '... Checking'}
-                  </span>
-                  <button
-                    className="p-0.5 hover:bg-gray-200 dark:hover:bg-gray-700 rounded transition-colors disabled:opacity-50"
-                    onClick={handleRefreshPricing}
-                    disabled={isRefreshing}
-                    title="Refresh pricing data"
-                    type="button"
-                  >
-                    <RefreshCw
-                      size={8}
-                      className={`text-textSubtle hover:text-textStandard ${isRefreshing ? 'animate-spin-fast' : ''}`}
-                    />
-                  </button>
-                </div>
-              </div>
-
-              {lastFetchTime && (
-                <div className="flex items-center justify-between text-xs mb-2 px-4">
-                  <span className="text-textSubtle">Last updated:</span>
-                  <span className="text-textSubtle">{lastFetchTime.toLocaleTimeString()}</span>
-                </div>
-              )}
-
-              {pricingStatus === 'error' && (
-                <p className="text-xs text-red-600 dark:text-red-400 px-4">
-                  Unable to fetch pricing data. Costs will not be displayed.
-                </p>
-              )}
-            </>
-          )}
        </CardContent>
      </Card>