chore: Update logic of runtime url rewrite

2024-05-12 21:14:27 +05:30 · 2024-05-12 21:14:27 +05:30 · 62ffe8346e
commit 62ffe8346e
parent bd84b90e5f
11 changed files with 148 additions and 38 deletions
--- a/bun.lockb
+++ b/bun.lockb
--- a/docs/connection-issue.md
+++ b/docs/connection-issue.md
@ -0,0 +1 @@
+# Ollama Connection Issues
--- a/src/components/Common/AdvanceOllamaSettings.tsx
+++ b/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>
+  )
+}
--- a/src/components/Option/Settings/ollama.tsx
+++ b/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>

--- a/src/libs/runtime.ts
+++ b/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 =
-                `${url.protocol}//${url.hostname}`
+              details.requestHeaders[i].value = origin
            }
          }
          return { requestHeaders: details.requestHeaders }
--- a/src/loader/html.ts
+++ b/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()

--- a/src/services/app.ts
+++ b/src/services/app.ts
@ -1,20 +1,39 @@
 import { Storage } from "@plasmohq/storage"
 const storage = new Storage()

-export const isUrlRewriteEnabled = async () => {
-  const enabled = await storage.get("urlRewriteEnabled")
-  return enabled === "true"
-}
+const DEFAULT_URL_REWRITE_URL = "http://127.0.0.1:11434"

+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
+  }
+}
--- a/src/services/ollama.ts
+++ b/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))
 }

--- a/src/web/search-engines/duckduckgo.ts
+++ b/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)
--- a/src/web/search-engines/google.ts
+++ b/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()
--- a/src/web/search-engines/sogou.ts
+++ b/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()