Add providers selector

server.js

@@ -8,6 +8,7 @@ import { InferenceClient } from "@huggingface/inference";
 import bodyParser from "body-parser";
 
 import checkUser from "./middlewares/checkUser.js";
+import { PROVIDERS } from "./utils/providers.js";
 
 // Load environment variables from .env file
 dotenv.config();
@@ -174,7 +175,7 @@ Check out the configuration reference at https://huggingface.co/docs/hub/spaces-
 });
 
 app.post("/api/ask-ai", async (req, res) => {
-  const { prompt, html, previousPrompt } = req.body;
+  const { prompt, html, previousPrompt, provider } = req.body;
   if (!prompt) {
     return res.status(400).send({
       ok: false,
@@ -192,7 +193,6 @@ app.post("/api/ask-ai", async (req, res) => {
     "0.0.0.0";
 
   if (!hf_token) {
-    // Rate limit requests from the same IP address, to prevent abuse, free is limited to 2 requests per IP
     ipAddresses.set(ip, (ipAddresses.get(ip) || 0) + 1);
     if (ipAddresses.get(ip) > MAX_REQUESTS_PER_IP) {
       return res.status(429).send({
@@ -213,10 +213,26 @@ app.post("/api/ask-ai", async (req, res) => {
   const client = new InferenceClient(token);
   let completeResponse = "";
 
+  const selectedProvider =
+    PROVIDERS.find((providerItem) => providerItem.id === provider) ??
+    PROVIDERS[0];
+
+  let TOKENS_USED = prompt?.length;
+  if (previousPrompt) TOKENS_USED += previousPrompt.length;
+  if (html) TOKENS_USED += html.length;
+
+  if (TOKENS_USED >= selectedProvider.max_tokens) {
+    return res.status(400).send({
+      ok: false,
+      openSelectProvider: true,
+      message: `Context is too long. ${selectedProvider.name} allow ${selectedProvider.max_tokens} max tokens.`,
+    });
+  }
+
   try {
     const chatCompletion = client.chatCompletionStream({
       model: MODEL_ID,
-      provider: "hyperbolic",
+      provider: selectedProvider.id,
       messages: [
         {
           role: "system",
@@ -244,7 +260,7 @@ app.post("/api/ask-ai", async (req, res) => {
           content: prompt,
         },
       ],
-      max_tokens: 12_000,
+      max_tokens: selectedProvider.max_tokens,
     });
 
     while (true) {
@@ -254,25 +270,37 @@ app.post("/api/ask-ai", async (req, res) => {
       }
       const chunk = value.choices[0]?.delta?.content;
       if (chunk) {
-        res.write(chunk);
-        completeResponse += chunk;
-
-        // Break when HTML is complete
-        if (completeResponse.includes("</html>")) {
-          break;
+        if (provider !== "sambanova") {
+          res.write(chunk);
+          completeResponse += chunk;
+
+          if (completeResponse.includes("</html>")) {
+            break;
+          }
+        } else {
+          let newChunk = chunk;
+          if (chunk.includes("</html>")) {
+            // Replace everything after the last </html> tag with an empty string
+            newChunk = newChunk.replace(/<\/html>[\s\S]*/, "</html>");
+          }
+          completeResponse += newChunk;
+          res.write(newChunk);
+          if (newChunk.includes("</html>")) {
+            break;
+          }
         }
       }
     }
-
     // End the response stream
+    // return the total_tokens used to the client
     res.end();
   } catch (error) {
-    console.error("Error:", error);
-    // If we haven't sent a response yet, send an error
     if (!res.headersSent) {
       res.status(500).send({
         ok: false,
-        message: `You probably reached the MAX_TOKENS limit, context is too long. You can start a new conversation by refreshing the page.`,
+        // use generic error,
+        message:
+          "An error occurred while processing your request. Please try again or switch provider.",
       });
     } else {
       // Otherwise end the stream

src/components/App.tsx

@@ -7,10 +7,10 @@ import { toast } from "react-toastify";
 
 import Header from "./header/header";
 import DeployButton from "./deploy-button/deploy-button";
-import { defaultHTML } from "../utils/consts";
+import { defaultHTML } from "./../../utils/consts";
 import Tabs from "./tabs/tabs";
 import AskAI from "./ask-ai/ask-ai";
-import { Auth } from "../utils/types";
+import { Auth } from "./../../utils/types";
 import Preview from "./preview/preview";
 
 function App() {

src/components/ask-ai/ask-ai.tsx

@@ -5,8 +5,10 @@ import classNames from "classnames";
 import { toast } from "react-toastify";
 
 import Login from "../login/login";
-import { defaultHTML } from "../../utils/consts";
+import { defaultHTML } from "./../../../utils/consts";
 import SuccessSound from "./../../assets/success.mp3";
+import Settings from "../settings/settings";
+import { useLocalStorage } from "react-use";
 
 function AskAI({
   html,
@@ -25,12 +27,17 @@ function AskAI({
   const [prompt, setPrompt] = useState("");
   const [hasAsked, setHasAsked] = useState(false);
   const [previousPrompt, setPreviousPrompt] = useState("");
+  const [provider, setProvider] = useLocalStorage("provider", "fireworks-ai");
+  const [openProvider, setOpenProvider] = useState(false);
+  const [providerError, setProviderError] = useState("");
+
   const audio = new Audio(SuccessSound);
   audio.volume = 0.5;
 
   const callAi = async () => {
     if (isAiWorking || !prompt.trim()) return;
     setisAiWorking(true);
+    setProviderError("");
 
     let contentResponse = "";
     let lastRenderTime = 0;
@@ -39,6 +46,7 @@ function AskAI({
         method: "POST",
         body: JSON.stringify({
           prompt,
+          provider,
           ...(html === defaultHTML ? {} : { html }),
           ...(previousPrompt ? { previousPrompt } : {}),
         }),
@@ -51,8 +59,10 @@ function AskAI({
           const res = await request.json();
           if (res.openLogin) {
             setOpen(true);
+          } else if (res.openSelectProvider) {
+            setOpenProvider(true);
+            setProviderError(res.message);
           } else {
-            // don't show toast if it's a login error
             toast.error(res.message);
           }
           setisAiWorking(false);
@@ -130,7 +140,7 @@ function AskAI({
         <input
           type="text"
           disabled={isAiWorking}
-          className="w-full bg-transparent max-lg:text-sm outline-none pl-3 text-white placeholder:text-gray-500 font-code"
+          className="w-full bg-transparent max-lg:text-sm outline-none px-3 text-white placeholder:text-gray-500 font-code"
           placeholder={
             hasAsked ? "What do you want to ask AI next?" : "Ask AI anything..."
           }
@@ -142,13 +152,22 @@ function AskAI({
             }
           }}
         />
-        <button
-          disabled={isAiWorking}
-          className="relative overflow-hidden cursor-pointer flex-none flex items-center justify-center rounded-full text-sm font-semibold size-8 text-center bg-pink-500 hover:bg-pink-400 text-white shadow-sm dark:shadow-highlight/20 disabled:bg-gray-300 disabled:text-gray-500 disabled:cursor-not-allowed disabled:hover:bg-gray-300"
-          onClick={callAi}
-        >
-          <GrSend className="-translate-x-[1px]" />
-        </button>
+        <div className="flex items-center justify-end gap-2">
+          <Settings
+            provider={provider as string}
+            onChange={setProvider}
+            open={openProvider}
+            error={providerError}
+            onClose={setOpenProvider}
+          />
+          <button
+            disabled={isAiWorking}
+            className="relative overflow-hidden cursor-pointer flex-none flex items-center justify-center rounded-full text-sm font-semibold size-8 text-center bg-pink-500 hover:bg-pink-400 text-white shadow-sm dark:shadow-highlight/20 disabled:bg-gray-300 disabled:text-gray-500 disabled:cursor-not-allowed disabled:hover:bg-gray-300"
+            onClick={callAi}
+          >
+            <GrSend className="-translate-x-[1px]" />
+          </button>
+        </div>
       </div>
       <div
         className={classNames(

src/components/deploy-button/deploy-button.tsx

@@ -6,7 +6,7 @@ import { toast } from "react-toastify";
 import SpaceIcon from "@/assets/space.svg";
 import Loading from "../loading/loading";
 import Login from "../login/login";
-import { Auth } from "../../utils/types";
+import { Auth } from "./../../../utils/types";
 
 const MsgToast = ({ url }: { url: string }) => (
   <div className="w-full flex items-center justify-center gap-3">

src/components/login/login.tsx

@@ -1,5 +1,5 @@
 import { useLocalStorage } from "react-use";
-import { defaultHTML } from "../../utils/consts";
+import { defaultHTML } from "./../../../utils/consts";
 
 function Login({
   html,

src/components/settings/settings.tsx

@@ -1,17 +1,32 @@
-import { useState } from "react";
+/* eslint-disable @typescript-eslint/no-explicit-any */
 import classNames from "classnames";
-import Login from "../login/login";
 
-function Settings() {
-  const [open, setOpen] = useState(false);
+import { PiGearSixFill } from "react-icons/pi";
+// @ts-expect-error not needed
+import { PROVIDERS } from "./../../../utils/providers";
 
+function Settings({
+  open,
+  onClose,
+  provider,
+  error,
+  onChange,
+}: {
+  open: boolean;
+  provider: string;
+  error?: string;
+  onClose: React.Dispatch<React.SetStateAction<boolean>>;
+  onChange: (provider: string) => void;
+}) {
   return (
-    <div className="relative">
+    <div className="">
       <button
-        className="bg-gray-800/70 rounded-md text-xs text-gray-300 hover:brightness-125 px-3 py-1.5 font-medium cursor-pointer"
-        onClick={() => setOpen(!open)}
+        className="relative overflow-hidden cursor-pointer flex-none flex items-center justify-center rounded-full text-base font-semibold size-8 text-center bg-gray-800 hover:bg-gray-700 text-gray-100 shadow-sm dark:shadow-highlight/20"
+        onClick={() => {
+          onClose((prev) => !prev);
+        }}
       >
-        Settings
+        <PiGearSixFill />
       </button>
       <div
         className={classNames(
@@ -20,17 +35,73 @@ function Settings() {
             "opacity-0 pointer-events-none": !open,
           }
         )}
-        onClick={() => setOpen(false)}
+        onClick={() => onClose(false)}
       ></div>
       <div
         className={classNames(
-          "absolute top-[calc(100%+8px)] right-0 z-10 w-80 bg-white border border-gray-200 rounded-lg shadow-lg transition-all duration-75 overflow-hidden",
+          "absolute top-0 -translate-y-[calc(100%+16px)] right-0 z-10 w-96 bg-white border border-gray-200 rounded-lg shadow-lg transition-all duration-75 overflow-hidden",
           {
             "opacity-0 pointer-events-none": !open,
           }
         )}
       >
-        <Login />
+        <header className="flex items-center text-sm px-4 py-2 border-b border-gray-200 gap-2 bg-gray-100 font-semibold text-gray-700">
+          <span className="text-xs bg-blue-500/10 text-blue-500 rounded-full pl-1.5 pr-2.5 py-0.5 flex items-center justify-start gap-1.5">
+            Provider
+          </span>
+          Customize Settings
+        </header>
+        <main className="px-4 pt-3 pb-4 space-y-3">
+          {error !== "" && (
+            <p className="text-red-500 text-sm font-medium mb-2 flex items-center justify-between bg-red-500/10 p-2 rounded-md">
+              {error}
+            </p>
+          )}
+          <label className="block">
+            <p className="text-gray-500 text-sm font-medium mb-2 flex items-center justify-between">
+              Inference Provider
+            </p>
+            <div className="grid grid-cols-2 gap-1.5">
+              {PROVIDERS.map((item: any) => (
+                <div
+                  key={item.id}
+                  className={classNames(
+                    "text-gray-600 text-sm font-medium cursor-pointer border p-2 rounded-md flex items-center justify-start gap-2",
+                    {
+                      "bg-blue-500/10 border-blue-500/15 text-blue-500":
+                        item.id === provider,
+                      "hover:bg-gray-100 border-gray-100": item.id !== provider,
+                    }
+                  )}
+                  onClick={() => {
+                    onChange(item.id);
+                  }}
+                >
+                  <img
+                    src={`/providers/${item.id}.svg`}
+                    alt={item.name}
+                    className="size-5"
+                  />
+                  {item.name}
+                </div>
+              ))}
+            </div>
+            {/* <input
+              type="password"
+              autoComplete="off"
+              className="mr-2 border rounded-md px-3 py-1.5 border-gray-300 w-full text-sm"
+              placeholder="hf_******"
+              value={tokenStorage[0] as string}
+              onChange={(e) => {
+                if (e.target.value.length > 0) {
+                  tokenStorage[1](e.target.value);
+                } else {
+                  tokenStorage[2]();
+                }
+              }}
+            /> */}
+          </label>
+        </main>
       </div>
     </div>
   );

src/components/tabs/tabs.tsx

@@ -14,7 +14,7 @@ function Tabs({ children }: { children?: React.ReactNode }) {
       </div>
       <div className="flex items-center justify-end gap-3">
         <a
-          href="https://huggingface.co/deepseek-ai/DeepSeek-V3-0324?inference_provider=fireworks-ai"
+          href="https://huggingface.co/deepseek-ai/DeepSeek-V3-0324"
           target="_blank"
           className="text-[12px] text-gray-300 hover:brightness-120 flex items-center gap-1 font-code"
         >

tsconfig.app.json

@@ -22,5 +22,5 @@
     "noFallthroughCasesInSwitch": true,
     "noUncheckedSideEffectImports": true
   },
-  "include": ["src", "middleware"]
+  "include": ["src", "middleware", "utils/consts.ts", "utils/types.ts"]
 }

utils/providers.js

@@ -0,0 +1,22 @@
+export const PROVIDERS = [
+  {
+    name: "Fireworks AI",
+    id: "fireworks-ai",
+    max_tokens: 131_000,
+  },
+  {
+    name: "Nebius AI Studio",
+    id: "nebius",
+    max_tokens: 131_000,
+  },
+  {
+    name: "SambaNova",
+    id: "sambanova",
+    max_tokens: 8_000,
+  },
+  {
+    name: "Hyperbolic",
+    id: "hyperbolic",
+    max_tokens: 131_000,
+  },
+];