feat: add model management UI

This commit introduces a new UI for managing models within the OpenAI integration. This UI allows users to view, add, and delete OpenAI models associated with their OpenAI providers. It includes functionality to fetch and refresh model lists, as well as to search for specific models. These changes enhance the user experience by offering greater control over their OpenAI model interactions.

This commit also includes improvements to the existing OpenAI configuration UI, enabling users to seamlessly manage multiple OpenAI providers and associated models.

src/components/Option/Models/CustomModelsTable.tsx

@@ -0,0 +1,85 @@
+import { getAllCustomModels, deleteModel } from "@/db/models"
+import { useStorage } from "@plasmohq/storage/hook"
+import { useQuery, useQueryClient, useMutation } from "@tanstack/react-query"
+import { Skeleton, Table, Tooltip } from "antd"
+import { Trash2 } from "lucide-react"
+import { useTranslation } from "react-i18next"
+
+export const CustomModelsTable = () => {
+  const [selectedModel, setSelectedModel] = useStorage("selectedModel")
+
+  const { t } = useTranslation(["openai", "common"])
+
+
+  const queryClient = useQueryClient()
+
+  const { data, status } = useQuery({
+    queryKey: ["fetchCustomModels"],
+    queryFn: () => getAllCustomModels()
+  })
+
+  const { mutate: deleteCustomModel } = useMutation({
+    mutationFn: deleteModel,
+    onSuccess: () => {
+      queryClient.invalidateQueries({
+        queryKey: ["fetchCustomModels"]
+      })
+    }
+  })
+
+
+  return (
+    <div>
+      <div>
+        {status === "pending" && <Skeleton paragraph={{ rows: 8 }} />}
+
+        {status === "success" && (
+          <div className="overflow-x-auto">
+            <Table
+              columns={[
+                {
+                  title: t("manageModels.columns.name"),
+                  dataIndex: "name",
+                  key: "name"
+                },
+                {
+                  title: t("manageModels.columns.model_id"),
+                  dataIndex: "model_id",
+                  key: "model_id"
+                },
+                {
+                  title: t("manageModels.columns.provider"),
+                  dataIndex: "provider",
+                  render: (_, record) => record.provider.name
+                },
+                {
+                  title: t("manageModels.columns.actions"),
+                  render: (_, record) => (
+                    <Tooltip title={t("manageModels.tooltip.delete")}>
+                      <button
+                        onClick={() => {
+                          if (
+                            window.confirm(t("manageModels.confirm.delete"))
+                          ) {
+                            deleteCustomModel(record.id)
+                            if (selectedModel && selectedModel === record.id) {
+                              setSelectedModel(null)
+                            }
+                          }
+                        }}
+                        className="text-red-500 dark:text-red-400">
+                        <Trash2 className="w-5 h-5" />
+                      </button>
+                    </Tooltip>
+                  )
+                }
+              ]}
+              bordered
+              dataSource={data}
+            />
+          </div>
+        )}
+      </div>
+    </div>
+  )
+}

src/components/Option/Models/OllamaModelsTable.tsx

@@ -0,0 +1,199 @@
+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 dayjs from "dayjs"
+import relativeTime from "dayjs/plugin/relativeTime"
+import { useForm } from "@mantine/form"
+import {  RotateCcw, Trash2 } from "lucide-react"
+import { useTranslation } from "react-i18next"
+import { useStorage } from "@plasmohq/storage/hook"
+
+dayjs.extend(relativeTime)
+
+export const OllamaModelsTable = () => {
+  const queryClient = useQueryClient()
+  const { t } = useTranslation(["settings", "common"])
+  const [selectedModel, setSelectedModel] = useStorage("selectedModel")
+
+  const form = useForm({
+    initialValues: {
+      model: ""
+    }
+  })
+
+  const { data, status } = useQuery({
+    queryKey: ["fetchAllModels"],
+    queryFn: () => getAllModels({ returnEmpty: true })
+  })
+
+  const { mutate: deleteOllamaModel } = useMutation({
+    mutationFn: deleteModel,
+    onSuccess: () => {
+      queryClient.invalidateQueries({
+        queryKey: ["fetchAllModels"]
+      })
+      notification.success({
+        message: t("manageModels.notification.success"),
+        description: t("manageModels.notification.successDeleteDescription")
+      })
+    },
+    onError: (error) => {
+      notification.error({
+        message: "Error",
+        description: error?.message || t("manageModels.notification.someError")
+      })
+    }
+  })
+
+  const pullModel = async (modelName: string) => {
+    notification.info({
+      message: t("manageModels.notification.pullModel"),
+      description: t("manageModels.notification.pullModelDescription", {
+        modelName
+      })
+    })
+
+    form.reset()
+
+    browser.runtime.sendMessage({
+      type: "pull_model",
+      modelName
+    })
+
+    return true
+  }
+
+  const { mutate: pullOllamaModel } = useMutation({
+    mutationFn: pullModel
+  })
+
+  return (
+    <div>
+      <div>
+        {status === "pending" && <Skeleton paragraph={{ rows: 8 }} />}
+
+        {status === "success" && (
+          <div className="overflow-x-auto">
+            <Table
+              columns={[
+                {
+                  title: t("manageModels.columns.name"),
+                  dataIndex: "name",
+                  key: "name"
+                },
+                {
+                  title: t("manageModels.columns.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: t("manageModels.columns.modifiedAt"),
+                  dataIndex: "modified_at",
+                  key: "modified_at",
+                  render: (text: string) => dayjs(text).fromNow(true)
+                },
+                {
+                  title: t("manageModels.columns.size"),
+                  dataIndex: "size",
+                  key: "size",
+                  render: (text: number) => bytePerSecondFormatter(text)
+                },
+                {
+                  title: t("manageModels.columns.actions"),
+                  render: (_, record) => (
+                    <div className="flex gap-4">
+                      <Tooltip title={t("manageModels.tooltip.delete")}>
+                        <button
+                          onClick={() => {
+                            if (
+                              window.confirm(t("manageModels.confirm.delete"))
+                            ) {
+                              deleteOllamaModel(record.model)
+                              if (
+                                selectedModel &&
+                                selectedModel === record.model
+                              ) {
+                                setSelectedModel(null)
+                              }
+                            }
+                          }}
+                          className="text-red-500 dark:text-red-400">
+                          <Trash2 className="w-5 h-5" />
+                        </button>
+                      </Tooltip>
+                      <Tooltip title={t("manageModels.tooltip.repull")}>
+                        <button
+                          onClick={() => {
+                            if (
+                              window.confirm(t("manageModels.confirm.repull"))
+                            ) {
+                              pullOllamaModel(record.model)
+                            }
+                          }}
+                          className="text-gray-700 dark:text-gray-400">
+                          <RotateCcw className="w-5 h-5" />
+                        </button>
+                      </Tooltip>
+                    </div>
+                  )
+                }
+              ]}
+              expandable={{
+                expandedRowRender: (record) => (
+                  <Table
+                    pagination={false}
+                    columns={[
+                      {
+                        title: t("manageModels.expandedColumns.parentModel"),
+                        key: "parent_model",
+                        dataIndex: "parent_model"
+                      },
+                      {
+                        title: t("manageModels.expandedColumns.format"),
+                        key: "format",
+                        dataIndex: "format"
+                      },
+                      {
+                        title: t("manageModels.expandedColumns.family"),
+                        key: "family",
+                        dataIndex: "family"
+                      },
+                      {
+                        title: t("manageModels.expandedColumns.parameterSize"),
+                        key: "parameter_size",
+                        dataIndex: "parameter_size"
+                      },
+                      {
+                        title: t(
+                          "manageModels.expandedColumns.quantizationLevel"
+                        ),
+                        key: "quantization_level",
+                        dataIndex: "quantization_level"
+                      }
+                    ]}
+                    dataSource={[record.details]}
+                    locale={{
+                      emptyText: t("common:noData")
+                    }}
+                  />
+                ),
+                defaultExpandAllRows: false
+              }}
+              bordered
+              dataSource={data}
+              rowKey={(record) => `${record.model}-${record.digest}`}
+            />
+          </div>
+        )}
+      </div>
+    </div>
+  )
+}

src/components/Option/Models/index.tsx

@@ -1,22 +1,30 @@
 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 {
+  Skeleton,
+  Table,
+  Tag,
+  Tooltip,
+  notification,
+  Modal,
+  Input,
+  Segmented
+} from "antd"
 import dayjs from "dayjs"
 import relativeTime from "dayjs/plugin/relativeTime"
 import { useState } from "react"
 import { useForm } from "@mantine/form"
-import { Download, RotateCcw, Trash2 } from "lucide-react"
+import { Download } from "lucide-react"
 import { useTranslation } from "react-i18next"
-import { useStorage } from "@plasmohq/storage/hook"
+import { OllamaModelsTable } from "./OllamaModelsTable"
+import { CustomModelsTable } from "./CustomModelsTable"
 
 dayjs.extend(relativeTime)
 
 export const ModelsBody = () => {
-  const queryClient = useQueryClient()
   const [open, setOpen] = useState(false)
-  const { t } = useTranslation(["settings", "common"])
-  const [selectedModel, setSelectedModel] = useStorage("selectedModel")
+  const [segmented, setSegmented] = useState<string>("ollama")
+  
+  const { t } = useTranslation(["settings", "common", "openai"])
 
   const form = useForm({
     initialValues: {
@@ -24,30 +32,6 @@ export const ModelsBody = () => {
     }
   })
 
-  const { data, status } = useQuery({
-    queryKey: ["fetchAllModels"],
-    queryFn: () => getAllModels({ returnEmpty: true })
-  })
-
-  const { mutate: deleteOllamaModel } = useMutation({
-    mutationFn: deleteModel,
-    onSuccess: () => {
-      queryClient.invalidateQueries({
-        queryKey: ["fetchAllModels"]
-      })
-      notification.success({
-        message: t("manageModels.notification.success"),
-        description: t("manageModels.notification.successDeleteDescription")
-      })
-    },
-    onError: (error) => {
-      notification.error({
-        message: "Error",
-        description: error?.message || t("manageModels.notification.someError")
-      })
-    }
-  })
-
   const pullModel = async (modelName: string) => {
     notification.info({
       message: t("manageModels.notification.pullModel"),
@@ -86,130 +70,26 @@ export const ModelsBody = () => {
               </button>
             </div>
           </div>
-        </div>
-
-        {status === "pending" && <Skeleton paragraph={{ rows: 8 }} />}
-
-        {status === "success" && (
-          <div className="overflow-x-auto">
-            <Table
-              columns={[
+          <div className="flex items-center justify-end mt-3">
+            <Segmented
+              options={[
                 {
-                  title: t("manageModels.columns.name"),
-                  dataIndex: "name",
-                  key: "name"
+                  label: t("common:segmented.ollama"),
+                  value: "ollama"
                 },
                 {
-                  title: t("manageModels.columns.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: t("manageModels.columns.modifiedAt"),
-                  dataIndex: "modified_at",
-                  key: "modified_at",
-                  render: (text: string) => dayjs(text).fromNow(true)
-                },
-                {
-                  title: t("manageModels.columns.size"),
-                  dataIndex: "size",
-                  key: "size",
-                  render: (text: number) => bytePerSecondFormatter(text)
-                },
-                {
-                  title: t("manageModels.columns.actions"),
-                  render: (_, record) => (
-                    <div className="flex gap-4">
-                      <Tooltip title={t("manageModels.tooltip.delete")}>
-                        <button
-                          onClick={() => {
-                            if (
-                              window.confirm(t("manageModels.confirm.delete"))
-                            ) {
-                              deleteOllamaModel(record.model)
-                              if (
-                                selectedModel &&
-                                selectedModel === record.model
-                              ) {
-                                setSelectedModel(null)
-                              }
-                            }
-                          }}
-                          className="text-red-500 dark:text-red-400">
-                          <Trash2 className="w-5 h-5" />
-                        </button>
-                      </Tooltip>
-                      <Tooltip title={t("manageModels.tooltip.repull")}>
-                        <button
-                          onClick={() => {
-                            if (
-                              window.confirm(t("manageModels.confirm.repull"))
-                            ) {
-                              pullOllamaModel(record.model)
-                            }
-                          }}
-                          className="text-gray-700 dark:text-gray-400">
-                          <RotateCcw className="w-5 h-5" />
-                        </button>
-                      </Tooltip>
-                    </div>
-                  )
+                  label: t("common:segmented.custom"),
+                  value: "custom"
                 }
               ]}
-              expandable={{
-                expandedRowRender: (record) => (
-                  <Table
-                    pagination={false}
-                    columns={[
-                      {
-                        title: t("manageModels.expandedColumns.parentModel"),
-                        key: "parent_model",
-                        dataIndex: "parent_model"
-                      },
-                      {
-                        title: t("manageModels.expandedColumns.format"),
-                        key: "format",
-                        dataIndex: "format"
-                      },
-                      {
-                        title: t("manageModels.expandedColumns.family"),
-                        key: "family",
-                        dataIndex: "family"
-                      },
-                      {
-                        title: t("manageModels.expandedColumns.parameterSize"),
-                        key: "parameter_size",
-                        dataIndex: "parameter_size"
-                      },
-                      {
-                        title: t(
-                          "manageModels.expandedColumns.quantizationLevel"
-                        ),
-                        key: "quantization_level",
-                        dataIndex: "quantization_level"
-                      }
-                    ]}
-                    dataSource={[record.details]}
-                    locale={{
-                      emptyText: t("common:noData")
-                    }}
-                  />
-                ),
-                defaultExpandAllRows: false
+              onChange={(value) => {
+                setSegmented(value)
               }}
-              bordered
-              dataSource={data}
-              rowKey={(record) => `${record.model}-${record.digest}`}
             />
           </div>
-        )}
+        </div>
+
+        {segmented === "ollama" ? <OllamaModelsTable /> : <CustomModelsTable />}
       </div>
 
       <Modal

src/components/Option/Settings/openai-fetch-model.tsx

@@ -0,0 +1,132 @@
+import { getOpenAIConfigById } from "@/db/openai"
+import { getAllOpenAIModels } from "@/libs/openai"
+import { useMutation, useQuery } from "@tanstack/react-query"
+import { useTranslation } from "react-i18next"
+import { Checkbox, Input, Spin, message } from "antd"
+import { useState, useMemo } from "react"
+import { createManyModels } from "@/db/models"
+
+type Props = {
+  openaiId: string
+  setOpenModelModal: (openModelModal: boolean) => void
+}
+
+export const OpenAIFetchModel = ({ openaiId, setOpenModelModal }: Props) => {
+  const { t } = useTranslation(["openai"])
+  const [selectedModels, setSelectedModels] = useState<string[]>([])
+  const [searchTerm, setSearchTerm] = useState("")
+
+  const { data, status } = useQuery({
+    queryKey: ["openAIConfigs", openaiId],
+    queryFn: async () => {
+      const config = await getOpenAIConfigById(openaiId)
+      const models = await getAllOpenAIModels(config.baseUrl, config.apiKey)
+      return models
+    },
+    enabled: !!openaiId
+  })
+
+  const filteredModels = useMemo(() => {
+    return (
+      data?.filter((model) =>
+        (model.name ?? model.id)
+          .toLowerCase()
+          .includes(searchTerm.toLowerCase())
+      ) || []
+    )
+  }, [data, searchTerm])
+
+  const handleSelectAll = (checked: boolean) => {
+    if (checked) {
+      setSelectedModels(filteredModels.map((model) => model.id))
+    } else {
+      setSelectedModels([])
+    }
+  }
+
+  const handleModelSelect = (modelId: string, checked: boolean) => {
+    if (checked) {
+      setSelectedModels((prev) => [...prev, modelId])
+    } else {
+      setSelectedModels((prev) => prev.filter((id) => id !== modelId))
+    }
+  }
+
+  const onSave = async (models: string[]) => {
+    const payload = models.map((id) => ({
+      model_id: id,
+      name: filteredModels.find((model) => model.id === id)?.name ?? id,
+      provider_id: openaiId
+    }))
+
+    await createManyModels(payload)
+
+    return true
+  }
+
+  const { mutate: saveModels, isPending: isSaving } = useMutation({
+    mutationFn: onSave,
+    onSuccess: () => {
+      setOpenModelModal(false)
+      message.success(t("modal.model.success"))
+    }
+  })
+
+  const handleSave = () => {
+    saveModels(selectedModels)
+  }
+
+  if (status === "pending") {
+    return <Spin />
+  }
+
+  if (status === "error" || !data || data.length === 0) {
+    return <div>{t("noModelFound")}</div>
+  }
+
+  return (
+    <div className="space-y-4">
+      <p className="text-sm text-gray-500 dark:text-gray-400">
+        {t("modal.model.subheading")}
+      </p>
+      <Input
+        placeholder={t("searchModel")}
+        value={searchTerm}
+        onChange={(e) => setSearchTerm(e.target.value)}
+        className="w-full"
+      />
+      <div className="flex  justify-between">
+        <Checkbox
+          checked={selectedModels.length === filteredModels.length}
+          indeterminate={
+            selectedModels.length > 0 &&
+            selectedModels.length < filteredModels.length
+          }
+          onChange={(e) => handleSelectAll(e.target.checked)}>
+          {t("selectAll")}
+        </Checkbox>
+        <div className="inline-flex items-center px-2.5 py-0.5 rounded-full text-xs font-medium bg-gray-100 text-gray-800 dark:bg-gray-700 dark:text-gray-200">
+          {`${selectedModels?.length} / ${data?.length}`}
+        </div>
+      </div>
+      <div className="space-y-2 custom-scrollbar max-h-[300px] border overflow-y-auto dark:border-gray-600 rounded-md p-3">
+        <div className="grid grid-cols-2 gap-4 items-center">
+          {filteredModels.map((model) => (
+            <Checkbox
+              key={model.id}
+              checked={selectedModels.includes(model.id)}
+              onChange={(e) => handleModelSelect(model.id, e.target.checked)}>
+              {model?.name || model.id}
+            </Checkbox>
+          ))}
+        </div>
+      </div>
+      <button
+        onClick={handleSave}
+        disabled={isSaving}
+        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-gray-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">
+        {isSaving ? t("saving") : t("save")}
+      </button>
+    </div>
+  )
+}

src/components/Option/Settings/openai.tsx

@@ -8,7 +8,8 @@ import {
   updateOpenAIConfig
 } from "@/db/openai"
 import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"
-import { Pencil, Trash2, Plus } from "lucide-react"
+import { Pencil, Trash2, RotateCwIcon } from "lucide-react"
+import { OpenAIFetchModel } from "./openai-fetch-model"
 
 export const OpenAIApp = () => {
   const { t } = useTranslation("openai")
@@ -16,6 +17,8 @@ export const OpenAIApp = () => {
   const [editingConfig, setEditingConfig] = useState(null)
   const queryClient = useQueryClient()
   const [form] = Form.useForm()
+  const [openaiId, setOpenaiId] = useState<string | null>(null)
+  const [openModelModal, setOpenModelModal] = useState(false)
 
   const { data: configs, isLoading } = useQuery({
     queryKey: ["openAIConfigs"],
@@ -24,12 +27,14 @@ export const OpenAIApp = () => {
 
   const addMutation = useMutation({
     mutationFn: addOpenAICofig,
-    onSuccess: () => {
+    onSuccess: (data) => {
       queryClient.invalidateQueries({
         queryKey: ["openAIConfigs"]
       })
       setOpen(false)
       message.success(t("addSuccess"))
+      setOpenaiId(data)
+      setOpenModelModal(true)
     }
   })
 
@@ -129,6 +134,18 @@ export const OpenAIApp = () => {
                       <Pencil className="size-4" />
                     </button>
                   </Tooltip>
+
+                  <Tooltip title={t("refetch")}>
+                    <button
+                      className="text-gray-700 dark:text-gray-400"
+                      onClick={() => {
+                        setOpenModelModal(true)
+                        setOpenaiId(record.id)
+                      }}
+                      disabled={!record.id}>
+                      <RotateCwIcon className="size-4" />
+                    </button>
+                  </Tooltip>
                   <Tooltip title={t("delete")}>
                     <button
                       className="text-red-500 dark:text-red-400"
@@ -212,6 +229,19 @@ export const OpenAIApp = () => {
             </button>
           </Form>
         </Modal>
+
+        <Modal
+          open={openModelModal}
+          title={t("modal.model.title")}
+          footer={null}
+          onCancel={() => setOpenModelModal(false)}>
+          {openaiId ? (
+            <OpenAIFetchModel
+              openaiId={openaiId}
+              setOpenModelModal={setOpenModelModal}
+            />
+          ) : null}
+        </Modal>
       </div>
     </div>
   )