Merge pull request #14370 from daw/feat/add-azure-openai-embeddings-option

feat:Add Azure OpenAI embedding support

backend/open_webui/config.py

@@ -2184,6 +2184,27 @@ RAG_OPENAI_API_KEY = PersistentConfig(
     os.getenv("RAG_OPENAI_API_KEY", OPENAI_API_KEY),
 )
 
+RAG_AZURE_OPENAI_BASE_URL = PersistentConfig(
+    "RAG_AZURE_OPENAI_BASE_URL",
+    "rag.azure_openai.base_url",
+    os.getenv("RAG_AZURE_OPENAI_BASE_URL", ""),
+)
+RAG_AZURE_OPENAI_API_KEY = PersistentConfig(
+    "RAG_AZURE_OPENAI_API_KEY",
+    "rag.azure_openai.api_key",
+    os.getenv("RAG_AZURE_OPENAI_API_KEY", ""),
+)
+RAG_AZURE_OPENAI_DEPLOYMENT = PersistentConfig(
+    "RAG_AZURE_OPENAI_DEPLOYMENT",
+    "rag.azure_openai.deployment",
+    os.getenv("RAG_AZURE_OPENAI_DEPLOYMENT", ""),
+)
+RAG_AZURE_OPENAI_VERSION = PersistentConfig(
+    "RAG_AZURE_OPENAI_VERSION",
+    "rag.azure_openai.version",
+    os.getenv("RAG_AZURE_OPENAI_VERSION", ""),
+)
+
 RAG_OLLAMA_BASE_URL = PersistentConfig(
     "RAG_OLLAMA_BASE_URL",
     "rag.ollama.url",

backend/open_webui/main.py

@@ -207,6 +207,10 @@ from open_webui.config import (
     RAG_FILE_MAX_SIZE,
     RAG_OPENAI_API_BASE_URL,
     RAG_OPENAI_API_KEY,
+    RAG_AZURE_OPENAI_BASE_URL,
+    RAG_AZURE_OPENAI_API_KEY,
+    RAG_AZURE_OPENAI_DEPLOYMENT,
+    RAG_AZURE_OPENAI_VERSION,
     RAG_OLLAMA_BASE_URL,
     RAG_OLLAMA_API_KEY,
     CHUNK_OVERLAP,
@@ -717,6 +721,11 @@ app.state.config.RAG_TEMPLATE = RAG_TEMPLATE
 app.state.config.RAG_OPENAI_API_BASE_URL = RAG_OPENAI_API_BASE_URL
 app.state.config.RAG_OPENAI_API_KEY = RAG_OPENAI_API_KEY
 
+app.state.config.RAG_AZURE_OPENAI_BASE_URL = RAG_AZURE_OPENAI_BASE_URL
+app.state.config.RAG_AZURE_OPENAI_API_KEY = RAG_AZURE_OPENAI_API_KEY
+app.state.config.RAG_AZURE_OPENAI_DEPLOYMENT = RAG_AZURE_OPENAI_DEPLOYMENT
+app.state.config.RAG_AZURE_OPENAI_VERSION = RAG_AZURE_OPENAI_VERSION
+
 app.state.config.RAG_OLLAMA_BASE_URL = RAG_OLLAMA_BASE_URL
 app.state.config.RAG_OLLAMA_API_KEY = RAG_OLLAMA_API_KEY
 
@@ -811,14 +820,32 @@ app.state.EMBEDDING_FUNCTION = get_embedding_function(
     (
         app.state.config.RAG_OPENAI_API_BASE_URL
         if app.state.config.RAG_EMBEDDING_ENGINE == "openai"
-        else app.state.config.RAG_OLLAMA_BASE_URL
+        else (
+            app.state.config.RAG_OLLAMA_BASE_URL
+            if app.state.config.RAG_EMBEDDING_ENGINE == "ollama"
+            else app.state.config.RAG_AZURE_OPENAI_BASE_URL
+        )
     ),
     (
         app.state.config.RAG_OPENAI_API_KEY
         if app.state.config.RAG_EMBEDDING_ENGINE == "openai"
-        else app.state.config.RAG_OLLAMA_API_KEY
+        else (
+            app.state.config.RAG_OLLAMA_API_KEY
+            if app.state.config.RAG_EMBEDDING_ENGINE == "ollama"
+            else app.state.config.RAG_AZURE_OPENAI_API_KEY
+        )
     ),
     app.state.config.RAG_EMBEDDING_BATCH_SIZE,
+    (
+        app.state.config.RAG_AZURE_OPENAI_DEPLOYMENT
+        if app.state.config.RAG_EMBEDDING_ENGINE == "azure_openai"
+        else None
+    ),
+    (
+        app.state.config.RAG_AZURE_OPENAI_VERSION
+        if app.state.config.RAG_EMBEDDING_ENGINE == "azure_openai"
+        else None
+    ),
 )
 
 ########################################

backend/open_webui/retrieval/utils.py

@@ -5,6 +5,7 @@ from typing import Optional, Union
 import requests
 import hashlib
 from concurrent.futures import ThreadPoolExecutor
+import time
 
 from huggingface_hub import snapshot_download
 from langchain.retrievers import ContextualCompressionRetriever, EnsembleRetriever
@@ -400,12 +401,14 @@ def get_embedding_function(
     url,
     key,
     embedding_batch_size,
+    deployment=None,
+    version=None,
 ):
     if embedding_engine == "":
         return lambda query, prefix=None, user=None: embedding_function.encode(
             query, **({"prompt": prefix} if prefix else {})
         ).tolist()
-    elif embedding_engine in ["ollama", "openai"]:
+    elif embedding_engine in ["ollama", "openai", "azure_openai"]:
         func = lambda query, prefix=None, user=None: generate_embeddings(
             engine=embedding_engine,
             model=embedding_model,
@@ -414,6 +417,8 @@ def get_embedding_function(
             url=url,
             key=key,
             user=user,
+            deployment=deployment,
+            version=version,
         )
 
         def generate_multiple(query, prefix, user, func):
@@ -697,6 +702,61 @@ def generate_openai_batch_embeddings(
         return None
 
 
+def generate_azure_openai_batch_embeddings(
+    deployment: str,
+    texts: list[str],
+    url: str,
+    key: str = "",
+    model: str = "",
+    version: str = "",
+    prefix: str = None,
+    user: UserModel = None,
+) -> Optional[list[list[float]]]:
+    try:
+        log.debug(
+            f"generate_azure_openai_batch_embeddings:deployment {deployment} batch size: {len(texts)}"
+        )
+        json_data = {"input": texts, "model": model}
+        if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str):
+            json_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
+
+        url = f"{url}/openai/deployments/{deployment}/embeddings?api-version={version}"
+
+        for _ in range(5):
+            r = requests.post(
+                url,
+                headers={
+                    "Content-Type": "application/json",
+                    "api-key": key,
+                    **(
+                        {
+                            "X-OpenWebUI-User-Name": user.name,
+                            "X-OpenWebUI-User-Id": user.id,
+                            "X-OpenWebUI-User-Email": user.email,
+                            "X-OpenWebUI-User-Role": user.role,
+                        }
+                        if ENABLE_FORWARD_USER_INFO_HEADERS and user
+                        else {}
+                    ),
+                },
+                json=json_data,
+            )
+            if r.status_code == 429:
+                retry = float(r.headers.get("Retry-After", "1"))
+                time.sleep(retry)
+                continue
+            r.raise_for_status()
+            data = r.json()
+            if "data" in data:
+                return [elem["embedding"] for elem in data["data"]]
+            else:
+                raise Exception("Something went wrong :/")
+        return None
+    except Exception as e:
+        log.exception(f"Error generating azure openai batch embeddings: {e}")
+        return None
+
+
 def generate_ollama_batch_embeddings(
     model: str,
     texts: list[str],
@@ -794,6 +854,32 @@ def generate_embeddings(
                 model, [text], url, key, prefix, user
             )
         return embeddings[0] if isinstance(text, str) else embeddings
+    elif engine == "azure_openai":
+        deployment = kwargs.get("deployment", "")
+        version = kwargs.get("version", "")
+        if isinstance(text, list):
+            embeddings = generate_azure_openai_batch_embeddings(
+                deployment,
+                text,
+                url,
+                key,
+                model,
+                version,
+                prefix,
+                user,
+            )
+        else:
+            embeddings = generate_azure_openai_batch_embeddings(
+                deployment,
+                [text],
+                url,
+                key,
+                model,
+                version,
+                prefix,
+                user,
+            )
+        return embeddings[0] if isinstance(text, str) else embeddings
 
 
 import operator

backend/open_webui/routers/retrieval.py

@@ -239,6 +239,12 @@ async def get_embedding_config(request: Request, user=Depends(get_admin_user)):
             "url": request.app.state.config.RAG_OLLAMA_BASE_URL,
             "key": request.app.state.config.RAG_OLLAMA_API_KEY,
         },
+        "azure_openai_config": {
+            "url": request.app.state.config.RAG_AZURE_OPENAI_BASE_URL,
+            "key": request.app.state.config.RAG_AZURE_OPENAI_API_KEY,
+            "deployment": request.app.state.config.RAG_AZURE_OPENAI_DEPLOYMENT,
+            "version": request.app.state.config.RAG_AZURE_OPENAI_VERSION,
+        },
     }
 
 
@@ -252,9 +258,17 @@ class OllamaConfigForm(BaseModel):
     key: str
 
 
+class AzureOpenAIConfigForm(BaseModel):
+    url: str
+    key: str
+    deployment: str
+    version: str
+
+
 class EmbeddingModelUpdateForm(BaseModel):
     openai_config: Optional[OpenAIConfigForm] = None
     ollama_config: Optional[OllamaConfigForm] = None
+    azure_openai_config: Optional[AzureOpenAIConfigForm] = None
     embedding_engine: str
     embedding_model: str
     embedding_batch_size: Optional[int] = 1
@@ -271,7 +285,7 @@ async def update_embedding_config(
         request.app.state.config.RAG_EMBEDDING_ENGINE = form_data.embedding_engine
         request.app.state.config.RAG_EMBEDDING_MODEL = form_data.embedding_model
 
-        if request.app.state.config.RAG_EMBEDDING_ENGINE in ["ollama", "openai"]:
+        if request.app.state.config.RAG_EMBEDDING_ENGINE in ["ollama", "openai", "azure_openai"]:
             if form_data.openai_config is not None:
                 request.app.state.config.RAG_OPENAI_API_BASE_URL = (
                     form_data.openai_config.url
@@ -288,6 +302,20 @@ async def update_embedding_config(
                     form_data.ollama_config.key
                 )
 
+            if form_data.azure_openai_config is not None:
+                request.app.state.config.RAG_AZURE_OPENAI_BASE_URL = (
+                    form_data.azure_openai_config.url
+                )
+                request.app.state.config.RAG_AZURE_OPENAI_API_KEY = (
+                    form_data.azure_openai_config.key
+                )
+                request.app.state.config.RAG_AZURE_OPENAI_DEPLOYMENT = (
+                    form_data.azure_openai_config.deployment
+                )
+                request.app.state.config.RAG_AZURE_OPENAI_VERSION = (
+                    form_data.azure_openai_config.version
+                )
+
             request.app.state.config.RAG_EMBEDDING_BATCH_SIZE = (
                 form_data.embedding_batch_size
             )
@@ -304,14 +332,32 @@ async def update_embedding_config(
             (
                 request.app.state.config.RAG_OPENAI_API_BASE_URL
                 if request.app.state.config.RAG_EMBEDDING_ENGINE == "openai"
-                else request.app.state.config.RAG_OLLAMA_BASE_URL
+                else (
+                    request.app.state.config.RAG_OLLAMA_BASE_URL
+                    if request.app.state.config.RAG_EMBEDDING_ENGINE == "ollama"
+                    else request.app.state.config.RAG_AZURE_OPENAI_BASE_URL
+                )
             ),
             (
                 request.app.state.config.RAG_OPENAI_API_KEY
                 if request.app.state.config.RAG_EMBEDDING_ENGINE == "openai"
-                else request.app.state.config.RAG_OLLAMA_API_KEY
+                else (
+                    request.app.state.config.RAG_OLLAMA_API_KEY
+                    if request.app.state.config.RAG_EMBEDDING_ENGINE == "ollama"
+                    else request.app.state.config.RAG_AZURE_OPENAI_API_KEY
+                )
             ),
             request.app.state.config.RAG_EMBEDDING_BATCH_SIZE,
+            (
+                request.app.state.config.RAG_AZURE_OPENAI_DEPLOYMENT
+                if request.app.state.config.RAG_EMBEDDING_ENGINE == "azure_openai"
+                else None
+            ),
+            (
+                request.app.state.config.RAG_AZURE_OPENAI_VERSION
+                if request.app.state.config.RAG_EMBEDDING_ENGINE == "azure_openai"
+                else None
+            ),
         )
 
         return {
@@ -327,6 +373,12 @@ async def update_embedding_config(
                 "url": request.app.state.config.RAG_OLLAMA_BASE_URL,
                 "key": request.app.state.config.RAG_OLLAMA_API_KEY,
             },
+            "azure_openai_config": {
+                "url": request.app.state.config.RAG_AZURE_OPENAI_BASE_URL,
+                "key": request.app.state.config.RAG_AZURE_OPENAI_API_KEY,
+                "deployment": request.app.state.config.RAG_AZURE_OPENAI_DEPLOYMENT,
+                "version": request.app.state.config.RAG_AZURE_OPENAI_VERSION,
+            },
         }
     except Exception as e:
         log.exception(f"Problem updating embedding model: {e}")
@@ -1129,14 +1181,32 @@ def save_docs_to_vector_db(
             (
                 request.app.state.config.RAG_OPENAI_API_BASE_URL
                 if request.app.state.config.RAG_EMBEDDING_ENGINE == "openai"
-                else request.app.state.config.RAG_OLLAMA_BASE_URL
+                else (
+                    request.app.state.config.RAG_OLLAMA_BASE_URL
+                    if request.app.state.config.RAG_EMBEDDING_ENGINE == "ollama"
+                    else request.app.state.config.RAG_AZURE_OPENAI_BASE_URL
+                )
             ),
             (
                 request.app.state.config.RAG_OPENAI_API_KEY
                 if request.app.state.config.RAG_EMBEDDING_ENGINE == "openai"
-                else request.app.state.config.RAG_OLLAMA_API_KEY
+                else (
+                    request.app.state.config.RAG_OLLAMA_API_KEY
+                    if request.app.state.config.RAG_EMBEDDING_ENGINE == "ollama"
+                    else request.app.state.config.RAG_AZURE_OPENAI_API_KEY
+                )
             ),
             request.app.state.config.RAG_EMBEDDING_BATCH_SIZE,
+            (
+                request.app.state.config.RAG_AZURE_OPENAI_DEPLOYMENT
+                if request.app.state.config.RAG_EMBEDDING_ENGINE == "azure_openai"
+                else None
+            ),
+            (
+                request.app.state.config.RAG_AZURE_OPENAI_VERSION
+                if request.app.state.config.RAG_EMBEDDING_ENGINE == "azure_openai"
+                else None
+            ),
         )
 
         embeddings = embedding_function(