From 53d999a596a0f0c9e9a644260ef47b4269a7524e Mon Sep 17 00:00:00 2001
From: n4ze3m <n4ze3m@gmail.com>
Date: Sat, 14 Sep 2024 12:07:54 +0530
Subject: [PATCH] feat: Add system and quick prompts on side panel

---
 src/components/Common/PromptSelect.tsx   | 23 ++++++++++------
 src/components/Layouts/Header.tsx        |  8 ++++--
 src/components/Sidepanel/Chat/form.tsx   | 26 +++++++++++++++---
 src/components/Sidepanel/Chat/header.tsx | 23 ++++++++++++----
 src/hooks/useMessage.tsx                 | 35 ++++++++++++++++++++++--
 src/hooks/useMessageOption.tsx           |  9 +++---
 src/services/ollama.ts                   |  6 ++--
 src/store/index.tsx                      | 13 ++++++++-
 src/store/option.tsx                     |  3 ++
 9 files changed, 115 insertions(+), 31 deletions(-)

diff --git a/src/components/Common/PromptSelect.tsx b/src/components/Common/PromptSelect.tsx
index fadd8c0..dd493a9 100644
--- a/src/components/Common/PromptSelect.tsx
+++ b/src/components/Common/PromptSelect.tsx
@@ -4,20 +4,27 @@ import { BookIcon, ComputerIcon, ZapIcon } from "lucide-react"
 import React from "react"
 import { useTranslation } from "react-i18next"
 import { getAllPrompts } from "@/db"
-import { useMessageOption } from "@/hooks/useMessageOption"
 
-export const PromptSelect: React.FC = () => {
+type Props = {
+  setSelectedSystemPrompt: (promptId: string | undefined) => void
+  setSelectedQuickPrompt: (prompt: string | undefined) => void
+  selectedSystemPrompt: string | undefined
+  className?: string
+}
+
+export const PromptSelect: React.FC<Props> = ({
+  setSelectedQuickPrompt,
+  setSelectedSystemPrompt,
+  selectedSystemPrompt,
+  className = "dark:text-gray-300"
+}) => {
   const { t } = useTranslation("option")
-  const {
-    selectedSystemPrompt,
-    setSelectedQuickPrompt,
-    setSelectedSystemPrompt
-  } = useMessageOption()
 
   const { data } = useQuery({
     queryKey: ["getAllPromptsForSelect"],
     queryFn: getAllPrompts
   })
+
   const handlePromptChange = (value?: string) => {
     if (!value) {
       setSelectedSystemPrompt(undefined)
@@ -79,7 +86,7 @@ export const PromptSelect: React.FC = () => {
           placement={"topLeft"}
           trigger={["click"]}>
           <Tooltip title={t("selectAPrompt")}>
-            <button type="button" className="dark:text-gray-300">
+            <button type="button" className={className}>
               <BookIcon className="h-5 w-5" />
             </button>
           </Tooltip>
diff --git a/src/components/Layouts/Header.tsx b/src/components/Layouts/Header.tsx
index f7251b1..65fab8e 100644
--- a/src/components/Layouts/Header.tsx
+++ b/src/components/Layouts/Header.tsx
@@ -45,7 +45,7 @@ export const Header: React.FC<Props> = ({
     setSelectedQuickPrompt,
     setSelectedSystemPrompt,
     messages,
-    streaming
+    streaming,
   } = useMessageOption()
   const {
     data: models,
@@ -182,7 +182,11 @@ export const Header: React.FC<Props> = ({
           />
         </div>
         <div className="lg:hidden">
-          <PromptSelect />
+          <PromptSelect
+            selectedSystemPrompt={selectedSystemPrompt}
+            setSelectedSystemPrompt={setSelectedSystemPrompt}
+            setSelectedQuickPrompt={setSelectedQuickPrompt}
+          />
         </div>
         <SelectedKnowledge />
       </div>
diff --git a/src/components/Sidepanel/Chat/form.tsx b/src/components/Sidepanel/Chat/form.tsx
index c413377..e1cc3f0 100644
--- a/src/components/Sidepanel/Chat/form.tsx
+++ b/src/components/Sidepanel/Chat/form.tsx
@@ -24,7 +24,6 @@ export const SidepanelForm = ({ dropedFile }: Props) => {
   const { sendWhenEnter, setSendWhenEnter } = useWebUI()
   const [typing, setTyping] = React.useState<boolean>(false)
   const { t } = useTranslation(["playground", "common"])
-
   const form = useForm({
     initialValues: {
       message: "",
@@ -37,7 +36,7 @@ export const SidepanelForm = ({ dropedFile }: Props) => {
     resetTranscript,
     start: startListening,
     stop: stopSpeechRecognition,
-    supported: browserSupportsSpeechRecognition
+    supported: browserSupportsSpeechRecognition, 
   } = useSpeechRecognition()
 
   const stopListening = async () => {
@@ -118,12 +117,14 @@ export const SidepanelForm = ({ dropedFile }: Props) => {
     onSubmit,
     selectedModel,
     chatMode,
-    speechToTextLanguage,
     stopStreamingRequest,
     streaming,
     setChatMode,
     webSearch,
-    setWebSearch
+    setWebSearch,
+    selectedQuickPrompt,
+    setSelectedQuickPrompt,
+    speechToTextLanguage
   } = useMessage()
 
   React.useEffect(() => {
@@ -139,6 +140,23 @@ export const SidepanelForm = ({ dropedFile }: Props) => {
       form.setFieldValue("message", transcript)
     }
   }, [transcript])
+
+  React.useEffect(() => {
+    if (selectedQuickPrompt) {
+      const word = getVariable(selectedQuickPrompt)
+      form.setFieldValue("message", selectedQuickPrompt)
+      if (word) {
+        textareaRef.current?.focus()
+        const interval = setTimeout(() => {
+          textareaRef.current?.setSelectionRange(word.start, word.end)
+          setSelectedQuickPrompt(null)
+        }, 100)
+        return () => {
+          clearInterval(interval)
+        }
+      }
+    }
+  }, [selectedQuickPrompt])
   const { mutateAsync: sendMessage, isPending: isSending } = useMutation({
     mutationFn: onSubmit,
     onSuccess: () => {
diff --git a/src/components/Sidepanel/Chat/header.tsx b/src/components/Sidepanel/Chat/header.tsx
index 5970a7c..95e1374 100644
--- a/src/components/Sidepanel/Chat/header.tsx
+++ b/src/components/Sidepanel/Chat/header.tsx
@@ -7,13 +7,22 @@ import { useTranslation } from "react-i18next"
 import { CurrentChatModelSettings } from "@/components/Common/Settings/CurrentChatModelSettings"
 import React from "react"
 import { useStorage } from "@plasmohq/storage/hook"
+import { PromptSelect } from "@/components/Common/PromptSelect"
 export const SidepanelHeader = () => {
   const [hideCurrentChatModelSettings] = useStorage(
     "hideCurrentChatModelSettings",
     false
   )
 
-  const { clearChat, isEmbedding, messages, streaming } = useMessage()
+  const {
+    clearChat,
+    isEmbedding,
+    messages,
+    streaming,
+    selectedSystemPrompt,
+    setSelectedSystemPrompt,
+    setSelectedQuickPrompt
+  } = useMessage()
   const { t } = useTranslation(["sidepanel", "common"])
   const [openModelSettings, setOpenModelSettings] = React.useState(false)
 
@@ -44,11 +53,13 @@ export const SidepanelHeader = () => {
             <EraserIcon className="h-5 w-5 text-gray-500 dark:text-gray-400" />
           </button>
         )}
-        {/* <Tooltip title={t("tooltip.history")}>
-          <Link to="/history">
-            <HistoryIcon className="h-5 w-5 text-gray-500 dark:text-gray-400" />
-          </Link>
-        </Tooltip> */}
+        <PromptSelect
+          selectedSystemPrompt={selectedSystemPrompt}
+          setSelectedSystemPrompt={setSelectedSystemPrompt}
+          setSelectedQuickPrompt={setSelectedQuickPrompt}
+          className="text-gray-500 dark:text-gray-400 hover:text-gray-600 dark:hover:text-gray-300 transition-colors"
+        />
+
         {!hideCurrentChatModelSettings && (
           <Tooltip title={t("common:currentChatModelSettings")}>
             <button
diff --git a/src/hooks/useMessage.tsx b/src/hooks/useMessage.tsx
index 7d1d209..4dcac4e 100644
--- a/src/hooks/useMessage.tsx
+++ b/src/hooks/useMessage.tsx
@@ -17,6 +17,7 @@ import { ChatHistory } from "@/store/option"
 import {
   deleteChatForEdit,
   generateID,
+  getPromptById,
   removeMessageUsingHistoryId,
   updateMessageByIndex
 } from "@/db"
@@ -75,9 +76,18 @@ export const useMessage = () => {
     setIsEmbedding,
     isEmbedding,
     currentURL,
-    setCurrentURL
+    setCurrentURL,
+    selectedQuickPrompt,
+    setSelectedQuickPrompt,
+    selectedSystemPrompt,
+    setSelectedSystemPrompt
   } = useStoreMessage()
 
+  const [speechToTextLanguage, setSpeechToTextLanguage] = useStorage(
+    "speechToTextLanguage",
+    "en-US"
+  )
+
   const [keepTrackOfEmbedding, setKeepTrackOfEmbedding] = React.useState<{
     [key: string]: MemoryVectorStore
   }>({})
@@ -488,6 +498,7 @@ export const useMessage = () => {
 
     try {
       const prompt = await systemPromptForNonRag()
+      const selectedPrompt = await getPromptById(selectedSystemPrompt)
 
       let humanMessage = new HumanMessage({
         content: [
@@ -514,7 +525,7 @@ export const useMessage = () => {
 
       const applicationChatHistory = generateHistory(history)
 
-      if (prompt) {
+      if (prompt && !selectedPrompt) {
         applicationChatHistory.unshift(
           new SystemMessage({
             content: [
@@ -526,6 +537,18 @@ export const useMessage = () => {
           })
         )
       }
+      if (selectedPrompt) {
+        applicationChatHistory.unshift(
+          new SystemMessage({
+            content: [
+              {
+                text: selectedPrompt.content,
+                type: "text"
+              }
+            ]
+          })
+        )
+      }
 
       const chunks = await ollama.stream(
         [...applicationChatHistory, humanMessage],
@@ -1231,6 +1254,12 @@ export const useMessage = () => {
     regenerateLastMessage,
     webSearch,
     setWebSearch,
-    isSearchingInternet
+    isSearchingInternet,
+    selectedQuickPrompt,
+    setSelectedQuickPrompt,
+    selectedSystemPrompt,
+    setSelectedSystemPrompt,
+    speechToTextLanguage,
+    setSpeechToTextLanguage
   }
 }
diff --git a/src/hooks/useMessageOption.tsx b/src/hooks/useMessageOption.tsx
index 7dd6d6d..4e633f8 100644
--- a/src/hooks/useMessageOption.tsx
+++ b/src/hooks/useMessageOption.tsx
@@ -55,8 +55,6 @@ export const useMessageOption = () => {
     setIsProcessing,
     chatMode,
     setChatMode,
-    speechToTextLanguage,
-    setSpeechToTextLanguage,
     webSearch,
     setWebSearch,
     isSearchingInternet,
@@ -70,7 +68,10 @@ export const useMessageOption = () => {
   } = useStoreMessageOption()
   const currentChatModelSettings = useStoreChatModelSettings()
   const [selectedModel, setSelectedModel] = useStorage("selectedModel")
-
+  const [ speechToTextLanguage, setSpeechToTextLanguage ] = useStorage(
+    "speechToTextLanguage",
+    "en-US"
+  )
   const { ttsEnabled } = useWebUI()
 
   const { t } = useTranslation("option")
@@ -411,8 +412,6 @@ export const useMessageOption = () => {
       const prompt = await systemPromptForNonRagOption()
       const selectedPrompt = await getPromptById(selectedSystemPrompt)
 
-      // message = message.trim().replaceAll("\n", " ")
-
       let humanMessage = new HumanMessage({
         content: [
           {
diff --git a/src/services/ollama.ts b/src/services/ollama.ts
index 7ec4f69..ee58b0e 100644
--- a/src/services/ollama.ts
+++ b/src/services/ollama.ts
@@ -329,13 +329,15 @@ export const saveForRag = async (
   chunkSize: number,
   overlap: number,
   totalFilePerKB: number,
-  noOfRetrievedDocs: number
+  noOfRetrievedDocs?: number
 ) => {
   await setDefaultEmbeddingModelForRag(model)
   await setDefaultEmbeddingChunkSize(chunkSize)
   await setDefaultEmbeddingChunkOverlap(overlap)
   await setTotalFilePerKB(totalFilePerKB)
-  await setNoOfRetrievedDocs(noOfRetrievedDocs)
+  if(noOfRetrievedDocs) {
+    await setNoOfRetrievedDocs(noOfRetrievedDocs)
+  }
 }
 
 export const getWebSearchPrompt = async () => {
diff --git a/src/store/index.tsx b/src/store/index.tsx
index c1929d3..02bf2cb 100644
--- a/src/store/index.tsx
+++ b/src/store/index.tsx
@@ -40,6 +40,11 @@ type State = {
   setSpeechToTextLanguage: (speechToTextLanguage: string) => void
   currentURL: string
   setCurrentURL: (currentURL: string) => void
+  selectedSystemPrompt: string | null
+  setSelectedSystemPrompt: (selectedSystemPrompt: string) => void
+
+  selectedQuickPrompt: string | null
+  setSelectedQuickPrompt: (selectedQuickPrompt: string) => void
 }
 
 export const useStoreMessage = create<State>((set) => ({
@@ -68,5 +73,11 @@ export const useStoreMessage = create<State>((set) => ({
   setSpeechToTextLanguage: (speechToTextLanguage) =>
     set({ speechToTextLanguage }),
   currentURL: "",
-  setCurrentURL: (currentURL) => set({ currentURL })
+  setCurrentURL: (currentURL) => set({ currentURL }),
+
+  selectedSystemPrompt: null,
+  setSelectedSystemPrompt: (selectedSystemPrompt) =>
+    set({ selectedSystemPrompt }),
+  selectedQuickPrompt: null,
+  setSelectedQuickPrompt: (selectedQuickPrompt) => set({ selectedQuickPrompt })
 }))
diff --git a/src/store/option.tsx b/src/store/option.tsx
index f3ff42e..8a55501 100644
--- a/src/store/option.tsx
+++ b/src/store/option.tsx
@@ -62,6 +62,9 @@ type State = {
 
   selectedKnowledge: Knowledge | null
   setSelectedKnowledge: (selectedKnowledge: Knowledge) => void
+
+  setSpeechToTextLanguage: (language: string) => void
+  speechToTextLanguage: string
 }
 
 export const useStoreMessageOption = create<State>((set) => ({