From 88e7e7521b46360ff7a819b910a83a47173d1097 Mon Sep 17 00:00:00 2001
From: n4ze3m <n4ze3m@gmail.com>
Date: Wed, 28 Feb 2024 22:50:38 +0530
Subject: [PATCH] Update OptionLayout component and add OptionPrompt route

---
 src/components/Option/Layout.tsx       |  23 ++-
 src/components/Option/Models/index.tsx |   2 +-
 src/components/Option/Prompt/index.tsx | 273 +++++++++++++++++++++++++
 src/libs/db.ts                         |  89 ++++++++
 src/routes/index.tsx                   |   2 +
 src/routes/option-prompt.tsx           |  10 +
 src/web/local-google.ts                |   9 +-
 7 files changed, 401 insertions(+), 7 deletions(-)
 create mode 100644 src/components/Option/Prompt/index.tsx
 create mode 100644 src/routes/option-prompt.tsx

diff --git a/src/components/Option/Layout.tsx b/src/components/Option/Layout.tsx
index d85a1a3..1fb6b7c 100644
--- a/src/components/Option/Layout.tsx
+++ b/src/components/Option/Layout.tsx
@@ -4,11 +4,17 @@ import { useLocation, NavLink } from "react-router-dom"
 import { Sidebar } from "./Sidebar"
 import { Drawer, Layout, Modal, Select, Tooltip } from "antd"
 import { useQuery } from "@tanstack/react-query"
-import {  getAllModels } from "~services/ollama"
+import { getAllModels } from "~services/ollama"
 import { useMessageOption } from "~hooks/useMessageOption"
 import { Settings } from "./Settings"
-import { BrainCircuit, ChevronLeft, CogIcon, GithubIcon, PanelLeftIcon, SquarePen } from "lucide-react"
-
+import {
+  Book,
+  BrainCircuit,
+  ChevronLeft,
+  CogIcon,
+  PanelLeftIcon,
+  SquarePen
+} from "lucide-react"
 
 export default function OptionLayout({
   children
@@ -83,14 +89,21 @@ export default function OptionLayout({
           </div>
         </div>
         <div className="flex gap-4 items-center">
-          <Tooltip title="Github Repository">
+          <Tooltip title="Manage Prompts">
+            <NavLink
+              to="/prompts"
+              className="!text-gray-500 dark:text-gray-400 hover:text-gray-600 dark:hover:text-gray-300 transition-colors">
+              <Book className="w-6 h-6" />
+            </NavLink>
+          </Tooltip>
+          {/* <Tooltip title="Github Repository">
             <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">
               <GithubIcon className="w-6 h-6" />
             </a>
-          </Tooltip>
+          </Tooltip> */}
           <Tooltip title="Manage Ollama Models">
             <NavLink
               to="/models"
diff --git a/src/components/Option/Models/index.tsx b/src/components/Option/Models/index.tsx
index b637d7f..a7c6116 100644
--- a/src/components/Option/Models/index.tsx
+++ b/src/components/Option/Models/index.tsx
@@ -75,7 +75,7 @@ export const ModelsBody = () => {
             <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">
+                className="inline-flex items-center rounded-md border border-transparent bg-black px-2 py-2 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>
diff --git a/src/components/Option/Prompt/index.tsx b/src/components/Option/Prompt/index.tsx
new file mode 100644
index 0000000..7e328b3
--- /dev/null
+++ b/src/components/Option/Prompt/index.tsx
@@ -0,0 +1,273 @@
+import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"
+import {
+  Skeleton,
+  Table,
+  Tag,
+  Tooltip,
+  notification,
+  Modal,
+  Input,
+  Form,
+  Switch
+} from "antd"
+import { Trash2, Pen } from "lucide-react"
+import { useState } from "react"
+import {
+  deletePromptById,
+  getAllPrompts,
+  savePrompt,
+  updatePrompt
+} from "~libs/db"
+
+export const PromptBody = () => {
+  const queryClient = useQueryClient()
+  const [open, setOpen] = useState(false)
+  const [openEdit, setOpenEdit] = useState(false)
+  const [editId, setEditId] = useState("")
+  const [createForm] = Form.useForm()
+  const [editForm] = Form.useForm()
+
+  const { data, status } = useQuery({
+    queryKey: ["fetchAllPrompts"],
+    queryFn: getAllPrompts
+  })
+
+  const { mutate: deletePrompt } = useMutation({
+    mutationFn: deletePromptById,
+    onSuccess: () => {
+      queryClient.invalidateQueries({
+        queryKey: ["fetchAllPrompts"]
+      })
+      notification.success({
+        message: "Model Deleted",
+        description: "Model has been deleted successfully"
+      })
+    },
+    onError: (error) => {
+      notification.error({
+        message: "Error",
+        description: error?.message || "Something went wrong"
+      })
+    }
+  })
+
+  const { mutate: savePromptMutation, isPending: savePromptLoading } =
+    useMutation({
+      mutationFn: savePrompt,
+      onSuccess: () => {
+        queryClient.invalidateQueries({
+          queryKey: ["fetchAllPrompts"]
+        })
+        setOpen(false)
+        createForm.resetFields()
+        notification.success({
+          message: "Prompt Added",
+          description: "Prompt has been added successfully"
+        })
+      },
+      onError: (error) => {
+        notification.error({
+          message: "Error",
+          description: error?.message || "Something went wrong"
+        })
+      }
+    })
+
+  const { mutate: updatePromptMutation, isPending: isUpdatingPrompt } =
+    useMutation({
+      mutationFn: async (data: any) => {
+        return await updatePrompt({
+          ...data,
+          id: editId
+        })
+      },
+      onSuccess: () => {
+        queryClient.invalidateQueries({
+          queryKey: ["fetchAllPrompts"]
+        })
+        setOpenEdit(false)
+        editForm.resetFields()
+        notification.success({
+          message: "Prompt Updated",
+          description: "Prompt has been updated successfully"
+        })
+      },
+      onError: (error) => {
+        notification.error({
+          message: "Error",
+          description: error?.message || "Something went wrong"
+        })
+      }
+    })
+
+  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">
+        <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-2 py-2 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 Prompt
+              </button>
+            </div>
+          </div>
+        </div>
+
+        {status === "pending" && <Skeleton paragraph={{ rows: 8 }} />}
+
+        {status === "success" && (
+          <Table
+            columns={[
+              {
+                title: "Title",
+                dataIndex: "title",
+                key: "title"
+              },
+              {
+                title: "Prompt",
+                dataIndex: "content",
+                key: "content"
+              },
+              {
+                title: "Is System Prompt",
+                dataIndex: "is_system",
+                key: "is_system",
+                render: (is_system) => (
+                  <Tag color={is_system ? "green" : "blue"}>
+                    {is_system ? "Yes" : "No"}
+                  </Tag>
+                )
+              },
+              {
+                title: "Action",
+                render: (_, record) => (
+                  <div className="flex gap-4">
+                    <Tooltip title="Delete Prompt">
+                      <button
+                        onClick={() => {
+                          if (
+                            window.confirm(
+                              "Are you sure you want to delete this prompt? This action cannot be undone."
+                            )
+                          ) {
+                            deletePrompt(record.id)
+                          }
+                        }}
+                        className="text-red-500 dark:text-red-400">
+                        <Trash2 className="w-5 h-5" />
+                      </button>
+                    </Tooltip>
+                    <Tooltip title="Edit Prompt">
+                      <button
+                        onClick={() => {
+                          setEditId(record.id)
+                          editForm.setFieldsValue(record)
+                          setOpenEdit(true)
+                        }}
+                        className="text-gray-500 dark:text-gray-400">
+                        <Pen className="w-5 h-5" />
+                      </button>
+                    </Tooltip>
+                  </div>
+                )
+              }
+            ]}
+            bordered
+            dataSource={data}
+            rowKey={(record) => record.id}
+          />
+        )}
+      </div>
+
+      <Modal
+        title="Add New Prompt"
+        open={open}
+        onCancel={() => setOpen(false)}
+        footer={null}>
+        <Form
+          onFinish={(values) => savePromptMutation(values)}
+          layout="vertical"
+          form={createForm}>
+          <Form.Item
+            name="title"
+            label="Title"
+            rules={[{ required: true, message: "Title is required" }]}>
+            <Input placeholder="My Awesome Prompt" />
+          </Form.Item>
+
+          <Form.Item
+            name="content"
+            label="Prompt"
+            rules={[{ required: true, message: "Prompt is required" }]}
+            help="You can use {key} as variable in your prompt.">
+            <Input.TextArea
+              placeholder="Your prompt goes here..."
+              autoSize={{ minRows: 3, maxRows: 10 }}
+            />
+          </Form.Item>
+
+          <Form.Item
+            name="is_system"
+            label="Is System Prompt"
+            valuePropName="checked">
+            <Switch />
+          </Form.Item>
+
+          <Form.Item>
+            <button
+              disabled={savePromptLoading}
+              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 ">
+              {savePromptLoading ? "Adding Prompt..." : "Add Prompt"}
+            </button>
+          </Form.Item>
+        </Form>
+      </Modal>
+
+      <Modal
+        title="Update Prompt"
+        open={openEdit}
+        onCancel={() => setOpenEdit(false)}
+        footer={null}>
+        <Form
+          onFinish={(values) => updatePromptMutation(values)}
+          layout="vertical"
+          form={editForm}>
+          <Form.Item
+            name="title"
+            label="Title"
+            rules={[{ required: true, message: "Title is required" }]}>
+            <Input placeholder="My Awesome Prompt" />
+          </Form.Item>
+
+          <Form.Item
+            name="content"
+            label="Prompt"
+            rules={[{ required: true, message: "Prompt is required" }]}
+            help="You can use {key} as variable in your prompt.">
+            <Input.TextArea
+              placeholder="Your prompt goes here..."
+              autoSize={{ minRows: 3, maxRows: 10 }}
+            />
+          </Form.Item>
+
+          <Form.Item
+            name="is_system"
+            label="Is System Prompt"
+            valuePropName="checked">
+            <Switch />
+          </Form.Item>
+
+          <Form.Item>
+            <button
+              disabled={isUpdatingPrompt}
+              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 ">
+              {isUpdatingPrompt ? "Updating Prompt..." : "Update Prompt"}
+            </button>
+          </Form.Item>
+        </Form>
+      </Modal>
+    </div>
+  )
+}
diff --git a/src/libs/db.ts b/src/libs/db.ts
index ff61f26..75d53a6 100644
--- a/src/libs/db.ts
+++ b/src/libs/db.ts
@@ -31,10 +31,22 @@ type Message = {
   createdAt: number
 }
 
+
+type Prompt = {
+  id: string
+  title: string
+  content: string
+  is_system: boolean
+  createdBy?: string
+  createdAt: number
+}
+
 type MessageHistory = Message[]
 
 type ChatHistory = HistoryInfo[]
 
+type Prompts = Prompt[]
+
 export class PageAssitDatabase {
   db: chrome.storage.StorageArea
 
@@ -102,8 +114,49 @@ export class PageAssitDatabase {
   async deleteMessage(history_id: string) {
     await this.db.remove(history_id)
   }
+
+
+  async getAllPrompts(): Promise<Prompts> {
+    return new Promise((resolve, reject) => {
+      this.db.get("prompts", (result) => {
+        resolve(result.prompts || [])
+      })
+    })
+  }
+
+  async addPrompt(prompt: Prompt) {
+    const prompts = await this.getAllPrompts()
+    const newPrompts = [prompt, ...prompts]
+    this.db.set({ prompts: newPrompts })
+  }
+
+  async deletePrompt(id: string) {
+    const prompts = await this.getAllPrompts()
+    const newPrompts = prompts.filter((prompt) => prompt.id !== id)
+    this.db.set({ prompts: newPrompts })
+  }
+
+  async updatePrompt(id: string, title: string, content: string, is_system: boolean) {
+    const prompts = await this.getAllPrompts()
+    const newPrompts = prompts.map((prompt) => {
+      if (prompt.id === id) {
+        prompt.title = title
+        prompt.content = content
+        prompt.is_system = is_system
+      }
+      return prompt
+    })
+    this.db.set({ prompts: newPrompts })
+  }
+
+  async getPromptById(id: string) {
+    const prompts = await this.getAllPrompts()
+    return prompts.find((prompt) => prompt.id === id)
+  }
+
 }
 
+
 const generateID = () => {
   return "pa_xxxx-xxxx-xxx-xxxx".replace(/[x]/g, () => {
     const r = Math.floor(Math.random() * 16)
@@ -188,3 +241,39 @@ export const removeMessageUsingHistoryId = async (history_id: string) => {
   const newChatHistory = chatHistory.slice(0, -1)
   await db.db.set({ [history_id]: newChatHistory })
 }
+
+
+export const getAllPrompts = async () => {
+  const db = new PageAssitDatabase()
+  return await db.getAllPrompts()
+}
+
+
+export const savePrompt = async ({ content, title, is_system = false }: { title: string, content: string, is_system: boolean }) => {
+  const db = new PageAssitDatabase()
+  const id = generateID()
+  const createdAt = Date.now()
+  const prompt = { id, title, content, is_system, createdAt }
+  await db.addPrompt(prompt)
+  return prompt
+}
+
+
+export const deletePromptById = async (id: string) => {
+  const db = new PageAssitDatabase()
+  await db.deletePrompt(id)
+  return id
+}
+
+
+export const updatePrompt = async ({ content, id, title, is_system }: { id: string, title: string, content: string, is_system: boolean }) => {
+  const db = new PageAssitDatabase()
+  await db.updatePrompt(id, title, content, is_system)
+  return id
+}
+
+
+export const getPromptById = async (id: string) => {
+  const db = new PageAssitDatabase()
+  return await db.getPromptById(id)
+}
\ No newline at end of file
diff --git a/src/routes/index.tsx b/src/routes/index.tsx
index f932315..4abef46 100644
--- a/src/routes/index.tsx
+++ b/src/routes/index.tsx
@@ -4,6 +4,7 @@ import { useDarkMode } from "~hooks/useDarkmode"
 import { SidepanelSettings } from "./sidepanel-settings"
 import { OptionIndex } from "./option-index"
 import { OptionModal } from "./option-model"
+import { OptionPrompt } from "./option-prompt"
 
 export const OptionRouting = () => {
   const { mode } = useDarkMode()
@@ -13,6 +14,7 @@ export const OptionRouting = () => {
       <Routes>
         <Route path="/" element={<OptionIndex />} />
         <Route path="/models" element={<OptionModal />} />
+        <Route path="/prompts" element={<OptionPrompt />} />
       </Routes>
     </div>
   )
diff --git a/src/routes/option-prompt.tsx b/src/routes/option-prompt.tsx
new file mode 100644
index 0000000..2fbd433
--- /dev/null
+++ b/src/routes/option-prompt.tsx
@@ -0,0 +1,10 @@
+import OptionLayout from "~components/Option/Layout"
+import { PromptBody } from "~components/Option/Prompt"
+
+export const OptionPrompt = () => {
+  return (
+    <OptionLayout>
+      <PromptBody />
+    </OptionLayout>
+  )
+}
diff --git a/src/web/local-google.ts b/src/web/local-google.ts
index cb2193a..e76475a 100644
--- a/src/web/local-google.ts
+++ b/src/web/local-google.ts
@@ -40,7 +40,14 @@ export const localGoogleSearch = async (query: string) => {
     (result) => {
       const title = result.querySelector("h3")?.textContent
       const link = result.querySelector("a")?.getAttribute("href")
-      return { title, link }
+      let content = result.querySelector("div[data-sncf='2']")?.textContent 
+      if(content === "") {
+        content = result.querySelector("div[data-sncf='1']")?.textContent
+        if(content === "") {
+          content = result.querySelector("div[data-sncf='3']")?.textContent
+        }
+      }
+      return { title, link, content }
     }
   )
   const filteredSearchResults = searchResults