Update src/App.tsx

src/App.tsx

@@ -1,12 +1,7 @@
-import React, { useState, useEffect } from 'react';
-import {
-  Card,
-  CardContent,
-  CardHeader,
-  CardTitle,
-} from '@/components/ui/card';
-import { Checkbox } from '@/components/ui/checkbox';
-import { Input } from '@/components/ui/input';
+import React, { useState, useEffect } from "react";
+import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
+import { Checkbox } from "@/components/ui/checkbox";
+import { Input } from "@/components/ui/input";
 import {
   Table,
   TableBody,
@@ -14,16 +9,17 @@ import {
   TableHead,
   TableHeader,
   TableRow,
-} from '@/components/ui/table';
-import { MultiSelect } from '@/components/ui/multi-select';
+} from "@/components/ui/table";
+import { MultiSelect } from "@/components/ui/multi-select";
 import {
   Collapsible,
   CollapsibleContent,
   CollapsibleTrigger,
-} from '@/components/ui/collapsible';
-import { Button } from '@/components/ui/button';
-import { ChevronDown, ChevronRight } from 'lucide-react';
-import { mockData } from './lib/data'; // Assuming you have this file for mock data
+} from "@/components/ui/collapsible";
+import { Button } from "@/components/ui/button";
+import { ChevronDown, ChevronRight } from "lucide-react";
+import { mockData } from "./lib/data"; // Assuming you have this file for mock data
+import { Switch } from "@/components/ui/switch";
 
 interface FlattenedModel extends Model {
   provider: string;
@@ -50,7 +46,8 @@ const App: React.FC = () => {
   const [selectedProviders, setSelectedProviders] = useState<string[]>([]);
   const [selectedModels, setSelectedModels] = useState<string[]>([]);
   const [expandedProviders, setExpandedProviders] = useState<string[]>([]);
-  const [tokenCalculation, setTokenCalculation] = useState<string>('million');
+  const [tokenCalculation, setTokenCalculation] = useState<string>("million");
+  const [linkProviderModel, setLinkProviderModel] = useState<boolean>(false);
 
   const [sortConfig, setSortConfig] = useState<{
     key: keyof FlattenedModel;
@@ -63,23 +60,23 @@ const App: React.FC = () => {
 
   const calculatePrice = (price: number, tokens: number): number => {
     let multiplier = 1;
-    if (tokenCalculation === 'thousand') {
+    if (tokenCalculation === "thousand") {
       multiplier = 1e-3;
-    } else if (tokenCalculation === 'unit') {
+    } else if (tokenCalculation === "unit") {
       multiplier = 1e-6;
-    } else if (tokenCalculation === 'billion') {
+    } else if (tokenCalculation === "billion") {
       multiplier = 1e3;
     }
     return price * tokens * multiplier;
   };
 
-
-
   const calculateComparison = (
     modelPrice: number,
     comparisonPrice: number
   ): string => {
-    return (((modelPrice - comparisonPrice) / comparisonPrice) * 100).toFixed(2);
+    return (((modelPrice - comparisonPrice) / comparisonPrice) * 100).toFixed(
+      2
+    );
   };
 
   const flattenData = (data: Provider[]) => {
@@ -92,20 +89,39 @@ const App: React.FC = () => {
     );
   };
 
-  const filteredData = data
-    .filter(
-      (provider) =>
-        selectedProviders.length === 0 ||
-        selectedProviders.includes(provider.provider)
-    )
-    .map((provider) => ({
-      ...provider,
-      models: provider.models.filter(
-        (model) =>
-          selectedModels.length === 0 || selectedModels.includes(model.name)
-      ),
-    }))
-    .filter((provider) => provider.models.length > 0);
+  const filteredData =
+    selectedProviders.length === 0 &&
+    selectedModels.length === 0 &&
+    !linkProviderModel
+      ? data.map((provider) => ({
+          ...provider,
+          models: provider.models,
+        }))
+      : data
+          .filter(
+            (provider) =>
+              selectedProviders.length === 0 ||
+              selectedProviders.includes(provider.provider)
+          )
+          .map((provider) => ({
+            ...provider,
+            models: provider.models.filter((model) => {
+              // If linking is enabled and no models are selected, filter by provider
+              if (linkProviderModel && selectedModels.length === 0)
+                return selectedProviders.includes(provider.provider);
+
+              // If no models are selected and linking is off, show all models from selected providers (or all if no providers selected)
+              if (!linkProviderModel && selectedModels.length === 0)
+                return (
+                  selectedProviders.length === 0 ||
+                  selectedProviders.includes(provider.provider)
+                );
+
+              // Otherwise, only show selected models
+              return selectedModels.includes(model.name);
+            }),
+          }))
+          .filter((provider) => provider.models.length > 0);
 
   const sortedFlattenedData = React.useMemo(() => {
     let sortableData: FlattenedModel[] = flattenData(filteredData);
@@ -114,12 +130,12 @@ const App: React.FC = () => {
         const aValue = a[sortConfig.key];
         const bValue = b[sortConfig.key];
 
-        if (typeof aValue === 'string' && typeof bValue === 'string') {
-          return sortConfig.direction === 'ascending'
+        if (typeof aValue === "string" && typeof bValue === "string") {
+          return sortConfig.direction === "ascending"
             ? aValue.localeCompare(bValue)
             : bValue.localeCompare(aValue);
-        } else if (typeof aValue === 'number' && typeof bValue === 'number') {
-          return sortConfig.direction === 'ascending'
+        } else if (typeof aValue === "number" && typeof bValue === "number") {
+          return sortConfig.direction === "ascending"
             ? aValue - bValue
             : bValue - aValue;
         } else {
@@ -131,13 +147,13 @@ const App: React.FC = () => {
   }, [filteredData, sortConfig]);
 
   const requestSort = (key: keyof FlattenedModel) => {
-    let direction = 'ascending';
+    let direction = "ascending";
     if (
       sortConfig &&
       sortConfig.key === key &&
-      sortConfig.direction === 'ascending'
+      sortConfig.direction === "ascending"
     ) {
-      direction = 'descending';
+      direction = "descending";
     }
     setSortConfig({ key, direction });
   };
@@ -150,6 +166,53 @@ const App: React.FC = () => {
     );
   };
 
+  const getModelsForSelectedProviders = () => {
+    if (!linkProviderModel) {
+      return data
+        .flatMap((provider) =>
+          provider.models.map((model) => ({
+            label: model.name,
+            value: model.name,
+            provider: provider.provider,
+          }))
+        )
+        .reduce(
+          (
+            acc: { label: string; value: string; provider: string }[],
+            curr: { label: string; value: string; provider: string }
+          ) => {
+            if (!acc.find((m) => m.value === curr.value)) {
+              acc.push(curr);
+            }
+            return acc;
+          },
+          []
+        );
+    }
+
+    return data
+      .filter((provider) => selectedProviders.includes(provider.provider))
+      .flatMap((provider) =>
+        provider.models.map((model) => ({
+          label: model.name,
+          value: model.name,
+          provider: provider.provider,
+        }))
+      )
+      .reduce(
+        (
+          acc: { label: string; value: string; provider: string }[],
+          curr: { label: string; value: string; provider: string }
+        ) => {
+          if (!acc.find((m) => m.value === curr.value)) {
+            acc.push(curr);
+          }
+          return acc;
+        },
+        []
+      );
+  };
+
   return (
     <Card className="w-full max-w-6xl mx-auto">
       <CardHeader>
@@ -158,9 +221,16 @@ const App: React.FC = () => {
       <CardContent>
         <div className="mb-4">
           <p className="italic text-sm text-muted-foreground mb-4">
-            <a href="https://huggingface.co/spaces/philschmid/llm-pricing" className='underline'>This is a fork of philschmid tool: philschmid/llm-pricing</a>
+            <a
+              href="https://huggingface.co/spaces/philschmid/llm-pricing"
+              className="underline"
+            >
+              This is a fork of philschmid tool: philschmid/llm-pricing
+            </a>
           </p>
-          <h3 className="text-lg font-semibold mb-2">Select Comparison Models</h3>
+          <h3 className="text-lg font-semibold mb-2">
+            Select Comparison Models
+          </h3>
           <div className="grid grid-cols-2 md:grid-cols-3 lg:grid-cols-4 gap-4">
             {data.map((provider) => (
               <Collapsible
@@ -198,7 +268,8 @@ const App: React.FC = () => {
                           } else {
                             setComparisonModels((prev) =>
                               prev.filter(
-                                (m) => m !== `${provider.provider}:${model.name}`
+                                (m) =>
+                                  m !== `${provider.provider}:${model.name}`
                               )
                             );
                           }
@@ -224,7 +295,6 @@ const App: React.FC = () => {
               htmlFor="inputTokens"
               className="block text-sm font-medium text-gray-700"
             >
-
               Input Tokens ({tokenCalculation})
             </label>
             <Input
@@ -270,72 +340,75 @@ const App: React.FC = () => {
               <option value="thousand">Thousand Tokens</option>
               <option value="unit">Unit Tokens</option>
             </select>
-
           </div>
-
         </div>
 
-
         <p className="italic text-sm text-muted-foreground mb-4">
           Note: If you use Amazon Bedrock or Azure prices for Anthropic, Cohere
           or OpenAI should be the same.
         </p>
+        <div className="flex items-center space-x-2 mb-4">
+          <Switch
+            id="linkProviderModel"
+            checked={linkProviderModel}
+            onCheckedChange={setLinkProviderModel}
+          />
+          <label htmlFor="linkProviderModel" className="text-sm">
+            Link Provider and Model
+          </label>
+        </div>
 
         <Table>
           <TableHeader>
             <TableRow>
               <TableHead>
-                <button type="button" onClick={() => requestSort('provider')}>
-                  Provider{' '}
-                  {sortConfig?.key === 'provider' ? (
-                    sortConfig.direction === 'ascending' ? (
-                      '▲'
-                    ) : (
-                      '▼'
-                    )
-                  ) : null}
+                <button type="button" onClick={() => requestSort("provider")}>
+                  Provider{" "}
+                  {sortConfig?.key === "provider"
+                    ? sortConfig.direction === "ascending"
+                      ? "▲"
+                      : "▼"
+                    : null}
                 </button>
               </TableHead>
               <TableHead>
-                <button type="button" onClick={() => requestSort('name')}>
-                  Model{' '}
-                  {sortConfig?.key === 'name' ? (
-                    sortConfig.direction === 'ascending' ? (
-                      '▲'
-                    ) : (
-                      '▼'
-                    )
-                  ) : null}
+                <button type="button" onClick={() => requestSort("name")}>
+                  Model{" "}
+                  {sortConfig?.key === "name"
+                    ? sortConfig.direction === "ascending"
+                      ? "▲"
+                      : "▼"
+                    : null}
                 </button>
               </TableHead>
 
               <TableHead>
-                <button type="button" onClick={() => requestSort('inputPrice')}>
+                <button type="button" onClick={() => requestSort("inputPrice")}>
                   Input Price (million tokens)
-                  {sortConfig?.key === 'inputPrice' ? (
-                    sortConfig.direction === 'ascending' ? (
-                      '▲'
-                    ) : (
-                      '▼'
-                    )
-                  ) : null}
+                  {sortConfig?.key === "inputPrice"
+                    ? sortConfig.direction === "ascending"
+                      ? "▲"
+                      : "▼"
+                    : null}
                 </button>
               </TableHead>
               <TableHead>
-                <button type="button" onClick={() => requestSort('outputPrice')}>
+                <button
+                  type="button"
+                  onClick={() => requestSort("outputPrice")}
+                >
                   Output Price (million tokens)
-                  {sortConfig?.key === 'outputPrice' ? (
-                    sortConfig.direction === 'ascending' ? (
-                      '▲'
-                    ) : (
-                      '▼'
-                    )
-                  ) : null}
+                  {sortConfig?.key === "outputPrice"
+                    ? sortConfig.direction === "ascending"
+                      ? "▲"
+                      : "▼"
+                    : null}
                 </button>
               </TableHead>
 
-
-              <TableHead>Total Price (per {tokenCalculation} tokens){' '}</TableHead>
+              <TableHead>
+                Total Price (per {tokenCalculation} tokens){" "}
+              </TableHead>
               {comparisonModels.map((model) => (
                 <TableHead key={model} colSpan={2}>
                   Compared to {model}
@@ -357,23 +430,7 @@ const App: React.FC = () => {
               </TableHead>
               <TableHead>
                 <MultiSelect
-                  options={
-                    data
-                      .flatMap((provider) => provider.models)
-                      .map((model) => ({ label: model.name, value: model.name }))
-                      .reduce(
-                        (
-                          acc: { label: string; value: string }[],
-                          curr: { label: string; value: string }
-                        ) => {
-                          if (!acc.find((m) => m.value === curr.value)) {
-                            acc.push(curr);
-                          }
-                          return acc;
-                        },
-                        []
-                      ) || []
-                  }
+                  options={getModelsForSelectedProviders()}
                   defaultValue={selectedModels}
                   onValueChange={setSelectedModels}
                 />
@@ -391,7 +448,7 @@ const App: React.FC = () => {
             {sortedFlattenedData.map((item) => (
               <TableRow key={`${item.provider}-${item.name}`}>
                 <TableCell>
-                  {' '}
+                  {" "}
                   <a href={item.uri} className="underline">
                     {item.provider}
                   </a>
@@ -409,65 +466,66 @@ const App: React.FC = () => {
                   ).toFixed(2)}
                 </TableCell>
 
-
                 {comparisonModels.flatMap((comparisonModel) => {
                   const [comparisonProvider, comparisonModelName] =
-                    comparisonModel.split(':');
+                    comparisonModel.split(":");
                   const comparisonModelData = data
                     .find((p) => p.provider === comparisonProvider)
                     ?.models.find((m) => m.name === comparisonModelName)!;
                   return [
                     <TableCell
                       key={`${comparisonModel}-input`}
-                      className={`${parseFloat(
-                        calculateComparison(
-                          item.inputPrice,
-                          comparisonModelData.inputPrice
-                        )
-                      ) < 0
-                        ? 'bg-green-100'
-                        : parseFloat(
+                      className={`${
+                        parseFloat(
                           calculateComparison(
                             item.inputPrice,
                             comparisonModelData.inputPrice
                           )
-                        ) > 0
-                          ? 'bg-red-100'
-                          : ''
-                        }`}
+                        ) < 0
+                          ? "bg-green-100"
+                          : parseFloat(
+                              calculateComparison(
+                                item.inputPrice,
+                                comparisonModelData.inputPrice
+                              )
+                            ) > 0
+                          ? "bg-red-100"
+                          : ""
+                      }`}
                     >
                       {`${item.provider}:${item.name}` === comparisonModel
-                        ? '0.00%'
+                        ? "0.00%"
                         : `${calculateComparison(
-                          item.inputPrice,
-                          comparisonModelData.inputPrice
-                        )}%`}
+                            item.inputPrice,
+                            comparisonModelData.inputPrice
+                          )}%`}
                     </TableCell>,
                     <TableCell
                       key={`${comparisonModel}-output`}
-                      className={`${parseFloat(
-                        calculateComparison(
-                          item.outputPrice,
-                          comparisonModelData.outputPrice
-                        )
-                      ) < 0
-                        ? 'bg-green-100'
-                        : parseFloat(
+                      className={`${
+                        parseFloat(
                           calculateComparison(
                             item.outputPrice,
                             comparisonModelData.outputPrice
                           )
-                        ) > 0
-                          ? 'bg-red-100'
-                          : ''
-                        }`}
+                        ) < 0
+                          ? "bg-green-100"
+                          : parseFloat(
+                              calculateComparison(
+                                item.outputPrice,
+                                comparisonModelData.outputPrice
+                              )
+                            ) > 0
+                          ? "bg-red-100"
+                          : ""
+                      }`}
                     >
                       {`${item.provider}:${item.name}` === comparisonModel
-                        ? '0.00%'
+                        ? "0.00%"
                         : `${calculateComparison(
-                          item.outputPrice,
-                          comparisonModelData.outputPrice
-                        )}%`}
+                            item.outputPrice,
+                            comparisonModelData.outputPrice
+                          )}%`}
                     </TableCell>,
                   ];
                 })}