add hf model support

components/Playground/index.tsx

@@ -36,7 +36,7 @@ const Playground = () => {
   const [prevMarkdownCode, setPrevMarkdownCode] = useState(markdownCode);
   const [isSystemInputVisible, setSystemInputVisible] = useState(false);
   const [isModelInputVisible, setModelInputVisible] = useState(false);
-  const [selectedModel, setSelectedModel] = useState("GPT-4");
+  const [aiProvider, setAIProvider] = useState<string>("openai");
 
   const [systemMessage, setSystemMessage] = useState(
     DEFAULT_PROMPT
@@ -48,8 +48,8 @@ const Playground = () => {
 
   const modifiedHandleSubmit = async (e: FormEvent<HTMLFormElement>, chatRequestOptions?: ChatRequestOptions) => {
     e.preventDefault();
-    // Simply append the user message
-    await handleSubmit(e, chatRequestOptions);
+    // Pass the aiProvider in chatRequestOptions
+    await handleSubmit(e, { ...chatRequestOptions, aiProvider } as any);
   };
 
   useEffect(() => {
@@ -67,7 +67,8 @@ const Playground = () => {
 
   const { append, messages, input, setInput, handleSubmit, ...rest } = useChat({
     body: {
-      systemMessage: systemMessage
+      systemMessage: systemMessage,
+      aiProvider: aiProvider
     },
     onError: (error) => {
       console.error(error);
@@ -149,11 +150,12 @@ const Playground = () => {
             {isModelInputVisible && (
               <div className="my-4">
                 <select 
-                  value={selectedModel}
-                  onChange={(e) => setSelectedModel(e.target.value)}
-                  className="border p-2 rounded-md shadow-sm w-full bg-transparent text-gray-700"
-                  >
-                  <option value="GPT-4">GPT-4</option>
+                  value={aiProvider}
+                  onChange={(e) => setAIProvider(e.target.value)}
+                  className="border p-2 rounded-md shadow-sm w-full bg-transparent text-gray-700 mt-2"
+                >
+                  <option value="openai">OpenAI</option>
+                  <option value="huggingface">Hugging Face</option>
                 </select>
               </div>
             )}

package-lock.json

@@ -14,6 +14,7 @@
         "@emotion/react": "^11.10.4",
         "@graphql-codegen/cli": "^2.6.2",
         "@graphql-codegen/typescript-react-apollo": "^3.3.3",
+        "@huggingface/inference": "^2.6.4",
         "@monaco-editor/react": "^4.2.0",
         "@reduxjs/toolkit": "^1.6.0",
         "@types/apollo-upload-client": "^17.0.1",
@@ -2234,6 +2235,14 @@
         "graphql": "^0.8.0 || ^0.9.0 || ^0.10.0 || ^0.11.0 || ^0.12.0 || ^0.13.0 || ^14.0.0 || ^15.0.0 || ^16.0.0"
       }
     },
+    "node_modules/@huggingface/inference": {
+      "version": "2.6.4",
+      "resolved": "https://registry.npmjs.org/@huggingface/inference/-/inference-2.6.4.tgz",
+      "integrity": "sha512-Xna7arltBSBoKaH3diGi3sYvkExgJMd/pF4T6vl2YbmDccbr1G/X5EPZ2048p+YgrJYG1jTYFCtY6Dr3HvJaow==",
+      "engines": {
+        "node": ">=18"
+      }
+    },
     "node_modules/@jridgewell/gen-mapping": {
       "version": "0.3.3",
       "resolved": "https://registry.npmjs.org/@jridgewell/gen-mapping/-/gen-mapping-0.3.3.tgz",

package.json

@@ -16,6 +16,7 @@
     "@emotion/react": "^11.10.4",
     "@graphql-codegen/cli": "^2.6.2",
     "@graphql-codegen/typescript-react-apollo": "^3.3.3",
+    "@huggingface/inference": "^2.6.4",
     "@monaco-editor/react": "^4.2.0",
     "@reduxjs/toolkit": "^1.6.0",
     "@types/apollo-upload-client": "^17.0.1",

pages/api/chat/index.ts

@@ -1,15 +1,22 @@
 import { OpenAIStream, StreamingTextResponse } from "ai";
 import { Configuration, OpenAIApi } from "openai-edge";
+import { HfInference } from '@huggingface/inference';
+import { HuggingFaceStream } from 'ai';
+import { experimental_buildLlama2Prompt } from 'ai/prompts'
 
-const config = new Configuration({
+// Configurations for OpenAI
+const openaiConfig = new Configuration({
     apiKey: process.env.OPENAI_API_KEY,
 });
-const openai = new OpenAIApi(config);
+const openai = new OpenAIApi(openaiConfig);
 
-export const runtime = "edge";
+// Create a new HuggingFace Inference instance
+const Hf = new HfInference(process.env.HUGGINGFACE_API_KEY);
+
+export const runtime = 'edge';
 
 export default async function(req: Request) {
-    let { messages, systemMessage } = await req.json();
+    let { messages, aiProvider = 'openai', systemMessage } = await req.json();
 
     // Prepend the system message if it's not already there
     if (messages.length === 0 || messages[0].role !== "system") {
@@ -19,12 +26,31 @@ export default async function(req: Request) {
         }, ...messages];
     }
 
-    const response = await openai.createChatCompletion({
-        model: 'gpt-4',
-        stream: true,
-        messages
-    })
+    if (aiProvider === 'openai') {
+        const response = await openai.createChatCompletion({
+            model: 'gpt-4',
+            stream: true,
+            messages
+        });
+
+        const stream = OpenAIStream(response);
+        return new StreamingTextResponse(stream);
+    } else if (aiProvider === 'huggingface') {
+        const response = Hf.textGenerationStream({
+            // @ts-ignore
+            model: 'meta-llama/Llama-2-7b-chat-hf',
+            inputs: experimental_buildLlama2Prompt(messages),
+            parameters: {
+              max_new_tokens: 500,
+              repetition_penalty: 1,
+              truncate: 4000,
+              return_full_text: false
+            }
+          })
 
-    const stream = OpenAIStream(response);
-    return new StreamingTextResponse(stream);
+        const stream = HuggingFaceStream(response);
+        return new StreamingTextResponse(stream);
+    } else {
+        throw new Error(`Unsupported AI provider: ${aiProvider}`);
+    }
 }