From b86332f9ecd0971bee5087b5cfadcd11117555ea Mon Sep 17 00:00:00 2001
From: n4ze3m <n4ze3m@gmail.com>
Date: Fri, 5 Apr 2024 20:59:54 +0530
Subject: [PATCH] Update Ollama system prompt in ollama.ts

---
 src/components/Option/Settings/prompt.tsx | 88 ++++++++++++++---------
 src/services/ollama.ts                    | 18 +++--
 2 files changed, 61 insertions(+), 45 deletions(-)

diff --git a/src/components/Option/Settings/prompt.tsx b/src/components/Option/Settings/prompt.tsx
index f0cbbc2..499011c 100644
--- a/src/components/Option/Settings/prompt.tsx
+++ b/src/components/Option/Settings/prompt.tsx
@@ -1,5 +1,5 @@
 import { useQuery, useQueryClient } from "@tanstack/react-query"
-import { Skeleton, Radio, Form, Alert } from "antd"
+import { Skeleton, Radio, Form, Input } from "antd"
 import React from "react"
 import { useTranslation } from "react-i18next"
 import { SaveButton } from "~/components/Common/SaveButton"
@@ -8,15 +8,15 @@ import {
   setSystemPromptForNonRagOption,
   systemPromptForNonRagOption,
   geWebSearchFollowUpPrompt,
-  setWebPrompts
+  setWebPrompts,
+  promptForRag,
+  setPromptForRag
 } from "~/services/ollama"
 
 export const SettingPrompt = () => {
   const { t } = useTranslation("settings")
 
-  const [selectedValue, setSelectedValue] = React.useState<"normal" | "web">(
-    "web"
-  )
+  const [selectedValue, setSelectedValue] = React.useState<"web" | "rag">("rag")
 
   const queryClient = useQueryClient()
 
@@ -25,7 +25,7 @@ export const SettingPrompt = () => {
     queryFn: async () => {
       const [prompt, webSearchPrompt, webSearchFollowUpPrompt] =
         await Promise.all([
-          systemPromptForNonRagOption(),
+          promptForRag(),
           getWebSearchPrompt(),
           geWebSearchFollowUpPrompt()
         ])
@@ -48,46 +48,60 @@ export const SettingPrompt = () => {
             <Radio.Group
               defaultValue={selectedValue}
               onChange={(e) => setSelectedValue(e.target.value)}>
-              <Radio.Button value="normal">
-                {t("ollamaSettings.settings.prompt.option1")}
-              </Radio.Button>
+              <Radio.Button value="rag">RAG</Radio.Button>
               <Radio.Button value="web">
                 {t("ollamaSettings.settings.prompt.option2")}
               </Radio.Button>
             </Radio.Group>
           </div>
 
-          {selectedValue === "normal" && (
+          {selectedValue === "rag" && (
             <Form
               layout="vertical"
               onFinish={(values) => {
-                setSystemPromptForNonRagOption(values?.prompt || "")
+                // setSystemPromptForNonRagOption(values?.prompt || "")
+                setPromptForRag(
+                  values?.systemPrompt || "",
+                  values?.questionPrompt || ""
+                )
                 queryClient.invalidateQueries({
                   queryKey: ["fetchOllaPrompt"]
                 })
               }}
               initialValues={{
-                prompt: data.prompt
+                systemPrompt: data.prompt.ragPrompt,
+                questionPrompt: data.prompt.ragQuestionPrompt
               }}>
-              <Form.Item>
-                <Alert
-                  message={t("ollamaSettings.settings.prompt.alert")}
-                  type="warning"
-                  showIcon
-                  closable
+              <Form.Item
+                label={t("managePrompts.systemPrompt")}
+                name="systemPrompt"
+                rules={[
+                  {
+                    required: true,
+                    message: "Enter a prompt."
+                  }
+                ]}>
+                <Input.TextArea
+                  value={data.webSearchPrompt}
+                  rows={5}
+                  placeholder="Enter a prompt."
                 />
               </Form.Item>
               <Form.Item
-                label={t("ollamaSettings.settings.prompt.systemPrompt")}
-                name="prompt">
-                <textarea
-                  value={data.prompt}
+                label={t("managePrompts.questionPrompt")}
+                name="questionPrompt"
+                rules={[
+                  {
+                    required: true,
+                    message: "Enter a follow up prompt."
+                  }
+                ]}>
+                <Input.TextArea
+                  value={data.webSearchFollowUpPrompt}
                   rows={5}
-                  id="ollamaPrompt"
                   placeholder={t(
-                    "ollamaSettings.settings.prompt.systemPromptPlaceholder"
+                    "ollamaSettings.settings.prompt.webSearchFollowUpPromptPlaceholder"
                   )}
-                  className="w-full p-2 border border-gray-300 rounded-md dark:bg-[#262626] dark:text-gray-100"
                 />
               </Form.Item>
               <Form.Item>
@@ -126,32 +140,36 @@ export const SettingPrompt = () => {
                     )
                   }
                 ]}>
-                <textarea
+                <Input.TextArea
                   value={data.webSearchPrompt}
                   rows={5}
-                  id="ollamaWebSearchPrompt"
                   placeholder={t(
                     "ollamaSettings.settings.prompt.webSearchPromptPlaceholder"
                   )}
-                  className="w-full p-2 border border-gray-300 rounded-md dark:bg-[#262626] dark:text-gray-100"
                 />
               </Form.Item>
               <Form.Item
-                label={t("ollamaSettings.settings.prompt.webSearchFollowUpPrompt")}
+                label={t(
+                  "ollamaSettings.settings.prompt.webSearchFollowUpPrompt"
+                )}
                 name="webSearchFollowUpPrompt"
-                help={t("ollamaSettings.settings.prompt.webSearchFollowUpPromptHelp")}
+                help={t(
+                  "ollamaSettings.settings.prompt.webSearchFollowUpPromptHelp"
+                )}
                 rules={[
                   {
                     required: true,
-                    message: t("ollamaSettings.settings.prompt.webSearchFollowUpPromptError")
+                    message: t(
+                      "ollamaSettings.settings.prompt.webSearchFollowUpPromptError"
+                    )
                   }
                 ]}>
-                <textarea
+                <Input.TextArea
                   value={data.webSearchFollowUpPrompt}
                   rows={5}
-                  id="ollamaWebSearchFollowUpPrompt"
-                  placeholder={t("ollamaSettings.settings.prompt.webSearchFollowUpPromptPlaceholder")}
-                  className="w-full p-2 border border-gray-300 rounded-md dark:bg-[#262626] dark:text-gray-100"
+                  placeholder={t(
+                    "ollamaSettings.settings.prompt.webSearchFollowUpPromptPlaceholder"
+                  )}
                 />
               </Form.Item>
               <Form.Item>
diff --git a/src/services/ollama.ts b/src/services/ollama.ts
index e8e65b1..0e785d0 100644
--- a/src/services/ollama.ts
+++ b/src/services/ollama.ts
@@ -11,8 +11,7 @@ const DEFAULT_PAGE_SHARE_URL = "https://pageassist.xyz"
 const DEFAULT_RAG_QUESTION_PROMPT =
   "Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question.   Chat History: {chat_history} Follow Up Input: {question} Standalone question:"
 
-const DEFAUTL_RAG_SYSTEM_PROMPT = `You are a helpful AI assistant. Use the following pieces of context to answer the question at the end. If you don't know the answer, just say you don't know. DO NOT try to make up an answer. If the question is not related to the context, politely respond that you are tuned to only answer questions that are related to the context.  {context}  Question: {question} Helpful answer in markdown:`
-
+const DEFAUTL_RAG_SYSTEM_PROMPT = `You are a helpful AI assistant. Use the following pieces of context to answer the question at the end. If you don't know the answer, just say you don't know. DO NOT try to make up an answer. If the question is not related to the context, politely respond that you are tuned to only answer questions that are related to the context.  {context}  Question: {question} Helpful answer:`
 
 const DEFAULT_WEBSEARCH_PROMP = `You are a helpful assistant that can answer any questions. You can use the following search results in case you want to answer questions about anything in real-time. The current date and time are {current_date_time}.  
 
@@ -61,7 +60,11 @@ export const isOllamaRunning = async () => {
   }
 }
 
-export const getAllModels = async ({ returnEmpty = false }: { returnEmpty?: boolean }) => {
+export const getAllModels = async ({
+  returnEmpty = false
+}: {
+  returnEmpty?: boolean
+}) => {
   try {
     const baseUrl = await getOllamaURL()
     const response = await fetch(`${cleanUrl(baseUrl)}/api/tags`)
@@ -264,7 +267,6 @@ export const saveForRag = async (
   await setDefaultEmbeddingChunkOverlap(overlap)
 }
 
-
 export const getWebSearchPrompt = async () => {
   const prompt = await storage.get("webSearchPrompt")
   if (!prompt || prompt.length === 0) {
@@ -280,23 +282,20 @@ export const setWebSearchPrompt = async (prompt: string) => {
 export const geWebSearchFollowUpPrompt = async () => {
   const prompt = await storage.get("webSearchFollowUpPrompt")
   if (!prompt || prompt.length === 0) {
-    return DEFAULT_RAG_QUESTION_PROMPT;
+    return DEFAULT_RAG_QUESTION_PROMPT
   }
   return prompt
 }
 
-
 export const setWebSearchFollowUpPrompt = async (prompt: string) => {
   await storage.set("webSearchFollowUpPrompt", prompt)
 }
 
-
 export const setWebPrompts = async (prompt: string, followUpPrompt: string) => {
   await setWebSearchPrompt(prompt)
   await setWebSearchFollowUpPrompt(followUpPrompt)
 }
 
-
 export const getPageShareUrl = async () => {
   const pageShareUrl = await storage.get("pageShareUrl")
   if (!pageShareUrl || pageShareUrl.length === 0) {
@@ -305,7 +304,6 @@ export const getPageShareUrl = async () => {
   return pageShareUrl
 }
 
-
 export const setPageShareUrl = async (pageShareUrl: string) => {
   await storage.set("pageShareUrl", pageShareUrl)
-}
\ No newline at end of file
+}