Sync models from remote for FireworksAI (#4475)

resolves #4474

Author: timothycarambat

frontend/src/components/LLMSelection/FireworksAiOptions/index.jsx

@@ -2,6 +2,11 @@ import System from "@/models/system";
 import { useState, useEffect } from "react";
 
 export default function FireworksAiOptions({ settings }) {
+  const [inputValue, setInputValue] = useState(settings?.FireworksAiLLMApiKey);
+  const [fireworksAiApiKey, setFireworksAiApiKey] = useState(
+    settings?.FireworksAiLLMApiKey
+  );
+
   return (
     <div className="flex gap-[36px] mt-1.5">
       <div className="flex flex-col w-60">
@@ -17,22 +22,27 @@ export default function FireworksAiOptions({ settings }) {
           required={true}
           autoComplete="off"
           spellCheck={false}
+          onChange={(e) => setInputValue(e.target.value)}
+          onBlur={() => setFireworksAiApiKey(inputValue)}
         />
       </div>
       {!settings?.credentialsOnly && (
-        <FireworksAiModelSelection settings={settings} />
+        <FireworksAiModelSelection
+          apiKey={fireworksAiApiKey}
+          settings={settings}
+        />
       )}
     </div>
   );
 }
-function FireworksAiModelSelection({ settings }) {
+function FireworksAiModelSelection({ apiKey, settings }) {
   const [groupedModels, setGroupedModels] = useState({});
   const [loading, setLoading] = useState(true);
 
   useEffect(() => {
     async function findCustomModels() {
       setLoading(true);
-      const { models } = await System.customModels("fireworksai");
+      const { models } = await System.customModels("fireworksai", apiKey);
 
       if (models?.length > 0) {
         const modelsByOrganization = models.reduce((acc, model) => {
@@ -47,7 +57,7 @@ function FireworksAiModelSelection({ settings }) {
       setLoading(false);
     }
     findCustomModels();
-  }, []);
+  }, [apiKey]);
 
   if (loading || Object.keys(groupedModels).length === 0) {
     return (

server/storage/models/.gitignore

@@ -11,4 +11,5 @@ tesseract
 ppio
 context-windows/*
 MintplexLabs
-cometapi
+cometapi
+fireworks

server/utils/AiProviders/fireworksAi/index.js

@@ -1,3 +1,6 @@
+const fs = require("fs");
+const path = require("path");
+const { safeJsonParse } = require("../../http");
 const { NativeEmbedder } = require("../../EmbeddingEngines/native");
 const {
   LLMPerformanceMonitor,
@@ -6,13 +9,16 @@ const {
   handleDefaultStreamResponseV2,
 } = require("../../helpers/chat/responses");
 
-function fireworksAiModels() {
-  const { MODELS } = require("./models.js");
-  return MODELS || {};
-}
+const cacheFolder = path.resolve(
+  process.env.STORAGE_DIR
+    ? path.resolve(process.env.STORAGE_DIR, "models", "fireworks")
+    : path.resolve(__dirname, `../../../storage/models/fireworks`)
+);
 
 class FireworksAiLLM {
   constructor(embedder = null, modelPreference = null) {
+    this.className = "FireworksAiLLM";
+
     if (!process.env.FIREWORKS_AI_LLM_API_KEY)
       throw new Error("No FireworksAI API key was set.");
     const { OpenAI: OpenAIApi } = require("openai");
@@ -29,6 +35,51 @@ class FireworksAiLLM {
 
     this.embedder = !embedder ? new NativeEmbedder() : embedder;
     this.defaultTemp = 0.7;
+
+    if (!fs.existsSync(cacheFolder))
+      fs.mkdirSync(cacheFolder, { recursive: true });
+    this.cacheModelPath = path.resolve(cacheFolder, "models.json");
+    this.cacheAtPath = path.resolve(cacheFolder, ".cached_at");
+  }
+
+  log(text, ...args) {
+    console.log(`\x1b[36m[${this.className}]\x1b[0m ${text}`, ...args);
+  }
+
+  // This checks if the .cached_at file has a timestamp that is more than 1Week (in millis)
+  // from the current date. If it is, then we will refetch the API so that all the models are up
+  // to date.
+  #cacheIsStale() {
+    const MAX_STALE = 6.048e8; // 1 Week in MS
+    if (!fs.existsSync(this.cacheAtPath)) return true;
+    const now = Number(new Date());
+    const timestampMs = Number(fs.readFileSync(this.cacheAtPath));
+    return now - timestampMs > MAX_STALE;
+  }
+
+  // This function fetches the models from the ApiPie API and caches them locally.
+  // We do this because the ApiPie API has a lot of models, and we need to get the proper token context window
+  // for each model and this is a constructor property - so we can really only get it if this cache exists.
+  // We used to have this as a chore, but given there is an API to get the info - this makes little sense.
+  // This might slow down the first request, but we need the proper token context window
+  // for each model and this is a constructor property - so we can really only get it if this cache exists.
+  async #syncModels() {
+    if (fs.existsSync(this.cacheModelPath) && !this.#cacheIsStale())
+      return false;
+
+    this.log(
+      "Model cache is not present or stale. Fetching from FireworksAI API."
+    );
+    await fireworksAiModels();
+    return;
+  }
+
+  models() {
+    if (!fs.existsSync(this.cacheModelPath)) return {};
+    return safeJsonParse(
+      fs.readFileSync(this.cacheModelPath, { encoding: "utf-8" }),
+      {}
+    );
   }
 
   #appendContext(contextTexts = []) {
@@ -43,28 +94,31 @@ class FireworksAiLLM {
     );
   }
 
-  allModelInformation() {
-    return fireworksAiModels();
-  }
-
   streamingEnabled() {
     return "streamGetChatCompletion" in this;
   }
 
   static promptWindowLimit(modelName) {
-    const availableModels = fireworksAiModels();
+    const cacheModelPath = path.resolve(cacheFolder, "models.json");
+    const availableModels = fs.existsSync(cacheModelPath)
+      ? safeJsonParse(
+          fs.readFileSync(cacheModelPath, { encoding: "utf-8" }),
+          {}
+        )
+      : {};
     return availableModels[modelName]?.maxLength || 4096;
   }
 
   // Ensure the user set a value for the token limit
   // and if undefined - assume 4096 window.
   promptWindowLimit() {
-    const availableModels = this.allModelInformation();
+    const availableModels = this.models();
     return availableModels[this.model]?.maxLength || 4096;
   }
 
   async isValidChatCompletionModel(model = "") {
-    const availableModels = this.allModelInformation();
+    await this.#syncModels();
+    const availableModels = this.models();
     return availableModels.hasOwnProperty(model);
   }
 
@@ -151,6 +205,63 @@ class FireworksAiLLM {
   }
 }
 
+async function fireworksAiModels(providedApiKey = null) {
+  const apiKey = providedApiKey || process.env.FIREWORKS_AI_LLM_API_KEY || null;
+  const { OpenAI: OpenAIApi } = require("openai");
+  const client = new OpenAIApi({
+    baseURL: "https://api.fireworks.ai/inference/v1",
+    apiKey: apiKey,
+  });
+
+  return await client.models
+    .list()
+    .then((res) => res.data)
+    .then((models = []) => {
+      const validModels = {};
+      models.forEach((model) => {
+        // There are many models - the ones without a context length are not chat models
+        if (!model.hasOwnProperty("context_length")) return;
+
+        validModels[model.id] = {
+          id: model.id,
+          name: model.id.split("/").pop(),
+          organization: model.owned_by,
+          subtype: model.type,
+          maxLength: model.context_length ?? 4096,
+        };
+      });
+
+      if (Object.keys(validModels).length === 0) {
+        console.log("fireworksAi: No models found");
+        return {};
+      }
+
+      // Cache all response information
+      if (!fs.existsSync(cacheFolder))
+        fs.mkdirSync(cacheFolder, { recursive: true });
+      fs.writeFileSync(
+        path.resolve(cacheFolder, "models.json"),
+        JSON.stringify(validModels),
+        {
+          encoding: "utf-8",
+        }
+      );
+      fs.writeFileSync(
+        path.resolve(cacheFolder, ".cached_at"),
+        String(Number(new Date())),
+        {
+          encoding: "utf-8",
+        }
+      );
+
+      return validModels;
+    })
+    .catch((e) => {
+      console.error(e);
+      return {};
+    });
+}
+
 module.exports = {
   FireworksAiLLM,
   fireworksAiModels,