add support for arbitrary HF endpoints

src/lib/agents/llm.ts

@@ -1,8 +1,9 @@
 import { get } from "svelte/store";
-import { OPENAI_API_KEY } from "$lib/store";
+import { OPENAI_API_KEY, HF_ENDPOINT } from "$lib/store";
 import { Configuration, OpenAIApi } from "openai";
 import { getInference } from "./getInference";
 import type { LLM } from "$lib/types";
+import type { TextGenerationOutput } from "@huggingface/inference";
 
 async function OpenAILLMCall(prompt: string) {
   const openai = new OpenAIApi(
@@ -21,19 +22,30 @@ async function OpenAILLMCall(prompt: string) {
 }
 
 async function HFLLMCall(prompt: string) {
-  const hf = getInference();
-
-  const formattedPrompt = "<|user|>" + prompt + "<|end|>";
-
-  const generatedOutput = await hf.textGeneration({
-    inputs: formattedPrompt,
-    model: "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
-    parameters: {
-      max_new_tokens: 512,
-    },
-  });
-
-  const text = generatedOutput.generated_text.slice(formattedPrompt.length);
+  const formattedPrompt = "<|user|>" + prompt + "<|end|><|assistant|>";
+
+  let output: TextGenerationOutput;
+
+  if (!!get(HF_ENDPOINT)) {
+    output = await getInference()
+      .endpoint(get(HF_ENDPOINT))
+      .textGeneration({
+        inputs: formattedPrompt,
+        parameters: {
+          max_new_tokens: 1400,
+        },
+      });
+  } else {
+    output = await getInference().textGeneration({
+      inputs: formattedPrompt,
+      model: "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
+      parameters: {
+        max_new_tokens: 900,
+      },
+    });
+  }
+
+  const text = output.generated_text.slice(formattedPrompt.length);
 
   return text;
 }

src/lib/components/ApiKeyModal.svelte

@@ -1,9 +1,9 @@
 <script lang="ts">
   import { onMount } from "svelte";
-  import { HF_ACCESS_TOKEN, OPENAI_API_KEY } from "$lib/store";
+  import { HF_ACCESS_TOKEN, OPENAI_API_KEY, HF_ENDPOINT } from "$lib/store";
   import { goto } from "$app/navigation";
 
-  let dialogElement: HTMLDialogElement;
+  export let dialogElement: HTMLDialogElement;
 
   onMount(() => {
     if ($HF_ACCESS_TOKEN === "") {
@@ -22,6 +22,7 @@
     } else {
       localStorage.setItem("HF_ACCESS_TOKEN", $HF_ACCESS_TOKEN);
       localStorage.setItem("OPENAI_API_KEY", $OPENAI_API_KEY);
+      localStorage.setItem("HF_ENDPOINT", $HF_ENDPOINT);
       goto("/");
     }
   }}
@@ -42,8 +43,20 @@
           bind:value={$HF_ACCESS_TOKEN}
         />
       </form>
+      <form aria-label="HF ENDPOINT" class="form-control">
+        <label for="hf_endpoint" class="label-text pb-2">HF Endpoint</label>
+        <input
+          class="input input-primary"
+          name="hf_endpoint"
+          type="text"
+          placeholder="Leave empty to use the default endpoint"
+          bind:value={$HF_ENDPOINT}
+        />
+      </form>
+
+      <div class="divider my-0" />
       <form aria-label="OPENAI API" class="form-control">
-        <p class="py-4">
+        <p class="pb-4">
           Optionally you can add your OpenAI key to use it as your LLM.
         </p>
 

src/lib/store.ts

@@ -10,7 +10,12 @@ export const OPENAI_API_KEY = writable(
   (browser && localStorage.OPENAI_API_KEY) || ""
 );
 
+export const HF_ENDPOINT = writable(
+  (browser && localStorage.HF_ENDPOINT) || ""
+);
+
 if (browser) {
   HF_ACCESS_TOKEN.subscribe((value) => (localStorage.content = value));
   OPENAI_API_KEY.subscribe((value) => (localStorage.content = value));
+  HF_ENDPOINT.subscribe((value) => (localStorage.content = value));
 }

src/routes/+layout.svelte

@@ -4,6 +4,5 @@
 </script>
 
 <div class="p-4 max-w-xl mx-auto">
-  <ApiKeyModal />
   <slot />
 </div>

src/routes/+page.svelte

@@ -11,6 +11,7 @@
   import type { LLM } from "$lib/types";
   import { HFLLM } from "$lib/agents/llm";
   import { OPENAI_API_KEY } from "$lib/store";
+  import ApiKeyModal from "$lib/components/ApiKeyModal.svelte";
 
   let prompt =
     "Draw a picture of a cat wearing a top hat. Then caption the picture and read it out loud.";
@@ -36,14 +37,15 @@
 
   const onRun = async (code: string) => {
     messages = [];
+    const callback = (message: string, data: string | Blob | undefined) => {
+      messages = [...messages, { message, data }];
+    };
 
     const wrapperEval = await evalBuilder(
       code,
       tools.filter((el) => selectedTools.includes(el.name)),
       files,
-      (message, data) => {
-        messages = [...messages, { message, data }];
-      }
+      callback
     );
 
     isLoading = true;
@@ -59,12 +61,21 @@
     }
     isLoading = false;
   };
+
+  let dialogElement: HTMLDialogElement;
 </script>
 
+<ApiKeyModal bind:dialogElement />
+
 <div class="flex flex-col space-y-4 max-w-xl">
-  <div class="flex flex-row justify-around">
+  <div class="flex flex-row">
     <LogoHuggingFaceBorderless classNames="text-4xl" />
-    <h1 class="text-3xl font-semibold w-fit mx-auto">Agents.js</h1>
+    <h1 class="text-3xl font-semibold mx-auto">Agents.js</h1>
+    <button
+      class="btn btn-ghost"
+      on:click={() => dialogElement.showModal()}
+      on:keydown={() => dialogElement.showModal()}>API keys</button
+    >
   </div>
 
   {#if $OPENAI_API_KEY !== ""}