Merge pull request #2785 from cheahjs/feat/openai-embeddings-batch

feat: add RAG_EMBEDDING_OPENAI_BATCH_SIZE to batch multiple embeddings

backend/apps/rag/main.py

@@ -78,6 +78,7 @@ from utils.misc import (
 from utils.utils import get_current_user, get_admin_user
 
 from config import (
+    AppConfig,
     ENV,
     SRC_LOG_LEVELS,
     UPLOAD_DIR,
@@ -114,7 +115,7 @@ from config import (
     SERPER_API_KEY,
     RAG_WEB_SEARCH_RESULT_COUNT,
     RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
-    AppConfig,
+    RAG_EMBEDDING_OPENAI_BATCH_SIZE,
 )
 
 from constants import ERROR_MESSAGES
@@ -139,6 +140,7 @@ app.state.config.CHUNK_OVERLAP = CHUNK_OVERLAP
 
 app.state.config.RAG_EMBEDDING_ENGINE = RAG_EMBEDDING_ENGINE
 app.state.config.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
+app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE = RAG_EMBEDDING_OPENAI_BATCH_SIZE
 app.state.config.RAG_RERANKING_MODEL = RAG_RERANKING_MODEL
 app.state.config.RAG_TEMPLATE = RAG_TEMPLATE
 
@@ -212,6 +214,7 @@ app.state.EMBEDDING_FUNCTION = get_embedding_function(
     app.state.sentence_transformer_ef,
     app.state.config.OPENAI_API_KEY,
     app.state.config.OPENAI_API_BASE_URL,
+    app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
 )
 
 origins = ["*"]
@@ -248,6 +251,7 @@ async def get_status():
         "embedding_engine": app.state.config.RAG_EMBEDDING_ENGINE,
         "embedding_model": app.state.config.RAG_EMBEDDING_MODEL,
         "reranking_model": app.state.config.RAG_RERANKING_MODEL,
+        "openai_batch_size": app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
     }
 
 
@@ -260,6 +264,7 @@ async def get_embedding_config(user=Depends(get_admin_user)):
         "openai_config": {
             "url": app.state.config.OPENAI_API_BASE_URL,
             "key": app.state.config.OPENAI_API_KEY,
+            "batch_size": app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
         },
     }
 
@@ -275,6 +280,7 @@ async def get_reraanking_config(user=Depends(get_admin_user)):
 class OpenAIConfigForm(BaseModel):
     url: str
     key: str
+    batch_size: Optional[int] = None
 
 
 class EmbeddingModelUpdateForm(BaseModel):
@@ -295,9 +301,14 @@ async def update_embedding_config(
         app.state.config.RAG_EMBEDDING_MODEL = form_data.embedding_model
 
         if app.state.config.RAG_EMBEDDING_ENGINE in ["ollama", "openai"]:
-            if form_data.openai_config != None:
+            if form_data.openai_config is not None:
                 app.state.config.OPENAI_API_BASE_URL = form_data.openai_config.url
                 app.state.config.OPENAI_API_KEY = form_data.openai_config.key
+                app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE = (
+                    form_data.openai_config.batch_size
+                    if form_data.openai_config.batch_size
+                    else 1
+                )
 
         update_embedding_model(app.state.config.RAG_EMBEDDING_MODEL)
 
@@ -307,6 +318,7 @@ async def update_embedding_config(
             app.state.sentence_transformer_ef,
             app.state.config.OPENAI_API_KEY,
             app.state.config.OPENAI_API_BASE_URL,
+            app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
         )
 
         return {
@@ -316,6 +328,7 @@ async def update_embedding_config(
             "openai_config": {
                 "url": app.state.config.OPENAI_API_BASE_URL,
                 "key": app.state.config.OPENAI_API_KEY,
+                "batch_size": app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
             },
         }
     except Exception as e:
@@ -881,6 +894,7 @@ def store_docs_in_vector_db(docs, collection_name, overwrite: bool = False) -> b
             app.state.sentence_transformer_ef,
             app.state.config.OPENAI_API_KEY,
             app.state.config.OPENAI_API_BASE_URL,
+            app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
         )
 
         embedding_texts = list(map(lambda x: x.replace("\n", " "), texts))

backend/apps/rag/utils.py

@@ -2,7 +2,7 @@ import os
 import logging
 import requests
 
-from typing import List
+from typing import List, Union
 
 from apps.ollama.main import (
     generate_ollama_embeddings,
@@ -21,17 +21,7 @@ from langchain.retrievers import (
 from typing import Optional
 
 
-from config import (
-    SRC_LOG_LEVELS,
-    CHROMA_CLIENT,
-    SEARXNG_QUERY_URL,
-    GOOGLE_PSE_API_KEY,
-    GOOGLE_PSE_ENGINE_ID,
-    BRAVE_SEARCH_API_KEY,
-    SERPSTACK_API_KEY,
-    SERPSTACK_HTTPS,
-    SERPER_API_KEY,
-)
+from config import SRC_LOG_LEVELS, CHROMA_CLIENT
 
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["RAG"])
@@ -209,6 +199,7 @@ def get_embedding_function(
     embedding_function,
     openai_key,
     openai_url,
+    batch_size,
 ):
     if embedding_engine == "":
         return lambda query: embedding_function.encode(query).tolist()
@@ -232,6 +223,12 @@ def get_embedding_function(
 
         def generate_multiple(query, f):
             if isinstance(query, list):
+                if embedding_engine == "openai":
+                    embeddings = []
+                    for i in range(0, len(query), batch_size):
+                        embeddings.extend(f(query[i : i + batch_size]))
+                    return embeddings
+                else:
                     return [f(q) for q in query]
             else:
                 return f(query)
@@ -413,8 +410,22 @@ def get_model_path(model: str, update_model: bool = False):
 
 
 def generate_openai_embeddings(
-    model: str, text: str, key: str, url: str = "https://api.openai.com/v1"
+    model: str,
+    text: Union[str, list[str]],
+    key: str,
+    url: str = "https://api.openai.com/v1",
 ):
+    if isinstance(text, list):
+        embeddings = generate_openai_batch_embeddings(model, text, key, url)
+    else:
+        embeddings = generate_openai_batch_embeddings(model, [text], key, url)
+
+    return embeddings[0] if isinstance(text, str) else embeddings
+
+
+def generate_openai_batch_embeddings(
+    model: str, texts: list[str], key: str, url: str = "https://api.openai.com/v1"
+) -> Optional[list[list[float]]]:
     try:
         r = requests.post(
             f"{url}/embeddings",
@@ -422,12 +433,12 @@ def generate_openai_embeddings(
                 "Content-Type": "application/json",
                 "Authorization": f"Bearer {key}",
             },
-            json={"input": text, "model": model},
+            json={"input": texts, "model": model},
         )
         r.raise_for_status()
         data = r.json()
         if "data" in data:
-            return data["data"][0]["embedding"]
+            return [elem["embedding"] for elem in data["data"]]
         else:
             raise "Something went wrong :/"
     except Exception as e:

backend/config.py

@@ -683,6 +683,12 @@ RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE = (
     os.environ.get("RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE", "").lower() == "true"
 )
 
+RAG_EMBEDDING_OPENAI_BATCH_SIZE = PersistentConfig(
+    "RAG_EMBEDDING_OPENAI_BATCH_SIZE",
+    "rag.embedding_openai_batch_size",
+    os.environ.get("RAG_EMBEDDING_OPENAI_BATCH_SIZE", 1),
+)
+
 RAG_RERANKING_MODEL = PersistentConfig(
     "RAG_RERANKING_MODEL",
     "rag.reranking_model",

src/lib/apis/rag/index.ts

@@ -415,6 +415,7 @@ export const getEmbeddingConfig = async (token: string) => {
 type OpenAIConfigForm = {
 	key: string;
 	url: string;
+	batch_size: number;
 };
 
 type EmbeddingModelUpdateForm = {

src/lib/components/documents/Settings/General.svelte

@@ -31,6 +31,7 @@
 
 	let OpenAIKey = '';
 	let OpenAIUrl = '';
+	let OpenAIBatchSize = 1;
 
 	let querySettings = {
 		template: '',
@@ -92,7 +93,8 @@
 				? {
 						openai_config: {
 							key: OpenAIKey,
-							url: OpenAIUrl
+							url: OpenAIUrl,
+							batch_size: OpenAIBatchSize
 						}
 				  }
 				: {})
@@ -159,6 +161,7 @@
 
 			OpenAIKey = embeddingConfig.openai_config.key;
 			OpenAIUrl = embeddingConfig.openai_config.url;
+			OpenAIBatchSize = embeddingConfig.openai_config.batch_size ?? 1;
 		}
 	};
 
@@ -282,6 +285,30 @@
 						required
 					/>
 				</div>
+				<div class="flex mt-0.5 space-x-2">
+					<div class=" self-center text-xs font-medium">{$i18n.t('Embedding Batch Size')}</div>
+					<div class=" flex-1">
+						<input
+							id="steps-range"
+							type="range"
+							min="1"
+							max="2048"
+							step="1"
+							bind:value={OpenAIBatchSize}
+							class="w-full h-2 rounded-lg appearance-none cursor-pointer dark:bg-gray-700"
+						/>
+					</div>
+					<div class="">
+						<input
+							bind:value={OpenAIBatchSize}
+							type="number"
+							class=" bg-transparent text-center w-14"
+							min="-2"
+							max="16000"
+							step="1"
+						/>
+					</div>
+				</div>
 			{/if}
 
 			<div class=" flex w-full justify-between">

src/lib/i18n/locales/ar-BH/translation.json

@@ -170,6 +170,7 @@
 	"Edit Doc": "تعديل الملف",
 	"Edit User": "تعديل المستخدم",
 	"Email": "البريد",
+	"Embedding Batch Size": "",
 	"Embedding Model": "نموذج التضمين",
 	"Embedding Model Engine": "تضمين محرك النموذج",
 	"Embedding model set to \"{{embedding_model}}\"": "تم تعيين نموذج التضمين على \"{{embedding_model}}\"",

src/lib/i18n/locales/bg-BG/translation.json

@@ -170,6 +170,7 @@
 	"Edit Doc": "Редактиране на документ",
 	"Edit User": "Редактиране на потребител",
 	"Email": "Имейл",
+	"Embedding Batch Size": "",
 	"Embedding Model": "Модел за вграждане",
 	"Embedding Model Engine": "Модел за вграждане",
 	"Embedding model set to \"{{embedding_model}}\"": "Модел за вграждане е настроен на \"{{embedding_model}}\"",

src/lib/i18n/locales/bn-BD/translation.json

@@ -170,6 +170,7 @@
 	"Edit Doc": "ডকুমেন্ট এডিট করুন",
 	"Edit User": "ইউজার এডিট করুন",
 	"Email": "ইমেইল",
+	"Embedding Batch Size": "",
 	"Embedding Model": "ইমেজ ইমেবডিং মডেল",
 	"Embedding Model Engine": "ইমেজ ইমেবডিং মডেল ইঞ্জিন",
 	"Embedding model set to \"{{embedding_model}}\"": "ইমেজ ইমেবডিং মডেল সেট করা হয়েছে - \"{{embedding_model}}\"",

src/lib/i18n/locales/ca-ES/translation.json

@@ -170,6 +170,7 @@
 	"Edit Doc": "Edita Document",
 	"Edit User": "Edita Usuari",
 	"Email": "Correu electrònic",
+	"Embedding Batch Size": "",
 	"Embedding Model": "Model d'embutiment",
 	"Embedding Model Engine": "Motor de model d'embutiment",
 	"Embedding model set to \"{{embedding_model}}\"": "Model d'embutiment configurat a \"{{embedding_model}}\"",

src/lib/i18n/locales/ceb-PH/translation.json

@@ -170,6 +170,7 @@
 	"Edit Doc": "I-edit ang dokumento",
 	"Edit User": "I-edit ang tiggamit",
 	"Email": "E-mail",
+	"Embedding Batch Size": "",
 	"Embedding Model": "",
 	"Embedding Model Engine": "",
 	"Embedding model set to \"{{embedding_model}}\"": "",

src/lib/i18n/locales/de-DE/translation.json

@@ -170,6 +170,7 @@
 	"Edit Doc": "Dokument bearbeiten",
 	"Edit User": "Benutzer bearbeiten",
 	"Email": "E-Mail",
+	"Embedding Batch Size": "",
 	"Embedding Model": "Embedding-Modell",
 	"Embedding Model Engine": "Embedding-Modell-Engine",
 	"Embedding model set to \"{{embedding_model}}\"": "Embedding-Modell auf \"{{embedding_model}}\" gesetzt",

src/lib/i18n/locales/dg-DG/translation.json

@@ -170,6 +170,7 @@
 	"Edit Doc": "Edit Doge",
 	"Edit User": "Edit Wowser",
 	"Email": "Email",
+	"Embedding Batch Size": "",
 	"Embedding Model": "",
 	"Embedding Model Engine": "",
 	"Embedding model set to \"{{embedding_model}}\"": "",

src/lib/i18n/locales/en-GB/translation.json

@@ -170,6 +170,7 @@
 	"Edit Doc": "",
 	"Edit User": "",
 	"Email": "",
+	"Embedding Batch Size": "",
 	"Embedding Model": "",
 	"Embedding Model Engine": "",
 	"Embedding model set to \"{{embedding_model}}\"": "",

src/lib/i18n/locales/en-US/translation.json

@@ -170,6 +170,7 @@
 	"Edit Doc": "",
 	"Edit User": "",
 	"Email": "",
+	"Embedding Batch Size": "",
 	"Embedding Model": "",
 	"Embedding Model Engine": "",
 	"Embedding model set to \"{{embedding_model}}\"": "",

src/lib/i18n/locales/es-ES/translation.json

@@ -170,6 +170,7 @@
 	"Edit Doc": "Editar Documento",
 	"Edit User": "Editar Usuario",
 	"Email": "Email",
+	"Embedding Batch Size": "",
 	"Embedding Model": "Modelo de Embedding",
 	"Embedding Model Engine": "Motor de Modelo de Embedding",
 	"Embedding model set to \"{{embedding_model}}\"": "Modelo de Embedding configurado a \"{{embedding_model}}\"",

src/lib/i18n/locales/fa-IR/translation.json

@@ -170,6 +170,7 @@
 	"Edit Doc": "ویرایش سند",
 	"Edit User": "ویرایش کاربر",
 	"Email": "ایمیل",
+	"Embedding Batch Size": "",
 	"Embedding Model": "مدل پیدائش",
 	"Embedding Model Engine": "محرک مدل پیدائش",
 	"Embedding model set to \"{{embedding_model}}\"": "مدل پیدائش را به \"{{embedding_model}}\" تنظیم کنید",

src/lib/i18n/locales/fi-FI/translation.json

@@ -170,6 +170,7 @@
 	"Edit Doc": "Muokkaa asiakirjaa",
 	"Edit User": "Muokkaa käyttäjää",
 	"Email": "Sähköposti",
+	"Embedding Batch Size": "",
 	"Embedding Model": "Upotusmalli",
 	"Embedding Model Engine": "Upotusmallin moottori",
 	"Embedding model set to \"{{embedding_model}}\"": "\"{{embedding_model}}\" valittu upotusmalliksi",

src/lib/i18n/locales/fr-CA/translation.json

@@ -170,6 +170,7 @@
 	"Edit Doc": "Éditer le document",
 	"Edit User": "Éditer l'utilisateur",
 	"Email": "Email",
+	"Embedding Batch Size": "",
 	"Embedding Model": "Modèle d'embedding",
 	"Embedding Model Engine": "Moteur du modèle d'embedding",
 	"Embedding model set to \"{{embedding_model}}\"": "Modèle d'embedding défini sur \"{{embedding_model}}\"",

src/lib/i18n/locales/fr-FR/translation.json

@@ -170,6 +170,7 @@
 	"Edit Doc": "Éditer le document",
 	"Edit User": "Éditer l'utilisateur",
 	"Email": "Email",
+	"Embedding Batch Size": "",
 	"Embedding Model": "Modèle pour l'Embedding",
 	"Embedding Model Engine": "Moteur du Modèle d'Embedding",
 	"Embedding model set to \"{{embedding_model}}\"": "Modèle d'embedding défini sur \"{{embedding_model}}\"",

src/lib/i18n/locales/he-IL/translation.json

@@ -170,6 +170,7 @@
 	"Edit Doc": "ערוך מסמך",
 	"Edit User": "ערוך משתמש",
 	"Email": "דוא\"ל",
+	"Embedding Batch Size": "",
 	"Embedding Model": "מודל הטמעה",
 	"Embedding Model Engine": "מנוע מודל הטמעה",
 	"Embedding model set to \"{{embedding_model}}\"": "מודל ההטמעה הוגדר ל-\"{{embedding_model}}\"",