From 88ad1fcab72cd99c22ee325c04a8590f8d4ff053 Mon Sep 17 00:00:00 2001
From: n4ze3m <n4ze3m@gmail.com>
Date: Wed, 8 May 2024 10:44:51 +0530
Subject: [PATCH] feat: Add ModelSelect component to SidepanelForm

---
 src/components/Common/ModelSelect.tsx   | 63 +++++++++++++++++++++++++
 src/components/Sidepanel/Chat/empty.tsx |  7 ++-
 src/components/Sidepanel/Chat/form.tsx  |  2 +
 3 files changed, 70 insertions(+), 2 deletions(-)
 create mode 100644 src/components/Common/ModelSelect.tsx

diff --git a/src/components/Common/ModelSelect.tsx b/src/components/Common/ModelSelect.tsx
new file mode 100644
index 0000000..22160b8
--- /dev/null
+++ b/src/components/Common/ModelSelect.tsx
@@ -0,0 +1,63 @@
+import { useQuery } from "@tanstack/react-query"
+import { Dropdown, Tooltip } from "antd"
+import { LucideBrain } from "lucide-react"
+import React from "react"
+import { useTranslation } from "react-i18next"
+import { OllamaIcon } from "../Icons/Ollama"
+import { fetchChatModels } from "@/services/ollama"
+import { useMessage } from "@/hooks/useMessage"
+
+export const ModelSelect: React.FC = () => {
+  const { t } = useTranslation("common")
+  const { setSelectedModel, selectedModel } = useMessage()
+  const { data } = useQuery({
+    queryKey: ["getAllModelsForSelect"],
+    queryFn: async () => {
+      const models = await fetchChatModels({ returnEmpty: false })
+      return models
+    }
+  })
+
+  return (
+    <>
+      {data && data.length > 0 && (
+        <Dropdown
+          menu={{
+            items:
+              data?.map((d) => ({
+                key: d.name,
+                label: (
+                  <div className="w-52 gap-2 text-lg truncate inline-flex line-clamp-3  items-center  dark:border-gray-700">
+                    <div>
+                      <OllamaIcon className="h-6 w-6 text-gray-400" />
+                    </div>
+                    {d.name}
+                  </div>
+                ),
+                onClick: () => {
+                  if (selectedModel === d.name) {
+                    setSelectedModel(null)
+                  } else {
+                    setSelectedModel(d.name)
+                  }
+                }
+              })) || [],
+            style: {
+              maxHeight: 500,
+              overflowY: "scroll"
+            },
+            className: "no-scrollbar",
+            activeKey: selectedModel
+          }}
+          placement={"topLeft"}
+          trigger={["click"]}>
+          <Tooltip title={t("selectAModel")}>
+            <button type="button" className="dark:text-gray-300">
+              <LucideBrain className="h-5 w-5"/>
+            </button>
+          </Tooltip>
+        </Dropdown>
+      )}
+    </>
+  )
+}
diff --git a/src/components/Sidepanel/Chat/empty.tsx b/src/components/Sidepanel/Chat/empty.tsx
index 41305b8..499e320 100644
--- a/src/components/Sidepanel/Chat/empty.tsx
+++ b/src/components/Sidepanel/Chat/empty.tsx
@@ -1,4 +1,4 @@
-import { useQuery } from "@tanstack/react-query"
+import { useQuery, useQueryClient } from "@tanstack/react-query"
 import { Select } from "antd"
 import { RotateCcw } from "lucide-react"
 import { useEffect, useState } from "react"
@@ -15,6 +15,7 @@ import {
 export const EmptySidePanel = () => {
   const [ollamaURL, setOllamaURL] = useState<string>("")
   const { t } = useTranslation(["playground", "common"])
+  const queryClient = useQueryClient()
   const {
     data: ollamaInfo,
     status: ollamaStatus,
@@ -26,7 +27,9 @@ export const EmptySidePanel = () => {
       const ollamaURL = await getOllamaURL()
       const isOk = await isOllamaRunning()
       const models = await fetchChatModels({ returnEmpty: false })
-
+      queryClient.invalidateQueries({
+        queryKey: ["getAllModelsForSelect"]
+      })
       return {
         isOk,
         models,
diff --git a/src/components/Sidepanel/Chat/form.tsx b/src/components/Sidepanel/Chat/form.tsx
index b8b3fc9..45ec340 100644
--- a/src/components/Sidepanel/Chat/form.tsx
+++ b/src/components/Sidepanel/Chat/form.tsx
@@ -10,6 +10,7 @@ import { useWebUI } from "~/store/webui"
 import { defaultEmbeddingModelForRag } from "~/services/ollama"
 import { ImageIcon, MicIcon, StopCircleIcon, X } from "lucide-react"
 import { useTranslation } from "react-i18next"
+import { ModelSelect } from "@/components/Common/ModelSelect"
 
 type Props = {
   dropedFile: File | undefined
@@ -186,6 +187,7 @@ export const SidepanelForm = ({ dropedFile }: Props) => {
                 {...form.getInputProps("message")}
               />
               <div className="flex mt-4 justify-end gap-3">
+                <ModelSelect />
                 <Tooltip title={t("tooltip.speechToText")}>
                   <button
                     type="button"