Integrate AI-Mask

README.md

@@ -289,4 +289,4 @@ We are working on a guide for contributing.
 
 ## Contact
 
-Message Mckay on [Twitter/X](https://twitter.com/mckaywrigley)
+Message Mckay on [Twitter/X](https://twitter.com/mckaywrigley)

app/[locale]/setup/page.tsx

@@ -21,6 +21,7 @@ import {
   SETUP_STEP_COUNT,
   StepContainer
 } from "../../../components/setup/step-container"
+import { AIMaskStep } from "@/components/setup/ai-mask-step"
 
 export default function SetupPage() {
   const {
@@ -179,8 +180,23 @@ export default function SetupPage() {
           </StepContainer>
         )
 
-      // API Step
+      // AI Mask Step
       case 2:
+        return (
+          <StepContainer
+            stepDescription="Enter API keys for each service you'd like to use."
+            stepNum={currentStep}
+            stepTitle="Install AI-Mask (optional)"
+            onShouldProceed={handleShouldProceed}
+            showNextButton={true}
+            showBackButton={true}
+          >
+            <AIMaskStep />
+          </StepContainer>
+        )
+
+      // API Step
+      case 3:
         return (
           <StepContainer
             stepDescription="Enter API keys for each service you'd like to use."
@@ -226,7 +242,7 @@ export default function SetupPage() {
         )
 
       // Finish Step
-      case 3:
+      case 4:
         return (
           <StepContainer
             stepDescription="You are all set up!"

components/chat/chat-helpers/index.ts

@@ -19,6 +19,7 @@ import {
   LLM,
   MessageImage
 } from "@/types"
+import { AIMaskClient, ChatCompletionParams } from "@ai-mask/sdk"
 import React from "react"
 import { toast } from "sonner"
 import { v4 as uuidv4 } from "uuid"
@@ -144,6 +145,66 @@ export const createTempMessages = (
   }
 }
 
+export const handleAIMaskChat = async (
+  payload: ChatPayload,
+  profile: Tables<"profiles">,
+  chatSettings: ChatSettings,
+  tempAssistantMessage: ChatMessage,
+  isRegeneration: boolean,
+  newAbortController: AbortController,
+  setIsGenerating: React.Dispatch<React.SetStateAction<boolean>>,
+  setFirstTokenReceived: React.Dispatch<React.SetStateAction<boolean>>,
+  setChatMessages: React.Dispatch<React.SetStateAction<ChatMessage[]>>,
+  setToolInUse: React.Dispatch<React.SetStateAction<string>>
+) => {
+  // TODO persist client in a hook
+  const aiMaskClient = new AIMaskClient({ name: "chatbot-ui" })
+  const lastChatMessage = isRegeneration
+    ? payload.chatMessages[payload.chatMessages.length - 1]
+    : tempAssistantMessage
+
+  const messages = (await buildFinalMessages(
+    payload,
+    profile,
+    []
+  )) as ChatCompletionParams["messages"]
+  let fullText = ""
+
+  const response = await aiMaskClient.chat(
+    chatSettings.model,
+    {
+      messages,
+      temperature: chatSettings.temperature
+    },
+    contentToAdd => {
+      fullText += contentToAdd
+      setFirstTokenReceived(true)
+      setToolInUse("none")
+
+      setChatMessages(prev =>
+        prev.map(chatMessage => {
+          if (chatMessage.message.id === lastChatMessage.message.id) {
+            const updatedChatMessage: ChatMessage = {
+              message: {
+                ...chatMessage.message,
+                content: fullText
+              },
+              fileItems: chatMessage.fileItems
+            }
+
+            return updatedChatMessage
+          }
+
+          return chatMessage
+        })
+      )
+    }
+  )
+  aiMaskClient.dispose()
+  setIsGenerating(false)
+  return response
+}
+
 export const handleLocalChat = async (
   payload: ChatPayload,
   profile: Tables<"profiles">,

components/chat/chat-hooks/use-chat-handler.tsx

@@ -17,6 +17,7 @@ import {
   handleCreateMessages,
   handleHostedChat,
   handleLocalChat,
+  handleAIMaskChat,
   handleRetrieval,
   processResponse,
   validateChatSettings
@@ -39,6 +40,7 @@ export const useChatHandler = () => {
     setSelectedChat,
     setChats,
     setSelectedTools,
+    availableAIMaskModels,
     availableLocalModels,
     availableOpenRouterModels,
     abortController,
@@ -216,6 +218,7 @@ export const useChatHandler = () => {
         })),
         ...LLM_LIST,
         ...availableLocalModels,
+        ...availableAIMaskModels,
         ...availableOpenRouterModels
       ].find(llm => llm.modelId === chatSettings?.model)
 
@@ -320,6 +323,19 @@ export const useChatHandler = () => {
             setChatMessages,
             setToolInUse
           )
+        } else if (modelData!.provider === "ai-mask") {
+          generatedText = await handleAIMaskChat(
+            payload,
+            profile!,
+            chatSettings!,
+            tempAssistantChatMessage,
+            isRegeneration,
+            newAbortController,
+            setIsGenerating,
+            setFirstTokenReceived,
+            setChatMessages,
+            setToolInUse
+          )
         } else {
           generatedText = await handleHostedChat(
             payload,

components/chat/chat-settings.tsx

@@ -18,6 +18,7 @@ export const ChatSettings: FC<ChatSettingsProps> = ({}) => {
     setChatSettings,
     models,
     availableHostedModels,
+    availableAIMaskModels,
     availableLocalModels,
     availableOpenRouterModels
   } = useContext(ChatbotUIContext)
@@ -59,6 +60,7 @@ export const ChatSettings: FC<ChatSettingsProps> = ({}) => {
     })),
     ...availableHostedModels,
     ...availableLocalModels,
+    ...availableAIMaskModels,
     ...availableOpenRouterModels
   ]
 

components/icons/ai-mask-svg.tsx

@@ -0,0 +1,38 @@
+import { FC } from "react"
+
+interface AIMaskSVGProps {
+  height?: number
+  width?: number
+  className?: string
+}
+
+export const AIMaskSVG: FC<AIMaskSVGProps> = ({
+  height = 40,
+  width = 40,
+  className
+}) => {
+  return (
+    <svg
+      className={className}
+      width={width}
+      height={height}
+      viewBox="0 0 128 128"
+      fill="none"
+      xmlns="http://www.w3.org/2000/svg"
+      role="img"
+    >
+      <g
+        transform="translate(0.000000,128.000000) scale(0.100000,-0.100000)"
+        fill="#000000"
+        stroke="none"
+      >
+        <path d="M180 1113 c-57 -22 -66 -99 -26 -217 l23 -69 -38 -55 c-47 -66 -100 -200 -102 -256 -2 -37 0 -41 22 -41 21 0 27 10 52 89 27 80 83 196 95 196 3 0 12 -7 20 -15 12 -12 20 -12 36 -4 23 12 23 13 -26 130 -44 105 -56 176 -33 185 32 12 128 -44 199 -117 68 -71 99 -84 106 -45 2 10 -6 28 -18 41 -20 21 -20 23 -3 30 27 11 204 19 254 11 68 -10 76 -19 50 -51 -16 -21 -18 -31 -10 -41 20 -24 39 -15 95 45 87 95 205 157 220 117 9 -23 -14 -109 -52 -201 -38 -90 -40 -101 -19 -109 18 -6 45 2 45 15 0 23 18 4 46 -48 33 -62 57 -121 74 -190 9 -36 15 -43 36 -43 13 0 27 6 30 14 9 24 -57 204 -103 278 l-44 72 20 53 c24 62 36 142 27 177 -9 36 -57 59 -104 50 -45 -8 -83 -26 -144 -70 l-48 -34 -57 15 c-76 20 -239 19 -320 -1 l-63 -16 -35 31 c-63 55 -161 91 -205 74z" />
+        <path d="M407 736 c-81 -29 -189 -125 -263 -231 -34 -51 -35 -53 -21 -87 19 -45 72 -104 121 -135 38 -24 40 -25 73 -9 18 9 51 16 73 16 33 0 48 -8 91 -46 63 -56 103 -74 164 -74 55 0 94 19 159 76 40 36 55 44 88 44 21 0 52 -7 69 -16 16 -8 38 -12 47 -9 42 14 162 151 162 184 0 47 -158 220 -243 267 -93 51 -206 33 -244 -39 -41 -76 -13 -200 72 -318 l34 -48 -46 -40 c-39 -35 -53 -41 -90 -41 -51 0 -76 11 -120 53 l-32 30 24 29 c13 16 41 63 62 105 31 63 38 88 38 137 0 79 -19 115 -75 144 -50 25 -88 27 -143 8z m81 -138 c15 -15 15 -61 0 -76 -15 -15 -61 -15 -76 0 -18 18 -15 65 6 77 23 14 55 14 70 -1z m369 -17 c21 -43 -12 -79 -61 -67 -33 8 -42 40 -21 72 21 33 64 30 82 -5z" />
+        <path
+          d="M610 326 c-18 -22 5 -48 39 -44 20 2 26 9 26 28 0 19 -6 26 -26 28
+-14 2 -31 -4 -39 -12z"
+        />
+      </g>
+    </svg>
+  )
+}

components/messages/message.tsx

@@ -52,6 +52,7 @@ export const Message: FC<MessageProps> = ({
     setIsGenerating,
     firstTokenReceived,
     availableLocalModels,
+    availableAIMaskModels,
     availableOpenRouterModels,
     chatMessages,
     selectedAssistant,
@@ -137,6 +138,7 @@ export const Message: FC<MessageProps> = ({
     })),
     ...LLM_LIST,
     ...availableLocalModels,
+    ...availableAIMaskModels,
     ...availableOpenRouterModels
   ].find(llm => llm.modelId === message.model) as LLM
 

components/models/model-icon.tsx

@@ -10,6 +10,7 @@ import { FC, HTMLAttributes } from "react"
 import { AnthropicSVG } from "../icons/anthropic-svg"
 import { GoogleSVG } from "../icons/google-svg"
 import { OpenAISVG } from "../icons/openai-svg"
+import { AIMaskSVG } from "../icons/ai-mask-svg"
 
 interface ModelIconProps extends HTMLAttributes<HTMLDivElement> {
   provider: ModelProvider
@@ -76,6 +77,18 @@ export const ModelIcon: FC<ModelIconProps> = ({
           height={height}
         />
       )
+    case "ai-mask":
+      return (
+        <AIMaskSVG
+          className={cn(
+            "rounded-sm bg-[#fff] p-1 text-black",
+            props.className,
+            theme === "dark" ? "bg-white" : "border-[1px] border-black"
+          )}
+          width={width}
+          height={height}
+        />
+      )
     case "google":
       return (
         <GoogleSVG

components/models/model-select.tsx

@@ -27,6 +27,7 @@ export const ModelSelect: FC<ModelSelectProps> = ({
     models,
     availableHostedModels,
     availableLocalModels,
+    availableAIMaskModels,
     availableOpenRouterModels
   } = useContext(ChatbotUIContext)
 
@@ -37,14 +38,6 @@ export const ModelSelect: FC<ModelSelectProps> = ({
   const [search, setSearch] = useState("")
   const [tab, setTab] = useState<"hosted" | "local">("hosted")
 
-  useEffect(() => {
-    if (isOpen) {
-      setTimeout(() => {
-        inputRef.current?.focus()
-      }, 100) // FIX: hacky
-    }
-  }, [isOpen])
-
   const handleSelectModel = (modelId: LLMID) => {
     onSelectModel(modelId)
     setIsOpen(false)
@@ -60,10 +53,13 @@ export const ModelSelect: FC<ModelSelectProps> = ({
       imageInput: false
     })),
     ...availableHostedModels,
+    ...availableAIMaskModels,
     ...availableLocalModels,
     ...availableOpenRouterModels
   ]
 
+  // TODO better differenciation between ollama/aimask/hosted
+
   const groupedModels = allModels.reduce<Record<string, LLM[]>>(
     (groups, model) => {
       const key = model.provider
@@ -80,6 +76,21 @@ export const ModelSelect: FC<ModelSelectProps> = ({
     model => model.modelId === selectedModelId
   )
 
+  useEffect(() => {
+    if (isOpen) {
+      setTimeout(() => {
+        inputRef.current?.focus()
+      }, 100) // FIX: hacky
+      if (
+        selectedModel &&
+        (selectedModel.provider === "ollama" ||
+          selectedModel.provider === "ai-mask")
+      )
+        setTab("local")
+      else setTab("hosted")
+    }
+  }, [isOpen, selectedModel])
+
   if (!profile) return null
 
   return (
@@ -133,11 +144,11 @@ export const ModelSelect: FC<ModelSelectProps> = ({
         align="start"
       >
         <Tabs value={tab} onValueChange={(value: any) => setTab(value)}>
-          {availableLocalModels.length > 0 && (
+          {(availableLocalModels.length > 0 ||
+            availableAIMaskModels.length > 0) && (
             <TabsList defaultValue="hosted" className="grid grid-cols-2">
-              <TabsTrigger value="hosted">Hosted</TabsTrigger>
-
               <TabsTrigger value="local">Local</TabsTrigger>
+              <TabsTrigger value="hosted">Hosted</TabsTrigger>
             </TabsList>
           )}
         </Tabs>
@@ -154,9 +165,14 @@ export const ModelSelect: FC<ModelSelectProps> = ({
           {Object.entries(groupedModels).map(([provider, models]) => {
             const filteredModels = models
               .filter(model => {
-                if (tab === "hosted") return model.provider !== "ollama"
-                if (tab === "local") return model.provider === "ollama"
-                if (tab === "openrouter") return model.provider === "openrouter"
+                if (tab === "local")
+                  return (
+                    model.provider === "ollama" || model.provider === "ai-mask"
+                  )
+                if (tab === "hosted")
+                  return (
+                    model.provider !== "ollama" && model.provider !== "ai-mask"
+                  )
               })
               .filter(model =>
                 model.modelName.toLowerCase().includes(search.toLowerCase())

components/setup/ai-mask-step.tsx

@@ -0,0 +1,28 @@
+import { FC, useMemo } from "react"
+import { AIMaskSVG } from "../icons/ai-mask-svg"
+import { Button } from "../ui/button"
+import { AIMaskClient } from "@ai-mask/sdk"
+
+export const AIMaskStep: FC = () => {
+  const installed = useMemo(() => AIMaskClient.isExtensionAvailable(), [])
+  return (
+    <div>
+      <div className="flex space-x-4">
+        <AIMaskSVG height={40} width={40} className="bg-white" />
+        <a
+          href="https://chromewebstore.google.com/detail/lkfaajachdpegnlpikpdajccldcgfdde"
+          target="_blank"
+          className=""
+        >
+          <Button disabled={installed}>
+            {installed ? "Already installed" : "Install AI-Mask extension"}
+          </Button>
+        </a>
+      </div>
+      <p className="mt-4">
+        AI-Mask is an extension which allows you to execute AI models right in
+        your browser. Fully local, private and free!
+      </p>
+    </div>
+  )
+}

components/setup/step-container.tsx

@@ -9,7 +9,7 @@ import {
 } from "@/components/ui/card"
 import { FC, useRef } from "react"
 
-export const SETUP_STEP_COUNT = 3
+export const SETUP_STEP_COUNT = 4
 
 interface StepContainerProps {
   stepDescription: string