chore: Update logic of runtime url rewrite

docs/connection-issue.md

@@ -0,0 +1 @@
+# Ollama Connection Issues

src/components/Common/AdvanceOllamaSettings.tsx

@@ -0,0 +1,49 @@
+import { useStorage } from "@plasmohq/storage/hook"
+import { Input, Switch } from "antd"
+import { useTranslation } from "react-i18next"
+
+export const AdvanceOllamaSettings = () => {
+  const [urlRewriteEnabled, setUrlRewriteEnabled] = useStorage(
+    "urlRewriteEnabled",
+    true
+  )
+
+  const [rewriteUrl, setRewriteUrl] = useStorage(
+    "rewriteUrl",
+    "http://127.0.0.1:11434"
+  )
+  const { t } = useTranslation("settings")
+
+  return (
+    <div className="space-y-4">
+      <div className="flex sm:flex-row flex-col space-y-4 sm:space-y-0 sm:justify-between">
+        <span className="text-gray-500 dark:text-neutral-50 ">
+          {t("generalSettings.settings.advanced.urlRewriteEnabled.label")}
+        </span>
+        <div>
+          <Switch
+            className="mt-4 sm:mt-0"
+            checked={urlRewriteEnabled}
+            onChange={(checked) => setUrlRewriteEnabled(checked)}
+          />
+        </div>
+      </div>
+      <div className="flex flex-col space-y-4 sm:space-y-0 sm:justify-between">
+        <span className="text-gray-500 dark:text-neutral-50 mb-3">
+          {t("generalSettings.settings.advanced.urlRewriteEnabled.label")}
+        </span>
+        <div>
+          <Input
+            className="w-full"
+            value={rewriteUrl}
+            disabled={urlRewriteEnabled}
+            placeholder={t(
+              "generalSettings.settings.advanced.urlRewriteEnabled.placeholder"
+            )}
+            onChange={(e) => setRewriteUrl(e.target.value)}
+          />
+        </div>
+      </div>
+    </div>
+  )
+}

src/components/Option/Settings/ollama.tsx

@@ -1,5 +1,5 @@
 import { useMutation, useQuery } from "@tanstack/react-query"
-import { Form, InputNumber, Select, Skeleton } from "antd"
+import { Collapse, Form, InputNumber, Select, Skeleton } from "antd"
 import { useState } from "react"
 import { SaveButton } from "~/components/Common/SaveButton"
 import {
@@ -13,9 +13,12 @@ import {
 } from "~/services/ollama"
 import { SettingPrompt } from "./prompt"
 import { useTranslation } from "react-i18next"
+import { useStorage } from "@plasmohq/storage/hook"
+import { AdvanceOllamaSettings } from "@/components/Common/AdvanceOllamaSettings"
 
 export const SettingsOllama = () => {
   const [ollamaURL, setOllamaURL] = useState<string>("")
+
   const { t } = useTranslation("settings")
 
   const { data: ollamaInfo, status } = useQuery({
@@ -61,7 +64,7 @@ export const SettingsOllama = () => {
               </h2>
               <div className="border border-b border-gray-200 dark:border-gray-600 mt-3 mb-6"></div>
             </div>
-            <div>
+            <div className="mb-3">
               <label
                 htmlFor="ollamaURL"
                 className="text-sm font-medium dark:text-gray-200">
@@ -78,6 +81,21 @@ export const SettingsOllama = () => {
                 className="w-full p-2 border border-gray-300 rounded-md dark:bg-[#262626] dark:text-gray-100"
               />
             </div>
+            <Collapse
+              size="small"
+              items={[
+                {
+                  key: "1",
+                  label: (
+                    <h2 className="text-base font-semibold leading-7 text-gray-900 dark:text-white">
+                      {t("ollamaSettings.settings.advanced.label")}
+                    </h2>
+                  ),
+                  children: <AdvanceOllamaSettings />
+                }
+              ]}
+            />
+
             <div className="flex justify-end">
               <SaveButton
                 onClick={() => {
@@ -130,7 +148,9 @@ export const SettingsOllama = () => {
                       0
                   }
                   showSearch
-                  placeholder={t("ollamaSettings.settings.ragSettings.model.placeholder")}
+                  placeholder={t(
+                    "ollamaSettings.settings.ragSettings.model.placeholder"
+                  )}
                   style={{ width: "100%" }}
                   className="mt-4"
                   options={ollamaInfo.models?.map((model) => ({
@@ -144,26 +164,38 @@ export const SettingsOllama = () => {
                 name="chunkSize"
                 label={t("ollamaSettings.settings.ragSettings.chunkSize.label")}
                 rules={[
-                  { required: true, message: t("ollamaSettings.settings.ragSettings.chunkSize.required")
+                  {
+                    required: true,
+                    message: t(
+                      "ollamaSettings.settings.ragSettings.chunkSize.required"
+                    )
                   }
                 ]}>
                 <InputNumber
                   style={{ width: "100%" }}
-                  placeholder={t("ollamaSettings.settings.ragSettings.chunkSize.placeholder")}
+                  placeholder={t(
+                    "ollamaSettings.settings.ragSettings.chunkSize.placeholder"
+                  )}
                 />
               </Form.Item>
               <Form.Item
                 name="chunkOverlap"
-                label={t("ollamaSettings.settings.ragSettings.chunkOverlap.label")}
+                label={t(
+                  "ollamaSettings.settings.ragSettings.chunkOverlap.label"
+                )}
                 rules={[
                   {
                     required: true,
-                    message: t("ollamaSettings.settings.ragSettings.chunkOverlap.required")
+                    message: t(
+                      "ollamaSettings.settings.ragSettings.chunkOverlap.required"
+                    )
                   }
                 ]}>
                 <InputNumber
                   style={{ width: "100%" }}
-                  placeholder={t("ollamaSettings.settings.ragSettings.chunkOverlap.placeholder")}
+                  placeholder={t(
+                    "ollamaSettings.settings.ragSettings.chunkOverlap.placeholder"
+                  )}
                 />
               </Form.Item>
 

src/libs/runtime.ts

@@ -1,8 +1,15 @@
-export const chromeRunTime = async function (domain: string) {
+import { getAdvancedOllamaSettings } from "@/services/app"
+
+export const urlRewriteRuntime = async function (domain: string) {
   if (browser.runtime && browser.runtime.id) {
+    const { isEnableRewriteUrl, rewriteUrl } = await getAdvancedOllamaSettings()
     if (import.meta.env.BROWSER === "chrome") {
       const url = new URL(domain)
       const domains = [url.hostname]
+      let origin = `${url.protocol}//${url.hostname}`
+      if (!isEnableRewriteUrl && rewriteUrl) {
+        origin = rewriteUrl
+      }
       const rules = [
         {
           id: 1,
@@ -16,13 +23,12 @@ export const chromeRunTime = async function (domain: string) {
               {
                 header: "Origin",
                 operation: "set",
-                value: `${url.protocol}//${url.hostname}`
+                value: origin
               }
             ]
           }
         }
       ]
-
       await browser.declarativeNetRequest.updateDynamicRules({
         removeRuleIds: rules.map((r) => r.id),
         // @ts-ignore
@@ -35,10 +41,13 @@ export const chromeRunTime = async function (domain: string) {
       const domains = [`*://${url.hostname}/*`]
       browser.webRequest.onBeforeSendHeaders.addListener(
         (details) => {
+          let origin = `${url.protocol}//${url.hostname}`
+          if (!isEnableRewriteUrl && rewriteUrl) {
+            origin = rewriteUrl
+          }
           for (let i = 0; i < details.requestHeaders.length; i++) {
             if (details.requestHeaders[i].name === "Origin") {
-              details.requestHeaders[i].value =
+              details.requestHeaders[i].value = origin
-                `${url.protocol}//${url.hostname}`
             }
           }
           return { requestHeaders: details.requestHeaders }

src/loader/html.ts

@@ -1,7 +1,7 @@
 import { BaseDocumentLoader } from "langchain/document_loaders/base"
 import { Document } from "@langchain/core/documents"
 import { compile } from "html-to-text"
-import { chromeRunTime } from "~/libs/runtime"
+import { urlRewriteRuntime } from "~/libs/runtime"
 import { YtTranscript } from "yt-transcript"
 import { isWikipedia, parseWikipedia } from "@/parser/wiki"
 
@@ -102,7 +102,7 @@ export class PageAssistHtmlLoader
         }
       ]
     }
-    await chromeRunTime(this.url)
+    await urlRewriteRuntime(this.url)
     const fetchHTML = await fetch(this.url)
     let html = await fetchHTML.text()
 

src/services/app.ts

@@ -1,20 +1,39 @@
 import { Storage } from "@plasmohq/storage"
 const storage = new Storage()
 
-export const isUrlRewriteEnabled = async () => {
+const DEFAULT_URL_REWRITE_URL = "http://127.0.0.1:11434"
-  const enabled = await storage.get("urlRewriteEnabled")
-  return enabled === "true"
-}
 
+export const isUrlRewriteEnabled = async () => {
+  const enabled = await storage.get<boolean | undefined>("urlRewriteEnabled")
+  if (typeof enabled === "undefined") {
+    return true
+  }
+  return enabled
+}
 export const setUrlRewriteEnabled = async (enabled: boolean) => {
   await storage.set("urlRewriteEnabled", enabled ? "true" : "false")
 }
 
 export const getRewriteUrl = async () => {
   const rewriteUrl = await storage.get("rewriteUrl")
+  if (!rewriteUrl || rewriteUrl.trim() === "") {
+    return DEFAULT_URL_REWRITE_URL
+  }
   return rewriteUrl
 }
 
 export const setRewriteUrl = async (url: string) => {
   await storage.set("rewriteUrl", url)
 }
+
+export const getAdvancedOllamaSettings = async () => {
+  const [isEnableRewriteUrl, rewriteUrl] = await Promise.all([
+    isUrlRewriteEnabled(),
+    getRewriteUrl()
+  ])
+
+  return {
+    isEnableRewriteUrl,
+    rewriteUrl
+  }
+}

src/services/ollama.ts

@@ -1,6 +1,6 @@
 import { Storage } from "@plasmohq/storage"
 import { cleanUrl } from "../libs/clean-url"
-import { chromeRunTime } from "../libs/runtime"
+import { urlRewriteRuntime } from "../libs/runtime"
 
 const storage = new Storage()
 
@@ -22,10 +22,10 @@ Search results:
 export const getOllamaURL = async () => {
   const ollamaURL = await storage.get("ollamaURL")
   if (!ollamaURL || ollamaURL.length === 0) {
-    await chromeRunTime(DEFAULT_OLLAMA_URL)
+    await urlRewriteRuntime(DEFAULT_OLLAMA_URL)
     return DEFAULT_OLLAMA_URL
   }
-  await chromeRunTime(cleanUrl(ollamaURL))
+  await urlRewriteRuntime(cleanUrl(ollamaURL))
   return ollamaURL
 }
 
@@ -163,7 +163,7 @@ export const setOllamaURL = async (ollamaURL: string) => {
       "http://127.0.0.1:"
     )
   }
-  await chromeRunTime(cleanUrl(formattedUrl))
+  await urlRewriteRuntime(cleanUrl(formattedUrl))
   await storage.set("ollamaURL", cleanUrl(formattedUrl))
 }
 

src/web/search-engines/duckduckgo.ts

@@ -1,5 +1,5 @@
 import { cleanUrl } from "@/libs/clean-url"
-import { chromeRunTime } from "@/libs/runtime"
+import { urlRewriteRuntime } from "@/libs/runtime"
 import { PageAssistHtmlLoader } from "@/loader/html"
 import {
   defaultEmbeddingChunkOverlap,
@@ -18,7 +18,7 @@ import { RecursiveCharacterTextSplitter } from "langchain/text_splitter"
 import { MemoryVectorStore } from "langchain/vectorstores/memory"
 
 export const localDuckDuckGoSearch = async (query: string) => {
-  await chromeRunTime(cleanUrl("https://html.duckduckgo.com/html/?q=" + query))
+  await urlRewriteRuntime(cleanUrl("https://html.duckduckgo.com/html/?q=" + query))
 
   const abortController = new AbortController()
   setTimeout(() => abortController.abort(), 10000)

src/web/search-engines/google.ts

@@ -7,7 +7,7 @@ import type { Document } from "@langchain/core/documents"
 import { RecursiveCharacterTextSplitter } from "langchain/text_splitter"
 import { MemoryVectorStore } from "langchain/vectorstores/memory"
 import { cleanUrl } from "~/libs/clean-url"
-import { chromeRunTime } from "~/libs/runtime"
+import { urlRewriteRuntime } from "~/libs/runtime"
 import { PageAssistHtmlLoader } from "~/loader/html"
 import {
   defaultEmbeddingChunkOverlap,
@@ -18,7 +18,7 @@ import {
 
 
 export const localGoogleSearch = async (query: string) => {
-  await chromeRunTime(
+  await urlRewriteRuntime(
     cleanUrl("https://www.google.com/search?hl=en&q=" + query)
   )
   const abortController = new AbortController()

src/web/search-engines/sogou.ts

@@ -1,5 +1,5 @@
 import { cleanUrl } from "@/libs/clean-url"
-import { chromeRunTime } from "@/libs/runtime"
+import { urlRewriteRuntime } from "@/libs/runtime"
 import { PageAssistHtmlLoader } from "@/loader/html"
 import {
   defaultEmbeddingChunkOverlap,
@@ -25,7 +25,7 @@ const getCorrectTargeUrl = async (url: string) => {
   return matches?.[1] || ""
 }
 export const localSogouSearch = async (query: string) => {
-  await chromeRunTime(cleanUrl("https://www.sogou.com/web?query=" + query))
+  await urlRewriteRuntime(cleanUrl("https://www.sogou.com/web?query=" + query))
 
   const abortController = new AbortController()