From 768ff2e5558d759e3f5a0e3d1a171223014031f1 Mon Sep 17 00:00:00 2001
From: n4ze3m <n4ze3m@gmail.com>
Date: Sat, 12 Oct 2024 23:32:00 +0530
Subject: [PATCH] feat: Improve model selection and embedding

Refactor embedding models and their handling to improve performance and simplify the process.
Add a new model selection mechanism,  and enhance the UI for model selection, offering clearer and more user-friendly options for embedding models.
Refactor embeddings to use a common model for page assist and RAG, further improving performance and streamlining the workflow.
---
 src/components/Option/Settings/rag.tsx | 28 +++++++++++++-------
 src/hooks/useMessage.tsx               |  4 +--
 src/hooks/useMessageOption.tsx         |  4 +--
 src/libs/process-knowledge.ts          |  4 +--
 src/models/OAIEmbedding.ts             | 16 +-----------
 src/models/embedding.ts                | 36 ++++++++++++++++++++++++++
 src/models/index.ts                    |  2 +-
 src/services/ollama.ts                 | 24 ++++++++++++++++-
 src/utils/memory-embeddings.ts         |  2 +-
 src/web/search-engines/brave.ts        |  5 ++--
 src/web/search-engines/duckduckgo.ts   |  4 +--
 src/web/search-engines/google.ts       |  4 +--
 src/web/search-engines/sogou.ts        |  4 +--
 src/web/website/index.ts               |  4 +--
 14 files changed, 98 insertions(+), 43 deletions(-)
 create mode 100644 src/models/embedding.ts

diff --git a/src/components/Option/Settings/rag.tsx b/src/components/Option/Settings/rag.tsx
index fd5d576..8cc6bbf 100644
--- a/src/components/Option/Settings/rag.tsx
+++ b/src/components/Option/Settings/rag.tsx
@@ -5,13 +5,14 @@ import {
   defaultEmbeddingChunkOverlap,
   defaultEmbeddingChunkSize,
   defaultEmbeddingModelForRag,
-  getAllModels,
+  getEmbeddingModels,
   saveForRag
 } from "~/services/ollama"
 import { SettingPrompt } from "./prompt"
 import { useTranslation } from "react-i18next"
 import { getNoOfRetrievedDocs, getTotalFilePerKB } from "@/services/app"
 import { SidepanelRag } from "./sidepanel-rag"
+import { ProviderIcons } from "@/components/Common/ProviderIcon"
 
 export const RagSettings = () => {
   const { t } = useTranslation("settings")
@@ -29,7 +30,7 @@ export const RagSettings = () => {
         totalFilePerKB,
         noOfRetrievedDocs
       ] = await Promise.all([
-        getAllModels({ returnEmpty: true }),
+        getEmbeddingModels({ returnEmpty: true }),
         defaultEmbeddingChunkOverlap(),
         defaultEmbeddingChunkSize(),
         defaultEmbeddingModelForRag(),
@@ -113,18 +114,27 @@ export const RagSettings = () => {
                 ]}>
                 <Select
                   size="large"
-                  filterOption={(input, option) =>
-                    option!.label.toLowerCase().indexOf(input.toLowerCase()) >=
-                      0 ||
-                    option!.value.toLowerCase().indexOf(input.toLowerCase()) >=
-                      0
-                  }
                   showSearch
                   placeholder={t("rag.ragSettings.model.placeholder")}
                   style={{ width: "100%" }}
                   className="mt-4"
+                  filterOption={(input, option) =>
+                    option.label.key
+                      .toLowerCase()
+                      .indexOf(input.toLowerCase()) >= 0
+                  }
                   options={ollamaInfo.models?.map((model) => ({
-                    label: model.name,
+                    label: (
+                      <span
+                        key={model.model}
+                        className="flex flex-row gap-3 items-center truncate">
+                        <ProviderIcons
+                          provider={model?.provider}
+                          className="w-5 h-5"
+                        />
+                        <span className="truncate">{model.name}</span>
+                      </span>
+                    ),
                     value: model.model
                   }))}
                 />
diff --git a/src/hooks/useMessage.tsx b/src/hooks/useMessage.tsx
index d267dcf..18cceb5 100644
--- a/src/hooks/useMessage.tsx
+++ b/src/hooks/useMessage.tsx
@@ -26,7 +26,6 @@ import { notification } from "antd"
 import { useTranslation } from "react-i18next"
 import { usePageAssist } from "@/context"
 import { formatDocs } from "@/chain/chat-with-x"
-import { OllamaEmbeddingsPageAssist } from "@/models/OllamaEmbedding"
 import { useStorage } from "@plasmohq/storage/hook"
 import { useStoreChatModelSettings } from "@/store/model"
 import { getAllDefaultModelSettings } from "@/services/model-settings"
@@ -34,6 +33,7 @@ import { getSystemPromptForWeb } from "@/web/web"
 import { pageAssistModel } from "@/models"
 import { getPrompt } from "@/services/application"
 import { humanMessageFormatter } from "@/utils/human-message"
+import { pageAssistEmbeddingModel } from "@/models/embedding"
 
 export const useMessage = () => {
   const {
@@ -202,7 +202,7 @@ export const useMessage = () => {
     const ollamaUrl = await getOllamaURL()
     const embeddingModle = await defaultEmbeddingModelForRag()
 
-    const ollamaEmbedding = new OllamaEmbeddingsPageAssist({
+    const ollamaEmbedding = await pageAssistEmbeddingModel({
       model: embeddingModle || selectedModel,
       baseUrl: cleanUrl(ollamaUrl),
       signal: embeddingSignal,
diff --git a/src/hooks/useMessageOption.tsx b/src/hooks/useMessageOption.tsx
index b70b6dd..3b38353 100644
--- a/src/hooks/useMessageOption.tsx
+++ b/src/hooks/useMessageOption.tsx
@@ -24,7 +24,6 @@ import { generateHistory } from "@/utils/generate-history"
 import { useTranslation } from "react-i18next"
 import { saveMessageOnError, saveMessageOnSuccess } from "./chat-helper"
 import { usePageAssist } from "@/context"
-import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama"
 import { PageAssistVectorStore } from "@/libs/PageAssistVectorStore"
 import { formatDocs } from "@/chain/chat-with-x"
 import { useWebUI } from "@/store/webui"
@@ -34,6 +33,7 @@ import { getAllDefaultModelSettings } from "@/services/model-settings"
 import { pageAssistModel } from "@/models"
 import { getNoOfRetrievedDocs } from "@/services/app"
 import { humanMessageFormatter } from "@/utils/human-message"
+import { pageAssistEmbeddingModel } from "@/models/embedding"
 
 export const useMessageOption = () => {
   const {
@@ -628,7 +628,7 @@ export const useMessageOption = () => {
 
     const embeddingModle = await defaultEmbeddingModelForRag()
     const ollamaUrl = await getOllamaURL()
-    const ollamaEmbedding = new OllamaEmbeddings({
+    const ollamaEmbedding = await pageAssistEmbeddingModel({
       model: embeddingModle || selectedModel,
       baseUrl: cleanUrl(ollamaUrl),
       keepAlive:
diff --git a/src/libs/process-knowledge.ts b/src/libs/process-knowledge.ts
index 0288e3a..618fb42 100644
--- a/src/libs/process-knowledge.ts
+++ b/src/libs/process-knowledge.ts
@@ -5,7 +5,6 @@ import {
   defaultEmbeddingChunkSize,
   getOllamaURL
 } from "@/services/ollama"
-import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama"
 import { RecursiveCharacterTextSplitter } from "langchain/text_splitter"
 import { PageAssistVectorStore } from "./PageAssistVectorStore"
 import { PageAssisCSVUrlLoader } from "@/loader/csv"
@@ -13,6 +12,7 @@ import { PageAssisTXTUrlLoader } from "@/loader/txt"
 import { PageAssistDocxLoader } from "@/loader/docx"
 import { cleanUrl } from "./clean-url"
 import { sendEmbeddingCompleteNotification } from "./send-notification"
+import { pageAssistEmbeddingModel } from "@/models/embedding"
 
 
 export const processKnowledge = async (msg: any, id: string): Promise<void> => {
@@ -28,7 +28,7 @@ export const processKnowledge = async (msg: any, id: string): Promise<void> => {
 
     await updateKnowledgeStatus(id, "processing")
 
-    const ollamaEmbedding = new OllamaEmbeddings({
+    const ollamaEmbedding = await pageAssistEmbeddingModel({
       baseUrl: cleanUrl(ollamaUrl),
       model: knowledge.embedding_model
     })
diff --git a/src/models/OAIEmbedding.ts b/src/models/OAIEmbedding.ts
index b2df653..7617021 100644
--- a/src/models/OAIEmbedding.ts
+++ b/src/models/OAIEmbedding.ts
@@ -44,21 +44,6 @@ export interface OpenAIEmbeddingsParams extends EmbeddingsParams {
     signal?: AbortSignal
 }
 
-/**
- * Class for generating embeddings using the OpenAI API. Extends the
- * Embeddings class and implements OpenAIEmbeddingsParams and
- * AzureOpenAIInput.
- * @example
- * ```typescript
- * // Embed a query using OpenAIEmbeddings to generate embeddings for a given text
- * const model = new OpenAIEmbeddings();
- * const res = await model.embedQuery(
- *   "What would be a good company name for a company that makes colorful socks?",
- * );
- * console.log({ res });
- *
- * ```
- */
 export class OAIEmbedding
     extends Embeddings
     implements OpenAIEmbeddingsParams {
@@ -96,6 +81,7 @@ export class OAIEmbedding
     protected client: OpenAIClient
 
     protected clientConfig: ClientOptions
+    
     signal?: AbortSignal
 
     constructor(
diff --git a/src/models/embedding.ts b/src/models/embedding.ts
new file mode 100644
index 0000000..03eb663
--- /dev/null
+++ b/src/models/embedding.ts
@@ -0,0 +1,36 @@
+import { getModelInfo, isCustomModel } from "@/db/models"
+import { OllamaEmbeddingsPageAssist } from "./OllamaEmbedding"
+import { OAIEmbedding } from "./OAIEmbedding"
+import { getOpenAIConfigById } from "@/db/openai"
+
+type EmbeddingModel = {
+    model: string
+    baseUrl: string
+    signal?: AbortSignal
+    keepAlive?: string
+}
+
+export const pageAssistEmbeddingModel = async ({ baseUrl, model, keepAlive, signal }: EmbeddingModel) => {
+    const isCustom = isCustomModel(model)
+    if (isCustom) {
+        const modelInfo = await getModelInfo(model)
+        const providerInfo = await getOpenAIConfigById(modelInfo.provider_id)
+        return new OAIEmbedding({
+            modelName: modelInfo.model_id,
+            model: modelInfo.model_id,
+            signal,
+            openAIApiKey: providerInfo.apiKey || "temp",
+            configuration: {
+                apiKey: providerInfo.apiKey || "temp",
+                baseURL: providerInfo.baseUrl || "",
+            }
+        }) as any
+    }
+
+    return new OllamaEmbeddingsPageAssist({
+        model,
+        baseUrl,
+        keepAlive,
+        signal
+    })
+}
\ No newline at end of file
diff --git a/src/models/index.ts b/src/models/index.ts
index 4df2419..d459e66 100644
--- a/src/models/index.ts
+++ b/src/models/index.ts
@@ -40,7 +40,7 @@ export const pageAssistModel = async ({
   if (isCustom) {
     const modelInfo = await getModelInfo(model)
     const providerInfo = await getOpenAIConfigById(modelInfo.provider_id)
-    console.log(modelInfo, providerInfo)
+
     return new ChatOpenAI({
       modelName: modelInfo.model_id,
       openAIApiKey: providerInfo.apiKey || "temp",
diff --git a/src/services/ollama.ts b/src/services/ollama.ts
index 5aff8c3..f6096ce 100644
--- a/src/services/ollama.ts
+++ b/src/services/ollama.ts
@@ -133,6 +133,28 @@ export const getAllModels = async ({
   }
 }
 
+export const getEmbeddingModels = async ({ returnEmpty }: {
+  returnEmpty?: boolean
+}) => {
+  try {
+    const ollamaModels = await getAllModels({ returnEmpty })
+    const customModels = await ollamaFormatAllCustomModels()
+
+    return [
+      ...ollamaModels.map((model) => {
+        return {
+          ...model,
+          provider: "ollama"
+        }
+      }),
+      ...customModels
+    ]
+  } catch (e) {
+    console.error(e)
+    return []
+  }
+}
+
 export const deleteModel = async (model: string) => {
   const baseUrl = await getOllamaURL()
   const response = await fetcher(`${cleanUrl(baseUrl)}/api/delete`, {
@@ -341,7 +363,7 @@ export const saveForRag = async (
   await setDefaultEmbeddingChunkSize(chunkSize)
   await setDefaultEmbeddingChunkOverlap(overlap)
   await setTotalFilePerKB(totalFilePerKB)
-  if(noOfRetrievedDocs) {
+  if (noOfRetrievedDocs) {
     await setNoOfRetrievedDocs(noOfRetrievedDocs)
   }
 }
diff --git a/src/utils/memory-embeddings.ts b/src/utils/memory-embeddings.ts
index f58f4b9..8d0decc 100644
--- a/src/utils/memory-embeddings.ts
+++ b/src/utils/memory-embeddings.ts
@@ -1,7 +1,7 @@
 import { PageAssistHtmlLoader } from "~/loader/html"
 import { RecursiveCharacterTextSplitter } from "langchain/text_splitter"
 import { MemoryVectorStore } from "langchain/vectorstores/memory"
-import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama"
+
 import {
   defaultEmbeddingChunkOverlap,
   defaultEmbeddingChunkSize
diff --git a/src/web/search-engines/brave.ts b/src/web/search-engines/brave.ts
index b387fbf..b795d8b 100644
--- a/src/web/search-engines/brave.ts
+++ b/src/web/search-engines/brave.ts
@@ -1,6 +1,7 @@
 import { cleanUrl } from "@/libs/clean-url"
 import { urlRewriteRuntime } from "@/libs/runtime"
 import { PageAssistHtmlLoader } from "@/loader/html"
+import { pageAssistEmbeddingModel } from "@/models/embedding"
 import {
     defaultEmbeddingChunkOverlap,
     defaultEmbeddingChunkSize,
@@ -11,7 +12,7 @@ import {
     getIsSimpleInternetSearch,
     totalSearchResults
 } from "@/services/search"
-import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama"
+
 import type { Document } from "@langchain/core/documents"
 import * as cheerio from "cheerio"
 import { RecursiveCharacterTextSplitter } from "langchain/text_splitter"
@@ -81,7 +82,7 @@ export const webBraveSearch = async (query: string) => {
     const ollamaUrl = await getOllamaURL()
 
     const embeddingModle = await defaultEmbeddingModelForRag()
-    const ollamaEmbedding = new OllamaEmbeddings({
+    const ollamaEmbedding = await pageAssistEmbeddingModel({
         model: embeddingModle || "",
         baseUrl: cleanUrl(ollamaUrl)
     })
diff --git a/src/web/search-engines/duckduckgo.ts b/src/web/search-engines/duckduckgo.ts
index 374ce60..e368500 100644
--- a/src/web/search-engines/duckduckgo.ts
+++ b/src/web/search-engines/duckduckgo.ts
@@ -1,6 +1,7 @@
 import { cleanUrl } from "@/libs/clean-url"
 import { urlRewriteRuntime } from "@/libs/runtime"
 import { PageAssistHtmlLoader } from "@/loader/html"
+import { pageAssistEmbeddingModel } from "@/models/embedding"
 import {
   defaultEmbeddingChunkOverlap,
   defaultEmbeddingChunkSize,
@@ -11,7 +12,6 @@ import {
   getIsSimpleInternetSearch,
   totalSearchResults
 } from "@/services/search"
-import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama"
 import type { Document } from "@langchain/core/documents"
 import * as cheerio from "cheerio"
 import { RecursiveCharacterTextSplitter } from "langchain/text_splitter"
@@ -85,7 +85,7 @@ export const webDuckDuckGoSearch = async (query: string) => {
   const ollamaUrl = await getOllamaURL()
 
   const embeddingModle = await defaultEmbeddingModelForRag()
-  const ollamaEmbedding = new OllamaEmbeddings({
+  const ollamaEmbedding = await pageAssistEmbeddingModel({
     model: embeddingModle || "",
     baseUrl: cleanUrl(ollamaUrl)
   })
diff --git a/src/web/search-engines/google.ts b/src/web/search-engines/google.ts
index 8f0a16a..51dd9a3 100644
--- a/src/web/search-engines/google.ts
+++ b/src/web/search-engines/google.ts
@@ -1,8 +1,8 @@
+import { pageAssistEmbeddingModel } from "@/models/embedding"
 import {
   getIsSimpleInternetSearch,
   totalSearchResults
 } from "@/services/search"
-import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama"
 import type { Document } from "@langchain/core/documents"
 import { RecursiveCharacterTextSplitter } from "langchain/text_splitter"
 import { MemoryVectorStore } from "langchain/vectorstores/memory"
@@ -84,7 +84,7 @@ export const webGoogleSearch = async (query: string) => {
   const ollamaUrl = await getOllamaURL()
 
   const embeddingModle = await defaultEmbeddingModelForRag()
-  const ollamaEmbedding = new OllamaEmbeddings({
+  const ollamaEmbedding = await pageAssistEmbeddingModel({
     model: embeddingModle || "",
     baseUrl: cleanUrl(ollamaUrl)
   })
diff --git a/src/web/search-engines/sogou.ts b/src/web/search-engines/sogou.ts
index 274d03f..d1a6090 100644
--- a/src/web/search-engines/sogou.ts
+++ b/src/web/search-engines/sogou.ts
@@ -1,6 +1,7 @@
 import { cleanUrl } from "@/libs/clean-url"
 import { urlRewriteRuntime } from "@/libs/runtime"
 import { PageAssistHtmlLoader } from "@/loader/html"
+import { pageAssistEmbeddingModel } from "@/models/embedding"
 import {
   defaultEmbeddingChunkOverlap,
   defaultEmbeddingChunkSize,
@@ -11,7 +12,6 @@ import {
   getIsSimpleInternetSearch,
   totalSearchResults
 } from "@/services/search"
-import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama"
 import type { Document } from "@langchain/core/documents"
 import * as cheerio from "cheerio"
 import { RecursiveCharacterTextSplitter } from "langchain/text_splitter"
@@ -99,7 +99,7 @@ export const webSogouSearch = async (query: string) => {
   const ollamaUrl = await getOllamaURL()
 
   const embeddingModle = await defaultEmbeddingModelForRag()
-  const ollamaEmbedding = new OllamaEmbeddings({
+  const ollamaEmbedding = await pageAssistEmbeddingModel({
     model: embeddingModle || "",
     baseUrl: cleanUrl(ollamaUrl)
   })
diff --git a/src/web/website/index.ts b/src/web/website/index.ts
index 31ed5f9..817fdb7 100644
--- a/src/web/website/index.ts
+++ b/src/web/website/index.ts
@@ -1,7 +1,7 @@
 import { cleanUrl } from "@/libs/clean-url"
 import { PageAssistHtmlLoader } from "@/loader/html"
+import { pageAssistEmbeddingModel } from "@/models/embedding"
 import { defaultEmbeddingChunkOverlap, defaultEmbeddingChunkSize, defaultEmbeddingModelForRag, getOllamaURL } from "@/services/ollama"
-import { OllamaEmbeddings } from "@langchain/community/embeddings/ollama"
 import { RecursiveCharacterTextSplitter } from "langchain/text_splitter"
 import { MemoryVectorStore } from "langchain/vectorstores/memory"
 
@@ -15,7 +15,7 @@ export const processSingleWebsite = async (url: string, query: string) => {
     const ollamaUrl = await getOllamaURL()
 
     const embeddingModle = await defaultEmbeddingModelForRag()
-    const ollamaEmbedding = new OllamaEmbeddings({
+    const ollamaEmbedding = await pageAssistEmbeddingModel({
         model: embeddingModle || "",
         baseUrl: cleanUrl(ollamaUrl)
     })