Support separate embedding base URL

surfsense_backend/app/config/__init__.py

@@ -9,6 +9,11 @@ from chonkie import AutoEmbeddings, CodeChunker, RecursiveChunker
 from dotenv import load_dotenv
 from rerankers import Reranker
 
+from app.config.embedding_settings import (
+    build_embedding_kwargs,
+    resolve_embedding_base_url,
+)
+
 # Get the base directory of the project
 BASE_DIR = Path(__file__).resolve().parent.parent.parent
 
@@ -876,16 +881,13 @@ class Config:
 
     # Chonkie Configuration | Edit this to your needs
     EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL")
+    EMBEDDING_BASE_URL = resolve_embedding_base_url()
     # Azure OpenAI credentials from environment variables
     AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
     AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
 
-    # Pass Azure credentials to embeddings when using Azure OpenAI
-    embedding_kwargs = {}
-    if AZURE_OPENAI_ENDPOINT:
-        embedding_kwargs["azure_endpoint"] = AZURE_OPENAI_ENDPOINT
-    if AZURE_OPENAI_API_KEY:
-        embedding_kwargs["azure_api_key"] = AZURE_OPENAI_API_KEY
+    # Pass provider-specific settings to embeddings when supported.
+    embedding_kwargs = build_embedding_kwargs(embedding_model=EMBEDDING_MODEL)
 
     embedding_model_instance = AutoEmbeddings.get_embeddings(
         EMBEDDING_MODEL,

surfsense_backend/app/config/embedding_settings.py

@@ -0,0 +1,48 @@
+import os
+from collections.abc import Mapping
+
+EMBEDDING_BASE_URL_ENV = "EMBEDDING_BASE_URL"
+OLLAMA_EMBEDDING_BASE_URL_ENV = "OLLAMA_EMBEDDING_BASE_URL"
+
+
+def _clean_env_value(value: str | None) -> str | None:
+    if value is None:
+        return None
+    stripped = value.strip()
+    return stripped or None
+
+
+def resolve_embedding_base_url(environ: Mapping[str, str] | None = None) -> str | None:
+    """Return the configured embedding endpoint, if any."""
+    environ = os.environ if environ is None else environ
+    return _clean_env_value(environ.get(EMBEDDING_BASE_URL_ENV)) or _clean_env_value(
+        environ.get(OLLAMA_EMBEDDING_BASE_URL_ENV)
+    )
+
+
+def _supports_embedding_api_base(embedding_model: str | None) -> bool:
+    return (embedding_model or "").startswith("litellm://")
+
+
+def build_embedding_kwargs(
+    environ: Mapping[str, str] | None = None,
+    *,
+    embedding_model: str | None = None,
+) -> dict[str, str]:
+    """Build keyword arguments for Chonkie's embedding provider."""
+    environ = os.environ if environ is None else environ
+
+    embedding_kwargs: dict[str, str] = {}
+    embedding_base_url = resolve_embedding_base_url(environ)
+    if embedding_base_url and _supports_embedding_api_base(embedding_model):
+        embedding_kwargs["api_base"] = embedding_base_url
+
+    azure_openai_endpoint = _clean_env_value(environ.get("AZURE_OPENAI_ENDPOINT"))
+    azure_openai_api_key = _clean_env_value(environ.get("AZURE_OPENAI_API_KEY"))
+
+    if azure_openai_endpoint:
+        embedding_kwargs["azure_endpoint"] = azure_openai_endpoint
+    if azure_openai_api_key:
+        embedding_kwargs["azure_api_key"] = azure_openai_api_key
+
+    return embedding_kwargs