manage model

package.json

@@ -21,6 +21,7 @@
     "@tanstack/react-query": "^5.17.19",
     "antd": "^5.13.3",
     "axios": "^1.6.7",
+    "dayjs": "^1.11.10",
     "html-to-text": "^9.0.5",
     "langchain": "^0.1.9",
     "lucide-react": "^0.323.0",

src/components/Common/Playground/Message.tsx

@@ -26,7 +26,7 @@ export const PlaygroundMessage = (props: Props) => {
 
   return (
     <div
-      className={`group w-full text-gray-800 dark:text-gray-100 border-b border-black/10 dark:border-gray-900/50 `}>
+      className={`group w-full text-gray-800 dark:text-gray-100`}>
       <div className="text-base gap-4 md:gap-6 md:max-w-2xl lg:max-w-xl xl:max-w-3xl flex lg:px-0 m-auto w-full">
         <div className="flex flex-row gap-4 md:gap-6 md:max-w-2xl lg:max-w-xl xl:max-w-3xl p-4 md:py-6 lg:px-0 m-auto w-full">
           <div className="w-8 flex flex-col relative items-end">

src/components/Option/Layout.tsx

@@ -93,14 +93,14 @@ export default function OptionLayout({
             <a
               href="https://github.com/n4ze3m/page-assist"
               target="_blank"
-              className="text-gray-500 dark:text-gray-400 hover:text-gray-600 dark:hover:text-gray-300 transition-colors">
+              className="!text-gray-500 dark:text-gray-400 hover:text-gray-600 dark:hover:text-gray-300 transition-colors">
               <GithubIcon className="w-6 h-6" />
             </a>
           </Tooltip>
           <Tooltip title="Manage Ollama Models">
             <NavLink
               to="/models"
-              className="text-gray-500 dark:text-gray-400 hover:text-gray-600 dark:hover:text-gray-300 transition-colors">
+              className="!text-gray-500 dark:text-gray-400 hover:text-gray-600 dark:hover:text-gray-300 transition-colors">
               <BrainCircuit className="w-6 h-6" />
             </NavLink>
           </Tooltip>

src/components/Option/Models/index.tsx

@@ -0,0 +1,223 @@
+import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"
+import { Skeleton, Table, Tag, Tooltip, notification, Modal, Input } from "antd"
+import { bytePerSecondFormatter } from "~libs/byte-formater"
+import { deleteModel, getAllModels } from "~services/ollama"
+import { Trash, RotateCcw, Download } from "lucide-react"
+import dayjs from "dayjs"
+import relativeTime from "dayjs/plugin/relativeTime"
+import { useState } from "react"
+import { useForm } from "@mantine/form"
+
+dayjs.extend(relativeTime)
+
+export const ModelsBody = () => {
+  const queryClient = useQueryClient()
+  const [open, setOpen] = useState(false)
+
+  const form = useForm({
+    initialValues: {
+      model: ""
+    }
+  })
+
+  const { data, status } = useQuery({
+    queryKey: ["fetchAllModels"],
+    queryFn: getAllModels
+  })
+
+  const { mutate: deleteOllamaModel } = useMutation({
+    mutationFn: deleteModel,
+    onSuccess: () => {
+      queryClient.invalidateQueries({
+        queryKey: ["fetchAllModels"]
+      })
+      notification.success({
+        message: "Model Deleted",
+        description: "Model has been deleted successfully"
+      })
+    },
+    onError: (error) => {
+      notification.error({
+        message: "Error",
+        description: error?.message || "Something went wrong"
+      })
+    }
+  })
+
+  const pullModel = async (modelName: string) => {
+    notification.info({
+      message: "Pulling Model",
+      description: `Pulling ${modelName} model. For more details, check the extension icon.`
+    })
+
+    setOpen(false)
+
+    form.reset()
+
+    chrome.runtime.sendMessage({
+      type: "pull_model",
+      modelName
+    })
+
+    return true
+  }
+
+  const { mutate: pullOllamaModel } = useMutation({
+    mutationFn: pullModel
+  })
+
+  return (
+    <div className="z-10 min-h-screen">
+      <div className="mt-16 mx-auto py-10 max-w-7xl px-3 sm:px-6 lg:px-8">
+        {/* Add new model button */}
+        <div className="mb-6">
+          <div className="-ml-4 -mt-2 flex flex-wrap items-center justify-end sm:flex-nowrap">
+            <div className="ml-4 mt-2 flex-shrink-0">
+              <button
+                onClick={() => setOpen(true)}
+                className="inline-flex items-center rounded-md border border-transparent bg-black px-3 py-3 text-md font-medium leading-4 text-white shadow-sm hover:bg-gray-800 focus:outline-none focus:ring-2 focus:ring-indigo-500 focus:ring-offset-2 dark:bg-white dark:text-gray-800 dark:hover:bg-gray-100 dark:focus:ring-gray-500 dark:focus:ring-offset-gray-100 disabled:opacity-50">
+                Add New Model
+              </button>
+            </div>
+          </div>
+        </div>
+
+        {status === "pending" && <Skeleton paragraph={{ rows: 8 }} />}
+
+        {status === "success" && (
+          <Table
+            columns={[
+              {
+                title: "Name",
+                dataIndex: "name",
+                key: "name"
+              },
+              {
+                title: "Digest",
+                dataIndex: "digest",
+                key: "digest",
+                render: (text: string) => (
+                  <Tooltip title={text}>
+                    <Tag
+                      className="cursor-pointer"
+                      color="blue">{`${text?.slice(0, 5)}...${text?.slice(-4)}`}</Tag>
+                  </Tooltip>
+                )
+              },
+              {
+                title: "Modified",
+                dataIndex: "modified_at",
+                key: "modified_at",
+                render: (text: string) => dayjs(text).fromNow(true)
+              },
+              {
+                title: "Size",
+                dataIndex: "size",
+                key: "size",
+                render: (text: number) => bytePerSecondFormatter(text)
+              },
+              {
+                title: "Action",
+                render: (_, record) => (
+                  <div className="flex gap-4">
+                    <Tooltip title="Delete Model">
+                      <button
+                        onClick={() => {
+                          if (
+                            window.confirm(
+                              "Are you sure you want to delete this model?"
+                            )
+                          ) {
+                            deleteOllamaModel(record.model)
+                          }
+                        }}
+                        className="text-red-500 dark:text-red-400">
+                        <Trash className="w-5 h-5" />
+                      </button>
+                    </Tooltip>
+                    <Tooltip title="Re-Pull Model">
+                      <button
+                        onClick={() => {
+                          if (
+                            window.confirm(
+                              "Are you sure you want to re-pull this model?"
+                            )
+                          ) {
+                            pullOllamaModel(record.model)
+                          }
+                        }}
+                        className="text-gray-500 dark:text-gray-400">
+                        <RotateCcw className="w-5 h-5" />
+                      </button>
+                    </Tooltip>
+                  </div>
+                )
+              }
+            ]}
+            expandable={{
+              expandedRowRender: (record) => (
+                <Table
+                  pagination={false}
+                  columns={[
+                    {
+                      title: "Parent Model",
+                      key: "parent_model",
+                      dataIndex: "parent_model"
+                    },
+                    {
+                      title: "Format",
+                      key: "format",
+                      dataIndex: "format"
+                    },
+                    {
+                      title: "Family",
+                      key: "family",
+                      dataIndex: "family"
+                    },
+                    {
+                      title: "Parameter Size",
+                      key: "parameter_size",
+                      dataIndex: "parameter_size"
+                    },
+                    {
+                      title: "Quantization Level",
+                      key: "quantization_level",
+                      dataIndex: "quantization_level"
+                    }
+                  ]}
+                  dataSource={[record.details]}
+                />
+              ),
+              defaultExpandAllRows: false
+            }}
+            bordered
+            dataSource={data}
+            rowKey={(record) => `${record.model}-${record.digest}`}
+          />
+        )}
+      </div>
+
+      <Modal
+        footer={null}
+        open={open}
+        title="Add New Model"
+        onCancel={() => setOpen(false)}>
+        <form
+          onSubmit={form.onSubmit((values) => pullOllamaModel(values.model))}>
+          <Input
+            {...form.getInputProps("model")}
+            placeholder="Enter model name"
+            size="large"
+          />
+
+          <button
+            type="submit"
+            className="inline-flex justify-center w-full text-center mt-4 items-center rounded-md border border-transparent bg-black px-2 py-2 text-sm font-medium leading-4 text-white shadow-sm hover:bg-indigo-700 focus:outline-none focus:ring-2 focus:ring-indigo-500 focus:ring-offset-2 dark:bg-white dark:text-gray-800 dark:hover:bg-gray-100 dark:focus:ring-gray-500 dark:focus:ring-offset-gray-100 disabled:opacity-50 ">
+            <Download className="w-5 h-5 mr-3" />
+            Pull Model
+          </button>
+        </form>
+      </Modal>
+    </div>
+  )
+}

src/components/Option/Playground/PlaygroundForm.tsx

@@ -7,7 +7,7 @@ import XMarkIcon from "@heroicons/react/24/outline/XMarkIcon"
 import { toBase64 } from "~libs/to-base64"
 import { useMessageOption } from "~hooks/useMessageOption"
 import { Tooltip } from "antd"
-import { MicIcon, MicOffIcon } from "lucide-react"
+import { MicIcon, StopCircleIcon } from "lucide-react"
 import { Image } from "antd"
 import { useSpeechRecognition } from "~hooks/useSpeechRecognition"
 
@@ -60,8 +60,13 @@ export const PlaygroundForm = ({ dropedFile }: Props) => {
 
   useDynamicTextareaSize(textareaRef, form.values.message, 300)
 
-  const { onSubmit, selectedModel, chatMode, speechToTextLanguage } =
-    useMessageOption()
+  const {
+    onSubmit,
+    selectedModel,
+    chatMode,
+    speechToTextLanguage,
+    stopStreamingRequest
+  } = useMessageOption()
 
   const { isListening, start, stop, transcript } = useSpeechRecognition()
 
@@ -208,23 +213,34 @@ export const PlaygroundForm = ({ dropedFile }: Props) => {
                     <PhotoIcon className="h-5 w-5" />
                   </button>
                 </Tooltip>
-                <button
-                  disabled={isSending || form.values.message.length === 0}
-                  className="inline-flex items-center rounded-md border border-transparent bg-black px-2 py-2 text-sm font-medium leading-4 text-white shadow-sm hover:bg-indigo-700 focus:outline-none focus:ring-2 focus:ring-indigo-500 focus:ring-offset-2 dark:bg-white dark:text-gray-800 dark:hover:bg-gray-100 dark:focus:ring-gray-500 dark:focus:ring-offset-gray-100 disabled:opacity-50 ">
-                  <svg
-                    xmlns="http://www.w3.org/2000/svg"
-                    fill="none"
-                    stroke="currentColor"
-                    strokeLinecap="round"
-                    strokeLinejoin="round"
-                    strokeWidth="2"
-                    className="h-4 w-4 mr-2"
-                    viewBox="0 0 24 24">
-                    <path d="M9 10L4 15 9 20"></path>
-                    <path d="M20 4v7a4 4 0 01-4 4H4"></path>
-                  </svg>
-                  Send
-                </button>
+                {!isSending ? (
+                  <button
+                    disabled={isSending || form.values.message.length === 0}
+                    className="inline-flex items-center rounded-md border border-transparent bg-black px-2 py-2 text-sm font-medium leading-4 text-white shadow-sm hover:bg-gray-800 focus:outline-none focus:ring-2 focus:ring-indigo-500 focus:ring-offset-2 dark:bg-white dark:text-gray-800 dark:hover:bg-gray-100 dark:focus:ring-gray-500 dark:focus:ring-offset-gray-100 disabled:opacity-50 ">
+                    <svg
+                      xmlns="http://www.w3.org/2000/svg"
+                      fill="none"
+                      stroke="currentColor"
+                      strokeLinecap="round"
+                      strokeLinejoin="round"
+                      strokeWidth="2"
+                      className="h-4 w-4 mr-2"
+                      viewBox="0 0 24 24">
+                      <path d="M9 10L4 15 9 20"></path>
+                      <path d="M20 4v7a4 4 0 01-4 4H4"></path>
+                    </svg>
+                    Send
+                  </button>
+                ) : (
+                  <Tooltip title="Stop Streaming">
+                    <button
+                      type="button"
+                      onClick={stopStreamingRequest}
+                      className="text-gray-800 dark:text-gray-300">
+                      <StopCircleIcon className="h-6 w-6" />
+                    </button>
+                  </Tooltip>
+                )}
               </div>
             </div>
           </form>

src/components/Option/Playground/PlaygroundMessage.tsx

@@ -25,8 +25,7 @@ export const PlaygroundMessage = (props: Props) => {
   }, [isBtnPressed])
 
   return (
-    <div
-      className={`group w-full text-gray-800 dark:text-gray-100 border-b border-black/10 dark:border-gray-900/50 `}>
+    <div className={`group w-full text-gray-800 dark:text-gray-100 `}>
       <div className="text-base gap-4 md:gap-6 md:max-w-2xl lg:max-w-xl xl:max-w-3xl flex lg:px-0 m-auto w-full">
         <div className="flex flex-row gap-4 md:gap-6 md:max-w-2xl lg:max-w-xl xl:max-w-3xl p-4 md:py-6 lg:px-0 m-auto w-full">
           <div className="w-8 flex flex-col relative items-end">

src/hooks/useMessageOption.tsx

@@ -12,6 +12,7 @@ import {
 import { useStoreMessageOption } from "~store/option"
 import { saveHistory, saveMessage } from "~libs/db"
 import { useNavigate } from "react-router-dom"
+import { notification } from "antd"
 
 export type BotResponse = {
   bot: {
@@ -253,22 +254,58 @@ export const useMessageOption = () => {
 
       setIsProcessing(false)
     } catch (e) {
+      console.log(e)
+      
+      if (e?.name === "AbortError") {
+        newMessage[appendingIndex].message = newMessage[
+          appendingIndex
+        ].message.slice(0, -1)
+
+        setHistory([
+          ...history,
+          {
+            role: "user",
+            content: message,
+            image
+          },
+          {
+            role: "assistant",
+            content: newMessage[appendingIndex].message
+          }
+        ])
+
+        if (historyId) {
+          await saveMessage(historyId, selectedModel, "user", message, [image])
+          await saveMessage(
+            historyId,
+            selectedModel,
+            "assistant",
+            newMessage[appendingIndex].message,
+            []
+          )
+        } else {
+          const newHistoryId = await saveHistory(message)
+          await saveMessage(newHistoryId.id, selectedModel, "user", message, [
+            image
+          ])
+          await saveMessage(
+            newHistoryId.id,
+            selectedModel,
+            "assistant",
+            newMessage[appendingIndex].message,
+            []
+          )
+          setHistoryId(newHistoryId.id)
+        }
+      } else {
+        notification.error({
+          message: "Error",
+          description: e?.message || "Something went wrong"
+        })
+      }
+
       setIsProcessing(false)
       setStreaming(false)
-
-      setMessages([
-        ...messages,
-        {
-          isBot: true,
-          name: selectedModel,
-          message: `Something went wrong. Check out the following logs:
-        \`\`\`
-        ${e?.message}
-        \`\`\`
-        `,
-          sources: []
-        }
-      ])
     }
   }
 

src/libs/byte-formater.ts

@@ -0,0 +1,23 @@
+const UNITS = [
+  "byte",
+  "kilobyte",
+  "megabyte",
+  "gigabyte",
+  "terabyte",
+  "petabyte"
+]
+
+const getValueAndUnit = (n: number) => {
+  const i = n == 0 ? 0 : Math.floor(Math.log(n) / Math.log(1024))
+  const value = n / Math.pow(1024, i)
+  return { value, unit: UNITS[i] }
+}
+
+export const bytePerSecondFormatter = (n: number) => {
+  const { unit, value } = getValueAndUnit(n)
+  return new Intl.NumberFormat("en", {
+    notation: "compact",
+    style: "unit",
+    unit
+  }).format(value)
+}

src/routes/option-model.tsx

@@ -1,9 +1,10 @@
 import OptionLayout from "~components/Option/Layout"
+import { ModelsBody } from "~components/Option/Models"
 
 export const OptionModal = () => {
   return (
     <OptionLayout>
-        yo
+      <ModelsBody />
     </OptionLayout>
   )
 }

src/services/ollama.ts

@@ -53,6 +53,47 @@ export const isOllamaRunning = async () => {
   }
 }
 
+export const getAllModels = async () => {
+  const baseUrl = await getOllamaURL()
+  const response = await fetch(`${cleanUrl(baseUrl)}/api/tags`)
+  if (!response.ok) {
+    throw new Error(response.statusText)
+  }
+  const json = await response.json()
+
+  return json.models as {
+    name: string
+    model: string
+    modified_at: string
+    size: number
+    digest: string
+    details: {
+      parent_model: string
+      format: string
+      family: string
+      families: string[]
+      parameter_size: string
+      quantization_level: string
+    }
+  }[]
+}
+
+export const deleteModel= async (model: string) => {
+  const baseUrl = await getOllamaURL()
+  const response = await fetch(`${cleanUrl(baseUrl)}/api/delete`, {
+    method: "DELETE",
+    headers: {
+      "Content-Type": "application/json"
+    },
+    body: JSON.stringify({ name: model })
+  })
+
+  if (!response.ok) {
+    throw new Error(response.statusText)
+  }
+  return response.json()
+}
+
 export const fetchModels = async () => {
   try {
     const baseUrl = await getOllamaURL()
@@ -65,6 +106,17 @@ export const fetchModels = async () => {
     return json.models as {
       name: string
       model: string
+      modified_at: string
+      size: number
+      digest: string
+      details: {
+        parent_model: string
+        format: string
+        family: string
+        families: string[]
+        parameter_size: string
+        quantization_level: string
+      }
     }[]
   } catch (e) {
     console.error(e)