Merge branch 'dev' into vector-search-branch

This commit is contained in:
hiwylee
2025-08-01 04:23:38 +09:00
committed by GitHub
367 changed files with 23637 additions and 9204 deletions
+133 -17
View File
@@ -13,12 +13,15 @@ from urllib.parse import urlparse
import requests
from pydantic import BaseModel
from sqlalchemy import JSON, Column, DateTime, Integer, func
from authlib.integrations.starlette_client import OAuth
from open_webui.env import (
DATA_DIR,
DATABASE_URL,
ENV,
REDIS_URL,
REDIS_KEY_PREFIX,
REDIS_SENTINEL_HOSTS,
REDIS_SENTINEL_PORT,
FRONTEND_BUILD_DIR,
@@ -211,11 +214,16 @@ class PersistentConfig(Generic[T]):
class AppConfig:
_state: dict[str, PersistentConfig]
_redis: Optional[redis.Redis] = None
_redis_key_prefix: str
def __init__(
self, redis_url: Optional[str] = None, redis_sentinels: Optional[list] = []
self,
redis_url: Optional[str] = None,
redis_sentinels: Optional[list] = [],
redis_key_prefix: str = "open-webui",
):
super().__setattr__("_state", {})
super().__setattr__("_redis_key_prefix", redis_key_prefix)
if redis_url:
super().__setattr__(
"_redis",
@@ -230,7 +238,7 @@ class AppConfig:
self._state[key].save()
if self._redis:
redis_key = f"open-webui:config:{key}"
redis_key = f"{self._redis_key_prefix}:config:{key}"
self._redis.set(redis_key, json.dumps(self._state[key].value))
def __getattr__(self, key):
@@ -239,7 +247,7 @@ class AppConfig:
# If Redis is available, check for an updated value
if self._redis:
redis_key = f"open-webui:config:{key}"
redis_key = f"{self._redis_key_prefix}:config:{key}"
redis_value = self._redis.get(redis_key)
if redis_value is not None:
@@ -431,6 +439,18 @@ OAUTH_SCOPES = PersistentConfig(
os.environ.get("OAUTH_SCOPES", "openid email profile"),
)
OAUTH_TIMEOUT = PersistentConfig(
"OAUTH_TIMEOUT",
"oauth.oidc.oauth_timeout",
os.environ.get("OAUTH_TIMEOUT", ""),
)
OAUTH_TOKEN_ENDPOINT_AUTH_METHOD = PersistentConfig(
"OAUTH_TOKEN_ENDPOINT_AUTH_METHOD",
"oauth.oidc.token_endpoint_auth_method",
os.environ.get("OAUTH_TOKEN_ENDPOINT_AUTH_METHOD", None),
)
OAUTH_CODE_CHALLENGE_METHOD = PersistentConfig(
"OAUTH_CODE_CHALLENGE_METHOD",
"oauth.oidc.code_challenge_method",
@@ -534,13 +554,20 @@ def load_oauth_providers():
OAUTH_PROVIDERS.clear()
if GOOGLE_CLIENT_ID.value and GOOGLE_CLIENT_SECRET.value:
def google_oauth_register(client):
def google_oauth_register(client: OAuth):
client.register(
name="google",
client_id=GOOGLE_CLIENT_ID.value,
client_secret=GOOGLE_CLIENT_SECRET.value,
server_metadata_url="https://accounts.google.com/.well-known/openid-configuration",
client_kwargs={"scope": GOOGLE_OAUTH_SCOPE.value},
client_kwargs={
"scope": GOOGLE_OAUTH_SCOPE.value,
**(
{"timeout": int(OAUTH_TIMEOUT.value)}
if OAUTH_TIMEOUT.value
else {}
),
},
redirect_uri=GOOGLE_REDIRECT_URI.value,
)
@@ -555,7 +582,7 @@ def load_oauth_providers():
and MICROSOFT_CLIENT_TENANT_ID.value
):
def microsoft_oauth_register(client):
def microsoft_oauth_register(client: OAuth):
client.register(
name="microsoft",
client_id=MICROSOFT_CLIENT_ID.value,
@@ -563,6 +590,11 @@ def load_oauth_providers():
server_metadata_url=f"{MICROSOFT_CLIENT_LOGIN_BASE_URL.value}/{MICROSOFT_CLIENT_TENANT_ID.value}/v2.0/.well-known/openid-configuration?appid={MICROSOFT_CLIENT_ID.value}",
client_kwargs={
"scope": MICROSOFT_OAUTH_SCOPE.value,
**(
{"timeout": int(OAUTH_TIMEOUT.value)}
if OAUTH_TIMEOUT.value
else {}
),
},
redirect_uri=MICROSOFT_REDIRECT_URI.value,
)
@@ -575,7 +607,7 @@ def load_oauth_providers():
if GITHUB_CLIENT_ID.value and GITHUB_CLIENT_SECRET.value:
def github_oauth_register(client):
def github_oauth_register(client: OAuth):
client.register(
name="github",
client_id=GITHUB_CLIENT_ID.value,
@@ -584,7 +616,14 @@ def load_oauth_providers():
authorize_url="https://github.com/login/oauth/authorize",
api_base_url="https://api.github.com",
userinfo_endpoint="https://api.github.com/user",
client_kwargs={"scope": GITHUB_CLIENT_SCOPE.value},
client_kwargs={
"scope": GITHUB_CLIENT_SCOPE.value,
**(
{"timeout": int(OAUTH_TIMEOUT.value)}
if OAUTH_TIMEOUT.value
else {}
),
},
redirect_uri=GITHUB_CLIENT_REDIRECT_URI.value,
)
@@ -600,9 +639,19 @@ def load_oauth_providers():
and OPENID_PROVIDER_URL.value
):
def oidc_oauth_register(client):
def oidc_oauth_register(client: OAuth):
client_kwargs = {
"scope": OAUTH_SCOPES.value,
**(
{
"token_endpoint_auth_method": OAUTH_TOKEN_ENDPOINT_AUTH_METHOD.value
}
if OAUTH_TOKEN_ENDPOINT_AUTH_METHOD.value
else {}
),
**(
{"timeout": int(OAUTH_TIMEOUT.value)} if OAUTH_TIMEOUT.value else {}
),
}
if (
@@ -640,6 +689,17 @@ load_oauth_providers()
STATIC_DIR = Path(os.getenv("STATIC_DIR", OPEN_WEBUI_DIR / "static")).resolve()
try:
if STATIC_DIR.exists():
for item in STATIC_DIR.iterdir():
if item.is_file() or item.is_symlink():
try:
item.unlink()
except Exception as e:
pass
except Exception as e:
pass
for file_path in (FRONTEND_BUILD_DIR / "static").glob("**/*"):
if file_path.is_file():
target_path = STATIC_DIR / file_path.relative_to(
@@ -719,12 +779,6 @@ if CUSTOM_NAME:
pass
####################################
# LICENSE_KEY
####################################
LICENSE_KEY = os.environ.get("LICENSE_KEY", "")
####################################
# STORAGE PROVIDER
####################################
@@ -895,6 +949,18 @@ except Exception:
pass
OPENAI_API_BASE_URL = "https://api.openai.com/v1"
####################################
# MODELS
####################################
ENABLE_BASE_MODELS_CACHE = PersistentConfig(
"ENABLE_BASE_MODELS_CACHE",
"models.base_models_cache",
os.environ.get("ENABLE_BASE_MODELS_CACHE", "False").lower() == "true",
)
####################################
# TOOL_SERVERS
####################################
@@ -1077,10 +1143,18 @@ USER_PERMISSIONS_CHAT_CONTROLS = (
os.environ.get("USER_PERMISSIONS_CHAT_CONTROLS", "True").lower() == "true"
)
USER_PERMISSIONS_CHAT_VALVES = (
os.environ.get("USER_PERMISSIONS_CHAT_VALVES", "True").lower() == "true"
)
USER_PERMISSIONS_CHAT_SYSTEM_PROMPT = (
os.environ.get("USER_PERMISSIONS_CHAT_SYSTEM_PROMPT", "True").lower() == "true"
)
USER_PERMISSIONS_CHAT_PARAMS = (
os.environ.get("USER_PERMISSIONS_CHAT_PARAMS", "True").lower() == "true"
)
USER_PERMISSIONS_CHAT_FILE_UPLOAD = (
os.environ.get("USER_PERMISSIONS_CHAT_FILE_UPLOAD", "True").lower() == "true"
)
@@ -1166,7 +1240,9 @@ DEFAULT_USER_PERMISSIONS = {
},
"chat": {
"controls": USER_PERMISSIONS_CHAT_CONTROLS,
"valves": USER_PERMISSIONS_CHAT_VALVES,
"system_prompt": USER_PERMISSIONS_CHAT_SYSTEM_PROMPT,
"params": USER_PERMISSIONS_CHAT_PARAMS,
"file_upload": USER_PERMISSIONS_CHAT_FILE_UPLOAD,
"delete": USER_PERMISSIONS_CHAT_DELETE,
"edit": USER_PERMISSIONS_CHAT_EDIT,
@@ -1794,11 +1870,12 @@ MILVUS_IVF_FLAT_NLIST = int(os.environ.get("MILVUS_IVF_FLAT_NLIST", "128"))
QDRANT_URI = os.environ.get("QDRANT_URI", None)
QDRANT_API_KEY = os.environ.get("QDRANT_API_KEY", None)
QDRANT_ON_DISK = os.environ.get("QDRANT_ON_DISK", "false").lower() == "true"
QDRANT_PREFER_GRPC = os.environ.get("QDRANT_PREFER_GRPC", "False").lower() == "true"
QDRANT_PREFER_GRPC = os.environ.get("QDRANT_PREFER_GRPC", "false").lower() == "true"
QDRANT_GRPC_PORT = int(os.environ.get("QDRANT_GRPC_PORT", "6334"))
ENABLE_QDRANT_MULTITENANCY_MODE = (
os.environ.get("ENABLE_QDRANT_MULTITENANCY_MODE", "false").lower() == "true"
os.environ.get("ENABLE_QDRANT_MULTITENANCY_MODE", "true").lower() == "true"
)
QDRANT_COLLECTION_PREFIX = os.environ.get("QDRANT_COLLECTION_PREFIX", "open-webui")
# OpenSearch
OPENSEARCH_URI = os.environ.get("OPENSEARCH_URI", "https://localhost:9200")
@@ -1837,6 +1914,45 @@ if PGVECTOR_PGCRYPTO and not PGVECTOR_PGCRYPTO_KEY:
"PGVECTOR_PGCRYPTO is enabled but PGVECTOR_PGCRYPTO_KEY is not set. Please provide a valid key."
)
PGVECTOR_POOL_SIZE = os.environ.get("PGVECTOR_POOL_SIZE", None)
if PGVECTOR_POOL_SIZE != None:
try:
PGVECTOR_POOL_SIZE = int(PGVECTOR_POOL_SIZE)
except Exception:
PGVECTOR_POOL_SIZE = None
PGVECTOR_POOL_MAX_OVERFLOW = os.environ.get("PGVECTOR_POOL_MAX_OVERFLOW", 0)
if PGVECTOR_POOL_MAX_OVERFLOW == "":
PGVECTOR_POOL_MAX_OVERFLOW = 0
else:
try:
PGVECTOR_POOL_MAX_OVERFLOW = int(PGVECTOR_POOL_MAX_OVERFLOW)
except Exception:
PGVECTOR_POOL_MAX_OVERFLOW = 0
PGVECTOR_POOL_TIMEOUT = os.environ.get("PGVECTOR_POOL_TIMEOUT", 30)
if PGVECTOR_POOL_TIMEOUT == "":
PGVECTOR_POOL_TIMEOUT = 30
else:
try:
PGVECTOR_POOL_TIMEOUT = int(PGVECTOR_POOL_TIMEOUT)
except Exception:
PGVECTOR_POOL_TIMEOUT = 30
PGVECTOR_POOL_RECYCLE = os.environ.get("PGVECTOR_POOL_RECYCLE", 3600)
if PGVECTOR_POOL_RECYCLE == "":
PGVECTOR_POOL_RECYCLE = 3600
else:
try:
PGVECTOR_POOL_RECYCLE = int(PGVECTOR_POOL_RECYCLE)
except Exception:
PGVECTOR_POOL_RECYCLE = 3600
# Pinecone
PINECONE_API_KEY = os.environ.get("PINECONE_API_KEY", None)
PINECONE_ENVIRONMENT = os.environ.get("PINECONE_ENVIRONMENT", None)
+121 -7
View File
@@ -7,6 +7,7 @@ import sys
import shutil
from uuid import uuid4
from pathlib import Path
from cryptography.hazmat.primitives import serialization
import markdown
from bs4 import BeautifulSoup
@@ -199,6 +200,7 @@ CHANGELOG = changelog_json
SAFE_MODE = os.environ.get("SAFE_MODE", "false").lower() == "true"
####################################
# ENABLE_FORWARD_USER_INFO_HEADERS
####################################
@@ -266,21 +268,40 @@ else:
DATABASE_URL = os.environ.get("DATABASE_URL", f"sqlite:///{DATA_DIR}/webui.db")
DATABASE_TYPE = os.environ.get("DATABASE_TYPE")
DATABASE_USER = os.environ.get("DATABASE_USER")
DATABASE_PASSWORD = os.environ.get("DATABASE_PASSWORD")
DATABASE_CRED = ""
if DATABASE_USER:
DATABASE_CRED += f"{DATABASE_USER}"
if DATABASE_PASSWORD:
DATABASE_CRED += f":{DATABASE_PASSWORD}"
DB_VARS = {
"db_type": DATABASE_TYPE,
"db_cred": DATABASE_CRED,
"db_host": os.environ.get("DATABASE_HOST"),
"db_port": os.environ.get("DATABASE_PORT"),
"db_name": os.environ.get("DATABASE_NAME"),
}
if all(DB_VARS.values()):
DATABASE_URL = f"{DB_VARS['db_type']}://{DB_VARS['db_cred']}@{DB_VARS['db_host']}:{DB_VARS['db_port']}/{DB_VARS['db_name']}"
# Replace the postgres:// with postgresql://
if "postgres://" in DATABASE_URL:
DATABASE_URL = DATABASE_URL.replace("postgres://", "postgresql://")
DATABASE_SCHEMA = os.environ.get("DATABASE_SCHEMA", None)
DATABASE_POOL_SIZE = os.environ.get("DATABASE_POOL_SIZE", 0)
DATABASE_POOL_SIZE = os.environ.get("DATABASE_POOL_SIZE", None)
if DATABASE_POOL_SIZE == "":
DATABASE_POOL_SIZE = 0
else:
if DATABASE_POOL_SIZE != None:
try:
DATABASE_POOL_SIZE = int(DATABASE_POOL_SIZE)
except Exception:
DATABASE_POOL_SIZE = 0
DATABASE_POOL_SIZE = None
DATABASE_POOL_MAX_OVERFLOW = os.environ.get("DATABASE_POOL_MAX_OVERFLOW", 0)
@@ -325,9 +346,19 @@ ENABLE_REALTIME_CHAT_SAVE = (
####################################
REDIS_URL = os.environ.get("REDIS_URL", "")
REDIS_KEY_PREFIX = os.environ.get("REDIS_KEY_PREFIX", "open-webui")
REDIS_SENTINEL_HOSTS = os.environ.get("REDIS_SENTINEL_HOSTS", "")
REDIS_SENTINEL_PORT = os.environ.get("REDIS_SENTINEL_PORT", "26379")
# Maximum number of retries for Redis operations when using Sentinel fail-over
REDIS_SENTINEL_MAX_RETRY_COUNT = os.environ.get("REDIS_SENTINEL_MAX_RETRY_COUNT", "2")
try:
REDIS_SENTINEL_MAX_RETRY_COUNT = int(REDIS_SENTINEL_MAX_RETRY_COUNT)
if REDIS_SENTINEL_MAX_RETRY_COUNT < 1:
REDIS_SENTINEL_MAX_RETRY_COUNT = 2
except ValueError:
REDIS_SENTINEL_MAX_RETRY_COUNT = 2
####################################
# UVICORN WORKERS
####################################
@@ -347,6 +378,10 @@ except ValueError:
####################################
WEBUI_AUTH = os.environ.get("WEBUI_AUTH", "True").lower() == "true"
ENABLE_SIGNUP_PASSWORD_CONFIRMATION = (
os.environ.get("ENABLE_SIGNUP_PASSWORD_CONFIRMATION", "False").lower() == "true"
)
WEBUI_AUTH_TRUSTED_EMAIL_HEADER = os.environ.get(
"WEBUI_AUTH_TRUSTED_EMAIL_HEADER", None
)
@@ -396,14 +431,71 @@ WEBUI_AUTH_COOKIE_SECURE = (
if WEBUI_AUTH and WEBUI_SECRET_KEY == "":
raise ValueError(ERROR_MESSAGES.ENV_VAR_NOT_FOUND)
ENABLE_COMPRESSION_MIDDLEWARE = (
os.environ.get("ENABLE_COMPRESSION_MIDDLEWARE", "True").lower() == "true"
)
####################################
# LICENSE_KEY
####################################
LICENSE_KEY = os.environ.get("LICENSE_KEY", "")
LICENSE_BLOB = None
LICENSE_BLOB_PATH = os.environ.get("LICENSE_BLOB_PATH", DATA_DIR / "l.data")
if LICENSE_BLOB_PATH and os.path.exists(LICENSE_BLOB_PATH):
with open(LICENSE_BLOB_PATH, "rb") as f:
LICENSE_BLOB = f.read()
LICENSE_PUBLIC_KEY = os.environ.get("LICENSE_PUBLIC_KEY", "")
pk = None
if LICENSE_PUBLIC_KEY:
pk = serialization.load_pem_public_key(
f"""
-----BEGIN PUBLIC KEY-----
{LICENSE_PUBLIC_KEY}
-----END PUBLIC KEY-----
""".encode(
"utf-8"
)
)
####################################
# MODELS
####################################
MODELS_CACHE_TTL = os.environ.get("MODELS_CACHE_TTL", "1")
if MODELS_CACHE_TTL == "":
MODELS_CACHE_TTL = None
else:
try:
MODELS_CACHE_TTL = int(MODELS_CACHE_TTL)
except Exception:
MODELS_CACHE_TTL = 1
####################################
# WEBSOCKET SUPPORT
####################################
ENABLE_WEBSOCKET_SUPPORT = (
os.environ.get("ENABLE_WEBSOCKET_SUPPORT", "True").lower() == "true"
)
WEBSOCKET_MANAGER = os.environ.get("WEBSOCKET_MANAGER", "")
WEBSOCKET_REDIS_URL = os.environ.get("WEBSOCKET_REDIS_URL", REDIS_URL)
WEBSOCKET_REDIS_LOCK_TIMEOUT = os.environ.get("WEBSOCKET_REDIS_LOCK_TIMEOUT", 60)
websocket_redis_lock_timeout = os.environ.get("WEBSOCKET_REDIS_LOCK_TIMEOUT", "60")
try:
WEBSOCKET_REDIS_LOCK_TIMEOUT = int(websocket_redis_lock_timeout)
except ValueError:
WEBSOCKET_REDIS_LOCK_TIMEOUT = 60
WEBSOCKET_SENTINEL_HOSTS = os.environ.get("WEBSOCKET_SENTINEL_HOSTS", "")
@@ -506,11 +598,14 @@ else:
# OFFLINE_MODE
####################################
ENABLE_VERSION_UPDATE_CHECK = (
os.environ.get("ENABLE_VERSION_UPDATE_CHECK", "true").lower() == "true"
)
OFFLINE_MODE = os.environ.get("OFFLINE_MODE", "false").lower() == "true"
if OFFLINE_MODE:
os.environ["HF_HUB_OFFLINE"] = "1"
ENABLE_VERSION_UPDATE_CHECK = False
####################################
# AUDIT LOGGING
@@ -519,6 +614,14 @@ if OFFLINE_MODE:
AUDIT_LOGS_FILE_PATH = f"{DATA_DIR}/audit.log"
# Maximum size of a file before rotating into a new log file
AUDIT_LOG_FILE_ROTATION_SIZE = os.getenv("AUDIT_LOG_FILE_ROTATION_SIZE", "10MB")
# Comma separated list of logger names to use for audit logging
# Default is "uvicorn.access" which is the access log for Uvicorn
# You can add more logger names to this list if you want to capture more logs
AUDIT_UVICORN_LOGGER_NAMES = os.getenv(
"AUDIT_UVICORN_LOGGER_NAMES", "uvicorn.access"
).split(",")
# METADATA | REQUEST | REQUEST_RESPONSE
AUDIT_LOG_LEVEL = os.getenv("AUDIT_LOG_LEVEL", "NONE").upper()
try:
@@ -543,6 +646,9 @@ ENABLE_OTEL_METRICS = os.environ.get("ENABLE_OTEL_METRICS", "False").lower() ==
OTEL_EXPORTER_OTLP_ENDPOINT = os.environ.get(
"OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317"
)
OTEL_EXPORTER_OTLP_INSECURE = (
os.environ.get("OTEL_EXPORTER_OTLP_INSECURE", "False").lower() == "true"
)
OTEL_SERVICE_NAME = os.environ.get("OTEL_SERVICE_NAME", "open-webui")
OTEL_RESOURCE_ATTRIBUTES = os.environ.get(
"OTEL_RESOURCE_ATTRIBUTES", ""
@@ -550,6 +656,14 @@ OTEL_RESOURCE_ATTRIBUTES = os.environ.get(
OTEL_TRACES_SAMPLER = os.environ.get(
"OTEL_TRACES_SAMPLER", "parentbased_always_on"
).lower()
OTEL_BASIC_AUTH_USERNAME = os.environ.get("OTEL_BASIC_AUTH_USERNAME", "")
OTEL_BASIC_AUTH_PASSWORD = os.environ.get("OTEL_BASIC_AUTH_PASSWORD", "")
OTEL_OTLP_SPAN_EXPORTER = os.environ.get(
"OTEL_OTLP_SPAN_EXPORTER", "grpc"
).lower() # grpc or http
####################################
# TOOLS/FUNCTIONS PIP OPTIONS
+19 -13
View File
@@ -62,6 +62,9 @@ def handle_peewee_migration(DATABASE_URL):
except Exception as e:
log.error(f"Failed to initialize the database connection: {e}")
log.warning(
"Hint: If your database password contains special characters, you may need to URL-encode it."
)
raise
finally:
# Properly closing the database connection
@@ -81,20 +84,23 @@ if "sqlite" in SQLALCHEMY_DATABASE_URL:
SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
)
else:
if DATABASE_POOL_SIZE > 0:
engine = create_engine(
SQLALCHEMY_DATABASE_URL,
pool_size=DATABASE_POOL_SIZE,
max_overflow=DATABASE_POOL_MAX_OVERFLOW,
pool_timeout=DATABASE_POOL_TIMEOUT,
pool_recycle=DATABASE_POOL_RECYCLE,
pool_pre_ping=True,
poolclass=QueuePool,
)
if isinstance(DATABASE_POOL_SIZE, int):
if DATABASE_POOL_SIZE > 0:
engine = create_engine(
SQLALCHEMY_DATABASE_URL,
pool_size=DATABASE_POOL_SIZE,
max_overflow=DATABASE_POOL_MAX_OVERFLOW,
pool_timeout=DATABASE_POOL_TIMEOUT,
pool_recycle=DATABASE_POOL_RECYCLE,
pool_pre_ping=True,
poolclass=QueuePool,
)
else:
engine = create_engine(
SQLALCHEMY_DATABASE_URL, pool_pre_ping=True, poolclass=NullPool
)
else:
engine = create_engine(
SQLALCHEMY_DATABASE_URL, pool_pre_ping=True, poolclass=NullPool
)
engine = create_engine(SQLALCHEMY_DATABASE_URL, pool_pre_ping=True)
SessionLocal = sessionmaker(
+93 -22
View File
@@ -36,7 +36,6 @@ from fastapi import (
applications,
BackgroundTasks,
)
from fastapi.openapi.docs import get_swagger_ui_html
from fastapi.middleware.cors import CORSMiddleware
@@ -49,6 +48,7 @@ from starlette.exceptions import HTTPException as StarletteHTTPException
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.middleware.sessions import SessionMiddleware
from starlette.responses import Response, StreamingResponse
from starlette.datastructures import Headers
from open_webui.utils import logger
@@ -89,6 +89,7 @@ from open_webui.routers import (
from open_webui.routers.retrieval import (
get_embedding_function,
get_reranking_function,
get_ef,
get_rf,
)
@@ -101,7 +102,6 @@ from open_webui.models.users import UserModel, Users
from open_webui.models.chats import Chats
from open_webui.config import (
LICENSE_KEY,
# Ollama
ENABLE_OLLAMA_API,
OLLAMA_BASE_URLS,
@@ -116,6 +116,8 @@ from open_webui.config import (
OPENAI_API_CONFIGS,
# Direct Connections
ENABLE_DIRECT_CONNECTIONS,
# Model list
ENABLE_BASE_MODELS_CACHE,
# Thread pool size for FastAPI/AnyIO
THREAD_POOL_SIZE,
# Tool Server Configs
@@ -392,10 +394,12 @@ from open_webui.config import (
reset_config,
)
from open_webui.env import (
LICENSE_KEY,
AUDIT_EXCLUDED_PATHS,
AUDIT_LOG_LEVEL,
CHANGELOG,
REDIS_URL,
REDIS_KEY_PREFIX,
REDIS_SENTINEL_HOSTS,
REDIS_SENTINEL_PORT,
GLOBAL_LOG_LEVEL,
@@ -408,13 +412,15 @@ from open_webui.env import (
WEBUI_SECRET_KEY,
WEBUI_SESSION_COOKIE_SAME_SITE,
WEBUI_SESSION_COOKIE_SECURE,
ENABLE_SIGNUP_PASSWORD_CONFIRMATION,
WEBUI_AUTH_TRUSTED_EMAIL_HEADER,
WEBUI_AUTH_TRUSTED_NAME_HEADER,
WEBUI_AUTH_SIGNOUT_REDIRECT_URL,
ENABLE_COMPRESSION_MIDDLEWARE,
ENABLE_WEBSOCKET_SUPPORT,
BYPASS_MODEL_ACCESS_CONTROL,
RESET_CONFIG_ON_START,
OFFLINE_MODE,
ENABLE_VERSION_UPDATE_CHECK,
ENABLE_OTEL,
EXTERNAL_PWA_MANIFEST_URL,
AIOHTTP_CLIENT_SESSION_SSL,
@@ -449,7 +455,7 @@ from open_webui.utils.redis import get_redis_connection
from open_webui.tasks import (
redis_task_command_listener,
list_task_ids_by_chat_id,
list_task_ids_by_item_id,
stop_task,
list_tasks,
) # Import from tasks.py
@@ -533,6 +539,27 @@ async def lifespan(app: FastAPI):
asyncio.create_task(periodic_usage_pool_cleanup())
if app.state.config.ENABLE_BASE_MODELS_CACHE:
await get_all_models(
Request(
# Creating a mock request object to pass to get_all_models
{
"type": "http",
"asgi.version": "3.0",
"asgi.spec_version": "2.0",
"method": "GET",
"path": "/internal",
"query_string": b"",
"headers": Headers({}).raw,
"client": ("127.0.0.1", 12345),
"server": ("127.0.0.1", 80),
"scheme": "http",
"app": app,
}
),
None,
)
yield
if hasattr(app.state, "redis_task_command_listener"):
@@ -553,6 +580,7 @@ app.state.instance_id = None
app.state.config = AppConfig(
redis_url=REDIS_URL,
redis_sentinels=get_sentinels_from_env(REDIS_SENTINEL_HOSTS, REDIS_SENTINEL_PORT),
redis_key_prefix=REDIS_KEY_PREFIX,
)
app.state.redis = None
@@ -615,6 +643,15 @@ app.state.TOOL_SERVERS = []
app.state.config.ENABLE_DIRECT_CONNECTIONS = ENABLE_DIRECT_CONNECTIONS
########################################
#
# MODELS
#
########################################
app.state.config.ENABLE_BASE_MODELS_CACHE = ENABLE_BASE_MODELS_CACHE
app.state.BASE_MODELS = []
########################################
#
# WEBUI
@@ -843,6 +880,7 @@ app.state.config.FIRECRAWL_API_KEY = FIRECRAWL_API_KEY
app.state.config.TAVILY_EXTRACT_DEPTH = TAVILY_EXTRACT_DEPTH
app.state.EMBEDDING_FUNCTION = None
app.state.RERANKING_FUNCTION = None
app.state.ef = None
app.state.rf = None
@@ -871,8 +909,8 @@ except Exception as e:
app.state.EMBEDDING_FUNCTION = get_embedding_function(
app.state.config.RAG_EMBEDDING_ENGINE,
app.state.config.RAG_EMBEDDING_MODEL,
app.state.ef,
(
embedding_function=app.state.ef,
url=(
app.state.config.RAG_OPENAI_API_BASE_URL
if app.state.config.RAG_EMBEDDING_ENGINE == "openai"
else (
@@ -881,7 +919,7 @@ app.state.EMBEDDING_FUNCTION = get_embedding_function(
else app.state.config.RAG_AZURE_OPENAI_BASE_URL
)
),
(
key=(
app.state.config.RAG_OPENAI_API_KEY
if app.state.config.RAG_EMBEDDING_ENGINE == "openai"
else (
@@ -890,7 +928,7 @@ app.state.EMBEDDING_FUNCTION = get_embedding_function(
else app.state.config.RAG_AZURE_OPENAI_API_KEY
)
),
app.state.config.RAG_EMBEDDING_BATCH_SIZE,
embedding_batch_size=app.state.config.RAG_EMBEDDING_BATCH_SIZE,
azure_api_version=(
app.state.config.RAG_AZURE_OPENAI_API_VERSION
if app.state.config.RAG_EMBEDDING_ENGINE == "azure_openai"
@@ -898,6 +936,12 @@ app.state.EMBEDDING_FUNCTION = get_embedding_function(
),
)
app.state.RERANKING_FUNCTION = get_reranking_function(
app.state.config.RAG_RERANKING_ENGINE,
app.state.config.RAG_RERANKING_MODEL,
reranking_function=app.state.rf,
)
########################################
#
# CODE EXECUTION
@@ -1072,7 +1116,9 @@ class RedirectMiddleware(BaseHTTPMiddleware):
# Add the middleware to the app
app.add_middleware(CompressMiddleware)
if ENABLE_COMPRESSION_MIDDLEWARE:
app.add_middleware(CompressMiddleware)
app.add_middleware(RedirectMiddleware)
app.add_middleware(SecurityHeadersMiddleware)
@@ -1188,7 +1234,9 @@ if audit_level != AuditLevel.NONE:
@app.get("/api/models")
async def get_models(request: Request, user=Depends(get_verified_user)):
async def get_models(
request: Request, refresh: bool = False, user=Depends(get_verified_user)
):
def get_filtered_models(models, user):
filtered_models = []
for model in models:
@@ -1212,7 +1260,7 @@ async def get_models(request: Request, user=Depends(get_verified_user)):
return filtered_models
all_models = await get_all_models(request, user=user)
all_models = await get_all_models(request, refresh=refresh, user=user)
models = []
for model in all_models:
@@ -1249,7 +1297,7 @@ async def get_models(request: Request, user=Depends(get_verified_user)):
models = get_filtered_models(models, user)
log.debug(
f"/api/models returned filtered models accessible to the user: {json.dumps([model['id'] for model in models])}"
f"/api/models returned filtered models accessible to the user: {json.dumps([model.get('id') for model in models])}"
)
return {"data": models}
@@ -1357,7 +1405,6 @@ async def chat_completion(
form_data, metadata, events = await process_chat_payload(
request, form_data, user, metadata, model
)
except Exception as e:
log.debug(f"Error processing chat payload: {e}")
if metadata.get("chat_id") and metadata.get("message_id"):
@@ -1377,6 +1424,14 @@ async def chat_completion(
try:
response = await chat_completion_handler(request, form_data, user)
if metadata.get("chat_id") and metadata.get("message_id"):
Chats.upsert_message_to_chat_by_id_and_message_id(
metadata["chat_id"],
metadata["message_id"],
{
"model": model_id,
},
)
return await process_chat_response(
request, response, form_data, user, metadata, model, events, tasks
@@ -1447,7 +1502,7 @@ async def stop_task_endpoint(
request: Request, task_id: str, user=Depends(get_verified_user)
):
try:
result = await stop_task(request, task_id)
result = await stop_task(request.app.state.redis, task_id)
return result
except ValueError as e:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e))
@@ -1455,7 +1510,7 @@ async def stop_task_endpoint(
@app.get("/api/tasks")
async def list_tasks_endpoint(request: Request, user=Depends(get_verified_user)):
return {"tasks": await list_tasks(request)}
return {"tasks": await list_tasks(request.app.state.redis)}
@app.get("/api/tasks/chat/{chat_id}")
@@ -1466,9 +1521,9 @@ async def list_tasks_by_chat_id_endpoint(
if chat is None or chat.user_id != user.id:
return {"task_ids": []}
task_ids = await list_task_ids_by_chat_id(request, chat_id)
task_ids = await list_task_ids_by_item_id(request.app.state.redis, chat_id)
print(f"Task IDs for chat {chat_id}: {task_ids}")
log.debug(f"Task IDs for chat {chat_id}: {task_ids}")
return {"task_ids": task_ids}
@@ -1516,11 +1571,13 @@ async def get_app_config(request: Request):
"features": {
"auth": WEBUI_AUTH,
"auth_trusted_header": bool(app.state.AUTH_TRUSTED_EMAIL_HEADER),
"enable_signup_password_confirmation": ENABLE_SIGNUP_PASSWORD_CONFIRMATION,
"enable_ldap": app.state.config.ENABLE_LDAP,
"enable_api_key": app.state.config.ENABLE_API_KEY,
"enable_signup": app.state.config.ENABLE_SIGNUP,
"enable_login_form": app.state.config.ENABLE_LOGIN_FORM,
"enable_websocket": ENABLE_WEBSOCKET_SUPPORT,
"enable_version_update_check": ENABLE_VERSION_UPDATE_CHECK,
**(
{
"enable_direct_connections": app.state.config.ENABLE_DIRECT_CONNECTIONS,
@@ -1593,8 +1650,23 @@ async def get_app_config(request: Request):
else {}
),
}
if user is not None
else {}
if user is not None and (user.role in ["admin", "user"])
else {
**(
{
"metadata": {
"login_footer": app.state.LICENSE_METADATA.get(
"login_footer", ""
),
"auth_logo_position": app.state.LICENSE_METADATA.get(
"auth_logo_position", ""
),
}
}
if app.state.LICENSE_METADATA
else {}
)
}
),
}
@@ -1626,9 +1698,9 @@ async def get_app_version():
@app.get("/api/version/updates")
async def get_app_latest_release_version(user=Depends(get_verified_user)):
if OFFLINE_MODE:
if not ENABLE_VERSION_UPDATE_CHECK:
log.debug(
f"Offline mode is enabled, returning current version as latest version"
f"Version update check is disabled, returning current version as latest version"
)
return {"current": VERSION, "latest": VERSION}
try:
@@ -1709,7 +1781,6 @@ async def get_manifest_json():
"start_url": "/",
"display": "standalone",
"background_color": "#343541",
"orientation": "any",
"icons": [
{
"src": "/static/logo.png",
@@ -0,0 +1,23 @@
"""Update folder table data
Revision ID: d31026856c01
Revises: 9f0c9cd09105
Create Date: 2025-07-13 03:00:00.000000
"""
from alembic import op
import sqlalchemy as sa
revision = "d31026856c01"
down_revision = "9f0c9cd09105"
branch_labels = None
depends_on = None
def upgrade():
op.add_column("folder", sa.Column("data", sa.JSON(), nullable=True))
def downgrade():
op.drop_column("folder", "data")
+43 -32
View File
@@ -12,6 +12,7 @@ from pydantic import BaseModel, ConfigDict
from sqlalchemy import BigInteger, Boolean, Column, String, Text, JSON
from sqlalchemy import or_, func, select, and_, text
from sqlalchemy.sql import exists
from sqlalchemy.sql.expression import bindparam
####################
# Chat DB Schema
@@ -66,12 +67,14 @@ class ChatModel(BaseModel):
class ChatForm(BaseModel):
chat: dict
folder_id: Optional[str] = None
class ChatImportForm(ChatForm):
meta: Optional[dict] = {}
pinned: Optional[bool] = False
folder_id: Optional[str] = None
created_at: Optional[int] = None
updated_at: Optional[int] = None
class ChatTitleMessagesForm(BaseModel):
@@ -118,6 +121,7 @@ class ChatTable:
else "New Chat"
),
"chat": form_data.chat,
"folder_id": form_data.folder_id,
"created_at": int(time.time()),
"updated_at": int(time.time()),
}
@@ -147,8 +151,16 @@ class ChatTable:
"meta": form_data.meta,
"pinned": form_data.pinned,
"folder_id": form_data.folder_id,
"created_at": int(time.time()),
"updated_at": int(time.time()),
"created_at": (
form_data.created_at
if form_data.created_at
else int(time.time())
),
"updated_at": (
form_data.updated_at
if form_data.updated_at
else int(time.time())
),
}
)
@@ -232,6 +244,10 @@ class ChatTable:
if chat is None:
return None
# Sanitize message content for null characters before upserting
if isinstance(message.get("content"), str):
message["content"] = message["content"].replace("\x00", "")
chat = chat.chat
history = chat.get("history", {})
@@ -580,7 +596,7 @@ class ChatTable:
"""
Filters chats based on a search query using Python, allowing pagination using skip and limit.
"""
search_text = search_text.lower().strip()
search_text = search_text.replace("\u0000", "").lower().strip()
if not search_text:
return self.get_chat_list_by_user_id(
@@ -614,21 +630,18 @@ class ChatTable:
dialect_name = db.bind.dialect.name
if dialect_name == "sqlite":
# SQLite case: using JSON1 extension for JSON searching
sqlite_content_sql = (
"EXISTS ("
" SELECT 1 "
" FROM json_each(Chat.chat, '$.messages') AS message "
" WHERE LOWER(message.value->>'content') LIKE '%' || :content_key || '%'"
")"
)
sqlite_content_clause = text(sqlite_content_sql)
query = query.filter(
(
Chat.title.ilike(
f"%{search_text}%"
) # Case-insensitive search in title
| text(
"""
EXISTS (
SELECT 1
FROM json_each(Chat.chat, '$.messages') AS message
WHERE LOWER(message.value->>'content') LIKE '%' || :search_text || '%'
)
"""
)
).params(search_text=search_text)
or_(
Chat.title.ilike(bindparam("title_key")), sqlite_content_clause
).params(title_key=f"%{search_text}%", content_key=search_text)
)
# Check if there are any tags to filter, it should have all the tags
@@ -663,21 +676,19 @@ class ChatTable:
elif dialect_name == "postgresql":
# PostgreSQL relies on proper JSON query for search
postgres_content_sql = (
"EXISTS ("
" SELECT 1 "
" FROM json_array_elements(Chat.chat->'messages') AS message "
" WHERE LOWER(message->>'content') LIKE '%' || :content_key || '%'"
")"
)
postgres_content_clause = text(postgres_content_sql)
query = query.filter(
(
Chat.title.ilike(
f"%{search_text}%"
) # Case-insensitive search in title
| text(
"""
EXISTS (
SELECT 1
FROM json_array_elements(Chat.chat->'messages') AS message
WHERE LOWER(message->>'content') LIKE '%' || :search_text || '%'
)
"""
)
).params(search_text=search_text)
or_(
Chat.title.ilike(bindparam("title_key")),
postgres_content_clause,
).params(title_key=f"%{search_text}%", content_key=search_text)
)
# Check if there are any tags to filter, it should have all the tags
+22 -7
View File
@@ -29,6 +29,7 @@ class Folder(Base):
name = Column(Text)
items = Column(JSON, nullable=True)
meta = Column(JSON, nullable=True)
data = Column(JSON, nullable=True)
is_expanded = Column(Boolean, default=False)
created_at = Column(BigInteger)
updated_at = Column(BigInteger)
@@ -41,6 +42,7 @@ class FolderModel(BaseModel):
name: str
items: Optional[dict] = None
meta: Optional[dict] = None
data: Optional[dict] = None
is_expanded: bool = False
created_at: int
updated_at: int
@@ -55,12 +57,13 @@ class FolderModel(BaseModel):
class FolderForm(BaseModel):
name: str
data: Optional[dict] = None
model_config = ConfigDict(extra="allow")
class FolderTable:
def insert_new_folder(
self, user_id: str, name: str, parent_id: Optional[str] = None
self, user_id: str, form_data: FolderForm, parent_id: Optional[str] = None
) -> Optional[FolderModel]:
with get_db() as db:
id = str(uuid.uuid4())
@@ -68,7 +71,7 @@ class FolderTable:
**{
"id": id,
"user_id": user_id,
"name": name,
**(form_data.model_dump(exclude_unset=True) or {}),
"parent_id": parent_id,
"created_at": int(time.time()),
"updated_at": int(time.time()),
@@ -187,8 +190,8 @@ class FolderTable:
log.error(f"update_folder: {e}")
return
def update_folder_name_by_id_and_user_id(
self, id: str, user_id: str, name: str
def update_folder_by_id_and_user_id(
self, id: str, user_id: str, form_data: FolderForm
) -> Optional[FolderModel]:
try:
with get_db() as db:
@@ -197,16 +200,28 @@ class FolderTable:
if not folder:
return None
form_data = form_data.model_dump(exclude_unset=True)
existing_folder = (
db.query(Folder)
.filter_by(name=name, parent_id=folder.parent_id, user_id=user_id)
.filter_by(
name=form_data.get("name"),
parent_id=folder.parent_id,
user_id=user_id,
)
.first()
)
if existing_folder:
if existing_folder and existing_folder.id != id:
return None
folder.name = name
folder.name = form_data.get("name", folder.name)
if "data" in form_data:
folder.data = {
**(folder.data or {}),
**form_data["data"],
}
folder.updated_at = int(time.time())
db.commit()
+53 -1
View File
@@ -83,10 +83,14 @@ class GroupForm(BaseModel):
permissions: Optional[dict] = None
class GroupUpdateForm(GroupForm):
class UserIdsForm(BaseModel):
user_ids: Optional[list[str]] = None
class GroupUpdateForm(GroupForm, UserIdsForm):
pass
class GroupTable:
def insert_new_group(
self, user_id: str, form_data: GroupForm
@@ -275,5 +279,53 @@ class GroupTable:
log.exception(e)
return False
def add_users_to_group(
self, id: str, user_ids: Optional[list[str]] = None
) -> Optional[GroupModel]:
try:
with get_db() as db:
group = db.query(Group).filter_by(id=id).first()
if not group:
return None
if not group.user_ids:
group.user_ids = []
for user_id in user_ids:
if user_id not in group.user_ids:
group.user_ids.append(user_id)
group.updated_at = int(time.time())
db.commit()
db.refresh(group)
return GroupModel.model_validate(group)
except Exception as e:
log.exception(e)
return None
def remove_users_from_group(
self, id: str, user_ids: Optional[list[str]] = None
) -> Optional[GroupModel]:
try:
with get_db() as db:
group = db.query(Group).filter_by(id=id).first()
if not group:
return None
if not group.user_ids:
return GroupModel.model_validate(group)
for user_id in user_ids:
if user_id in group.user_ids:
group.user_ids.remove(user_id)
group.updated_at = int(time.time())
db.commit()
db.refresh(group)
return GroupModel.model_validate(group)
except Exception as e:
log.exception(e)
return None
Groups = GroupTable()
+13 -4
View File
@@ -71,9 +71,13 @@ class MemoriesTable:
) -> Optional[MemoryModel]:
with get_db() as db:
try:
db.query(Memory).filter_by(id=id, user_id=user_id).update(
{"content": content, "updated_at": int(time.time())}
)
memory = db.get(Memory, id)
if not memory or memory.user_id != user_id:
return None
memory.content = content
memory.updated_at = int(time.time())
db.commit()
return self.get_memory_by_id(id)
except Exception:
@@ -127,7 +131,12 @@ class MemoriesTable:
def delete_memory_by_id_and_user_id(self, id: str, user_id: str) -> bool:
with get_db() as db:
try:
db.query(Memory).filter_by(id=id, user_id=user_id).delete()
memory = db.get(Memory, id)
if not memory or memory.user_id != user_id:
return None
# Delete the memory
db.delete(memory)
db.commit()
return True
+44
View File
@@ -269,5 +269,49 @@ class ModelsTable:
except Exception:
return False
def sync_models(self, user_id: str, models: list[ModelModel]) -> list[ModelModel]:
try:
with get_db() as db:
# Get existing models
existing_models = db.query(Model).all()
existing_ids = {model.id for model in existing_models}
# Prepare a set of new model IDs
new_model_ids = {model.id for model in models}
# Update or insert models
for model in models:
if model.id in existing_ids:
db.query(Model).filter_by(id=model.id).update(
{
**model.model_dump(),
"user_id": user_id,
"updated_at": int(time.time()),
}
)
else:
new_model = Model(
**{
**model.model_dump(),
"user_id": user_id,
"updated_at": int(time.time()),
}
)
db.add(new_model)
# Remove models that are no longer present
for model in existing_models:
if model.id not in new_model_ids:
db.delete(model)
db.commit()
return [
ModelModel.model_validate(model) for model in db.query(Model).all()
]
except Exception as e:
log.exception(f"Error syncing models for user {user_id}: {e}")
return []
Models = ModelsTable()
+22 -5
View File
@@ -62,6 +62,13 @@ class NoteForm(BaseModel):
access_control: Optional[dict] = None
class NoteUpdateForm(BaseModel):
title: Optional[str] = None
data: Optional[dict] = None
meta: Optional[dict] = None
access_control: Optional[dict] = None
class NoteUserResponse(NoteModel):
user: Optional[UserResponse] = None
@@ -110,16 +117,26 @@ class NoteTable:
note = db.query(Note).filter(Note.id == id).first()
return NoteModel.model_validate(note) if note else None
def update_note_by_id(self, id: str, form_data: NoteForm) -> Optional[NoteModel]:
def update_note_by_id(
self, id: str, form_data: NoteUpdateForm
) -> Optional[NoteModel]:
with get_db() as db:
note = db.query(Note).filter(Note.id == id).first()
if not note:
return None
note.title = form_data.title
note.data = form_data.data
note.meta = form_data.meta
note.access_control = form_data.access_control
form_data = form_data.model_dump(exclude_unset=True)
if "title" in form_data:
note.title = form_data["title"]
if "data" in form_data:
note.data = {**note.data, **form_data["data"]}
if "meta" in form_data:
note.meta = {**note.meta, **form_data["meta"]}
if "access_control" in form_data:
note.access_control = form_data["access_control"]
note.updated_at = int(time.time_ns())
db.commit()
+12
View File
@@ -74,6 +74,18 @@ class UserListResponse(BaseModel):
total: int
class UserInfoResponse(BaseModel):
id: str
name: str
email: str
role: str
class UserInfoListResponse(BaseModel):
users: list[UserInfoResponse]
total: int
class UserResponse(BaseModel):
id: str
name: str
+8 -3
View File
@@ -14,7 +14,7 @@ from langchain_community.document_loaders import (
TextLoader,
UnstructuredEPubLoader,
UnstructuredExcelLoader,
UnstructuredMarkdownLoader,
UnstructuredODTLoader,
UnstructuredPowerPointLoader,
UnstructuredRSTLoader,
UnstructuredXMLLoader,
@@ -181,7 +181,7 @@ class DoclingLoader:
if lang.strip()
]
endpoint = f"{self.url}/v1alpha/convert/file"
endpoint = f"{self.url}/v1/convert/file"
r = requests.post(endpoint, files=files, data=params)
if r.ok:
@@ -226,7 +226,10 @@ class Loader:
def _is_text_file(self, file_ext: str, file_content_type: str) -> bool:
return file_ext in known_source_ext or (
file_content_type and file_content_type.find("text/") >= 0
file_content_type
and file_content_type.find("text/") >= 0
# Avoid text/html files being detected as text
and not file_content_type.find("html") >= 0
)
def _get_loader(self, filename: str, file_content_type: str, file_path: str):
@@ -389,6 +392,8 @@ class Loader:
loader = UnstructuredPowerPointLoader(file_path)
elif file_ext == "msg":
loader = OutlookMessageLoader(file_path)
elif file_ext == "odt":
loader = UnstructuredODTLoader(file_path)
elif self._is_text_file(file_ext, file_content_type):
loader = TextLoader(file_path, autodetect_encoding=True)
else:
@@ -507,6 +507,7 @@ class MistralLoader:
timeout=timeout,
headers={"User-Agent": "OpenWebUI-MistralLoader/2.0"},
raise_for_status=False, # We handle status codes manually
trust_env=True,
) as session:
yield session
@@ -1,8 +1,10 @@
import logging
import requests
from typing import Optional, List, Tuple
from urllib.parse import quote
from open_webui.env import SRC_LOG_LEVELS
from open_webui.env import ENABLE_FORWARD_USER_INFO_HEADERS, SRC_LOG_LEVELS
from open_webui.retrieval.models.base_reranker import BaseReranker
@@ -21,7 +23,9 @@ class ExternalReranker(BaseReranker):
self.url = url
self.model = model
def predict(self, sentences: List[Tuple[str, str]]) -> Optional[List[float]]:
def predict(
self, sentences: List[Tuple[str, str]], user=None
) -> Optional[List[float]]:
query = sentences[0][0]
docs = [i[1] for i in sentences]
@@ -41,6 +45,16 @@ class ExternalReranker(BaseReranker):
headers={
"Content-Type": "application/json",
"Authorization": f"Bearer {self.api_key}",
**(
{
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"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=payload,
)
+196 -123
View File
@@ -7,6 +7,7 @@ import hashlib
from concurrent.futures import ThreadPoolExecutor
import time
from urllib.parse import quote
from huggingface_hub import snapshot_download
from langchain.retrievers import ContextualCompressionRetriever, EnsembleRetriever
from langchain_community.retrievers import BM25Retriever
@@ -17,8 +18,11 @@ from open_webui.retrieval.vector.factory import VECTOR_DB_CLIENT
from open_webui.models.users import UserModel
from open_webui.models.files import Files
from open_webui.models.knowledge import Knowledges
from open_webui.models.notes import Notes
from open_webui.retrieval.vector.main import GetResult
from open_webui.utils.access_control import has_access
from open_webui.env import (
@@ -441,9 +445,20 @@ def get_embedding_function(
raise ValueError(f"Unknown embedding engine: {embedding_engine}")
def get_sources_from_files(
def get_reranking_function(reranking_engine, reranking_model, reranking_function):
if reranking_function is None:
return None
if reranking_engine == "external":
return lambda sentences, user=None: reranking_function.predict(
sentences, user=user
)
else:
return lambda sentences, user=None: reranking_function.predict(sentences)
def get_sources_from_items(
request,
files,
items,
queries,
embedding_function,
k,
@@ -453,159 +468,217 @@ def get_sources_from_files(
hybrid_bm25_weight,
hybrid_search,
full_context=False,
user: Optional[UserModel] = None,
):
log.debug(
f"files: {files} {queries} {embedding_function} {reranking_function} {full_context}"
f"items: {items} {queries} {embedding_function} {reranking_function} {full_context}"
)
extracted_collections = []
relevant_contexts = []
query_results = []
for file in files:
for item in items:
query_result = None
collection_names = []
context = None
if file.get("docs"):
# BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL
context = {
"documents": [[doc.get("content") for doc in file.get("docs")]],
"metadatas": [[doc.get("metadata") for doc in file.get("docs")]],
}
elif file.get("context") == "full":
# Manual Full Mode Toggle
context = {
"documents": [[file.get("file").get("data", {}).get("content")]],
"metadatas": [[{"file_id": file.get("id"), "name": file.get("name")}]],
}
elif (
file.get("type") != "web_search"
and request.app.state.config.BYPASS_EMBEDDING_AND_RETRIEVAL
):
# BYPASS_EMBEDDING_AND_RETRIEVAL
if file.get("type") == "collection":
file_ids = file.get("data", {}).get("file_ids", [])
if item.get("type") == "text":
# Raw Text
# Used during temporary chat file uploads
documents = []
metadatas = []
for file_id in file_ids:
file_object = Files.get_file_by_id(file_id)
if file_object:
documents.append(file_object.data.get("content", ""))
metadatas.append(
{
"file_id": file_id,
"name": file_object.filename,
"source": file_object.filename,
}
)
context = {
"documents": [documents],
"metadatas": [metadatas],
if item.get("file"):
# if item has file data, use it
query_result = {
"documents": [
[item.get("file", {}).get("data", {}).get("content")]
],
"metadatas": [
[item.get("file", {}).get("data", {}).get("meta", {})]
],
}
else:
# Fallback to item content
query_result = {
"documents": [[item.get("content")]],
"metadatas": [
[{"file_id": item.get("id"), "name": item.get("name")}]
],
}
elif file.get("id"):
file_object = Files.get_file_by_id(file.get("id"))
if file_object:
context = {
"documents": [[file_object.data.get("content", "")]],
elif item.get("type") == "note":
# Note Attached
note = Notes.get_note_by_id(item.get("id"))
if note and (
user.role == "admin"
or note.user_id == user.id
or has_access(user.id, "read", note.access_control)
):
# User has access to the note
query_result = {
"documents": [[note.data.get("content", {}).get("md", "")]],
"metadatas": [[{"file_id": note.id, "name": note.title}]],
}
elif item.get("type") == "file":
if (
item.get("context") == "full"
or request.app.state.config.BYPASS_EMBEDDING_AND_RETRIEVAL
):
if item.get("file", {}).get("data", {}).get("content", ""):
# Manual Full Mode Toggle
# Used from chat file modal, we can assume that the file content will be available from item.get("file").get("data", {}).get("content")
query_result = {
"documents": [
[item.get("file", {}).get("data", {}).get("content", "")]
],
"metadatas": [
[
{
"file_id": file.get("id"),
"name": file_object.filename,
"source": file_object.filename,
"file_id": item.get("id"),
"name": item.get("name"),
**item.get("file")
.get("data", {})
.get("metadata", {}),
}
]
],
}
elif file.get("file").get("data"):
context = {
"documents": [[file.get("file").get("data", {}).get("content")]],
"metadatas": [
[file.get("file").get("data", {}).get("metadata", {})]
],
}
else:
collection_names = []
if file.get("type") == "collection":
if file.get("legacy"):
collection_names = file.get("collection_names", [])
elif item.get("id"):
file_object = Files.get_file_by_id(item.get("id"))
if file_object:
query_result = {
"documents": [[file_object.data.get("content", "")]],
"metadatas": [
[
{
"file_id": item.get("id"),
"name": file_object.filename,
"source": file_object.filename,
}
]
],
}
else:
# Fallback to collection names
if item.get("legacy"):
collection_names.append(f"{item['id']}")
else:
collection_names.append(file["id"])
elif file.get("collection_name"):
collection_names.append(file["collection_name"])
elif file.get("id"):
if file.get("legacy"):
collection_names.append(f"{file['id']}")
else:
collection_names.append(f"file-{file['id']}")
collection_names.append(f"file-{item['id']}")
elif item.get("type") == "collection":
if (
item.get("context") == "full"
or request.app.state.config.BYPASS_EMBEDDING_AND_RETRIEVAL
):
# Manual Full Mode Toggle for Collection
knowledge_base = Knowledges.get_knowledge_by_id(item.get("id"))
if knowledge_base and (
user.role == "admin"
or has_access(user.id, "read", knowledge_base.access_control)
):
file_ids = knowledge_base.data.get("file_ids", [])
documents = []
metadatas = []
for file_id in file_ids:
file_object = Files.get_file_by_id(file_id)
if file_object:
documents.append(file_object.data.get("content", ""))
metadatas.append(
{
"file_id": file_id,
"name": file_object.filename,
"source": file_object.filename,
}
)
query_result = {
"documents": [documents],
"metadatas": [metadatas],
}
else:
# Fallback to collection names
if item.get("legacy"):
collection_names = item.get("collection_names", [])
else:
collection_names.append(item["id"])
elif item.get("docs"):
# BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL
query_result = {
"documents": [[doc.get("content") for doc in item.get("docs")]],
"metadatas": [[doc.get("metadata") for doc in item.get("docs")]],
}
elif item.get("collection_name"):
# Direct Collection Name
collection_names.append(item["collection_name"])
elif item.get("collection_names"):
# Collection Names List
collection_names.extend(item["collection_names"])
# If query_result is None
# Fallback to collection names and vector search the collections
if query_result is None and collection_names:
collection_names = set(collection_names).difference(extracted_collections)
if not collection_names:
log.debug(f"skipping {file} as it has already been extracted")
log.debug(f"skipping {item} as it has already been extracted")
continue
if full_context:
try:
context = get_all_items_from_collections(collection_names)
except Exception as e:
log.exception(e)
else:
try:
context = None
if file.get("type") == "text":
context = file["content"]
else:
if hybrid_search:
try:
context = query_collection_with_hybrid_search(
collection_names=collection_names,
queries=queries,
embedding_function=embedding_function,
k=k,
reranking_function=reranking_function,
k_reranker=k_reranker,
r=r,
hybrid_bm25_weight=hybrid_bm25_weight,
)
except Exception as e:
log.debug(
"Error when using hybrid search, using"
" non hybrid search as fallback."
)
if (not hybrid_search) or (context is None):
context = query_collection(
try:
if full_context:
query_result = get_all_items_from_collections(collection_names)
else:
query_result = None # Initialize to None
if hybrid_search:
try:
query_result = query_collection_with_hybrid_search(
collection_names=collection_names,
queries=queries,
embedding_function=embedding_function,
k=k,
reranking_function=reranking_function,
k_reranker=k_reranker,
r=r,
hybrid_bm25_weight=hybrid_bm25_weight,
)
except Exception as e:
log.exception(e)
except Exception as e:
log.debug(
"Error when using hybrid search, using non hybrid search as fallback."
)
# fallback to non-hybrid search
if not hybrid_search and query_result is None:
query_result = query_collection(
collection_names=collection_names,
queries=queries,
embedding_function=embedding_function,
k=k,
)
except Exception as e:
log.exception(e)
extracted_collections.extend(collection_names)
if context:
if "data" in file:
del file["data"]
relevant_contexts.append({**context, "file": file})
if query_result:
if "data" in item:
del item["data"]
query_results.append({**query_result, "file": item})
sources = []
for context in relevant_contexts:
for query_result in query_results:
try:
if "documents" in context:
if "metadatas" in context:
if "documents" in query_result:
if "metadatas" in query_result:
source = {
"source": context["file"],
"document": context["documents"][0],
"metadata": context["metadatas"][0],
"source": query_result["file"],
"document": query_result["documents"][0],
"metadata": query_result["metadatas"][0],
}
if "distances" in context and context["distances"]:
source["distances"] = context["distances"][0]
if "distances" in query_result and query_result["distances"]:
source["distances"] = query_result["distances"][0]
sources.append(source)
except Exception as e:
@@ -678,7 +751,7 @@ def generate_openai_batch_embeddings(
"Authorization": f"Bearer {key}",
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -727,7 +800,7 @@ def generate_azure_openai_batch_embeddings(
"api-key": key,
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -777,7 +850,7 @@ def generate_ollama_batch_embeddings(
"Authorization": f"Bearer {key}",
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -874,7 +947,7 @@ class RerankCompressor(BaseDocumentCompressor):
reranking = self.reranking_function is not None
if reranking:
scores = self.reranking_function.predict(
scores = self.reranking_function(
[(query, doc.page_content) for doc in documents]
)
else:
@@ -11,6 +11,8 @@ from open_webui.retrieval.vector.main import (
SearchResult,
GetResult,
)
from open_webui.retrieval.vector.utils import stringify_metadata
from open_webui.config import (
CHROMA_DATA_PATH,
CHROMA_HTTP_HOST,
@@ -144,7 +146,7 @@ class ChromaClient(VectorDBBase):
ids = [item["id"] for item in items]
documents = [item["text"] for item in items]
embeddings = [item["vector"] for item in items]
metadatas = [item["metadata"] for item in items]
metadatas = [stringify_metadata(item["metadata"]) for item in items]
for batch in create_batches(
api=self.client,
@@ -164,7 +166,7 @@ class ChromaClient(VectorDBBase):
ids = [item["id"] for item in items]
documents = [item["text"] for item in items]
embeddings = [item["vector"] for item in items]
metadatas = [item["metadata"] for item in items]
metadatas = [stringify_metadata(item["metadata"]) for item in items]
collection.upsert(
ids=ids, documents=documents, embeddings=embeddings, metadatas=metadatas
@@ -3,6 +3,8 @@ from pymilvus import FieldSchema, DataType
import json
import logging
from typing import Optional
from open_webui.retrieval.vector.utils import stringify_metadata
from open_webui.retrieval.vector.main import (
VectorDBBase,
VectorItem,
@@ -311,7 +313,7 @@ class MilvusClient(VectorDBBase):
"id": item["id"],
"vector": item["vector"],
"data": {"text": item["text"]},
"metadata": item["metadata"],
"metadata": stringify_metadata(item["metadata"]),
}
for item in items
],
@@ -347,7 +349,7 @@ class MilvusClient(VectorDBBase):
"id": item["id"],
"vector": item["vector"],
"data": {"text": item["text"]},
"metadata": item["metadata"],
"metadata": stringify_metadata(item["metadata"]),
}
for item in items
],
@@ -157,10 +157,10 @@ class OpenSearchClient(VectorDBBase):
for field, value in filter.items():
query_body["query"]["bool"]["filter"].append(
{"match": {"metadata." + str(field): value}}
{"term": {"metadata." + str(field) + ".keyword": value}}
)
size = limit if limit else 10
size = limit if limit else 10000
try:
result = self.client.search(
@@ -206,6 +206,7 @@ class OpenSearchClient(VectorDBBase):
for item in batch
]
bulk(self.client, actions)
self.client.indices.refresh(self._get_index_name(collection_name))
def upsert(self, collection_name: str, items: list[VectorItem]):
self._create_index_if_not_exists(
@@ -228,6 +229,7 @@ class OpenSearchClient(VectorDBBase):
for item in batch
]
bulk(self.client, actions)
self.client.indices.refresh(self._get_index_name(collection_name))
def delete(
self,
@@ -251,11 +253,12 @@ class OpenSearchClient(VectorDBBase):
}
for field, value in filter.items():
query_body["query"]["bool"]["filter"].append(
{"match": {"metadata." + str(field): value}}
{"term": {"metadata." + str(field) + ".keyword": value}}
)
self.client.delete_by_query(
index=self._get_index_name(collection_name), body=query_body
)
self.client.indices.refresh(self._get_index_name(collection_name))
def reset(self):
indices = self.client.indices.get(index=f"{self.index_prefix}_*")
@@ -18,7 +18,7 @@ from sqlalchemy import (
values,
)
from sqlalchemy.sql import true
from sqlalchemy.pool import NullPool
from sqlalchemy.pool import NullPool, QueuePool
from sqlalchemy.orm import declarative_base, scoped_session, sessionmaker
from sqlalchemy.dialects.postgresql import JSONB, array
@@ -26,6 +26,8 @@ from pgvector.sqlalchemy import Vector
from sqlalchemy.ext.mutable import MutableDict
from sqlalchemy.exc import NoSuchTableError
from open_webui.retrieval.vector.utils import stringify_metadata
from open_webui.retrieval.vector.main import (
VectorDBBase,
VectorItem,
@@ -37,6 +39,10 @@ from open_webui.config import (
PGVECTOR_INITIALIZE_MAX_VECTOR_LENGTH,
PGVECTOR_PGCRYPTO,
PGVECTOR_PGCRYPTO_KEY,
PGVECTOR_POOL_SIZE,
PGVECTOR_POOL_MAX_OVERFLOW,
PGVECTOR_POOL_TIMEOUT,
PGVECTOR_POOL_RECYCLE,
)
from open_webui.env import SRC_LOG_LEVELS
@@ -80,9 +86,24 @@ class PgvectorClient(VectorDBBase):
self.session = Session
else:
engine = create_engine(
PGVECTOR_DB_URL, pool_pre_ping=True, poolclass=NullPool
)
if isinstance(PGVECTOR_POOL_SIZE, int):
if PGVECTOR_POOL_SIZE > 0:
engine = create_engine(
PGVECTOR_DB_URL,
pool_size=PGVECTOR_POOL_SIZE,
max_overflow=PGVECTOR_POOL_MAX_OVERFLOW,
pool_timeout=PGVECTOR_POOL_TIMEOUT,
pool_recycle=PGVECTOR_POOL_RECYCLE,
pool_pre_ping=True,
poolclass=QueuePool,
)
else:
engine = create_engine(
PGVECTOR_DB_URL, pool_pre_ping=True, poolclass=NullPool
)
else:
engine = create_engine(PGVECTOR_DB_URL, pool_pre_ping=True)
SessionLocal = sessionmaker(
autocommit=False, autoflush=False, bind=engine, expire_on_commit=False
)
@@ -216,7 +237,7 @@ class PgvectorClient(VectorDBBase):
vector=vector,
collection_name=collection_name,
text=item["text"],
vmetadata=item["metadata"],
vmetadata=stringify_metadata(item["metadata"]),
)
new_items.append(new_chunk)
self.session.bulk_save_objects(new_items)
@@ -273,7 +294,7 @@ class PgvectorClient(VectorDBBase):
if existing:
existing.vector = vector
existing.text = item["text"]
existing.vmetadata = item["metadata"]
existing.vmetadata = stringify_metadata(item["metadata"])
existing.collection_name = (
collection_name # Update collection_name if necessary
)
@@ -283,7 +304,7 @@ class PgvectorClient(VectorDBBase):
vector=vector,
collection_name=collection_name,
text=item["text"],
vmetadata=item["metadata"],
vmetadata=stringify_metadata(item["metadata"]),
)
self.session.add(new_chunk)
self.session.commit()
@@ -18,6 +18,7 @@ from open_webui.config import (
QDRANT_ON_DISK,
QDRANT_GRPC_PORT,
QDRANT_PREFER_GRPC,
QDRANT_COLLECTION_PREFIX,
)
from open_webui.env import SRC_LOG_LEVELS
@@ -29,7 +30,7 @@ log.setLevel(SRC_LOG_LEVELS["RAG"])
class QdrantClient(VectorDBBase):
def __init__(self):
self.collection_prefix = "open-webui"
self.collection_prefix = QDRANT_COLLECTION_PREFIX
self.QDRANT_URI = QDRANT_URI
self.QDRANT_API_KEY = QDRANT_API_KEY
self.QDRANT_ON_DISK = QDRANT_ON_DISK
@@ -86,6 +87,25 @@ class QdrantClient(VectorDBBase):
),
)
# Create payload indexes for efficient filtering
self.client.create_payload_index(
collection_name=collection_name_with_prefix,
field_name="metadata.hash",
field_schema=models.KeywordIndexParams(
type=models.KeywordIndexType.KEYWORD,
is_tenant=False,
on_disk=self.QDRANT_ON_DISK,
),
)
self.client.create_payload_index(
collection_name=collection_name_with_prefix,
field_name="metadata.file_id",
field_schema=models.KeywordIndexParams(
type=models.KeywordIndexType.KEYWORD,
is_tenant=False,
on_disk=self.QDRANT_ON_DISK,
),
)
log.info(f"collection {collection_name_with_prefix} successfully created!")
def _create_collection_if_not_exists(self, collection_name, dimension):
@@ -1,5 +1,5 @@
import logging
from typing import Optional, Tuple
from typing import Optional, Tuple, List, Dict, Any
from urllib.parse import urlparse
import grpc
@@ -9,6 +9,7 @@ from open_webui.config import (
QDRANT_ON_DISK,
QDRANT_PREFER_GRPC,
QDRANT_URI,
QDRANT_COLLECTION_PREFIX,
)
from open_webui.env import SRC_LOG_LEVELS
from open_webui.retrieval.vector.main import (
@@ -23,14 +24,28 @@ from qdrant_client.http.models import PointStruct
from qdrant_client.models import models
NO_LIMIT = 999999999
TENANT_ID_FIELD = "tenant_id"
DEFAULT_DIMENSION = 384
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])
def _tenant_filter(tenant_id: str) -> models.FieldCondition:
return models.FieldCondition(
key=TENANT_ID_FIELD, match=models.MatchValue(value=tenant_id)
)
def _metadata_filter(key: str, value: Any) -> models.FieldCondition:
return models.FieldCondition(
key=f"metadata.{key}", match=models.MatchValue(value=value)
)
class QdrantClient(VectorDBBase):
def __init__(self):
self.collection_prefix = "open-webui"
self.collection_prefix = QDRANT_COLLECTION_PREFIX
self.QDRANT_URI = QDRANT_URI
self.QDRANT_API_KEY = QDRANT_API_KEY
self.QDRANT_ON_DISK = QDRANT_ON_DISK
@@ -38,24 +53,26 @@ class QdrantClient(VectorDBBase):
self.GRPC_PORT = QDRANT_GRPC_PORT
if not self.QDRANT_URI:
self.client = None
return
raise ValueError(
"QDRANT_URI is not set. Please configure it in the environment variables."
)
# Unified handling for either scheme
parsed = urlparse(self.QDRANT_URI)
host = parsed.hostname or self.QDRANT_URI
http_port = parsed.port or 6333 # default REST port
if self.PREFER_GRPC:
self.client = Qclient(
self.client = (
Qclient(
host=host,
port=http_port,
grpc_port=self.GRPC_PORT,
prefer_grpc=self.PREFER_GRPC,
api_key=self.QDRANT_API_KEY,
)
else:
self.client = Qclient(url=self.QDRANT_URI, api_key=self.QDRANT_API_KEY)
if self.PREFER_GRPC
else Qclient(url=self.QDRANT_URI, api_key=self.QDRANT_API_KEY)
)
# Main collection types for multi-tenancy
self.MEMORY_COLLECTION = f"{self.collection_prefix}_memories"
@@ -65,23 +82,13 @@ class QdrantClient(VectorDBBase):
self.HASH_BASED_COLLECTION = f"{self.collection_prefix}_hash-based"
def _result_to_get_result(self, points) -> GetResult:
ids = []
documents = []
metadatas = []
ids, documents, metadatas = [], [], []
for point in points:
payload = point.payload
ids.append(point.id)
documents.append(payload["text"])
metadatas.append(payload["metadata"])
return GetResult(
**{
"ids": [ids],
"documents": [documents],
"metadatas": [metadatas],
}
)
return GetResult(ids=[ids], documents=[documents], metadatas=[metadatas])
def _get_collection_and_tenant_id(self, collection_name: str) -> Tuple[str, str]:
"""
@@ -113,143 +120,47 @@ class QdrantClient(VectorDBBase):
else:
return self.KNOWLEDGE_COLLECTION, tenant_id
def _extract_error_message(self, exception):
"""
Extract error message from either HTTP or gRPC exceptions
Returns:
tuple: (status_code, error_message)
"""
# Check if it's an HTTP exception
if isinstance(exception, UnexpectedResponse):
try:
error_data = exception.structured()
error_msg = error_data.get("status", {}).get("error", "")
return exception.status_code, error_msg
except Exception as inner_e:
log.error(f"Failed to parse HTTP error: {inner_e}")
return exception.status_code, str(exception)
# Check if it's a gRPC exception
elif isinstance(exception, grpc.RpcError):
# Extract status code from gRPC error
status_code = None
if hasattr(exception, "code") and callable(exception.code):
status_code = exception.code().value[0]
# Extract error message
error_msg = str(exception)
if "details =" in error_msg:
# Parse the details line which contains the actual error message
try:
details_line = [
line.strip()
for line in error_msg.split("\n")
if "details =" in line
][0]
error_msg = details_line.split("details =")[1].strip(' "')
except (IndexError, AttributeError):
# Fall back to full message if parsing fails
pass
return status_code, error_msg
# For any other type of exception
return None, str(exception)
def _is_collection_not_found_error(self, exception):
"""
Check if the exception is due to collection not found, supporting both HTTP and gRPC
"""
status_code, error_msg = self._extract_error_message(exception)
# HTTP error (404)
if (
status_code == 404
and "Collection" in error_msg
and "doesn't exist" in error_msg
):
return True
# gRPC error (NOT_FOUND status)
if (
isinstance(exception, grpc.RpcError)
and exception.code() == grpc.StatusCode.NOT_FOUND
):
return True
return False
def _is_dimension_mismatch_error(self, exception):
"""
Check if the exception is due to dimension mismatch, supporting both HTTP and gRPC
"""
status_code, error_msg = self._extract_error_message(exception)
# Common patterns in both HTTP and gRPC
return (
"Vector dimension error" in error_msg
or "dimensions mismatch" in error_msg
or "invalid vector size" in error_msg
)
def _create_multi_tenant_collection_if_not_exists(
self, mt_collection_name: str, dimension: int = 384
def _create_multi_tenant_collection(
self, mt_collection_name: str, dimension: int = DEFAULT_DIMENSION
):
"""
Creates a collection with multi-tenancy configuration if it doesn't exist.
Default dimension is set to 384 which corresponds to 'sentence-transformers/all-MiniLM-L6-v2'.
When creating collections dynamically (insert/upsert), the actual vector dimensions will be used.
Creates a collection with multi-tenancy configuration and payload indexes for tenant_id and metadata fields.
"""
try:
# Try to create the collection directly - will fail if it already exists
self.client.create_collection(
collection_name=mt_collection_name,
vectors_config=models.VectorParams(
size=dimension,
distance=models.Distance.COSINE,
on_disk=self.QDRANT_ON_DISK,
),
hnsw_config=models.HnswConfigDiff(
payload_m=16, # Enable per-tenant indexing
m=0,
on_disk=self.QDRANT_ON_DISK,
),
)
self.client.create_collection(
collection_name=mt_collection_name,
vectors_config=models.VectorParams(
size=dimension,
distance=models.Distance.COSINE,
on_disk=self.QDRANT_ON_DISK,
),
)
log.info(
f"Multi-tenant collection {mt_collection_name} created with dimension {dimension}!"
)
# Create tenant ID payload index
self.client.create_payload_index(
collection_name=mt_collection_name,
field_name=TENANT_ID_FIELD,
field_schema=models.KeywordIndexParams(
type=models.KeywordIndexType.KEYWORD,
is_tenant=True,
on_disk=self.QDRANT_ON_DISK,
),
)
for field in ("metadata.hash", "metadata.file_id"):
self.client.create_payload_index(
collection_name=mt_collection_name,
field_name="tenant_id",
field_name=field,
field_schema=models.KeywordIndexParams(
type=models.KeywordIndexType.KEYWORD,
is_tenant=True,
on_disk=self.QDRANT_ON_DISK,
),
wait=True,
)
log.info(
f"Multi-tenant collection {mt_collection_name} created with dimension {dimension}!"
)
except (UnexpectedResponse, grpc.RpcError) as e:
# Check for the specific error indicating collection already exists
status_code, error_msg = self._extract_error_message(e)
# HTTP status code 409 or gRPC ALREADY_EXISTS
if (isinstance(e, UnexpectedResponse) and status_code == 409) or (
isinstance(e, grpc.RpcError)
and e.code() == grpc.StatusCode.ALREADY_EXISTS
):
if "already exists" in error_msg:
log.debug(f"Collection {mt_collection_name} already exists")
return
# If it's not an already exists error, re-raise
raise e
except Exception as e:
raise e
def _create_points(self, items: list[VectorItem], tenant_id: str):
def _create_points(
self, items: List[VectorItem], tenant_id: str
) -> List[PointStruct]:
"""
Create point structs from vector items with tenant ID.
"""
@@ -260,56 +171,42 @@ class QdrantClient(VectorDBBase):
payload={
"text": item["text"],
"metadata": item["metadata"],
"tenant_id": tenant_id,
TENANT_ID_FIELD: tenant_id,
},
)
for item in items
]
def _ensure_collection(
self, mt_collection_name: str, dimension: int = DEFAULT_DIMENSION
):
"""
Ensure the collection exists and payload indexes are created for tenant_id and metadata fields.
"""
if not self.client.collection_exists(collection_name=mt_collection_name):
self._create_multi_tenant_collection(mt_collection_name, dimension)
def has_collection(self, collection_name: str) -> bool:
"""
Check if a logical collection exists by checking for any points with the tenant ID.
"""
if not self.client:
return False
# Map to multi-tenant collection and tenant ID
mt_collection, tenant_id = self._get_collection_and_tenant_id(collection_name)
# Create tenant filter
tenant_filter = models.FieldCondition(
key="tenant_id", match=models.MatchValue(value=tenant_id)
)
try:
# Try directly querying - most of the time collection should exist
response = self.client.query_points(
collection_name=mt_collection,
query_filter=models.Filter(must=[tenant_filter]),
limit=1,
)
# Collection exists with this tenant ID if there are points
return len(response.points) > 0
except (UnexpectedResponse, grpc.RpcError) as e:
if self._is_collection_not_found_error(e):
log.debug(f"Collection {mt_collection} doesn't exist")
return False
else:
# For other API errors, log and return False
_, error_msg = self._extract_error_message(e)
log.warning(f"Unexpected Qdrant error: {error_msg}")
return False
except Exception as e:
# For any other errors, log and return False
log.debug(f"Error checking collection {mt_collection}: {e}")
if not self.client.collection_exists(collection_name=mt_collection):
return False
tenant_filter = _tenant_filter(tenant_id)
count_result = self.client.count(
collection_name=mt_collection,
count_filter=models.Filter(must=[tenant_filter]),
)
return count_result.count > 0
def delete(
self,
collection_name: str,
ids: Optional[list[str]] = None,
filter: Optional[dict] = None,
ids: Optional[List[str]] = None,
filter: Optional[Dict[str, Any]] = None,
):
"""
Delete vectors by ID or filter from a collection with tenant isolation.
@@ -317,189 +214,76 @@ class QdrantClient(VectorDBBase):
if not self.client:
return None
# Map to multi-tenant collection and tenant ID
mt_collection, tenant_id = self._get_collection_and_tenant_id(collection_name)
if not self.client.collection_exists(collection_name=mt_collection):
log.debug(f"Collection {mt_collection} doesn't exist, nothing to delete")
return None
# Create tenant filter
tenant_filter = models.FieldCondition(
key="tenant_id", match=models.MatchValue(value=tenant_id)
must_conditions = [_tenant_filter(tenant_id)]
should_conditions = []
if ids:
should_conditions = [_metadata_filter("id", id_value) for id_value in ids]
elif filter:
must_conditions += [_metadata_filter(k, v) for k, v in filter.items()]
return self.client.delete(
collection_name=mt_collection,
points_selector=models.FilterSelector(
filter=models.Filter(must=must_conditions, should=should_conditions)
),
)
must_conditions = [tenant_filter]
should_conditions = []
if ids:
for id_value in ids:
should_conditions.append(
models.FieldCondition(
key="metadata.id",
match=models.MatchValue(value=id_value),
),
)
elif filter:
for key, value in filter.items():
must_conditions.append(
models.FieldCondition(
key=f"metadata.{key}",
match=models.MatchValue(value=value),
),
)
try:
# Try to delete directly - most of the time collection should exist
update_result = self.client.delete(
collection_name=mt_collection,
points_selector=models.FilterSelector(
filter=models.Filter(must=must_conditions, should=should_conditions)
),
)
return update_result
except (UnexpectedResponse, grpc.RpcError) as e:
if self._is_collection_not_found_error(e):
log.debug(
f"Collection {mt_collection} doesn't exist, nothing to delete"
)
return None
else:
# For other API errors, log and re-raise
_, error_msg = self._extract_error_message(e)
log.warning(f"Unexpected Qdrant error: {error_msg}")
raise
except Exception as e:
# For non-Qdrant exceptions, re-raise
raise
def search(
self, collection_name: str, vectors: list[list[float | int]], limit: int
self, collection_name: str, vectors: List[List[float | int]], limit: int
) -> Optional[SearchResult]:
"""
Search for the nearest neighbor items based on the vectors with tenant isolation.
"""
if not self.client:
if not self.client or not vectors:
return None
# Map to multi-tenant collection and tenant ID
mt_collection, tenant_id = self._get_collection_and_tenant_id(collection_name)
# Get the vector dimension from the query vector
dimension = len(vectors[0]) if vectors and len(vectors) > 0 else None
try:
# Try the search operation directly - most of the time collection should exist
# Create tenant filter
tenant_filter = models.FieldCondition(
key="tenant_id", match=models.MatchValue(value=tenant_id)
)
# Ensure vector dimensions match the collection
collection_dim = self.client.get_collection(
mt_collection
).config.params.vectors.size
if collection_dim != dimension:
if collection_dim < dimension:
vectors = [vector[:collection_dim] for vector in vectors]
else:
vectors = [
vector + [0] * (collection_dim - dimension)
for vector in vectors
]
# Search with tenant filter
prefetch_query = models.Prefetch(
filter=models.Filter(must=[tenant_filter]),
limit=NO_LIMIT,
)
query_response = self.client.query_points(
collection_name=mt_collection,
query=vectors[0],
prefetch=prefetch_query,
limit=limit,
)
get_result = self._result_to_get_result(query_response.points)
return SearchResult(
ids=get_result.ids,
documents=get_result.documents,
metadatas=get_result.metadatas,
# qdrant distance is [-1, 1], normalize to [0, 1]
distances=[
[(point.score + 1.0) / 2.0 for point in query_response.points]
],
)
except (UnexpectedResponse, grpc.RpcError) as e:
if self._is_collection_not_found_error(e):
log.debug(
f"Collection {mt_collection} doesn't exist, search returns None"
)
return None
else:
# For other API errors, log and re-raise
_, error_msg = self._extract_error_message(e)
log.warning(f"Unexpected Qdrant error during search: {error_msg}")
raise
except Exception as e:
# For non-Qdrant exceptions, log and return None
log.exception(f"Error searching collection '{collection_name}': {e}")
if not self.client.collection_exists(collection_name=mt_collection):
log.debug(f"Collection {mt_collection} doesn't exist, search returns None")
return None
def query(self, collection_name: str, filter: dict, limit: Optional[int] = None):
tenant_filter = _tenant_filter(tenant_id)
query_response = self.client.query_points(
collection_name=mt_collection,
query=vectors[0],
limit=limit,
query_filter=models.Filter(must=[tenant_filter]),
)
get_result = self._result_to_get_result(query_response.points)
return SearchResult(
ids=get_result.ids,
documents=get_result.documents,
metadatas=get_result.metadatas,
distances=[[(point.score + 1.0) / 2.0 for point in query_response.points]],
)
def query(
self, collection_name: str, filter: Dict[str, Any], limit: Optional[int] = None
):
"""
Query points with filters and tenant isolation.
"""
if not self.client:
return None
# Map to multi-tenant collection and tenant ID
mt_collection, tenant_id = self._get_collection_and_tenant_id(collection_name)
# Set default limit if not provided
if not self.client.collection_exists(collection_name=mt_collection):
log.debug(f"Collection {mt_collection} doesn't exist, query returns None")
return None
if limit is None:
limit = NO_LIMIT
# Create tenant filter
tenant_filter = models.FieldCondition(
key="tenant_id", match=models.MatchValue(value=tenant_id)
)
# Create metadata filters
field_conditions = []
for key, value in filter.items():
field_conditions.append(
models.FieldCondition(
key=f"metadata.{key}", match=models.MatchValue(value=value)
)
)
# Combine tenant filter with metadata filters
tenant_filter = _tenant_filter(tenant_id)
field_conditions = [_metadata_filter(k, v) for k, v in filter.items()]
combined_filter = models.Filter(must=[tenant_filter, *field_conditions])
try:
# Try the query directly - most of the time collection should exist
points = self.client.query_points(
collection_name=mt_collection,
query_filter=combined_filter,
limit=limit,
)
return self._result_to_get_result(points.points)
except (UnexpectedResponse, grpc.RpcError) as e:
if self._is_collection_not_found_error(e):
log.debug(
f"Collection {mt_collection} doesn't exist, query returns None"
)
return None
else:
# For other API errors, log and re-raise
_, error_msg = self._extract_error_message(e)
log.warning(f"Unexpected Qdrant error during query: {error_msg}")
raise
except Exception as e:
# For non-Qdrant exceptions, log and re-raise
log.exception(f"Error querying collection '{collection_name}': {e}")
return None
points = self.client.query_points(
collection_name=mt_collection,
query_filter=combined_filter,
limit=limit,
)
return self._result_to_get_result(points.points)
def get(self, collection_name: str) -> Optional[GetResult]:
"""
@@ -507,169 +291,36 @@ class QdrantClient(VectorDBBase):
"""
if not self.client:
return None
# Map to multi-tenant collection and tenant ID
mt_collection, tenant_id = self._get_collection_and_tenant_id(collection_name)
# Create tenant filter
tenant_filter = models.FieldCondition(
key="tenant_id", match=models.MatchValue(value=tenant_id)
)
try:
# Try to get points directly - most of the time collection should exist
points = self.client.query_points(
collection_name=mt_collection,
query_filter=models.Filter(must=[tenant_filter]),
limit=NO_LIMIT,
)
return self._result_to_get_result(points.points)
except (UnexpectedResponse, grpc.RpcError) as e:
if self._is_collection_not_found_error(e):
log.debug(f"Collection {mt_collection} doesn't exist, get returns None")
return None
else:
# For other API errors, log and re-raise
_, error_msg = self._extract_error_message(e)
log.warning(f"Unexpected Qdrant error during get: {error_msg}")
raise
except Exception as e:
# For non-Qdrant exceptions, log and return None
log.exception(f"Error getting collection '{collection_name}': {e}")
if not self.client.collection_exists(collection_name=mt_collection):
log.debug(f"Collection {mt_collection} doesn't exist, get returns None")
return None
def _handle_operation_with_error_retry(
self, operation_name, mt_collection, points, dimension
):
"""
Private helper to handle common error cases for insert and upsert operations.
Args:
operation_name: 'insert' or 'upsert'
mt_collection: The multi-tenant collection name
points: The vector points to insert/upsert
dimension: The dimension of the vectors
Returns:
The operation result (for upsert) or None (for insert)
"""
try:
if operation_name == "insert":
self.client.upload_points(mt_collection, points)
return None
else: # upsert
return self.client.upsert(mt_collection, points)
except (UnexpectedResponse, grpc.RpcError) as e:
# Handle collection not found
if self._is_collection_not_found_error(e):
log.info(
f"Collection {mt_collection} doesn't exist. Creating it with dimension {dimension}."
)
# Create collection with correct dimensions from our vectors
self._create_multi_tenant_collection_if_not_exists(
mt_collection_name=mt_collection, dimension=dimension
)
# Try operation again - no need for dimension adjustment since we just created with correct dimensions
if operation_name == "insert":
self.client.upload_points(mt_collection, points)
return None
else: # upsert
return self.client.upsert(mt_collection, points)
# Handle dimension mismatch
elif self._is_dimension_mismatch_error(e):
# For dimension errors, the collection must exist, so get its configuration
mt_collection_info = self.client.get_collection(mt_collection)
existing_size = mt_collection_info.config.params.vectors.size
log.info(
f"Dimension mismatch: Collection {mt_collection} expects {existing_size}, got {dimension}"
)
if existing_size < dimension:
# Truncate vectors to fit
log.info(
f"Truncating vectors from {dimension} to {existing_size} dimensions"
)
points = [
PointStruct(
id=point.id,
vector=point.vector[:existing_size],
payload=point.payload,
)
for point in points
]
elif existing_size > dimension:
# Pad vectors with zeros
log.info(
f"Padding vectors from {dimension} to {existing_size} dimensions with zeros"
)
points = [
PointStruct(
id=point.id,
vector=point.vector
+ [0] * (existing_size - len(point.vector)),
payload=point.payload,
)
for point in points
]
# Try operation again with adjusted dimensions
if operation_name == "insert":
self.client.upload_points(mt_collection, points)
return None
else: # upsert
return self.client.upsert(mt_collection, points)
else:
# Not a known error we can handle, log and re-raise
_, error_msg = self._extract_error_message(e)
log.warning(f"Unhandled Qdrant error: {error_msg}")
raise
except Exception as e:
# For non-Qdrant exceptions, re-raise
raise
def insert(self, collection_name: str, items: list[VectorItem]):
"""
Insert items with tenant ID.
"""
if not self.client or not items:
return None
# Map to multi-tenant collection and tenant ID
mt_collection, tenant_id = self._get_collection_and_tenant_id(collection_name)
# Get dimensions from the actual vectors
dimension = len(items[0]["vector"]) if items else None
# Create points with tenant ID
points = self._create_points(items, tenant_id)
# Handle the operation with error retry
return self._handle_operation_with_error_retry(
"insert", mt_collection, points, dimension
tenant_filter = _tenant_filter(tenant_id)
points = self.client.query_points(
collection_name=mt_collection,
query_filter=models.Filter(must=[tenant_filter]),
limit=NO_LIMIT,
)
return self._result_to_get_result(points.points)
def upsert(self, collection_name: str, items: list[VectorItem]):
def upsert(self, collection_name: str, items: List[VectorItem]):
"""
Upsert items with tenant ID.
"""
if not self.client or not items:
return None
# Map to multi-tenant collection and tenant ID
mt_collection, tenant_id = self._get_collection_and_tenant_id(collection_name)
# Get dimensions from the actual vectors
dimension = len(items[0]["vector"]) if items else None
# Create points with tenant ID
dimension = len(items[0]["vector"])
self._ensure_collection(mt_collection, dimension)
points = self._create_points(items, tenant_id)
self.client.upload_points(mt_collection, points)
return None
# Handle the operation with error retry
return self._handle_operation_with_error_retry(
"upsert", mt_collection, points, dimension
)
def insert(self, collection_name: str, items: List[VectorItem]):
"""
Insert items with tenant ID.
"""
return self.upsert(collection_name, items)
def reset(self):
"""
@@ -677,11 +328,9 @@ class QdrantClient(VectorDBBase):
"""
if not self.client:
return None
collection_names = self.client.get_collections().collections
for collection_name in collection_names:
if collection_name.name.startswith(self.collection_prefix):
self.client.delete_collection(collection_name=collection_name.name)
for collection in self.client.get_collections().collections:
if collection.name.startswith(self.collection_prefix):
self.client.delete_collection(collection_name=collection.name)
def delete_collection(self, collection_name: str):
"""
@@ -689,24 +338,13 @@ class QdrantClient(VectorDBBase):
"""
if not self.client:
return None
# Map to multi-tenant collection and tenant ID
mt_collection, tenant_id = self._get_collection_and_tenant_id(collection_name)
tenant_filter = models.FieldCondition(
key="tenant_id", match=models.MatchValue(value=tenant_id)
)
field_conditions = [tenant_filter]
update_result = self.client.delete(
if not self.client.collection_exists(collection_name=mt_collection):
log.debug(f"Collection {mt_collection} doesn't exist, nothing to delete")
return None
self.client.delete(
collection_name=mt_collection,
points_selector=models.FilterSelector(
filter=models.Filter(must=field_conditions)
filter=models.Filter(must=[_tenant_filter(tenant_id)])
),
)
if self.client.get_collection(mt_collection).points_count == 0:
self.client.delete_collection(mt_collection)
return update_result
@@ -0,0 +1,745 @@
from open_webui.retrieval.vector.main import (
VectorDBBase,
VectorItem,
GetResult,
SearchResult,
)
from open_webui.config import S3_VECTOR_BUCKET_NAME, S3_VECTOR_REGION
from open_webui.env import SRC_LOG_LEVELS
from typing import List, Optional, Dict, Any, Union
import logging
import boto3
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])
class S3VectorClient(VectorDBBase):
"""
AWS S3 Vector integration for Open WebUI Knowledge.
"""
def __init__(self):
self.bucket_name = S3_VECTOR_BUCKET_NAME
self.region = S3_VECTOR_REGION
# Simple validation - log warnings instead of raising exceptions
if not self.bucket_name:
log.warning("S3_VECTOR_BUCKET_NAME not set - S3Vector will not work")
if not self.region:
log.warning("S3_VECTOR_REGION not set - S3Vector will not work")
if self.bucket_name and self.region:
try:
self.client = boto3.client("s3vectors", region_name=self.region)
log.info(
f"S3Vector client initialized for bucket '{self.bucket_name}' in region '{self.region}'"
)
except Exception as e:
log.error(f"Failed to initialize S3Vector client: {e}")
self.client = None
else:
self.client = None
def _create_index(
self,
index_name: str,
dimension: int,
data_type: str = "float32",
distance_metric: str = "cosine",
) -> None:
"""
Create a new index in the S3 vector bucket for the given collection if it does not exist.
"""
if self.has_collection(index_name):
log.debug(f"Index '{index_name}' already exists, skipping creation")
return
try:
self.client.create_index(
vectorBucketName=self.bucket_name,
indexName=index_name,
dataType=data_type,
dimension=dimension,
distanceMetric=distance_metric,
)
log.info(
f"Created S3 index: {index_name} (dim={dimension}, type={data_type}, metric={distance_metric})"
)
except Exception as e:
log.error(f"Error creating S3 index '{index_name}': {e}")
raise
def _filter_metadata(
self, metadata: Dict[str, Any], item_id: str
) -> Dict[str, Any]:
"""
Filter vector metadata keys to comply with S3 Vector API limit of 10 keys maximum.
"""
if not isinstance(metadata, dict) or len(metadata) <= 10:
return metadata
# Keep only the first 10 keys, prioritizing important ones based on actual Open WebUI metadata
important_keys = [
"text", # The actual document content
"file_id", # File ID
"source", # Document source file
"title", # Document title
"page", # Page number
"total_pages", # Total pages in document
"embedding_config", # Embedding configuration
"created_by", # User who created it
"name", # Document name
"hash", # Content hash
]
filtered_metadata = {}
# First, add important keys if they exist
for key in important_keys:
if key in metadata:
filtered_metadata[key] = metadata[key]
if len(filtered_metadata) >= 10:
break
# If we still have room, add other keys
if len(filtered_metadata) < 10:
for key, value in metadata.items():
if key not in filtered_metadata:
filtered_metadata[key] = value
if len(filtered_metadata) >= 10:
break
log.warning(
f"Metadata for key '{item_id}' had {len(metadata)} keys, limited to 10 keys"
)
return filtered_metadata
def has_collection(self, collection_name: str) -> bool:
"""
Check if a vector index (collection) exists in the S3 vector bucket.
"""
try:
response = self.client.list_indexes(vectorBucketName=self.bucket_name)
indexes = response.get("indexes", [])
return any(idx.get("indexName") == collection_name for idx in indexes)
except Exception as e:
log.error(f"Error listing indexes: {e}")
return False
def delete_collection(self, collection_name: str) -> None:
"""
Delete an entire S3 Vector index/collection.
"""
if not self.has_collection(collection_name):
log.warning(
f"Collection '{collection_name}' does not exist, nothing to delete"
)
return
try:
log.info(f"Deleting collection '{collection_name}'")
self.client.delete_index(
vectorBucketName=self.bucket_name, indexName=collection_name
)
log.info(f"Successfully deleted collection '{collection_name}'")
except Exception as e:
log.error(f"Error deleting collection '{collection_name}': {e}")
raise
def insert(self, collection_name: str, items: List[VectorItem]) -> None:
"""
Insert vector items into the S3 Vector index. Create index if it does not exist.
"""
if not items:
log.warning("No items to insert")
return
dimension = len(items[0]["vector"])
try:
if not self.has_collection(collection_name):
log.info(f"Index '{collection_name}' does not exist. Creating index.")
self._create_index(
index_name=collection_name,
dimension=dimension,
data_type="float32",
distance_metric="cosine",
)
# Prepare vectors for insertion
vectors = []
for item in items:
# Ensure vector data is in the correct format for S3 Vector API
vector_data = item["vector"]
if isinstance(vector_data, list):
# Convert list to float32 values as required by S3 Vector API
vector_data = [float(x) for x in vector_data]
# Prepare metadata, ensuring the text field is preserved
metadata = item.get("metadata", {}).copy()
# Add the text field to metadata so it's available for retrieval
metadata["text"] = item["text"]
# Filter metadata to comply with S3 Vector API limit of 10 keys
metadata = self._filter_metadata(metadata, item["id"])
vectors.append(
{
"key": item["id"],
"data": {"float32": vector_data},
"metadata": metadata,
}
)
# Insert vectors
self.client.put_vectors(
vectorBucketName=self.bucket_name,
indexName=collection_name,
vectors=vectors,
)
log.info(f"Inserted {len(vectors)} vectors into index '{collection_name}'.")
except Exception as e:
log.error(f"Error inserting vectors: {e}")
raise
def upsert(self, collection_name: str, items: List[VectorItem]) -> None:
"""
Insert or update vector items in the S3 Vector index. Create index if it does not exist.
"""
if not items:
log.warning("No items to upsert")
return
dimension = len(items[0]["vector"])
log.info(f"Upsert dimension: {dimension}")
try:
if not self.has_collection(collection_name):
log.info(
f"Index '{collection_name}' does not exist. Creating index for upsert."
)
self._create_index(
index_name=collection_name,
dimension=dimension,
data_type="float32",
distance_metric="cosine",
)
# Prepare vectors for upsert
vectors = []
for item in items:
# Ensure vector data is in the correct format for S3 Vector API
vector_data = item["vector"]
if isinstance(vector_data, list):
# Convert list to float32 values as required by S3 Vector API
vector_data = [float(x) for x in vector_data]
# Prepare metadata, ensuring the text field is preserved
metadata = item.get("metadata", {}).copy()
# Add the text field to metadata so it's available for retrieval
metadata["text"] = item["text"]
# Filter metadata to comply with S3 Vector API limit of 10 keys
metadata = self._filter_metadata(metadata, item["id"])
vectors.append(
{
"key": item["id"],
"data": {"float32": vector_data},
"metadata": metadata,
}
)
# Upsert vectors (using put_vectors for upsert semantics)
log.info(
f"Upserting {len(vectors)} vectors. First vector sample: key={vectors[0]['key']}, data_type={type(vectors[0]['data']['float32'])}, data_len={len(vectors[0]['data']['float32'])}"
)
self.client.put_vectors(
vectorBucketName=self.bucket_name,
indexName=collection_name,
vectors=vectors,
)
log.info(f"Upserted {len(vectors)} vectors into index '{collection_name}'.")
except Exception as e:
log.error(f"Error upserting vectors: {e}")
raise
def search(
self, collection_name: str, vectors: List[List[Union[float, int]]], limit: int
) -> Optional[SearchResult]:
"""
Search for similar vectors in a collection using multiple query vectors.
"""
if not self.has_collection(collection_name):
log.warning(f"Collection '{collection_name}' does not exist")
return None
if not vectors:
log.warning("No query vectors provided")
return None
try:
log.info(
f"Searching collection '{collection_name}' with {len(vectors)} query vectors, limit={limit}"
)
# Initialize result lists
all_ids = []
all_documents = []
all_metadatas = []
all_distances = []
# Process each query vector
for i, query_vector in enumerate(vectors):
log.debug(f"Processing query vector {i+1}/{len(vectors)}")
# Prepare the query vector in S3 Vector format
query_vector_dict = {"float32": [float(x) for x in query_vector]}
# Call S3 Vector query API
response = self.client.query_vectors(
vectorBucketName=self.bucket_name,
indexName=collection_name,
topK=limit,
queryVector=query_vector_dict,
returnMetadata=True,
returnDistance=True,
)
# Process results for this query
query_ids = []
query_documents = []
query_metadatas = []
query_distances = []
result_vectors = response.get("vectors", [])
for vector in result_vectors:
vector_id = vector.get("key")
vector_metadata = vector.get("metadata", {})
vector_distance = vector.get("distance", 0.0)
# Extract document text from metadata
document_text = ""
if isinstance(vector_metadata, dict):
# Get the text field first (highest priority)
document_text = vector_metadata.get("text")
if not document_text:
# Fallback to other possible text fields
document_text = (
vector_metadata.get("content")
or vector_metadata.get("document")
or vector_id
)
else:
document_text = vector_id
query_ids.append(vector_id)
query_documents.append(document_text)
query_metadatas.append(vector_metadata)
query_distances.append(vector_distance)
# Add this query's results to the overall results
all_ids.append(query_ids)
all_documents.append(query_documents)
all_metadatas.append(query_metadatas)
all_distances.append(query_distances)
log.info(f"Search completed. Found results for {len(all_ids)} queries")
# Return SearchResult format
return SearchResult(
ids=all_ids if all_ids else None,
documents=all_documents if all_documents else None,
metadatas=all_metadatas if all_metadatas else None,
distances=all_distances if all_distances else None,
)
except Exception as e:
log.error(f"Error searching collection '{collection_name}': {str(e)}")
# Handle specific AWS exceptions
if hasattr(e, "response") and "Error" in e.response:
error_code = e.response["Error"]["Code"]
if error_code == "NotFoundException":
log.warning(f"Collection '{collection_name}' not found")
return None
elif error_code == "ValidationException":
log.error(f"Invalid query vector dimensions or parameters")
return None
elif error_code == "AccessDeniedException":
log.error(
f"Access denied for collection '{collection_name}'. Check permissions."
)
return None
raise
def query(
self, collection_name: str, filter: Dict, limit: Optional[int] = None
) -> Optional[GetResult]:
"""
Query vectors from a collection using metadata filter.
"""
if not self.has_collection(collection_name):
log.warning(f"Collection '{collection_name}' does not exist")
return GetResult(ids=[[]], documents=[[]], metadatas=[[]])
if not filter:
log.warning("No filter provided, returning all vectors")
return self.get(collection_name)
try:
log.info(f"Querying collection '{collection_name}' with filter: {filter}")
# For S3 Vector, we need to use list_vectors and then filter results
# Since S3 Vector may not support complex server-side filtering,
# we'll retrieve all vectors and filter client-side
# Get all vectors first
all_vectors_result = self.get(collection_name)
if not all_vectors_result or not all_vectors_result.ids:
log.warning("No vectors found in collection")
return GetResult(ids=[[]], documents=[[]], metadatas=[[]])
# Extract the lists from the result
all_ids = all_vectors_result.ids[0] if all_vectors_result.ids else []
all_documents = (
all_vectors_result.documents[0] if all_vectors_result.documents else []
)
all_metadatas = (
all_vectors_result.metadatas[0] if all_vectors_result.metadatas else []
)
# Apply client-side filtering
filtered_ids = []
filtered_documents = []
filtered_metadatas = []
for i, metadata in enumerate(all_metadatas):
if self._matches_filter(metadata, filter):
if i < len(all_ids):
filtered_ids.append(all_ids[i])
if i < len(all_documents):
filtered_documents.append(all_documents[i])
filtered_metadatas.append(metadata)
# Apply limit if specified
if limit and len(filtered_ids) >= limit:
break
log.info(
f"Filter applied: {len(filtered_ids)} vectors match out of {len(all_ids)} total"
)
# Return GetResult format
if filtered_ids:
return GetResult(
ids=[filtered_ids],
documents=[filtered_documents],
metadatas=[filtered_metadatas],
)
else:
return GetResult(ids=[[]], documents=[[]], metadatas=[[]])
except Exception as e:
log.error(f"Error querying collection '{collection_name}': {str(e)}")
# Handle specific AWS exceptions
if hasattr(e, "response") and "Error" in e.response:
error_code = e.response["Error"]["Code"]
if error_code == "NotFoundException":
log.warning(f"Collection '{collection_name}' not found")
return GetResult(ids=[[]], documents=[[]], metadatas=[[]])
elif error_code == "AccessDeniedException":
log.error(
f"Access denied for collection '{collection_name}'. Check permissions."
)
return GetResult(ids=[[]], documents=[[]], metadatas=[[]])
raise
def get(self, collection_name: str) -> Optional[GetResult]:
"""
Retrieve all vectors from a collection.
"""
if not self.has_collection(collection_name):
log.warning(f"Collection '{collection_name}' does not exist")
return GetResult(ids=[[]], documents=[[]], metadatas=[[]])
try:
log.info(f"Retrieving all vectors from collection '{collection_name}'")
# Initialize result lists
all_ids = []
all_documents = []
all_metadatas = []
# Handle pagination
next_token = None
while True:
# Prepare request parameters
request_params = {
"vectorBucketName": self.bucket_name,
"indexName": collection_name,
"returnData": False, # Don't include vector data (not needed for get)
"returnMetadata": True, # Include metadata
"maxResults": 500, # Use reasonable page size
}
if next_token:
request_params["nextToken"] = next_token
# Call S3 Vector API
response = self.client.list_vectors(**request_params)
# Process vectors in this page
vectors = response.get("vectors", [])
for vector in vectors:
vector_id = vector.get("key")
vector_data = vector.get("data", {})
vector_metadata = vector.get("metadata", {})
# Extract the actual vector array
vector_array = vector_data.get("float32", [])
# For documents, we try to extract text from metadata or use the vector ID
document_text = ""
if isinstance(vector_metadata, dict):
# Get the text field first (highest priority)
document_text = vector_metadata.get("text")
if not document_text:
# Fallback to other possible text fields
document_text = (
vector_metadata.get("content")
or vector_metadata.get("document")
or vector_id
)
# Log the actual content for debugging
log.debug(
f"Document text preview (first 200 chars): {str(document_text)[:200]}"
)
else:
document_text = vector_id
all_ids.append(vector_id)
all_documents.append(document_text)
all_metadatas.append(vector_metadata)
# Check if there are more pages
next_token = response.get("nextToken")
if not next_token:
break
log.info(
f"Retrieved {len(all_ids)} vectors from collection '{collection_name}'"
)
# Return in GetResult format
# The Open WebUI GetResult expects lists of lists, so we wrap each list
if all_ids:
return GetResult(
ids=[all_ids], documents=[all_documents], metadatas=[all_metadatas]
)
else:
return GetResult(ids=[[]], documents=[[]], metadatas=[[]])
except Exception as e:
log.error(
f"Error retrieving vectors from collection '{collection_name}': {str(e)}"
)
# Handle specific AWS exceptions
if hasattr(e, "response") and "Error" in e.response:
error_code = e.response["Error"]["Code"]
if error_code == "NotFoundException":
log.warning(f"Collection '{collection_name}' not found")
return GetResult(ids=[[]], documents=[[]], metadatas=[[]])
elif error_code == "AccessDeniedException":
log.error(
f"Access denied for collection '{collection_name}'. Check permissions."
)
return GetResult(ids=[[]], documents=[[]], metadatas=[[]])
raise
def delete(
self,
collection_name: str,
ids: Optional[List[str]] = None,
filter: Optional[Dict] = None,
) -> None:
"""
Delete vectors by ID or filter from a collection.
"""
if not self.has_collection(collection_name):
log.warning(
f"Collection '{collection_name}' does not exist, nothing to delete"
)
return
# Check if this is a knowledge collection (not file-specific)
is_knowledge_collection = not collection_name.startswith("file-")
try:
if ids:
# Delete by specific vector IDs/keys
log.info(
f"Deleting {len(ids)} vectors by IDs from collection '{collection_name}'"
)
self.client.delete_vectors(
vectorBucketName=self.bucket_name,
indexName=collection_name,
keys=ids,
)
log.info(f"Deleted {len(ids)} vectors from index '{collection_name}'")
elif filter:
# Handle filter-based deletion
log.info(
f"Deleting vectors by filter from collection '{collection_name}': {filter}"
)
# If this is a knowledge collection and we have a file_id filter,
# also clean up the corresponding file-specific collection
if is_knowledge_collection and "file_id" in filter:
file_id = filter["file_id"]
file_collection_name = f"file-{file_id}"
if self.has_collection(file_collection_name):
log.info(
f"Found related file-specific collection '{file_collection_name}', deleting it to prevent duplicates"
)
self.delete_collection(file_collection_name)
# For the main collection, implement query-then-delete
# First, query to get IDs matching the filter
query_result = self.query(collection_name, filter)
if query_result and query_result.ids and query_result.ids[0]:
matching_ids = query_result.ids[0]
log.info(
f"Found {len(matching_ids)} vectors matching filter, deleting them"
)
# Delete the matching vectors by ID
self.client.delete_vectors(
vectorBucketName=self.bucket_name,
indexName=collection_name,
keys=matching_ids,
)
log.info(
f"Deleted {len(matching_ids)} vectors from index '{collection_name}' using filter"
)
else:
log.warning("No vectors found matching the filter criteria")
else:
log.warning("No IDs or filter provided for deletion")
except Exception as e:
log.error(
f"Error deleting vectors from collection '{collection_name}': {e}"
)
raise
def reset(self) -> None:
"""
Reset/clear all vector data. For S3 Vector, this deletes all indexes.
"""
try:
log.warning(
"Reset called - this will delete all vector indexes in the S3 bucket"
)
# List all indexes
response = self.client.list_indexes(vectorBucketName=self.bucket_name)
indexes = response.get("indexes", [])
if not indexes:
log.warning("No indexes found to delete")
return
# Delete all indexes
deleted_count = 0
for index in indexes:
index_name = index.get("indexName")
if index_name:
try:
self.client.delete_index(
vectorBucketName=self.bucket_name, indexName=index_name
)
deleted_count += 1
log.info(f"Deleted index: {index_name}")
except Exception as e:
log.error(f"Error deleting index '{index_name}': {e}")
log.info(f"Reset completed: deleted {deleted_count} indexes")
except Exception as e:
log.error(f"Error during reset: {e}")
raise
def _matches_filter(self, metadata: Dict[str, Any], filter: Dict[str, Any]) -> bool:
"""
Check if metadata matches the given filter conditions.
"""
if not isinstance(metadata, dict) or not isinstance(filter, dict):
return False
# Check each filter condition
for key, expected_value in filter.items():
# Handle special operators
if key.startswith("$"):
if key == "$and":
# All conditions must match
if not isinstance(expected_value, list):
continue
for condition in expected_value:
if not self._matches_filter(metadata, condition):
return False
elif key == "$or":
# At least one condition must match
if not isinstance(expected_value, list):
continue
any_match = False
for condition in expected_value:
if self._matches_filter(metadata, condition):
any_match = True
break
if not any_match:
return False
continue
# Get the actual value from metadata
actual_value = metadata.get(key)
# Handle different types of expected values
if isinstance(expected_value, dict):
# Handle comparison operators
for op, op_value in expected_value.items():
if op == "$eq":
if actual_value != op_value:
return False
elif op == "$ne":
if actual_value == op_value:
return False
elif op == "$in":
if (
not isinstance(op_value, list)
or actual_value not in op_value
):
return False
elif op == "$nin":
if isinstance(op_value, list) and actual_value in op_value:
return False
elif op == "$exists":
if bool(op_value) != (key in metadata):
return False
# Add more operators as needed
else:
# Simple equality check
if actual_value != expected_value:
return False
return True
@@ -30,6 +30,10 @@ class Vector:
from open_webui.retrieval.vector.dbs.pinecone import PineconeClient
return PineconeClient()
case VectorType.S3VECTOR:
from open_webui.retrieval.vector.dbs.s3vector import S3VectorClient
return S3VectorClient()
case VectorType.OPENSEARCH:
from open_webui.retrieval.vector.dbs.opensearch import OpenSearchClient
@@ -0,0 +1,14 @@
from datetime import datetime
def stringify_metadata(
metadata: dict[str, any],
) -> dict[str, any]:
for key, value in metadata.items():
if (
isinstance(value, datetime)
or isinstance(value, list)
or isinstance(value, dict)
):
metadata[key] = str(value)
return metadata
+3 -1
View File
@@ -36,7 +36,9 @@ def search_brave(
return [
SearchResult(
link=result["url"], title=result.get("title"), snippet=result.get("snippet")
link=result["url"],
title=result.get("title"),
snippet=result.get("description"),
)
for result in results[:count]
]
@@ -2,8 +2,8 @@ import logging
from typing import Optional
from open_webui.retrieval.web.main import SearchResult, get_filtered_results
from duckduckgo_search import DDGS
from duckduckgo_search.exceptions import RatelimitException
from ddgs import DDGS
from ddgs.exceptions import RatelimitException
from open_webui.env import SRC_LOG_LEVELS
log = logging.getLogger(__name__)
+37 -22
View File
@@ -15,6 +15,7 @@ import aiohttp
import aiofiles
import requests
import mimetypes
from urllib.parse import quote
from fastapi import (
Depends,
@@ -327,6 +328,7 @@ async def speech(request: Request, user=Depends(get_verified_user)):
log.exception(e)
raise HTTPException(status_code=400, detail="Invalid JSON payload")
r = None
if request.app.state.config.TTS_ENGINE == "openai":
payload["model"] = request.app.state.config.TTS_MODEL
@@ -335,7 +337,7 @@ async def speech(request: Request, user=Depends(get_verified_user)):
async with aiohttp.ClientSession(
timeout=timeout, trust_env=True
) as session:
async with session.post(
r = await session.post(
url=f"{request.app.state.config.TTS_OPENAI_API_BASE_URL}/audio/speech",
json=payload,
headers={
@@ -343,7 +345,7 @@ async def speech(request: Request, user=Depends(get_verified_user)):
"Authorization": f"Bearer {request.app.state.config.TTS_OPENAI_API_KEY}",
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -353,14 +355,15 @@ async def speech(request: Request, user=Depends(get_verified_user)):
),
},
ssl=AIOHTTP_CLIENT_SESSION_SSL,
) as r:
r.raise_for_status()
)
async with aiofiles.open(file_path, "wb") as f:
await f.write(await r.read())
r.raise_for_status()
async with aiofiles.open(file_body_path, "w") as f:
await f.write(json.dumps(payload))
async with aiofiles.open(file_path, "wb") as f:
await f.write(await r.read())
async with aiofiles.open(file_body_path, "w") as f:
await f.write(json.dumps(payload))
return FileResponse(file_path)
@@ -368,18 +371,22 @@ async def speech(request: Request, user=Depends(get_verified_user)):
log.exception(e)
detail = None
try:
if r.status != 200:
res = await r.json()
status_code = 500
detail = f"Open WebUI: Server Connection Error"
if r is not None:
status_code = r.status
try:
res = await r.json()
if "error" in res:
detail = f"External: {res['error'].get('message', '')}"
except Exception:
detail = f"External: {e}"
detail = f"External: {res['error']}"
except Exception:
detail = f"External: {e}"
raise HTTPException(
status_code=getattr(r, "status", 500) if r else 500,
detail=detail if detail else "Open WebUI: Server Connection Error",
status_code=status_code,
detail=detail,
)
elif request.app.state.config.TTS_ENGINE == "elevenlabs":
@@ -554,7 +561,11 @@ def transcription_handler(request, file_path, metadata):
file_path,
beam_size=5,
vad_filter=request.app.state.config.WHISPER_VAD_FILTER,
language=metadata.get("language") or WHISPER_LANGUAGE,
language=(
metadata.get("language", None)
if WHISPER_LANGUAGE == ""
else WHISPER_LANGUAGE
),
)
log.info(
"Detected language '%s' with probability %f"
@@ -919,14 +930,18 @@ def transcription(
):
log.info(f"file.content_type: {file.content_type}")
supported_content_types = request.app.state.config.STT_SUPPORTED_CONTENT_TYPES or [
"audio/*",
"video/webm",
]
stt_supported_content_types = getattr(
request.app.state.config, "STT_SUPPORTED_CONTENT_TYPES", []
)
if not any(
fnmatch(file.content_type, content_type)
for content_type in supported_content_types
for content_type in (
stt_supported_content_types
if stt_supported_content_types
and any(t.strip() for t in stt_supported_content_types)
else ["audio/*", "video/webm"]
)
):
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
+8 -2
View File
@@ -669,12 +669,13 @@ async def signup(request: Request, response: Response, form_data: SignupForm):
@router.get("/signout")
async def signout(request: Request, response: Response):
response.delete_cookie("token")
response.delete_cookie("oui-session")
if ENABLE_OAUTH_SIGNUP.value:
oauth_id_token = request.cookies.get("oauth_id_token")
if oauth_id_token:
try:
async with ClientSession() as session:
async with ClientSession(trust_env=True) as session:
async with session.get(OPENID_PROVIDER_URL.value) as resp:
if resp.status == 200:
openid_data = await resp.json()
@@ -686,7 +687,12 @@ async def signout(request: Request, response: Response):
status_code=200,
content={
"status": True,
"redirect_url": f"{logout_url}?id_token_hint={oauth_id_token}",
"redirect_url": f"{logout_url}?id_token_hint={oauth_id_token}"
+ (
f"&post_logout_redirect_uri={WEBUI_AUTH_SIGNOUT_REDIRECT_URL}"
if WEBUI_AUTH_SIGNOUT_REDIRECT_URL
else ""
),
},
headers=response.headers,
)
+24 -16
View File
@@ -40,10 +40,14 @@ router = APIRouter()
@router.get("/", response_model=list[ChannelModel])
async def get_channels(user=Depends(get_verified_user)):
return Channels.get_channels_by_user_id(user.id)
@router.get("/list", response_model=list[ChannelModel])
async def get_all_channels(user=Depends(get_verified_user)):
if user.role == "admin":
return Channels.get_channels()
else:
return Channels.get_channels_by_user_id(user.id)
return Channels.get_channels_by_user_id(user.id)
############################
@@ -430,13 +434,6 @@ async def update_message_by_id(
status_code=status.HTTP_404_NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND
)
if user.role != "admin" and not has_access(
user.id, type="read", access_control=channel.access_control
):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.DEFAULT()
)
message = Messages.get_message_by_id(message_id)
if not message:
raise HTTPException(
@@ -448,6 +445,15 @@ async def update_message_by_id(
status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT()
)
if (
user.role != "admin"
and message.user_id != user.id
and not has_access(user.id, type="read", access_control=channel.access_control)
):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.DEFAULT()
)
try:
message = Messages.update_message_by_id(message_id, form_data)
message = Messages.get_message_by_id(message_id)
@@ -637,13 +643,6 @@ async def delete_message_by_id(
status_code=status.HTTP_404_NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND
)
if user.role != "admin" and not has_access(
user.id, type="read", access_control=channel.access_control
):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.DEFAULT()
)
message = Messages.get_message_by_id(message_id)
if not message:
raise HTTPException(
@@ -655,6 +654,15 @@ async def delete_message_by_id(
status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT()
)
if (
user.role != "admin"
and message.user_id != user.id
and not has_access(user.id, type="read", access_control=channel.access_control)
):
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN, detail=ERROR_MESSAGES.DEFAULT()
)
try:
Messages.delete_message_by_id(message_id)
await sio.emit(
+18 -8
View File
@@ -39,13 +39,21 @@ router = APIRouter()
async def get_session_user_chat_list(
user=Depends(get_verified_user), page: Optional[int] = None
):
if page is not None:
limit = 60
skip = (page - 1) * limit
try:
if page is not None:
limit = 60
skip = (page - 1) * limit
return Chats.get_chat_title_id_list_by_user_id(user.id, skip=skip, limit=limit)
else:
return Chats.get_chat_title_id_list_by_user_id(user.id)
return Chats.get_chat_title_id_list_by_user_id(
user.id, skip=skip, limit=limit
)
else:
return Chats.get_chat_title_id_list_by_user_id(user.id)
except Exception as e:
log.exception(e)
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT()
)
############################
@@ -684,8 +692,10 @@ async def archive_chat_by_id(id: str, user=Depends(get_verified_user)):
@router.post("/{id}/share", response_model=Optional[ChatResponse])
async def share_chat_by_id(request: Request, id: str, user=Depends(get_verified_user)):
if not has_permission(
user.id, "chat.share", request.app.state.config.USER_PERMISSIONS
if (user.role != "admin") and (
not has_permission(
user.id, "chat.share", request.app.state.config.USER_PERMISSIONS
)
):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
+20 -9
View File
@@ -7,7 +7,11 @@ from open_webui.utils.auth import get_admin_user, get_verified_user
from open_webui.config import get_config, save_config
from open_webui.config import BannerModel
from open_webui.utils.tools import get_tool_server_data, get_tool_servers_data
from open_webui.utils.tools import (
get_tool_server_data,
get_tool_servers_data,
get_tool_server_url,
)
router = APIRouter()
@@ -39,32 +43,39 @@ async def export_config(user=Depends(get_admin_user)):
############################
# Direct Connections Config
# Connections Config
############################
class DirectConnectionsConfigForm(BaseModel):
class ConnectionsConfigForm(BaseModel):
ENABLE_DIRECT_CONNECTIONS: bool
ENABLE_BASE_MODELS_CACHE: bool
@router.get("/direct_connections", response_model=DirectConnectionsConfigForm)
async def get_direct_connections_config(request: Request, user=Depends(get_admin_user)):
@router.get("/connections", response_model=ConnectionsConfigForm)
async def get_connections_config(request: Request, user=Depends(get_admin_user)):
return {
"ENABLE_DIRECT_CONNECTIONS": request.app.state.config.ENABLE_DIRECT_CONNECTIONS,
"ENABLE_BASE_MODELS_CACHE": request.app.state.config.ENABLE_BASE_MODELS_CACHE,
}
@router.post("/direct_connections", response_model=DirectConnectionsConfigForm)
async def set_direct_connections_config(
@router.post("/connections", response_model=ConnectionsConfigForm)
async def set_connections_config(
request: Request,
form_data: DirectConnectionsConfigForm,
form_data: ConnectionsConfigForm,
user=Depends(get_admin_user),
):
request.app.state.config.ENABLE_DIRECT_CONNECTIONS = (
form_data.ENABLE_DIRECT_CONNECTIONS
)
request.app.state.config.ENABLE_BASE_MODELS_CACHE = (
form_data.ENABLE_BASE_MODELS_CACHE
)
return {
"ENABLE_DIRECT_CONNECTIONS": request.app.state.config.ENABLE_DIRECT_CONNECTIONS,
"ENABLE_BASE_MODELS_CACHE": request.app.state.config.ENABLE_BASE_MODELS_CACHE,
}
@@ -128,7 +139,7 @@ async def verify_tool_servers_config(
elif form_data.auth_type == "session":
token = request.state.token.credentials
url = f"{form_data.url}/{form_data.path}"
url = get_tool_server_url(form_data.url, form_data.path)
return await get_tool_server_data(token, url)
except Exception as e:
raise HTTPException(
+10 -4
View File
@@ -129,7 +129,10 @@ async def create_feedback(
@router.get("/feedback/{id}", response_model=FeedbackModel)
async def get_feedback_by_id(id: str, user=Depends(get_verified_user)):
feedback = Feedbacks.get_feedback_by_id_and_user_id(id=id, user_id=user.id)
if user.role == "admin":
feedback = Feedbacks.get_feedback_by_id(id=id)
else:
feedback = Feedbacks.get_feedback_by_id_and_user_id(id=id, user_id=user.id)
if not feedback:
raise HTTPException(
@@ -143,9 +146,12 @@ async def get_feedback_by_id(id: str, user=Depends(get_verified_user)):
async def update_feedback_by_id(
id: str, form_data: FeedbackForm, user=Depends(get_verified_user)
):
feedback = Feedbacks.update_feedback_by_id_and_user_id(
id=id, user_id=user.id, form_data=form_data
)
if user.role == "admin":
feedback = Feedbacks.update_feedback_by_id(id=id, form_data=form_data)
else:
feedback = Feedbacks.update_feedback_by_id_and_user_id(
id=id, user_id=user.id, form_data=form_data
)
if not feedback:
raise HTTPException(
+11 -8
View File
@@ -21,6 +21,7 @@ from fastapi import (
from fastapi.responses import FileResponse, StreamingResponse
from open_webui.constants import ERROR_MESSAGES
from open_webui.env import SRC_LOG_LEVELS
from open_webui.retrieval.vector.factory import VECTOR_DB_CLIENT
from open_webui.models.users import Users
from open_webui.models.files import (
@@ -155,17 +156,18 @@ def upload_file(
if process:
try:
if file.content_type:
stt_supported_content_types = (
request.app.state.config.STT_SUPPORTED_CONTENT_TYPES
or [
"audio/*",
"video/webm",
]
stt_supported_content_types = getattr(
request.app.state.config, "STT_SUPPORTED_CONTENT_TYPES", []
)
if any(
fnmatch(file.content_type, content_type)
for content_type in stt_supported_content_types
for content_type in (
stt_supported_content_types
if stt_supported_content_types
and any(t.strip() for t in stt_supported_content_types)
else ["audio/*", "video/webm"]
)
):
file_path = Storage.get_file(file_path)
result = transcribe(request, file_path, file_metadata)
@@ -285,6 +287,7 @@ async def delete_all_files(user=Depends(get_admin_user)):
if result:
try:
Storage.delete_all_files()
VECTOR_DB_CLIENT.reset()
except Exception as e:
log.exception(e)
log.error("Error deleting files")
@@ -602,12 +605,12 @@ async def delete_file_by_id(id: str, user=Depends(get_verified_user)):
or user.role == "admin"
or has_access_to_file(id, "write", user)
):
# We should add Chroma cleanup here
result = Files.delete_file_by_id(id)
if result:
try:
Storage.delete_file(file.path)
VECTOR_DB_CLIENT.delete(collection_name=f"file-{id}")
except Exception as e:
log.exception(e)
log.error("Error deleting files")
+4 -6
View File
@@ -49,7 +49,7 @@ async def get_folders(user=Depends(get_verified_user)):
**folder.model_dump(),
"items": {
"chats": [
{"title": chat.title, "id": chat.id}
{"title": chat.title, "id": chat.id, "updated_at": chat.updated_at}
for chat in Chats.get_chats_by_folder_id_and_user_id(
folder.id, user.id
)
@@ -78,7 +78,7 @@ def create_folder(form_data: FolderForm, user=Depends(get_verified_user)):
)
try:
folder = Folders.insert_new_folder(user.id, form_data.name)
folder = Folders.insert_new_folder(user.id, form_data)
return folder
except Exception as e:
log.exception(e)
@@ -120,16 +120,14 @@ async def update_folder_name_by_id(
existing_folder = Folders.get_folder_by_parent_id_and_user_id_and_name(
folder.parent_id, user.id, form_data.name
)
if existing_folder:
if existing_folder and existing_folder.id != id:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT("Folder already exists"),
)
try:
folder = Folders.update_folder_name_by_id_and_user_id(
id, user.id, form_data.name
)
folder = Folders.update_folder_by_id_and_user_id(id, user.id, form_data)
return folder
except Exception as e:
+18 -4
View File
@@ -105,7 +105,7 @@ async def load_function_from_url(
)
try:
async with aiohttp.ClientSession() as session:
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.get(
url, headers={"Content-Type": "application/json"}
) as resp:
@@ -131,15 +131,29 @@ async def load_function_from_url(
############################
class SyncFunctionsForm(FunctionForm):
class SyncFunctionsForm(BaseModel):
functions: list[FunctionModel] = []
@router.post("/sync", response_model=Optional[FunctionModel])
@router.post("/sync", response_model=list[FunctionModel])
async def sync_functions(
request: Request, form_data: SyncFunctionsForm, user=Depends(get_admin_user)
):
return Functions.sync_functions(user.id, form_data.functions)
try:
for function in form_data.functions:
function.content = replace_imports(function.content)
function_module, function_type, frontmatter = load_function_module_by_id(
function.id,
content=function.content,
)
return Functions.sync_functions(user.id, form_data.functions)
except Exception as e:
log.exception(f"Failed to load a function: {e}")
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT(e),
)
############################
+51
View File
@@ -9,6 +9,7 @@ from open_webui.models.groups import (
GroupForm,
GroupUpdateForm,
GroupResponse,
UserIdsForm,
)
from open_webui.config import CACHE_DIR
@@ -107,6 +108,56 @@ async def update_group_by_id(
)
############################
# AddUserToGroupByUserIdAndGroupId
############################
@router.post("/id/{id}/users/add", response_model=Optional[GroupResponse])
async def add_user_to_group(
id: str, form_data: UserIdsForm, user=Depends(get_admin_user)
):
try:
if form_data.user_ids:
form_data.user_ids = Users.get_valid_user_ids(form_data.user_ids)
group = Groups.add_users_to_group(id, form_data.user_ids)
if group:
return group
else:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT("Error adding users to group"),
)
except Exception as e:
log.exception(f"Error adding users to group {id}: {e}")
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT(e),
)
@router.post("/id/{id}/users/remove", response_model=Optional[GroupResponse])
async def remove_users_from_group(
id: str, form_data: UserIdsForm, user=Depends(get_admin_user)
):
try:
group = Groups.remove_users_from_group(id, form_data.user_ids)
if group:
return group
else:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT("Error removing users from group"),
)
except Exception as e:
log.exception(f"Error removing users from group {id}: {e}")
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.DEFAULT(e),
)
############################
# DeleteGroupById
############################
+19 -3
View File
@@ -8,6 +8,7 @@ import re
from pathlib import Path
from typing import Optional
from urllib.parse import quote
import requests
from fastapi import APIRouter, Depends, HTTPException, Request, UploadFile
from open_webui.config import CACHE_DIR
@@ -302,8 +303,16 @@ async def update_image_config(
):
set_image_model(request, form_data.MODEL)
if form_data.IMAGE_SIZE == "auto" and form_data.MODEL != "gpt-image-1":
raise HTTPException(
status_code=400,
detail=ERROR_MESSAGES.INCORRECT_FORMAT(
" (auto is only allowed with gpt-image-1)."
),
)
pattern = r"^\d+x\d+$"
if re.match(pattern, form_data.IMAGE_SIZE):
if form_data.IMAGE_SIZE == "auto" or re.match(pattern, form_data.IMAGE_SIZE):
request.app.state.config.IMAGE_SIZE = form_data.IMAGE_SIZE
else:
raise HTTPException(
@@ -471,7 +480,14 @@ async def image_generations(
form_data: GenerateImageForm,
user=Depends(get_verified_user),
):
width, height = tuple(map(int, request.app.state.config.IMAGE_SIZE.split("x")))
# if IMAGE_SIZE = 'auto', default WidthxHeight to the 512x512 default
# This is only relevant when the user has set IMAGE_SIZE to 'auto' with an
# image model other than gpt-image-1, which is warned about on settings save
width, height = (
tuple(map(int, request.app.state.config.IMAGE_SIZE.split("x")))
if "x" in request.app.state.config.IMAGE_SIZE
else (512, 512)
)
r = None
try:
@@ -483,7 +499,7 @@ async def image_generations(
headers["Content-Type"] = "application/json"
if ENABLE_FORWARD_USER_INFO_HEADERS:
headers["X-OpenWebUI-User-Name"] = user.name
headers["X-OpenWebUI-User-Name"] = quote(user.name, safe=" ")
headers["X-OpenWebUI-User-Id"] = user.id
headers["X-OpenWebUI-User-Email"] = user.email
headers["X-OpenWebUI-User-Role"] = user.role
+4
View File
@@ -82,6 +82,10 @@ class QueryMemoryForm(BaseModel):
async def query_memory(
request: Request, form_data: QueryMemoryForm, user=Depends(get_verified_user)
):
memories = Memories.get_memories_by_user_id(user.id)
if not memories:
raise HTTPException(status_code=404, detail="No memories found for user")
results = VECTOR_DB_CLIENT.search(
collection_name=f"user-memory-{user.id}",
vectors=[request.app.state.EMBEDDING_FUNCTION(form_data.content, user=user)],
+29 -1
View File
@@ -7,6 +7,8 @@ from open_webui.models.models import (
ModelUserResponse,
Models,
)
from pydantic import BaseModel
from open_webui.constants import ERROR_MESSAGES
from fastapi import APIRouter, Depends, HTTPException, Request, status
@@ -78,6 +80,32 @@ async def create_new_model(
)
############################
# ExportModels
############################
@router.get("/export", response_model=list[ModelModel])
async def export_models(user=Depends(get_admin_user)):
return Models.get_models()
############################
# SyncModels
############################
class SyncModelsForm(BaseModel):
models: list[ModelModel] = []
@router.post("/sync", response_model=list[ModelModel])
async def sync_models(
request: Request, form_data: SyncModelsForm, user=Depends(get_admin_user)
):
return Models.sync_models(user.id, form_data.models)
###########################
# GetModelById
###########################
@@ -102,7 +130,7 @@ async def get_model_by_id(id: str, user=Depends(get_verified_user)):
############################
# ToggelModelById
# ToggleModelById
############################
+19 -8
View File
@@ -6,6 +6,9 @@ from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, Request, status, BackgroundTasks
from pydantic import BaseModel
from open_webui.socket.main import sio
from open_webui.models.users import Users, UserResponse
from open_webui.models.notes import Notes, NoteModel, NoteForm, NoteUserResponse
@@ -51,7 +54,14 @@ async def get_notes(request: Request, user=Depends(get_verified_user)):
return notes
@router.get("/list", response_model=list[NoteUserResponse])
class NoteTitleIdResponse(BaseModel):
id: str
title: str
updated_at: int
created_at: int
@router.get("/list", response_model=list[NoteTitleIdResponse])
async def get_note_list(request: Request, user=Depends(get_verified_user)):
if user.role != "admin" and not has_permission(
@@ -63,13 +73,8 @@ async def get_note_list(request: Request, user=Depends(get_verified_user)):
)
notes = [
NoteUserResponse(
**{
**note.model_dump(),
"user": UserResponse(**Users.get_user_by_id(note.user_id).model_dump()),
}
)
for note in Notes.get_notes_by_user_id(user.id, "read")
NoteTitleIdResponse(**note.model_dump())
for note in Notes.get_notes_by_user_id(user.id, "write")
]
return notes
@@ -168,6 +173,12 @@ async def update_note_by_id(
try:
note = Notes.update_note_by_id(id, form_data)
await sio.emit(
"note-events",
note.model_dump(),
to=f"note:{note.id}",
)
return note
except Exception as e:
log.exception(e)
+60 -27
View File
@@ -16,6 +16,7 @@ from urllib.parse import urlparse
import aiohttp
from aiocache import cached
import requests
from urllib.parse import quote
from open_webui.models.chats import Chats
from open_webui.models.users import UserModel
@@ -58,6 +59,7 @@ from open_webui.config import (
from open_webui.env import (
ENV,
SRC_LOG_LEVELS,
MODELS_CACHE_TTL,
AIOHTTP_CLIENT_SESSION_SSL,
AIOHTTP_CLIENT_TIMEOUT,
AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST,
@@ -87,7 +89,7 @@ async def send_get_request(url, key=None, user: UserModel = None):
**({"Authorization": f"Bearer {key}"} if key else {}),
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -122,6 +124,7 @@ async def send_post_request(
key: Optional[str] = None,
content_type: Optional[str] = None,
user: UserModel = None,
metadata: Optional[dict] = None,
):
r = None
@@ -138,10 +141,15 @@ async def send_post_request(
**({"Authorization": f"Bearer {key}"} if key else {}),
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
**(
{"X-OpenWebUI-Chat-Id": metadata.get("chat_id")}
if metadata and metadata.get("chat_id")
else {}
),
}
if ENABLE_FORWARD_USER_INFO_HEADERS and user
else {}
@@ -182,7 +190,6 @@ async def send_post_request(
)
else:
res = await r.json()
await cleanup_response(r, session)
return res
except HTTPException as e:
@@ -194,6 +201,9 @@ async def send_post_request(
status_code=r.status if r else 500,
detail=detail if e else "Open WebUI: Server Connection Error",
)
finally:
if not stream:
await cleanup_response(r, session)
def get_api_key(idx, url, configs):
@@ -242,7 +252,7 @@ async def verify_connection(
**({"Authorization": f"Bearer {key}"} if key else {}),
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -329,7 +339,7 @@ def merge_ollama_models_lists(model_lists):
return list(merged_models.values())
@cached(ttl=1)
@cached(ttl=MODELS_CACHE_TTL)
async def get_all_models(request: Request, user: UserModel = None):
log.info("get_all_models()")
if request.app.state.config.ENABLE_OLLAMA_API:
@@ -462,7 +472,7 @@ async def get_ollama_tags(
**({"Authorization": f"Bearer {key}"} if key else {}),
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -634,7 +644,10 @@ async def get_ollama_versions(request: Request, url_idx: Optional[int] = None):
class ModelNameForm(BaseModel):
name: str
model: Optional[str] = None
model_config = ConfigDict(
extra="allow",
)
@router.post("/api/unload")
@@ -643,10 +656,12 @@ async def unload_model(
form_data: ModelNameForm,
user=Depends(get_admin_user),
):
model_name = form_data.name
form_data = form_data.model_dump(exclude_none=True)
model_name = form_data.get("model", form_data.get("name"))
if not model_name:
raise HTTPException(
status_code=400, detail="Missing 'name' of model to unload."
status_code=400, detail="Missing name of the model to unload."
)
# Refresh/load models if needed, get mapping from name to URLs
@@ -709,11 +724,14 @@ async def pull_model(
url_idx: int = 0,
user=Depends(get_admin_user),
):
form_data = form_data.model_dump(exclude_none=True)
form_data["model"] = form_data.get("model", form_data.get("name"))
url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
log.info(f"url: {url}")
# Admin should be able to pull models from any source
payload = {**form_data.model_dump(exclude_none=True), "insecure": True}
payload = {**form_data, "insecure": True}
return await send_post_request(
url=f"{url}/api/pull",
@@ -724,7 +742,7 @@ async def pull_model(
class PushModelForm(BaseModel):
name: str
model: str
insecure: Optional[bool] = None
stream: Optional[bool] = None
@@ -741,12 +759,12 @@ async def push_model(
await get_all_models(request, user=user)
models = request.app.state.OLLAMA_MODELS
if form_data.name in models:
url_idx = models[form_data.name]["urls"][0]
if form_data.model in models:
url_idx = models[form_data.model]["urls"][0]
else:
raise HTTPException(
status_code=400,
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.name),
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model),
)
url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
@@ -824,7 +842,7 @@ async def copy_model(
**({"Authorization": f"Bearer {key}"} if key else {}),
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -865,16 +883,21 @@ async def delete_model(
url_idx: Optional[int] = None,
user=Depends(get_admin_user),
):
form_data = form_data.model_dump(exclude_none=True)
form_data["model"] = form_data.get("model", form_data.get("name"))
model = form_data.get("model")
if url_idx is None:
await get_all_models(request, user=user)
models = request.app.state.OLLAMA_MODELS
if form_data.name in models:
url_idx = models[form_data.name]["urls"][0]
if model in models:
url_idx = models[model]["urls"][0]
else:
raise HTTPException(
status_code=400,
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.name),
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(model),
)
url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
@@ -884,13 +907,13 @@ async def delete_model(
r = requests.request(
method="DELETE",
url=f"{url}/api/delete",
data=form_data.model_dump_json(exclude_none=True).encode(),
data=json.dumps(form_data).encode(),
headers={
"Content-Type": "application/json",
**({"Authorization": f"Bearer {key}"} if key else {}),
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -926,16 +949,21 @@ async def delete_model(
async def show_model_info(
request: Request, form_data: ModelNameForm, user=Depends(get_verified_user)
):
form_data = form_data.model_dump(exclude_none=True)
form_data["model"] = form_data.get("model", form_data.get("name"))
await get_all_models(request, user=user)
models = request.app.state.OLLAMA_MODELS
if form_data.name not in models:
model = form_data.get("model")
if model not in models:
raise HTTPException(
status_code=400,
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.name),
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(model),
)
url_idx = random.choice(models[form_data.name]["urls"])
url_idx = random.choice(models[model]["urls"])
url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS)
@@ -949,7 +977,7 @@ async def show_model_info(
**({"Authorization": f"Bearer {key}"} if key else {}),
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -958,7 +986,7 @@ async def show_model_info(
else {}
),
},
data=form_data.model_dump_json(exclude_none=True).encode(),
data=json.dumps(form_data).encode(),
)
r.raise_for_status()
@@ -1036,7 +1064,7 @@ async def embed(
**({"Authorization": f"Bearer {key}"} if key else {}),
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -1123,7 +1151,7 @@ async def embeddings(
**({"Authorization": f"Bearer {key}"} if key else {}),
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -1343,6 +1371,7 @@ async def generate_chat_completion(
key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
content_type="application/x-ndjson",
user=user,
metadata=metadata,
)
@@ -1381,6 +1410,8 @@ async def generate_openai_completion(
url_idx: Optional[int] = None,
user=Depends(get_verified_user),
):
metadata = form_data.pop("metadata", None)
try:
form_data = OpenAICompletionForm(**form_data)
except Exception as e:
@@ -1446,6 +1477,7 @@ async def generate_openai_completion(
stream=payload.get("stream", False),
key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
user=user,
metadata=metadata,
)
@@ -1527,6 +1559,7 @@ async def generate_openai_chat_completion(
stream=payload.get("stream", False),
key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
user=user,
metadata=metadata,
)
+48 -36
View File
@@ -8,7 +8,7 @@ from typing import Literal, Optional, overload
import aiohttp
from aiocache import cached
import requests
from urllib.parse import quote
from fastapi import Depends, FastAPI, HTTPException, Request, APIRouter
from fastapi.middleware.cors import CORSMiddleware
@@ -21,6 +21,7 @@ from open_webui.config import (
CACHE_DIR,
)
from open_webui.env import (
MODELS_CACHE_TTL,
AIOHTTP_CLIENT_SESSION_SSL,
AIOHTTP_CLIENT_TIMEOUT,
AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST,
@@ -66,7 +67,7 @@ async def send_get_request(url, key=None, user: UserModel = None):
**({"Authorization": f"Bearer {key}"} if key else {}),
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -225,7 +226,7 @@ async def speech(request: Request, user=Depends(get_verified_user)):
),
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -361,7 +362,9 @@ async def get_all_models_responses(request: Request, user: UserModel) -> list:
response if isinstance(response, list) else response.get("data", [])
):
if prefix_id:
model["id"] = f"{prefix_id}.{model['id']}"
model["id"] = (
f"{prefix_id}.{model.get('id', model.get('name', ''))}"
)
if tags:
model["tags"] = tags
@@ -386,7 +389,7 @@ async def get_filtered_models(models, user):
return filtered_models
@cached(ttl=1)
@cached(ttl=MODELS_CACHE_TTL)
async def get_all_models(request: Request, user: UserModel) -> dict[str, list]:
log.info("get_all_models()")
@@ -478,7 +481,7 @@ async def get_models(
"Content-Type": "application/json",
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -573,7 +576,7 @@ async def verify_connection(
"Content-Type": "application/json",
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -633,13 +636,7 @@ async def verify_connection(
raise HTTPException(status_code=500, detail=error_detail)
def convert_to_azure_payload(
url,
payload: dict,
):
model = payload.get("model", "")
# Filter allowed parameters based on Azure OpenAI API
def get_azure_allowed_params(api_version: str) -> set[str]:
allowed_params = {
"messages",
"temperature",
@@ -669,6 +666,23 @@ def convert_to_azure_payload(
"max_completion_tokens",
}
try:
if api_version >= "2024-09-01-preview":
allowed_params.add("stream_options")
except ValueError:
log.debug(
f"Invalid API version {api_version} for Azure OpenAI. Defaulting to allowed parameters."
)
return allowed_params
def convert_to_azure_payload(url, payload: dict, api_version: str):
model = payload.get("model", "")
# Filter allowed parameters based on Azure OpenAI API
allowed_params = get_azure_allowed_params(api_version)
# Special handling for o-series models
if model.startswith("o") and model.endswith("-mini"):
# Convert max_tokens to max_completion_tokens for o-series models
@@ -806,10 +820,15 @@ async def generate_chat_completion(
),
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
**(
{"X-OpenWebUI-Chat-Id": metadata.get("chat_id")}
if metadata and metadata.get("chat_id")
else {}
),
}
if ENABLE_FORWARD_USER_INFO_HEADERS
else {}
@@ -817,8 +836,8 @@ async def generate_chat_completion(
}
if api_config.get("azure", False):
request_url, payload = convert_to_azure_payload(url, payload)
api_version = api_config.get("api_version", "") or "2023-03-15-preview"
api_version = api_config.get("api_version", "2023-03-15-preview")
request_url, payload = convert_to_azure_payload(url, payload, api_version)
headers["api-key"] = key
headers["api-version"] = api_version
request_url = f"{request_url}/chat/completions?api-version={api_version}"
@@ -881,10 +900,8 @@ async def generate_chat_completion(
detail=detail if detail else "Open WebUI: Server Connection Error",
)
finally:
if not streaming and session:
if r:
r.close()
await session.close()
if not streaming:
await cleanup_response(r, session)
async def embeddings(request: Request, form_data: dict, user):
@@ -924,7 +941,7 @@ async def embeddings(request: Request, form_data: dict, user):
"Content-Type": "application/json",
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -963,10 +980,8 @@ async def embeddings(request: Request, form_data: dict, user):
detail=detail if detail else "Open WebUI: Server Connection Error",
)
finally:
if not streaming and session:
if r:
r.close()
await session.close()
if not streaming:
await cleanup_response(r, session)
@router.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE"])
@@ -996,7 +1011,7 @@ async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
"Content-Type": "application/json",
**(
{
"X-OpenWebUI-User-Name": user.name,
"X-OpenWebUI-User-Name": quote(user.name, safe=" "),
"X-OpenWebUI-User-Id": user.id,
"X-OpenWebUI-User-Email": user.email,
"X-OpenWebUI-User-Role": user.role,
@@ -1007,16 +1022,15 @@ async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
}
if api_config.get("azure", False):
api_version = api_config.get("api_version", "2023-03-15-preview")
headers["api-key"] = key
headers["api-version"] = (
api_config.get("api_version", "") or "2023-03-15-preview"
)
headers["api-version"] = api_version
payload = json.loads(body)
url, payload = convert_to_azure_payload(url, payload)
url, payload = convert_to_azure_payload(url, payload, api_version)
body = json.dumps(payload).encode()
request_url = f"{url}/{path}?api-version={api_config.get('api_version', '2023-03-15-preview')}"
request_url = f"{url}/{path}?api-version={api_version}"
else:
headers["Authorization"] = f"Bearer {key}"
request_url = f"{url}/{path}"
@@ -1063,7 +1077,5 @@ async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
detail=detail if detail else "Open WebUI: Server Connection Error",
)
finally:
if not streaming and session:
if r:
r.close()
await session.close()
if not streaming:
await cleanup_response(r, session)
+100 -22
View File
@@ -29,6 +29,7 @@ import tiktoken
from langchain.text_splitter import RecursiveCharacterTextSplitter, TokenTextSplitter
from langchain_text_splitters import MarkdownHeaderTextSplitter
from langchain_core.documents import Document
from open_webui.models.files import FileModel, Files
@@ -69,6 +70,7 @@ from open_webui.retrieval.web.external import search_external
from open_webui.retrieval.utils import (
get_embedding_function,
get_reranking_function,
get_model_path,
query_collection,
query_collection_with_hybrid_search,
@@ -813,7 +815,11 @@ async def update_rag_config(
f"Updating reranking model: {request.app.state.config.RAG_RERANKING_MODEL} to {form_data.RAG_RERANKING_MODEL}"
)
try:
request.app.state.config.RAG_RERANKING_MODEL = form_data.RAG_RERANKING_MODEL
request.app.state.config.RAG_RERANKING_MODEL = (
form_data.RAG_RERANKING_MODEL
if form_data.RAG_RERANKING_MODEL is not None
else request.app.state.config.RAG_RERANKING_MODEL
)
try:
request.app.state.rf = get_rf(
@@ -823,6 +829,12 @@ async def update_rag_config(
request.app.state.config.RAG_EXTERNAL_RERANKER_API_KEY,
True,
)
request.app.state.RERANKING_FUNCTION = get_reranking_function(
request.app.state.config.RAG_RERANKING_ENGINE,
request.app.state.config.RAG_RERANKING_MODEL,
request.app.state.rf,
)
except Exception as e:
log.error(f"Error loading reranking model: {e}")
request.app.state.config.ENABLE_RAG_HYBRID_SEARCH = False
@@ -1146,6 +1158,7 @@ def save_docs_to_vector_db(
chunk_overlap=request.app.state.config.CHUNK_OVERLAP,
add_start_index=True,
)
docs = text_splitter.split_documents(docs)
elif request.app.state.config.TEXT_SPLITTER == "token":
log.info(
f"Using token text splitter: {request.app.state.config.TIKTOKEN_ENCODING_NAME}"
@@ -1158,11 +1171,56 @@ def save_docs_to_vector_db(
chunk_overlap=request.app.state.config.CHUNK_OVERLAP,
add_start_index=True,
)
docs = text_splitter.split_documents(docs)
elif request.app.state.config.TEXT_SPLITTER == "markdown_header":
log.info("Using markdown header text splitter")
# Define headers to split on - covering most common markdown header levels
headers_to_split_on = [
("#", "Header 1"),
("##", "Header 2"),
("###", "Header 3"),
("####", "Header 4"),
("#####", "Header 5"),
("######", "Header 6"),
]
markdown_splitter = MarkdownHeaderTextSplitter(
headers_to_split_on=headers_to_split_on,
strip_headers=False, # Keep headers in content for context
)
md_split_docs = []
for doc in docs:
md_header_splits = markdown_splitter.split_text(doc.page_content)
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=request.app.state.config.CHUNK_SIZE,
chunk_overlap=request.app.state.config.CHUNK_OVERLAP,
add_start_index=True,
)
md_header_splits = text_splitter.split_documents(md_header_splits)
# Convert back to Document objects, preserving original metadata
for split_chunk in md_header_splits:
headings_list = []
# Extract header values in order based on headers_to_split_on
for _, header_meta_key_name in headers_to_split_on:
if header_meta_key_name in split_chunk.metadata:
headings_list.append(
split_chunk.metadata[header_meta_key_name]
)
md_split_docs.append(
Document(
page_content=split_chunk.page_content,
metadata={**doc.metadata, "headings": headings_list},
)
)
docs = md_split_docs
else:
raise ValueError(ERROR_MESSAGES.DEFAULT("Invalid text splitter"))
docs = text_splitter.split_documents(docs)
if len(docs) == 0:
raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
@@ -1171,27 +1229,14 @@ def save_docs_to_vector_db(
{
**doc.metadata,
**(metadata if metadata else {}),
"embedding_config": json.dumps(
{
"engine": request.app.state.config.RAG_EMBEDDING_ENGINE,
"model": request.app.state.config.RAG_EMBEDDING_MODEL,
}
),
"embedding_config": {
"engine": request.app.state.config.RAG_EMBEDDING_ENGINE,
"model": request.app.state.config.RAG_EMBEDDING_MODEL,
},
}
for doc in docs
]
# ChromaDB does not like datetime formats
# for meta-data so convert them to string.
for metadata in metadatas:
for key, value in metadata.items():
if (
isinstance(value, datetime)
or isinstance(value, list)
or isinstance(value, dict)
):
metadata[key] = str(value)
try:
if VECTOR_DB_CLIENT.has_collection(collection_name=collection_name):
log.info(f"collection {collection_name} already exists")
@@ -1747,6 +1792,16 @@ def search_web(request: Request, engine: str, query: str) -> list[SearchResult]:
)
else:
raise Exception("No TAVILY_API_KEY found in environment variables")
elif engine == "exa":
if request.app.state.config.EXA_API_KEY:
return search_exa(
request.app.state.config.EXA_API_KEY,
query,
request.app.state.config.WEB_SEARCH_RESULT_COUNT,
request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
)
else:
raise Exception("No EXA_API_KEY found in environment variables")
elif engine == "searchapi":
if request.app.state.config.SEARCHAPI_API_KEY:
return search_searchapi(
@@ -1784,6 +1839,13 @@ def search_web(request: Request, engine: str, query: str) -> list[SearchResult]:
request.app.state.config.WEB_SEARCH_RESULT_COUNT,
request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
)
elif engine == "exa":
return search_exa(
request.app.state.config.EXA_API_KEY,
query,
request.app.state.config.WEB_SEARCH_RESULT_COUNT,
request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
)
elif engine == "perplexity":
return search_perplexity(
request.app.state.config.PERPLEXITY_API_KEY,
@@ -1978,7 +2040,15 @@ def query_doc_handler(
query, prefix=prefix, user=user
),
k=form_data.k if form_data.k else request.app.state.config.TOP_K,
reranking_function=request.app.state.rf,
reranking_function=(
(
lambda sentences: request.app.state.RERANKING_FUNCTION(
sentences, user=user
)
)
if request.app.state.RERANKING_FUNCTION
else None
),
k_reranker=form_data.k_reranker
or request.app.state.config.TOP_K_RERANKER,
r=(
@@ -2035,7 +2105,15 @@ def query_collection_handler(
query, prefix=prefix, user=user
),
k=form_data.k if form_data.k else request.app.state.config.TOP_K,
reranking_function=request.app.state.rf,
reranking_function=(
(
lambda sentences: request.app.state.RERANKING_FUNCTION(
sentences, user=user
)
)
if request.app.state.RERANKING_FUNCTION
else None
),
k_reranker=form_data.k_reranker
or request.app.state.config.TOP_K_RERANKER,
r=(
+1 -1
View File
@@ -695,11 +695,11 @@ async def generate_emoji(
"max_completion_tokens": 4,
}
),
"chat_id": form_data.get("chat_id", None),
"metadata": {
**(request.state.metadata if hasattr(request.state, "metadata") else {}),
"task": str(TASKS.EMOJI_GENERATION),
"task_body": form_data,
"chat_id": form_data.get("chat_id", None),
},
}
+1 -1
View File
@@ -153,7 +153,7 @@ async def load_tool_from_url(
)
try:
async with aiohttp.ClientSession() as session:
async with aiohttp.ClientSession(trust_env=True) as session:
async with session.get(
url, headers={"Content-Type": "application/json"}
) as resp:
+4 -1
View File
@@ -7,6 +7,7 @@ from open_webui.models.chats import Chats
from open_webui.models.users import (
UserModel,
UserListResponse,
UserInfoListResponse,
UserRoleUpdateForm,
Users,
UserSettings,
@@ -83,7 +84,7 @@ async def get_users(
return Users.get_users(filter=filter, skip=skip, limit=limit)
@router.get("/all", response_model=UserListResponse)
@router.get("/all", response_model=UserInfoListResponse)
async def get_all_users(
user=Depends(get_admin_user),
):
@@ -133,7 +134,9 @@ class SharingPermissions(BaseModel):
class ChatPermissions(BaseModel):
controls: bool = True
valves: bool = True
system_prompt: bool = True
params: bool = True
file_upload: bool = True
delete: bool = True
edit: bool = True
+328 -12
View File
@@ -1,13 +1,18 @@
import asyncio
import random
import socketio
import logging
import sys
import time
from typing import Dict, Set
from redis import asyncio as aioredis
import pycrdt as Y
from open_webui.models.users import Users, UserNameResponse
from open_webui.models.channels import Channels
from open_webui.models.chats import Chats
from open_webui.models.notes import Notes, NoteUpdateForm
from open_webui.utils.redis import (
get_sentinels_from_env,
get_sentinel_url_from_env,
@@ -20,9 +25,14 @@ from open_webui.env import (
WEBSOCKET_REDIS_LOCK_TIMEOUT,
WEBSOCKET_SENTINEL_PORT,
WEBSOCKET_SENTINEL_HOSTS,
REDIS_KEY_PREFIX,
)
from open_webui.utils.auth import decode_token
from open_webui.socket.utils import RedisDict, RedisLock
from open_webui.socket.utils import RedisDict, RedisLock, YdocManager
from open_webui.tasks import create_task, stop_item_tasks
from open_webui.utils.redis import get_redis_connection
from open_webui.utils.access_control import has_access, get_users_with_access
from open_webui.env import (
GLOBAL_LOG_LEVEL,
@@ -35,6 +45,8 @@ log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["SOCKET"])
REDIS = None
if WEBSOCKET_MANAGER == "redis":
if WEBSOCKET_SENTINEL_HOSTS:
mgr = socketio.AsyncRedisManager(
@@ -69,21 +81,29 @@ TIMEOUT_DURATION = 3
if WEBSOCKET_MANAGER == "redis":
log.debug("Using Redis to manage websockets.")
REDIS = get_redis_connection(
redis_url=WEBSOCKET_REDIS_URL,
redis_sentinels=get_sentinels_from_env(
WEBSOCKET_SENTINEL_HOSTS, WEBSOCKET_SENTINEL_PORT
),
async_mode=True,
)
redis_sentinels = get_sentinels_from_env(
WEBSOCKET_SENTINEL_HOSTS, WEBSOCKET_SENTINEL_PORT
)
SESSION_POOL = RedisDict(
"open-webui:session_pool",
f"{REDIS_KEY_PREFIX}:session_pool",
redis_url=WEBSOCKET_REDIS_URL,
redis_sentinels=redis_sentinels,
)
USER_POOL = RedisDict(
"open-webui:user_pool",
f"{REDIS_KEY_PREFIX}:user_pool",
redis_url=WEBSOCKET_REDIS_URL,
redis_sentinels=redis_sentinels,
)
USAGE_POOL = RedisDict(
"open-webui:usage_pool",
f"{REDIS_KEY_PREFIX}:usage_pool",
redis_url=WEBSOCKET_REDIS_URL,
redis_sentinels=redis_sentinels,
)
@@ -101,14 +121,37 @@ else:
SESSION_POOL = {}
USER_POOL = {}
USAGE_POOL = {}
aquire_func = release_func = renew_func = lambda: True
YDOC_MANAGER = YdocManager(
redis=REDIS,
redis_key_prefix=f"{REDIS_KEY_PREFIX}:ydoc:documents",
)
async def periodic_usage_pool_cleanup():
if not aquire_func():
log.debug("Usage pool cleanup lock already exists. Not running it.")
return
log.debug("Running periodic_usage_pool_cleanup")
max_retries = 2
retry_delay = random.uniform(
WEBSOCKET_REDIS_LOCK_TIMEOUT / 2, WEBSOCKET_REDIS_LOCK_TIMEOUT
)
for attempt in range(max_retries + 1):
if aquire_func():
break
else:
if attempt < max_retries:
log.debug(
f"Cleanup lock already exists. Retry {attempt + 1} after {retry_delay}s..."
)
await asyncio.sleep(retry_delay)
else:
log.warning(
"Failed to acquire cleanup lock after retries. Skipping cleanup."
)
return
log.debug("Running periodic_cleanup")
try:
while True:
if not renew_func():
@@ -169,16 +212,20 @@ def get_user_id_from_session_pool(sid):
return None
def get_user_ids_from_room(room):
def get_session_ids_from_room(room):
"""Get all session IDs from a specific room."""
active_session_ids = sio.manager.get_participants(
namespace="/",
room=room,
)
return [session_id[0] for session_id in active_session_ids]
def get_user_ids_from_room(room):
active_session_ids = get_session_ids_from_room(room)
active_user_ids = list(
set(
[SESSION_POOL.get(session_id[0])["id"] for session_id in active_session_ids]
)
set([SESSION_POOL.get(session_id)["id"] for session_id in active_session_ids])
)
return active_user_ids
@@ -270,6 +317,37 @@ async def join_channel(sid, data):
await sio.enter_room(sid, f"channel:{channel.id}")
@sio.on("join-note")
async def join_note(sid, data):
auth = data["auth"] if "auth" in data else None
if not auth or "token" not in auth:
return
token_data = decode_token(auth["token"])
if token_data is None or "id" not in token_data:
return
user = Users.get_user_by_id(token_data["id"])
if not user:
return
note = Notes.get_note_by_id(data["note_id"])
if not note:
log.error(f"Note {data['note_id']} not found for user {user.id}")
return
if (
user.role != "admin"
and user.id != note.user_id
and not has_access(user.id, type="read", access_control=note.access_control)
):
log.error(f"User {user.id} does not have access to note {data['note_id']}")
return
log.debug(f"Joining note {note.id} for user {user.id}")
await sio.enter_room(sid, f"note:{note.id}")
@sio.on("channel-events")
async def channel_events(sid, data):
room = f"channel:{data['channel_id']}"
@@ -298,6 +376,242 @@ async def channel_events(sid, data):
)
@sio.on("ydoc:document:join")
async def ydoc_document_join(sid, data):
"""Handle user joining a document"""
user = SESSION_POOL.get(sid)
try:
document_id = data["document_id"]
if document_id.startswith("note:"):
note_id = document_id.split(":")[1]
note = Notes.get_note_by_id(note_id)
if not note:
log.error(f"Note {note_id} not found")
return
if (
user.get("role") != "admin"
and user.get("id") != note.user_id
and not has_access(
user.get("id"), type="read", access_control=note.access_control
)
):
log.error(
f"User {user.get('id')} does not have access to note {note_id}"
)
return
user_id = data.get("user_id", sid)
user_name = data.get("user_name", "Anonymous")
user_color = data.get("user_color", "#000000")
log.info(f"User {user_id} joining document {document_id}")
await YDOC_MANAGER.add_user(document_id=document_id, user_id=sid)
# Join Socket.IO room
await sio.enter_room(sid, f"doc_{document_id}")
active_session_ids = get_session_ids_from_room(f"doc_{document_id}")
# Get the Yjs document state
ydoc = Y.Doc()
updates = await YDOC_MANAGER.get_updates(document_id)
for update in updates:
ydoc.apply_update(bytes(update))
# Encode the entire document state as an update
state_update = ydoc.get_update()
await sio.emit(
"ydoc:document:state",
{
"document_id": document_id,
"state": list(state_update), # Convert bytes to list for JSON
"sessions": active_session_ids,
},
room=sid,
)
# Notify other users about the new user
await sio.emit(
"ydoc:user:joined",
{
"document_id": document_id,
"user_id": user_id,
"user_name": user_name,
"user_color": user_color,
},
room=f"doc_{document_id}",
skip_sid=sid,
)
log.info(f"User {user_id} successfully joined document {document_id}")
except Exception as e:
log.error(f"Error in yjs_document_join: {e}")
await sio.emit("error", {"message": "Failed to join document"}, room=sid)
async def document_save_handler(document_id, data, user):
if document_id.startswith("note:"):
note_id = document_id.split(":")[1]
note = Notes.get_note_by_id(note_id)
if not note:
log.error(f"Note {note_id} not found")
return
if (
user.get("role") != "admin"
and user.get("id") != note.user_id
and not has_access(
user.get("id"), type="read", access_control=note.access_control
)
):
log.error(f"User {user.get('id')} does not have access to note {note_id}")
return
Notes.update_note_by_id(note_id, NoteUpdateForm(data=data))
@sio.on("ydoc:document:state")
async def yjs_document_state(sid, data):
"""Send the current state of the Yjs document to the user"""
try:
document_id = data["document_id"]
room = f"doc_{document_id}"
active_session_ids = get_session_ids_from_room(room)
if sid not in active_session_ids:
log.warning(f"Session {sid} not in room {room}. Cannot send state.")
return
if not await YDOC_MANAGER.document_exists(document_id):
log.warning(f"Document {document_id} not found")
return
# Get the Yjs document state
ydoc = Y.Doc()
updates = await YDOC_MANAGER.get_updates(document_id)
for update in updates:
ydoc.apply_update(bytes(update))
# Encode the entire document state as an update
state_update = ydoc.get_update()
await sio.emit(
"ydoc:document:state",
{
"document_id": document_id,
"state": list(state_update), # Convert bytes to list for JSON
"sessions": active_session_ids,
},
room=sid,
)
except Exception as e:
log.error(f"Error in yjs_document_state: {e}")
@sio.on("ydoc:document:update")
async def yjs_document_update(sid, data):
"""Handle Yjs document updates"""
try:
document_id = data["document_id"]
try:
await stop_item_tasks(REDIS, document_id)
except:
pass
user_id = data.get("user_id", sid)
update = data["update"] # List of bytes from frontend
await YDOC_MANAGER.append_to_updates(
document_id=document_id,
update=update, # Convert list of bytes to bytes
)
# Broadcast update to all other users in the document
await sio.emit(
"ydoc:document:update",
{
"document_id": document_id,
"user_id": user_id,
"update": update,
"socket_id": sid, # Add socket_id to match frontend filtering
},
room=f"doc_{document_id}",
skip_sid=sid,
)
async def debounced_save():
await asyncio.sleep(0.5)
await document_save_handler(
document_id, data.get("data", {}), SESSION_POOL.get(sid)
)
if data.get("data"):
await create_task(REDIS, debounced_save(), document_id)
except Exception as e:
log.error(f"Error in yjs_document_update: {e}")
@sio.on("ydoc:document:leave")
async def yjs_document_leave(sid, data):
"""Handle user leaving a document"""
try:
document_id = data["document_id"]
user_id = data.get("user_id", sid)
log.info(f"User {user_id} leaving document {document_id}")
# Remove user from the document
await YDOC_MANAGER.remove_user(document_id=document_id, user_id=sid)
# Leave Socket.IO room
await sio.leave_room(sid, f"doc_{document_id}")
# Notify other users
await sio.emit(
"ydoc:user:left",
{"document_id": document_id, "user_id": user_id},
room=f"doc_{document_id}",
)
if (
await YDOC_MANAGER.document_exists(document_id)
and len(await YDOC_MANAGER.get_users(document_id)) == 0
):
log.info(f"Cleaning up document {document_id} as no users are left")
await YDOC_MANAGER.clear_document(document_id)
except Exception as e:
log.error(f"Error in yjs_document_leave: {e}")
@sio.on("ydoc:awareness:update")
async def yjs_awareness_update(sid, data):
"""Handle awareness updates (cursors, selections, etc.)"""
try:
document_id = data["document_id"]
user_id = data.get("user_id", sid)
update = data["update"]
# Broadcast awareness update to all other users in the document
await sio.emit(
"ydoc:awareness:update",
{"document_id": document_id, "user_id": user_id, "update": update},
room=f"doc_{document_id}",
skip_sid=sid,
)
except Exception as e:
log.error(f"Error in yjs_awareness_update: {e}")
@sio.event
async def disconnect(sid):
if sid in SESSION_POOL:
@@ -309,6 +623,8 @@ async def disconnect(sid):
if len(USER_POOL[user_id]) == 0:
del USER_POOL[user_id]
await YDOC_MANAGER.remove_user_from_all_documents(sid)
else:
pass
# print(f"Unknown session ID {sid} disconnected")
+109
View File
@@ -1,6 +1,9 @@
import json
import uuid
from open_webui.utils.redis import get_redis_connection
from open_webui.env import REDIS_KEY_PREFIX
from typing import Optional, List, Tuple
import pycrdt as Y
class RedisLock:
@@ -89,3 +92,109 @@ class RedisDict:
if key not in self:
self[key] = default
return self[key]
class YdocManager:
def __init__(
self,
redis=None,
redis_key_prefix: str = f"{REDIS_KEY_PREFIX}:ydoc:documents",
):
self._updates = {}
self._users = {}
self._redis = redis
self._redis_key_prefix = redis_key_prefix
async def append_to_updates(self, document_id: str, update: bytes):
document_id = document_id.replace(":", "_")
if self._redis:
redis_key = f"{self._redis_key_prefix}:{document_id}:updates"
await self._redis.rpush(redis_key, json.dumps(list(update)))
else:
if document_id not in self._updates:
self._updates[document_id] = []
self._updates[document_id].append(update)
async def get_updates(self, document_id: str) -> List[bytes]:
document_id = document_id.replace(":", "_")
if self._redis:
redis_key = f"{self._redis_key_prefix}:{document_id}:updates"
updates = await self._redis.lrange(redis_key, 0, -1)
return [bytes(json.loads(update)) for update in updates]
else:
return self._updates.get(document_id, [])
async def document_exists(self, document_id: str) -> bool:
document_id = document_id.replace(":", "_")
if self._redis:
redis_key = f"{self._redis_key_prefix}:{document_id}:updates"
return await self._redis.exists(redis_key) > 0
else:
return document_id in self._updates
async def get_users(self, document_id: str) -> List[str]:
document_id = document_id.replace(":", "_")
if self._redis:
redis_key = f"{self._redis_key_prefix}:{document_id}:users"
users = await self._redis.smembers(redis_key)
return list(users)
else:
return self._users.get(document_id, [])
async def add_user(self, document_id: str, user_id: str):
document_id = document_id.replace(":", "_")
if self._redis:
redis_key = f"{self._redis_key_prefix}:{document_id}:users"
await self._redis.sadd(redis_key, user_id)
else:
if document_id not in self._users:
self._users[document_id] = set()
self._users[document_id].add(user_id)
async def remove_user(self, document_id: str, user_id: str):
document_id = document_id.replace(":", "_")
if self._redis:
redis_key = f"{self._redis_key_prefix}:{document_id}:users"
await self._redis.srem(redis_key, user_id)
else:
if document_id in self._users and user_id in self._users[document_id]:
self._users[document_id].remove(user_id)
async def remove_user_from_all_documents(self, user_id: str):
if self._redis:
keys = await self._redis.keys(f"{self._redis_key_prefix}:*")
for key in keys:
if key.endswith(":users"):
await self._redis.srem(key, user_id)
document_id = key.split(":")[-2]
if len(await self.get_users(document_id)) == 0:
await self.clear_document(document_id)
else:
for document_id in list(self._users.keys()):
if user_id in self._users[document_id]:
self._users[document_id].remove(user_id)
if not self._users[document_id]:
del self._users[document_id]
await self.clear_document(document_id)
async def clear_document(self, document_id: str):
document_id = document_id.replace(":", "_")
if self._redis:
redis_key = f"{self._redis_key_prefix}:{document_id}:updates"
await self._redis.delete(redis_key)
redis_users_key = f"{self._redis_key_prefix}:{document_id}:users"
await self._redis.delete(redis_users_key)
else:
if document_id in self._updates:
del self._updates[document_id]
if document_id in self._users:
del self._users[document_id]
+65 -48
View File
@@ -3,23 +3,25 @@ import asyncio
from typing import Dict
from uuid import uuid4
import json
import logging
from redis.asyncio import Redis
from fastapi import Request
from typing import Dict, List, Optional
from open_webui.env import SRC_LOG_LEVELS, REDIS_KEY_PREFIX
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["MAIN"])
# A dictionary to keep track of active tasks
tasks: Dict[str, asyncio.Task] = {}
chat_tasks = {}
item_tasks = {}
REDIS_TASKS_KEY = "open-webui:tasks"
REDIS_CHAT_TASKS_KEY = "open-webui:tasks:chat"
REDIS_PUBSUB_CHANNEL = "open-webui:tasks:commands"
def is_redis(request: Request) -> bool:
# Called everywhere a request is available to check Redis
return hasattr(request.app.state, "redis") and (request.app.state.redis is not None)
REDIS_TASKS_KEY = f"{REDIS_KEY_PREFIX}:tasks"
REDIS_ITEM_TASKS_KEY = f"{REDIS_KEY_PREFIX}:tasks:item"
REDIS_PUBSUB_CHANNEL = f"{REDIS_KEY_PREFIX}:tasks:commands"
async def redis_task_command_listener(app):
@@ -38,7 +40,7 @@ async def redis_task_command_listener(app):
if local_task:
local_task.cancel()
except Exception as e:
print(f"Error handling distributed task command: {e}")
log.exception(f"Error handling distributed task command: {e}")
### ------------------------------
@@ -46,21 +48,21 @@ async def redis_task_command_listener(app):
### ------------------------------
async def redis_save_task(redis: Redis, task_id: str, chat_id: Optional[str]):
async def redis_save_task(redis: Redis, task_id: str, item_id: Optional[str]):
pipe = redis.pipeline()
pipe.hset(REDIS_TASKS_KEY, task_id, chat_id or "")
if chat_id:
pipe.sadd(f"{REDIS_CHAT_TASKS_KEY}:{chat_id}", task_id)
pipe.hset(REDIS_TASKS_KEY, task_id, item_id or "")
if item_id:
pipe.sadd(f"{REDIS_ITEM_TASKS_KEY}:{item_id}", task_id)
await pipe.execute()
async def redis_cleanup_task(redis: Redis, task_id: str, chat_id: Optional[str]):
async def redis_cleanup_task(redis: Redis, task_id: str, item_id: Optional[str]):
pipe = redis.pipeline()
pipe.hdel(REDIS_TASKS_KEY, task_id)
if chat_id:
pipe.srem(f"{REDIS_CHAT_TASKS_KEY}:{chat_id}", task_id)
if (await pipe.scard(f"{REDIS_CHAT_TASKS_KEY}:{chat_id}").execute())[-1] == 0:
pipe.delete(f"{REDIS_CHAT_TASKS_KEY}:{chat_id}") # Remove if empty set
if item_id:
pipe.srem(f"{REDIS_ITEM_TASKS_KEY}:{item_id}", task_id)
if (await pipe.scard(f"{REDIS_ITEM_TASKS_KEY}:{item_id}").execute())[-1] == 0:
pipe.delete(f"{REDIS_ITEM_TASKS_KEY}:{item_id}") # Remove if empty set
await pipe.execute()
@@ -68,31 +70,31 @@ async def redis_list_tasks(redis: Redis) -> List[str]:
return list(await redis.hkeys(REDIS_TASKS_KEY))
async def redis_list_chat_tasks(redis: Redis, chat_id: str) -> List[str]:
return list(await redis.smembers(f"{REDIS_CHAT_TASKS_KEY}:{chat_id}"))
async def redis_list_item_tasks(redis: Redis, item_id: str) -> List[str]:
return list(await redis.smembers(f"{REDIS_ITEM_TASKS_KEY}:{item_id}"))
async def redis_send_command(redis: Redis, command: dict):
await redis.publish(REDIS_PUBSUB_CHANNEL, json.dumps(command))
async def cleanup_task(request, task_id: str, id=None):
async def cleanup_task(redis, task_id: str, id=None):
"""
Remove a completed or canceled task from the global `tasks` dictionary.
"""
if is_redis(request):
await redis_cleanup_task(request.app.state.redis, task_id, id)
if redis:
await redis_cleanup_task(redis, task_id, id)
tasks.pop(task_id, None) # Remove the task if it exists
# If an ID is provided, remove the task from the chat_tasks dictionary
if id and task_id in chat_tasks.get(id, []):
chat_tasks[id].remove(task_id)
if not chat_tasks[id]: # If no tasks left for this ID, remove the entry
chat_tasks.pop(id, None)
# If an ID is provided, remove the task from the item_tasks dictionary
if id and task_id in item_tasks.get(id, []):
item_tasks[id].remove(task_id)
if not item_tasks[id]: # If no tasks left for this ID, remove the entry
item_tasks.pop(id, None)
async def create_task(request, coroutine, id=None):
async def create_task(redis, coroutine, id=None):
"""
Create a new asyncio task and add it to the global task dictionary.
"""
@@ -101,48 +103,48 @@ async def create_task(request, coroutine, id=None):
# Add a done callback for cleanup
task.add_done_callback(
lambda t: asyncio.create_task(cleanup_task(request, task_id, id))
lambda t: asyncio.create_task(cleanup_task(redis, task_id, id))
)
tasks[task_id] = task
# If an ID is provided, associate the task with that ID
if chat_tasks.get(id):
chat_tasks[id].append(task_id)
if item_tasks.get(id):
item_tasks[id].append(task_id)
else:
chat_tasks[id] = [task_id]
item_tasks[id] = [task_id]
if is_redis(request):
await redis_save_task(request.app.state.redis, task_id, id)
if redis:
await redis_save_task(redis, task_id, id)
return task_id, task
async def list_tasks(request):
async def list_tasks(redis):
"""
List all currently active task IDs.
"""
if is_redis(request):
return await redis_list_tasks(request.app.state.redis)
if redis:
return await redis_list_tasks(redis)
return list(tasks.keys())
async def list_task_ids_by_chat_id(request, id):
async def list_task_ids_by_item_id(redis, id):
"""
List all tasks associated with a specific ID.
"""
if is_redis(request):
return await redis_list_chat_tasks(request.app.state.redis, id)
return chat_tasks.get(id, [])
if redis:
return await redis_list_item_tasks(redis, id)
return item_tasks.get(id, [])
async def stop_task(request, task_id: str):
async def stop_task(redis, task_id: str):
"""
Cancel a running task and remove it from the global task list.
"""
if is_redis(request):
if redis:
# PUBSUB: All instances check if they have this task, and stop if so.
await redis_send_command(
request.app.state.redis,
redis,
{
"action": "stop",
"task_id": task_id,
@@ -151,7 +153,7 @@ async def stop_task(request, task_id: str):
# Optionally check if task_id still in Redis a few moments later for feedback?
return {"status": True, "message": f"Stop signal sent for {task_id}"}
task = tasks.get(task_id)
task = tasks.pop(task_id)
if not task:
raise ValueError(f"Task with ID {task_id} not found.")
@@ -160,7 +162,22 @@ async def stop_task(request, task_id: str):
await task # Wait for the task to handle the cancellation
except asyncio.CancelledError:
# Task successfully canceled
tasks.pop(task_id, None) # Remove it from the dictionary
return {"status": True, "message": f"Task {task_id} successfully stopped."}
return {"status": False, "message": f"Failed to stop task {task_id}."}
async def stop_item_tasks(redis: Redis, item_id: str):
"""
Stop all tasks associated with a specific item ID.
"""
task_ids = await list_task_ids_by_item_id(redis, item_id)
if not task_ids:
return {"status": True, "message": f"No tasks found for item {item_id}."}
for task_id in task_ids:
result = await stop_task(redis, task_id)
if not result["status"]:
return result # Return the first failure
return {"status": True, "message": f"All tasks for item {item_id} stopped."}
+793
View File
@@ -0,0 +1,793 @@
import pytest
from unittest.mock import Mock, patch, AsyncMock
import redis
from open_webui.utils.redis import (
SentinelRedisProxy,
parse_redis_service_url,
get_redis_connection,
get_sentinels_from_env,
MAX_RETRY_COUNT,
)
import inspect
class TestSentinelRedisProxy:
"""Test Redis Sentinel failover functionality"""
def test_parse_redis_service_url_valid(self):
"""Test parsing valid Redis service URL"""
url = "redis://user:pass@mymaster:6379/0"
result = parse_redis_service_url(url)
assert result["username"] == "user"
assert result["password"] == "pass"
assert result["service"] == "mymaster"
assert result["port"] == 6379
assert result["db"] == 0
def test_parse_redis_service_url_defaults(self):
"""Test parsing Redis service URL with defaults"""
url = "redis://mymaster"
result = parse_redis_service_url(url)
assert result["username"] is None
assert result["password"] is None
assert result["service"] == "mymaster"
assert result["port"] == 6379
assert result["db"] == 0
def test_parse_redis_service_url_invalid_scheme(self):
"""Test parsing invalid URL scheme"""
with pytest.raises(ValueError, match="Invalid Redis URL scheme"):
parse_redis_service_url("http://invalid")
def test_get_sentinels_from_env(self):
"""Test parsing sentinel hosts from environment"""
hosts = "sentinel1,sentinel2,sentinel3"
port = "26379"
result = get_sentinels_from_env(hosts, port)
expected = [("sentinel1", 26379), ("sentinel2", 26379), ("sentinel3", 26379)]
assert result == expected
def test_get_sentinels_from_env_empty(self):
"""Test empty sentinel hosts"""
result = get_sentinels_from_env(None, "26379")
assert result == []
@patch("redis.sentinel.Sentinel")
def test_sentinel_redis_proxy_sync_success(self, mock_sentinel_class):
"""Test successful sync operation with SentinelRedisProxy"""
mock_sentinel = Mock()
mock_master = Mock()
mock_master.get.return_value = "test_value"
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=False)
# Test attribute access
get_method = proxy.__getattr__("get")
result = get_method("test_key")
assert result == "test_value"
mock_sentinel.master_for.assert_called_with("mymaster")
mock_master.get.assert_called_with("test_key")
@patch("redis.sentinel.Sentinel")
@pytest.mark.asyncio
async def test_sentinel_redis_proxy_async_success(self, mock_sentinel_class):
"""Test successful async operation with SentinelRedisProxy"""
mock_sentinel = Mock()
mock_master = Mock()
mock_master.get = AsyncMock(return_value="test_value")
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=True)
# Test async attribute access
get_method = proxy.__getattr__("get")
result = await get_method("test_key")
assert result == "test_value"
mock_sentinel.master_for.assert_called_with("mymaster")
mock_master.get.assert_called_with("test_key")
@patch("redis.sentinel.Sentinel")
def test_sentinel_redis_proxy_failover_retry(self, mock_sentinel_class):
"""Test retry mechanism during failover"""
mock_sentinel = Mock()
mock_master = Mock()
# First call fails, second succeeds
mock_master.get.side_effect = [
redis.exceptions.ConnectionError("Master down"),
"test_value",
]
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=False)
get_method = proxy.__getattr__("get")
result = get_method("test_key")
assert result == "test_value"
assert mock_master.get.call_count == 2
@patch("redis.sentinel.Sentinel")
def test_sentinel_redis_proxy_max_retries_exceeded(self, mock_sentinel_class):
"""Test failure after max retries exceeded"""
mock_sentinel = Mock()
mock_master = Mock()
# All calls fail
mock_master.get.side_effect = redis.exceptions.ConnectionError("Master down")
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=False)
get_method = proxy.__getattr__("get")
with pytest.raises(redis.exceptions.ConnectionError):
get_method("test_key")
assert mock_master.get.call_count == MAX_RETRY_COUNT
@patch("redis.sentinel.Sentinel")
def test_sentinel_redis_proxy_readonly_error_retry(self, mock_sentinel_class):
"""Test retry on ReadOnlyError"""
mock_sentinel = Mock()
mock_master = Mock()
# First call gets ReadOnlyError (old master), second succeeds (new master)
mock_master.get.side_effect = [
redis.exceptions.ReadOnlyError("Read only"),
"test_value",
]
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=False)
get_method = proxy.__getattr__("get")
result = get_method("test_key")
assert result == "test_value"
assert mock_master.get.call_count == 2
@patch("redis.sentinel.Sentinel")
def test_sentinel_redis_proxy_factory_methods(self, mock_sentinel_class):
"""Test factory methods are passed through directly"""
mock_sentinel = Mock()
mock_master = Mock()
mock_pipeline = Mock()
mock_master.pipeline.return_value = mock_pipeline
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=False)
# Factory methods should be passed through without wrapping
pipeline_method = proxy.__getattr__("pipeline")
result = pipeline_method()
assert result == mock_pipeline
mock_master.pipeline.assert_called_once()
@patch("redis.sentinel.Sentinel")
@patch("redis.from_url")
def test_get_redis_connection_with_sentinel(
self, mock_from_url, mock_sentinel_class
):
"""Test getting Redis connection with Sentinel"""
mock_sentinel = Mock()
mock_sentinel_class.return_value = mock_sentinel
sentinels = [("sentinel1", 26379), ("sentinel2", 26379)]
redis_url = "redis://user:pass@mymaster:6379/0"
result = get_redis_connection(
redis_url=redis_url, redis_sentinels=sentinels, async_mode=False
)
assert isinstance(result, SentinelRedisProxy)
mock_sentinel_class.assert_called_once()
mock_from_url.assert_not_called()
@patch("redis.Redis.from_url")
def test_get_redis_connection_without_sentinel(self, mock_from_url):
"""Test getting Redis connection without Sentinel"""
mock_redis = Mock()
mock_from_url.return_value = mock_redis
redis_url = "redis://localhost:6379/0"
result = get_redis_connection(
redis_url=redis_url, redis_sentinels=None, async_mode=False
)
assert result == mock_redis
mock_from_url.assert_called_once_with(redis_url, decode_responses=True)
@patch("redis.asyncio.from_url")
def test_get_redis_connection_without_sentinel_async(self, mock_from_url):
"""Test getting async Redis connection without Sentinel"""
mock_redis = Mock()
mock_from_url.return_value = mock_redis
redis_url = "redis://localhost:6379/0"
result = get_redis_connection(
redis_url=redis_url, redis_sentinels=None, async_mode=True
)
assert result == mock_redis
mock_from_url.assert_called_once_with(redis_url, decode_responses=True)
class TestSentinelRedisProxyCommands:
"""Test Redis commands through SentinelRedisProxy"""
@patch("redis.sentinel.Sentinel")
def test_hash_commands_sync(self, mock_sentinel_class):
"""Test Redis hash commands in sync mode"""
mock_sentinel = Mock()
mock_master = Mock()
# Mock hash command responses
mock_master.hset.return_value = 1
mock_master.hget.return_value = "test_value"
mock_master.hgetall.return_value = {"key1": "value1", "key2": "value2"}
mock_master.hdel.return_value = 1
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=False)
# Test hset
hset_method = proxy.__getattr__("hset")
result = hset_method("test_hash", "field1", "value1")
assert result == 1
mock_master.hset.assert_called_with("test_hash", "field1", "value1")
# Test hget
hget_method = proxy.__getattr__("hget")
result = hget_method("test_hash", "field1")
assert result == "test_value"
mock_master.hget.assert_called_with("test_hash", "field1")
# Test hgetall
hgetall_method = proxy.__getattr__("hgetall")
result = hgetall_method("test_hash")
assert result == {"key1": "value1", "key2": "value2"}
mock_master.hgetall.assert_called_with("test_hash")
# Test hdel
hdel_method = proxy.__getattr__("hdel")
result = hdel_method("test_hash", "field1")
assert result == 1
mock_master.hdel.assert_called_with("test_hash", "field1")
@patch("redis.sentinel.Sentinel")
@pytest.mark.asyncio
async def test_hash_commands_async(self, mock_sentinel_class):
"""Test Redis hash commands in async mode"""
mock_sentinel = Mock()
mock_master = Mock()
# Mock async hash command responses
mock_master.hset = AsyncMock(return_value=1)
mock_master.hget = AsyncMock(return_value="test_value")
mock_master.hgetall = AsyncMock(
return_value={"key1": "value1", "key2": "value2"}
)
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=True)
# Test hset
hset_method = proxy.__getattr__("hset")
result = await hset_method("test_hash", "field1", "value1")
assert result == 1
mock_master.hset.assert_called_with("test_hash", "field1", "value1")
# Test hget
hget_method = proxy.__getattr__("hget")
result = await hget_method("test_hash", "field1")
assert result == "test_value"
mock_master.hget.assert_called_with("test_hash", "field1")
# Test hgetall
hgetall_method = proxy.__getattr__("hgetall")
result = await hgetall_method("test_hash")
assert result == {"key1": "value1", "key2": "value2"}
mock_master.hgetall.assert_called_with("test_hash")
@patch("redis.sentinel.Sentinel")
def test_string_commands_sync(self, mock_sentinel_class):
"""Test Redis string commands in sync mode"""
mock_sentinel = Mock()
mock_master = Mock()
# Mock string command responses
mock_master.set.return_value = True
mock_master.get.return_value = "test_value"
mock_master.delete.return_value = 1
mock_master.exists.return_value = True
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=False)
# Test set
set_method = proxy.__getattr__("set")
result = set_method("test_key", "test_value")
assert result is True
mock_master.set.assert_called_with("test_key", "test_value")
# Test get
get_method = proxy.__getattr__("get")
result = get_method("test_key")
assert result == "test_value"
mock_master.get.assert_called_with("test_key")
# Test delete
delete_method = proxy.__getattr__("delete")
result = delete_method("test_key")
assert result == 1
mock_master.delete.assert_called_with("test_key")
# Test exists
exists_method = proxy.__getattr__("exists")
result = exists_method("test_key")
assert result is True
mock_master.exists.assert_called_with("test_key")
@patch("redis.sentinel.Sentinel")
def test_list_commands_sync(self, mock_sentinel_class):
"""Test Redis list commands in sync mode"""
mock_sentinel = Mock()
mock_master = Mock()
# Mock list command responses
mock_master.lpush.return_value = 1
mock_master.rpop.return_value = "test_value"
mock_master.llen.return_value = 5
mock_master.lrange.return_value = ["item1", "item2", "item3"]
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=False)
# Test lpush
lpush_method = proxy.__getattr__("lpush")
result = lpush_method("test_list", "item1")
assert result == 1
mock_master.lpush.assert_called_with("test_list", "item1")
# Test rpop
rpop_method = proxy.__getattr__("rpop")
result = rpop_method("test_list")
assert result == "test_value"
mock_master.rpop.assert_called_with("test_list")
# Test llen
llen_method = proxy.__getattr__("llen")
result = llen_method("test_list")
assert result == 5
mock_master.llen.assert_called_with("test_list")
# Test lrange
lrange_method = proxy.__getattr__("lrange")
result = lrange_method("test_list", 0, -1)
assert result == ["item1", "item2", "item3"]
mock_master.lrange.assert_called_with("test_list", 0, -1)
@patch("redis.sentinel.Sentinel")
def test_pubsub_commands_sync(self, mock_sentinel_class):
"""Test Redis pubsub commands in sync mode"""
mock_sentinel = Mock()
mock_master = Mock()
mock_pubsub = Mock()
# Mock pubsub responses
mock_master.pubsub.return_value = mock_pubsub
mock_master.publish.return_value = 1
mock_pubsub.subscribe.return_value = None
mock_pubsub.get_message.return_value = {"type": "message", "data": "test_data"}
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=False)
# Test pubsub (factory method - should pass through)
pubsub_method = proxy.__getattr__("pubsub")
result = pubsub_method()
assert result == mock_pubsub
mock_master.pubsub.assert_called_once()
# Test publish
publish_method = proxy.__getattr__("publish")
result = publish_method("test_channel", "test_message")
assert result == 1
mock_master.publish.assert_called_with("test_channel", "test_message")
@patch("redis.sentinel.Sentinel")
def test_pipeline_commands_sync(self, mock_sentinel_class):
"""Test Redis pipeline commands in sync mode"""
mock_sentinel = Mock()
mock_master = Mock()
mock_pipeline = Mock()
# Mock pipeline responses
mock_master.pipeline.return_value = mock_pipeline
mock_pipeline.set.return_value = mock_pipeline
mock_pipeline.get.return_value = mock_pipeline
mock_pipeline.execute.return_value = [True, "test_value"]
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=False)
# Test pipeline (factory method - should pass through)
pipeline_method = proxy.__getattr__("pipeline")
result = pipeline_method()
assert result == mock_pipeline
mock_master.pipeline.assert_called_once()
@patch("redis.sentinel.Sentinel")
def test_commands_with_failover_retry(self, mock_sentinel_class):
"""Test Redis commands with failover retry mechanism"""
mock_sentinel = Mock()
mock_master = Mock()
# First call fails with connection error, second succeeds
mock_master.hget.side_effect = [
redis.exceptions.ConnectionError("Connection failed"),
"recovered_value",
]
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=False)
# Test hget with retry
hget_method = proxy.__getattr__("hget")
result = hget_method("test_hash", "field1")
assert result == "recovered_value"
assert mock_master.hget.call_count == 2
# Verify both calls were made with same parameters
expected_calls = [(("test_hash", "field1"),), (("test_hash", "field1"),)]
actual_calls = [call.args for call in mock_master.hget.call_args_list]
assert actual_calls == expected_calls
@patch("redis.sentinel.Sentinel")
def test_commands_with_readonly_error_retry(self, mock_sentinel_class):
"""Test Redis commands with ReadOnlyError retry mechanism"""
mock_sentinel = Mock()
mock_master = Mock()
# First call fails with ReadOnlyError, second succeeds
mock_master.hset.side_effect = [
redis.exceptions.ReadOnlyError(
"READONLY You can't write against a read only replica"
),
1,
]
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=False)
# Test hset with retry
hset_method = proxy.__getattr__("hset")
result = hset_method("test_hash", "field1", "value1")
assert result == 1
assert mock_master.hset.call_count == 2
# Verify both calls were made with same parameters
expected_calls = [
(("test_hash", "field1", "value1"),),
(("test_hash", "field1", "value1"),),
]
actual_calls = [call.args for call in mock_master.hset.call_args_list]
assert actual_calls == expected_calls
@patch("redis.sentinel.Sentinel")
@pytest.mark.asyncio
async def test_async_commands_with_failover_retry(self, mock_sentinel_class):
"""Test async Redis commands with failover retry mechanism"""
mock_sentinel = Mock()
mock_master = Mock()
# First call fails with connection error, second succeeds
mock_master.hget = AsyncMock(
side_effect=[
redis.exceptions.ConnectionError("Connection failed"),
"recovered_value",
]
)
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=True)
# Test async hget with retry
hget_method = proxy.__getattr__("hget")
result = await hget_method("test_hash", "field1")
assert result == "recovered_value"
assert mock_master.hget.call_count == 2
# Verify both calls were made with same parameters
expected_calls = [(("test_hash", "field1"),), (("test_hash", "field1"),)]
actual_calls = [call.args for call in mock_master.hget.call_args_list]
assert actual_calls == expected_calls
class TestSentinelRedisProxyFactoryMethods:
"""Test Redis factory methods in async mode - these are special cases that remain sync"""
@patch("redis.sentinel.Sentinel")
@pytest.mark.asyncio
async def test_pubsub_factory_method_async(self, mock_sentinel_class):
"""Test pubsub factory method in async mode - should pass through without wrapping"""
mock_sentinel = Mock()
mock_master = Mock()
mock_pubsub = Mock()
# Mock pubsub factory method
mock_master.pubsub.return_value = mock_pubsub
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=True)
# Test pubsub factory method - should NOT be wrapped as async
pubsub_method = proxy.__getattr__("pubsub")
result = pubsub_method()
assert result == mock_pubsub
mock_master.pubsub.assert_called_once()
# Verify it's not wrapped as async (no await needed)
assert not inspect.iscoroutine(result)
@patch("redis.sentinel.Sentinel")
@pytest.mark.asyncio
async def test_pipeline_factory_method_async(self, mock_sentinel_class):
"""Test pipeline factory method in async mode - should pass through without wrapping"""
mock_sentinel = Mock()
mock_master = Mock()
mock_pipeline = Mock()
# Mock pipeline factory method
mock_master.pipeline.return_value = mock_pipeline
mock_pipeline.set.return_value = mock_pipeline
mock_pipeline.get.return_value = mock_pipeline
mock_pipeline.execute.return_value = [True, "test_value"]
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=True)
# Test pipeline factory method - should NOT be wrapped as async
pipeline_method = proxy.__getattr__("pipeline")
result = pipeline_method()
assert result == mock_pipeline
mock_master.pipeline.assert_called_once()
# Verify it's not wrapped as async (no await needed)
assert not inspect.iscoroutine(result)
# Test pipeline usage (these should also be sync)
pipeline_result = result.set("key", "value").get("key").execute()
assert pipeline_result == [True, "test_value"]
@patch("redis.sentinel.Sentinel")
@pytest.mark.asyncio
async def test_factory_methods_vs_regular_commands_async(self, mock_sentinel_class):
"""Test that factory methods behave differently from regular commands in async mode"""
mock_sentinel = Mock()
mock_master = Mock()
# Mock both factory method and regular command
mock_pubsub = Mock()
mock_master.pubsub.return_value = mock_pubsub
mock_master.get = AsyncMock(return_value="test_value")
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=True)
# Test factory method - should NOT be wrapped
pubsub_method = proxy.__getattr__("pubsub")
pubsub_result = pubsub_method()
# Test regular command - should be wrapped as async
get_method = proxy.__getattr__("get")
get_result = get_method("test_key")
# Factory method returns directly
assert pubsub_result == mock_pubsub
assert not inspect.iscoroutine(pubsub_result)
# Regular command returns coroutine
assert inspect.iscoroutine(get_result)
# Regular command needs await
actual_value = await get_result
assert actual_value == "test_value"
@patch("redis.sentinel.Sentinel")
@pytest.mark.asyncio
async def test_factory_methods_with_failover_async(self, mock_sentinel_class):
"""Test factory methods with failover in async mode"""
mock_sentinel = Mock()
mock_master = Mock()
# First call fails, second succeeds
mock_pubsub = Mock()
mock_master.pubsub.side_effect = [
redis.exceptions.ConnectionError("Connection failed"),
mock_pubsub,
]
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=True)
# Test pubsub factory method with failover
pubsub_method = proxy.__getattr__("pubsub")
result = pubsub_method()
assert result == mock_pubsub
assert mock_master.pubsub.call_count == 2 # Retry happened
# Verify it's still not wrapped as async after retry
assert not inspect.iscoroutine(result)
@patch("redis.sentinel.Sentinel")
@pytest.mark.asyncio
async def test_monitor_factory_method_async(self, mock_sentinel_class):
"""Test monitor factory method in async mode - should pass through without wrapping"""
mock_sentinel = Mock()
mock_master = Mock()
mock_monitor = Mock()
# Mock monitor factory method
mock_master.monitor.return_value = mock_monitor
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=True)
# Test monitor factory method - should NOT be wrapped as async
monitor_method = proxy.__getattr__("monitor")
result = monitor_method()
assert result == mock_monitor
mock_master.monitor.assert_called_once()
# Verify it's not wrapped as async (no await needed)
assert not inspect.iscoroutine(result)
@patch("redis.sentinel.Sentinel")
@pytest.mark.asyncio
async def test_client_factory_method_async(self, mock_sentinel_class):
"""Test client factory method in async mode - should pass through without wrapping"""
mock_sentinel = Mock()
mock_master = Mock()
mock_client = Mock()
# Mock client factory method
mock_master.client.return_value = mock_client
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=True)
# Test client factory method - should NOT be wrapped as async
client_method = proxy.__getattr__("client")
result = client_method()
assert result == mock_client
mock_master.client.assert_called_once()
# Verify it's not wrapped as async (no await needed)
assert not inspect.iscoroutine(result)
@patch("redis.sentinel.Sentinel")
@pytest.mark.asyncio
async def test_transaction_factory_method_async(self, mock_sentinel_class):
"""Test transaction factory method in async mode - should pass through without wrapping"""
mock_sentinel = Mock()
mock_master = Mock()
mock_transaction = Mock()
# Mock transaction factory method
mock_master.transaction.return_value = mock_transaction
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=True)
# Test transaction factory method - should NOT be wrapped as async
transaction_method = proxy.__getattr__("transaction")
result = transaction_method()
assert result == mock_transaction
mock_master.transaction.assert_called_once()
# Verify it's not wrapped as async (no await needed)
assert not inspect.iscoroutine(result)
@patch("redis.sentinel.Sentinel")
@pytest.mark.asyncio
async def test_all_factory_methods_async(self, mock_sentinel_class):
"""Test all factory methods in async mode - comprehensive test"""
mock_sentinel = Mock()
mock_master = Mock()
# Mock all factory methods
mock_objects = {
"pipeline": Mock(),
"pubsub": Mock(),
"monitor": Mock(),
"client": Mock(),
"transaction": Mock(),
}
for method_name, mock_obj in mock_objects.items():
setattr(mock_master, method_name, Mock(return_value=mock_obj))
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=True)
# Test all factory methods
for method_name, expected_obj in mock_objects.items():
method = proxy.__getattr__(method_name)
result = method()
assert result == expected_obj
assert not inspect.iscoroutine(result)
getattr(mock_master, method_name).assert_called_once()
# Reset mock for next iteration
getattr(mock_master, method_name).reset_mock()
@patch("redis.sentinel.Sentinel")
@pytest.mark.asyncio
async def test_mixed_factory_and_regular_commands_async(self, mock_sentinel_class):
"""Test using both factory methods and regular commands in async mode"""
mock_sentinel = Mock()
mock_master = Mock()
# Mock pipeline factory and regular commands
mock_pipeline = Mock()
mock_master.pipeline.return_value = mock_pipeline
mock_pipeline.set.return_value = mock_pipeline
mock_pipeline.get.return_value = mock_pipeline
mock_pipeline.execute.return_value = [True, "pipeline_value"]
mock_master.get = AsyncMock(return_value="regular_value")
mock_sentinel.master_for.return_value = mock_master
proxy = SentinelRedisProxy(mock_sentinel, "mymaster", async_mode=True)
# Use factory method (sync)
pipeline = proxy.__getattr__("pipeline")()
pipeline_result = pipeline.set("key1", "value1").get("key1").execute()
# Use regular command (async)
get_method = proxy.__getattr__("get")
regular_result = await get_method("key2")
# Verify both work correctly
assert pipeline_result == [True, "pipeline_value"]
assert regular_result == "regular_value"
# Verify calls
mock_master.pipeline.assert_called_once()
mock_master.get.assert_called_with("key2")
+73 -23
View File
@@ -8,6 +8,12 @@ import requests
import os
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.asymmetric import ed25519
from cryptography.hazmat.primitives import serialization
import json
from datetime import datetime, timedelta
import pytz
from pytz import UTC
@@ -18,7 +24,11 @@ from opentelemetry import trace
from open_webui.models.users import Users
from open_webui.constants import ERROR_MESSAGES
from open_webui.env import (
OFFLINE_MODE,
LICENSE_BLOB,
pk,
WEBUI_SECRET_KEY,
TRUSTED_SIGNATURE_KEY,
STATIC_DIR,
@@ -74,33 +84,73 @@ def override_static(path: str, content: str):
def get_license_data(app, key):
if key:
try:
res = requests.post(
"https://api.openwebui.com/api/v1/license/",
json={"key": key, "version": "1"},
timeout=5,
def data_handler(data):
for k, v in data.items():
if k == "resources":
for p, c in v.items():
globals().get("override_static", lambda a, b: None)(p, c)
elif k == "count":
setattr(app.state, "USER_COUNT", v)
elif k == "name":
setattr(app.state, "WEBUI_NAME", v)
elif k == "metadata":
setattr(app.state, "LICENSE_METADATA", v)
def handler(u):
res = requests.post(
f"{u}/api/v1/license/",
json={"key": key, "version": "1"},
timeout=5,
)
if getattr(res, "ok", False):
payload = getattr(res, "json", lambda: {})()
data_handler(payload)
return True
else:
log.error(
f"License: retrieval issue: {getattr(res, 'text', 'unknown error')}"
)
if getattr(res, "ok", False):
payload = getattr(res, "json", lambda: {})()
for k, v in payload.items():
if k == "resources":
for p, c in v.items():
globals().get("override_static", lambda a, b: None)(p, c)
elif k == "count":
setattr(app.state, "USER_COUNT", v)
elif k == "name":
setattr(app.state, "WEBUI_NAME", v)
elif k == "metadata":
setattr(app.state, "LICENSE_METADATA", v)
return True
else:
log.error(
f"License: retrieval issue: {getattr(res, 'text', 'unknown error')}"
)
if key:
us = [
"https://api.openwebui.com",
"https://licenses.api.openwebui.com",
]
try:
for u in us:
if handler(u):
return True
except Exception as ex:
log.exception(f"License: Uncaught Exception: {ex}")
try:
if LICENSE_BLOB:
nl = 12
kb = hashlib.sha256((key.replace("-", "").upper()).encode()).digest()
def nt(b):
return b[:nl], b[nl:]
lb = base64.b64decode(LICENSE_BLOB)
ln, lt = nt(lb)
aesgcm = AESGCM(kb)
p = json.loads(aesgcm.decrypt(ln, lt, None))
pk.verify(base64.b64decode(p["s"]), p["p"].encode())
pb = base64.b64decode(p["p"])
pn, pt = nt(pb)
data = json.loads(aesgcm.decrypt(pn, pt, None).decode())
if not data.get("exp") and data.get("exp") < datetime.now().date():
return False
data_handler(data)
return True
except Exception as e:
log.error(f"License: {e}")
return False
+1 -1
View File
@@ -419,7 +419,7 @@ async def chat_action(request: Request, action_id: str, form_data: dict, user: A
params[key] = value
if "__user__" in sig.parameters:
__user__ = (user.model_dump() if isinstance(user, UserModel) else {},)
__user__ = user.model_dump() if isinstance(user, UserModel) else {}
try:
if hasattr(function_module, "UserValves"):
+21 -4
View File
@@ -4,12 +4,16 @@ import sys
from typing import TYPE_CHECKING
from loguru import logger
from opentelemetry import trace
from open_webui.env import (
AUDIT_UVICORN_LOGGER_NAMES,
AUDIT_LOG_FILE_ROTATION_SIZE,
AUDIT_LOG_LEVEL,
AUDIT_LOGS_FILE_PATH,
GLOBAL_LOG_LEVEL,
ENABLE_OTEL,
)
@@ -58,9 +62,20 @@ class InterceptHandler(logging.Handler):
frame = frame.f_back
depth += 1
logger.opt(depth=depth, exception=record.exc_info).log(
level, record.getMessage()
)
logger.opt(depth=depth, exception=record.exc_info).bind(
**self._get_extras()
).log(level, record.getMessage())
def _get_extras(self):
if not ENABLE_OTEL:
return {}
extras = {}
context = trace.get_current_span().get_span_context()
if context.is_valid:
extras["trace_id"] = trace.format_trace_id(context.trace_id)
extras["span_id"] = trace.format_span_id(context.span_id)
return extras
def file_format(record: "Record"):
@@ -128,11 +143,13 @@ def start_logger():
logging.basicConfig(
handlers=[InterceptHandler()], level=GLOBAL_LOG_LEVEL, force=True
)
for uvicorn_logger_name in ["uvicorn", "uvicorn.error"]:
uvicorn_logger = logging.getLogger(uvicorn_logger_name)
uvicorn_logger.setLevel(GLOBAL_LOG_LEVEL)
uvicorn_logger.handlers = []
for uvicorn_logger_name in ["uvicorn.access"]:
for uvicorn_logger_name in AUDIT_UVICORN_LOGGER_NAMES:
uvicorn_logger = logging.getLogger(uvicorn_logger_name)
uvicorn_logger.setLevel(GLOBAL_LOG_LEVEL)
uvicorn_logger.handlers = [InterceptHandler()]
+182 -120
View File
@@ -23,6 +23,7 @@ from starlette.responses import Response, StreamingResponse
from open_webui.models.chats import Chats
from open_webui.models.folders import Folders
from open_webui.models.users import Users
from open_webui.socket.main import (
get_event_call,
@@ -56,7 +57,7 @@ from open_webui.models.users import UserModel
from open_webui.models.functions import Functions
from open_webui.models.models import Models
from open_webui.retrieval.utils import get_sources_from_files
from open_webui.retrieval.utils import get_sources_from_items
from open_webui.utils.chat import generate_chat_completion
@@ -82,6 +83,7 @@ from open_webui.utils.filter import (
process_filter_functions,
)
from open_webui.utils.code_interpreter import execute_code_jupyter
from open_webui.utils.payload import apply_model_system_prompt_to_body
from open_webui.tasks import create_task
@@ -248,30 +250,28 @@ async def chat_completion_tools_handler(
if tool_id
else f"{tool_function_name}"
)
if tool.get("metadata", {}).get("citation", False) or tool.get(
"direct", False
):
# Citation is enabled for this tool
sources.append(
{
"source": {
"name": (f"TOOL:{tool_name}"),
},
"document": [tool_result],
"metadata": [
{
"source": (f"TOOL:{tool_name}"),
"parameters": tool_function_params,
}
],
}
)
else:
# Citation is not enabled for this tool
body["messages"] = add_or_update_user_message(
f"\nTool `{tool_name}` Output: {tool_result}",
body["messages"],
)
# Citation is enabled for this tool
sources.append(
{
"source": {
"name": (f"TOOL:{tool_name}"),
},
"document": [tool_result],
"metadata": [
{
"source": (f"TOOL:{tool_name}"),
"parameters": tool_function_params,
}
],
"tool_result": True,
}
)
# Citation is not enabled for this tool
body["messages"] = add_or_update_user_message(
f"\nTool `{tool_name}` Output: {tool_result}",
body["messages"],
)
if (
tools[tool_function_name]
@@ -640,25 +640,34 @@ async def chat_completion_files_handler(
queries = [get_last_user_message(body["messages"])]
try:
# Offload get_sources_from_files to a separate thread
# Offload get_sources_from_items to a separate thread
loop = asyncio.get_running_loop()
with ThreadPoolExecutor() as executor:
sources = await loop.run_in_executor(
executor,
lambda: get_sources_from_files(
lambda: get_sources_from_items(
request=request,
files=files,
items=files,
queries=queries,
embedding_function=lambda query, prefix: request.app.state.EMBEDDING_FUNCTION(
query, prefix=prefix, user=user
),
k=request.app.state.config.TOP_K,
reranking_function=request.app.state.rf,
reranking_function=(
(
lambda sentences: request.app.state.RERANKING_FUNCTION(
sentences, user=user
)
)
if request.app.state.RERANKING_FUNCTION
else None
),
k_reranker=request.app.state.config.TOP_K_RERANKER,
r=request.app.state.config.RELEVANCE_THRESHOLD,
hybrid_bm25_weight=request.app.state.config.HYBRID_BM25_WEIGHT,
hybrid_search=request.app.state.config.ENABLE_RAG_HYBRID_SEARCH,
full_context=request.app.state.config.RAG_FULL_CONTEXT,
user=user,
),
)
except Exception as e:
@@ -718,6 +727,10 @@ def apply_params_to_form_data(form_data, model):
async def process_chat_payload(request, form_data, user, metadata, model):
# Pipeline Inlet -> Filter Inlet -> Chat Memory -> Chat Web Search -> Chat Image Generation
# -> Chat Code Interpreter (Form Data Update) -> (Default) Chat Tools Function Calling
# -> Chat Files
form_data = apply_params_to_form_data(form_data, model)
log.debug(f"form_data: {form_data}")
@@ -752,6 +765,29 @@ async def process_chat_payload(request, form_data, user, metadata, model):
events = []
sources = []
# Folder "Project" handling
# Check if the request has chat_id and is inside of a folder
chat_id = metadata.get("chat_id", None)
if chat_id and user:
chat = Chats.get_chat_by_id_and_user_id(chat_id, user.id)
if chat and chat.folder_id:
folder = Folders.get_folder_by_id_and_user_id(chat.folder_id, user.id)
if folder and folder.data:
if "system_prompt" in folder.data:
form_data = apply_model_system_prompt_to_body(
folder.data["system_prompt"],
form_data,
metadata,
user
)
if "files" in folder.data:
form_data["files"] = [
*folder.data["files"],
*form_data.get("files", []),
]
# Model "Knowledge" handling
user_message = get_last_user_message(form_data["messages"])
model_knowledge = model.get("info", {}).get("meta", {}).get("knowledge", False)
@@ -804,7 +840,6 @@ async def process_chat_payload(request, form_data, user, metadata, model):
raise e
try:
filter_functions = [
Functions.get_function_by_id(filter_id)
for filter_id in get_sorted_filter_ids(
@@ -912,7 +947,6 @@ async def process_chat_payload(request, form_data, user, metadata, model):
request, form_data, extra_params, user, models, tools_dict
)
sources.extend(flags.get("sources", []))
except Exception as e:
log.exception(e)
@@ -925,55 +959,59 @@ async def process_chat_payload(request, form_data, user, metadata, model):
# If context is not empty, insert it into the messages
if len(sources) > 0:
context_string = ""
citation_idx = {}
citation_idx_map = {}
for source in sources:
if "document" in source:
for doc_context, doc_meta in zip(
is_tool_result = source.get("tool_result", False)
if "document" in source and not is_tool_result:
for document_text, document_metadata in zip(
source["document"], source["metadata"]
):
source_name = source.get("source", {}).get("name", None)
citation_id = (
doc_meta.get("source", None)
source_id = (
document_metadata.get("source", None)
or source.get("source", {}).get("id", None)
or "N/A"
)
if citation_id not in citation_idx:
citation_idx[citation_id] = len(citation_idx) + 1
if source_id not in citation_idx_map:
citation_idx_map[source_id] = len(citation_idx_map) + 1
context_string += (
f'<source id="{citation_idx[citation_id]}"'
f'<source id="{citation_idx_map[source_id]}"'
+ (f' name="{source_name}"' if source_name else "")
+ f">{doc_context}</source>\n"
+ f">{document_text}</source>\n"
)
context_string = context_string.strip()
prompt = get_last_user_message(form_data["messages"])
prompt = get_last_user_message(form_data["messages"])
if prompt is None:
raise Exception("No user message found")
if (
request.app.state.config.RELEVANCE_THRESHOLD == 0
and context_string.strip() == ""
):
log.debug(
f"With a 0 relevancy threshold for RAG, the context cannot be empty"
)
# Workaround for Ollama 2.0+ system prompt issue
# TODO: replace with add_or_update_system_message
if model.get("owned_by") == "ollama":
form_data["messages"] = prepend_to_first_user_message_content(
rag_template(
request.app.state.config.RAG_TEMPLATE, context_string, prompt
),
form_data["messages"],
)
if context_string == "":
if request.app.state.config.RELEVANCE_THRESHOLD == 0:
log.debug(
f"With a 0 relevancy threshold for RAG, the context cannot be empty"
)
else:
form_data["messages"] = add_or_update_system_message(
rag_template(
request.app.state.config.RAG_TEMPLATE, context_string, prompt
),
form_data["messages"],
)
# Workaround for Ollama 2.0+ system prompt issue
# TODO: replace with add_or_update_system_message
if model.get("owned_by") == "ollama":
form_data["messages"] = prepend_to_first_user_message_content(
rag_template(
request.app.state.config.RAG_TEMPLATE, context_string, prompt
),
form_data["messages"],
)
else:
form_data["messages"] = add_or_update_system_message(
rag_template(
request.app.state.config.RAG_TEMPLATE, context_string, prompt
),
form_data["messages"],
)
# If there are citations, add them to the data_items
sources = [
@@ -1347,14 +1385,6 @@ async def process_chat_response(
task_id = str(uuid4()) # Create a unique task ID.
model_id = form_data.get("model", "")
Chats.upsert_message_to_chat_by_id_and_message_id(
metadata["chat_id"],
metadata["message_id"],
{
"model": model_id,
},
)
def split_content_and_whitespace(content):
content_stripped = content.rstrip()
original_whitespace = (
@@ -1370,7 +1400,7 @@ async def process_chat_response(
return len(backtick_segments) > 1 and len(backtick_segments) % 2 == 0
# Handle as a background task
async def post_response_handler(response, events):
async def response_handler(response, events):
def serialize_content_blocks(content_blocks, raw=False):
content = ""
@@ -1405,7 +1435,7 @@ async def process_chat_response(
break
if tool_result:
tool_calls_display_content = f'{tool_calls_display_content}\n<details type="tool_calls" done="true" id="{tool_call_id}" name="{tool_name}" arguments="{html.escape(json.dumps(tool_arguments))}" result="{html.escape(json.dumps(tool_result))}" files="{html.escape(json.dumps(tool_result_files)) if tool_result_files else ""}">\n<summary>Tool Executed</summary>\n</details>\n'
tool_calls_display_content = f'{tool_calls_display_content}\n<details type="tool_calls" done="true" id="{tool_call_id}" name="{tool_name}" arguments="{html.escape(json.dumps(tool_arguments))}" result="{html.escape(json.dumps(tool_result, ensure_ascii=False))}" files="{html.escape(json.dumps(tool_result_files)) if tool_result_files else ""}">\n<summary>Tool Executed</summary>\n</details>\n'
else:
tool_calls_display_content = f'{tool_calls_display_content}\n<details type="tool_calls" done="false" id="{tool_call_id}" name="{tool_name}" arguments="{html.escape(json.dumps(tool_arguments))}">\n<summary>Executing...</summary>\n</details>'
@@ -1438,12 +1468,12 @@ async def process_chat_response(
if reasoning_duration is not None:
if raw:
content = f'{content}\n<{block["start_tag"]}>{block["content"]}<{block["end_tag"]}>\n'
content = f'{content}\n{block["start_tag"]}{block["content"]}{block["end_tag"]}\n'
else:
content = f'{content}\n<details type="reasoning" done="true" duration="{reasoning_duration}">\n<summary>Thought for {reasoning_duration} seconds</summary>\n{reasoning_display_content}\n</details>\n'
else:
if raw:
content = f'{content}\n<{block["start_tag"]}>{block["content"]}<{block["end_tag"]}>\n'
content = f'{content}\n{block["start_tag"]}{block["content"]}{block["end_tag"]}\n'
else:
content = f'{content}\n<details type="reasoning" done="false">\n<summary>Thinking…</summary>\n{reasoning_display_content}\n</details>\n'
@@ -1540,8 +1570,16 @@ async def process_chat_response(
if content_blocks[-1]["type"] == "text":
for start_tag, end_tag in tags:
# Match start tag e.g., <tag> or <tag attr="value">
start_tag_pattern = rf"<{re.escape(start_tag)}(\s.*?)?>"
start_tag_pattern = rf"{re.escape(start_tag)}"
if start_tag.startswith("<") and start_tag.endswith(">"):
# Match start tag e.g., <tag> or <tag attr="value">
# remove both '<' and '>' from start_tag
# Match start tag with attributes
start_tag_pattern = (
rf"<{re.escape(start_tag[1:-1])}(\s.*?)?>"
)
match = re.search(start_tag_pattern, content)
if match:
attr_content = (
@@ -1592,8 +1630,13 @@ async def process_chat_response(
elif content_blocks[-1]["type"] == content_type:
start_tag = content_blocks[-1]["start_tag"]
end_tag = content_blocks[-1]["end_tag"]
# Match end tag e.g., </tag>
end_tag_pattern = rf"<{re.escape(end_tag)}>"
if end_tag.startswith("<") and end_tag.endswith(">"):
# Match end tag e.g., </tag>
end_tag_pattern = rf"{re.escape(end_tag)}"
else:
# Handle cases where end_tag is just a tag name
end_tag_pattern = rf"{re.escape(end_tag)}"
# Check if the content has the end tag
if re.search(end_tag_pattern, content):
@@ -1665,8 +1708,17 @@ async def process_chat_response(
)
# Clean processed content
start_tag_pattern = rf"{re.escape(start_tag)}"
if start_tag.startswith("<") and start_tag.endswith(">"):
# Match start tag e.g., <tag> or <tag attr="value">
# remove both '<' and '>' from start_tag
# Match start tag with attributes
start_tag_pattern = (
rf"<{re.escape(start_tag[1:-1])}(\s.*?)?>"
)
content = re.sub(
rf"<{re.escape(start_tag)}(.*?)>(.|\n)*?<{re.escape(end_tag)}>",
rf"{start_tag_pattern}(.|\n)*?{re.escape(end_tag)}",
"",
content,
flags=re.DOTALL,
@@ -1710,18 +1762,19 @@ async def process_chat_response(
)
reasoning_tags = [
("think", "/think"),
("thinking", "/thinking"),
("reason", "/reason"),
("reasoning", "/reasoning"),
("thought", "/thought"),
("Thought", "/Thought"),
("|begin_of_thought|", "|end_of_thought|"),
("<think>", "</think>"),
("<thinking>", "</thinking>"),
("<reason>", "</reason>"),
("<reasoning>", "</reasoning>"),
("<thought>", "</thought>"),
("<Thought>", "</Thought>"),
("<|begin_of_thought|>", "<|end_of_thought|>"),
("◁think▷", "◁/think▷"),
]
code_interpreter_tags = [("code_interpreter", "/code_interpreter")]
code_interpreter_tags = [("<code_interpreter>", "</code_interpreter>")]
solution_tags = [("|begin_of_solution|", "|end_of_solution|")]
solution_tags = [("<|begin_of_solution|>", "<|end_of_solution|>")]
try:
for event in events:
@@ -1741,7 +1794,7 @@ async def process_chat_response(
},
)
async def stream_body_handler(response):
async def stream_body_handler(response, form_data):
nonlocal content
nonlocal content_blocks
@@ -1770,7 +1823,7 @@ async def process_chat_response(
filter_functions=filter_functions,
filter_type="stream",
form_data=data,
extra_params=extra_params,
extra_params={"__body__": form_data, **extra_params},
)
if data:
@@ -2005,7 +2058,7 @@ async def process_chat_response(
if done:
pass
else:
log.debug("Error: ", e)
log.debug(f"Error: {e}")
continue
if content_blocks:
@@ -2032,7 +2085,7 @@ async def process_chat_response(
if response.background:
await response.background()
await stream_body_handler(response)
await stream_body_handler(response, form_data)
MAX_TOOL_CALL_RETRIES = 10
tool_call_retries = 0
@@ -2148,7 +2201,9 @@ async def process_chat_response(
if isinstance(tool_result, dict) or isinstance(
tool_result, list
):
tool_result = json.dumps(tool_result, indent=2)
tool_result = json.dumps(
tool_result, indent=2, ensure_ascii=False
)
results.append(
{
@@ -2181,22 +2236,24 @@ async def process_chat_response(
)
try:
new_form_data = {
"model": model_id,
"stream": True,
"tools": form_data["tools"],
"messages": [
*form_data["messages"],
*convert_content_blocks_to_messages(content_blocks),
],
}
res = await generate_chat_completion(
request,
{
"model": model_id,
"stream": True,
"tools": form_data["tools"],
"messages": [
*form_data["messages"],
*convert_content_blocks_to_messages(content_blocks),
],
},
new_form_data,
user,
)
if isinstance(res, StreamingResponse):
await stream_body_handler(res)
await stream_body_handler(res, new_form_data)
else:
break
except Exception as e:
@@ -2211,6 +2268,7 @@ async def process_chat_response(
content_blocks[-1]["type"] == "code_interpreter"
and retries < MAX_RETRIES
):
await event_emitter(
{
"type": "chat:completion",
@@ -2343,26 +2401,28 @@ async def process_chat_response(
)
try:
new_form_data = {
"model": model_id,
"stream": True,
"messages": [
*form_data["messages"],
{
"role": "assistant",
"content": serialize_content_blocks(
content_blocks, raw=True
),
},
],
}
res = await generate_chat_completion(
request,
{
"model": model_id,
"stream": True,
"messages": [
*form_data["messages"],
{
"role": "assistant",
"content": serialize_content_blocks(
content_blocks, raw=True
),
},
],
},
new_form_data,
user,
)
if isinstance(res, StreamingResponse):
await stream_body_handler(res)
await stream_body_handler(res, new_form_data)
else:
break
except Exception as e:
@@ -2427,9 +2487,11 @@ async def process_chat_response(
if response.background is not None:
await response.background()
# background_tasks.add_task(post_response_handler, response, events)
# background_tasks.add_task(response_handler, response, events)
task_id, _ = await create_task(
request, post_response_handler(response, events), id=metadata["chat_id"]
request.app.state.redis,
response_handler(response, events),
id=metadata["chat_id"],
)
return {"status": True, "task_id": task_id}
+14 -2
View File
@@ -76,8 +76,19 @@ async def get_all_base_models(request: Request, user: UserModel = None):
return function_models + openai_models + ollama_models
async def get_all_models(request, user: UserModel = None):
models = await get_all_base_models(request, user=user)
async def get_all_models(request, refresh: bool = False, user: UserModel = None):
if (
request.app.state.MODELS
and request.app.state.BASE_MODELS
and (request.app.state.config.ENABLE_BASE_MODELS_CACHE and not refresh)
):
base_models = request.app.state.BASE_MODELS
else:
base_models = await get_all_base_models(request, user=user)
request.app.state.BASE_MODELS = base_models
# deep copy the base models to avoid modifying the original list
models = [model.copy() for model in base_models]
# If there are no models, return an empty list
if len(models) == 0:
@@ -137,6 +148,7 @@ async def get_all_models(request, user: UserModel = None):
custom_models = Models.get_all_models()
for custom_model in custom_models:
if custom_model.base_model_id is None:
# Applied directly to a base model
for model in models:
if custom_model.id == model["id"] or (
model.get("owned_by") == "ollama"
+116 -10
View File
@@ -1,6 +1,97 @@
import socketio
import inspect
from urllib.parse import urlparse
from typing import Optional
import logging
import redis
from open_webui.env import REDIS_SENTINEL_MAX_RETRY_COUNT
log = logging.getLogger(__name__)
_CONNECTION_CACHE = {}
class SentinelRedisProxy:
def __init__(self, sentinel, service, *, async_mode: bool = True, **kw):
self._sentinel = sentinel
self._service = service
self._kw = kw
self._async_mode = async_mode
def _master(self):
return self._sentinel.master_for(self._service, **self._kw)
def __getattr__(self, item):
master = self._master()
orig_attr = getattr(master, item)
if not callable(orig_attr):
return orig_attr
FACTORY_METHODS = {"pipeline", "pubsub", "monitor", "client", "transaction"}
if item in FACTORY_METHODS:
return orig_attr
if self._async_mode:
async def _wrapped(*args, **kwargs):
for i in range(REDIS_SENTINEL_MAX_RETRY_COUNT):
try:
method = getattr(self._master(), item)
result = method(*args, **kwargs)
if inspect.iscoroutine(result):
return await result
return result
except (
redis.exceptions.ConnectionError,
redis.exceptions.ReadOnlyError,
) as e:
if i < REDIS_SENTINEL_MAX_RETRY_COUNT - 1:
log.debug(
"Redis sentinel fail-over (%s). Retry %s/%s",
type(e).__name__,
i + 1,
REDIS_SENTINEL_MAX_RETRY_COUNT,
)
continue
log.error(
"Redis operation failed after %s retries: %s",
REDIS_SENTINEL_MAX_RETRY_COUNT,
e,
)
raise e from e
return _wrapped
else:
def _wrapped(*args, **kwargs):
for i in range(REDIS_SENTINEL_MAX_RETRY_COUNT):
try:
method = getattr(self._master(), item)
return method(*args, **kwargs)
except (
redis.exceptions.ConnectionError,
redis.exceptions.ReadOnlyError,
) as e:
if i < REDIS_SENTINEL_MAX_RETRY_COUNT - 1:
log.debug(
"Redis sentinel fail-over (%s). Retry %s/%s",
type(e).__name__,
i + 1,
REDIS_SENTINEL_MAX_RETRY_COUNT,
)
continue
log.error(
"Redis operation failed after %s retries: %s",
REDIS_SENTINEL_MAX_RETRY_COUNT,
e,
)
raise e from e
return _wrapped
def parse_redis_service_url(redis_url):
@@ -20,6 +111,14 @@ def parse_redis_service_url(redis_url):
def get_redis_connection(
redis_url, redis_sentinels, async_mode=False, decode_responses=True
):
cache_key = (redis_url, tuple(redis_sentinels) if redis_sentinels else (), async_mode, decode_responses)
if cache_key in _CONNECTION_CACHE:
return _CONNECTION_CACHE[cache_key]
connection = None
if async_mode:
import redis.asyncio as redis
@@ -34,11 +133,13 @@ def get_redis_connection(
password=redis_config["password"],
decode_responses=decode_responses,
)
return sentinel.master_for(redis_config["service"])
connection = SentinelRedisProxy(
sentinel,
redis_config["service"],
async_mode=async_mode,
)
elif redis_url:
return redis.from_url(redis_url, decode_responses=decode_responses)
else:
return None
connection = redis.from_url(redis_url, decode_responses=decode_responses)
else:
import redis
@@ -52,11 +153,16 @@ def get_redis_connection(
password=redis_config["password"],
decode_responses=decode_responses,
)
return sentinel.master_for(redis_config["service"])
connection = SentinelRedisProxy(
sentinel,
redis_config["service"],
async_mode=async_mode,
)
elif redis_url:
return redis.Redis.from_url(redis_url, decode_responses=decode_responses)
else:
return None
connection = redis.Redis.from_url(redis_url, decode_responses=decode_responses)
_CONNECTION_CACHE[cache_key] = connection
return connection
def get_sentinels_from_env(sentinel_hosts_env, sentinel_port_env):
+5 -6
View File
@@ -6,18 +6,17 @@ from open_webui.utils.misc import (
)
def convert_ollama_tool_call_to_openai(tool_calls: dict) -> dict:
def convert_ollama_tool_call_to_openai(tool_calls: list) -> list:
openai_tool_calls = []
for tool_call in tool_calls:
function = tool_call.get("function", {})
openai_tool_call = {
"index": tool_call.get("index", 0),
"index": tool_call.get("index", function.get("index", 0)),
"id": tool_call.get("id", f"call_{str(uuid4())}"),
"type": "function",
"function": {
"name": tool_call.get("function", {}).get("name", ""),
"arguments": json.dumps(
tool_call.get("function", {}).get("arguments", {})
),
"name": function.get("name", ""),
"arguments": json.dumps(function.get("arguments", {})),
},
}
openai_tool_calls.append(openai_tool_call)
@@ -1,31 +0,0 @@
import threading
from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export import BatchSpanProcessor
class LazyBatchSpanProcessor(BatchSpanProcessor):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.done = True
with self.condition:
self.condition.notify_all()
self.worker_thread.join()
self.done = False
self.worker_thread = None
def on_end(self, span: ReadableSpan) -> None:
if self.worker_thread is None:
self.worker_thread = threading.Thread(
name=self.__class__.__name__, target=self.worker, daemon=True
)
self.worker_thread.start()
super().on_end(span)
def shutdown(self) -> None:
self.done = True
with self.condition:
self.condition.notify_all()
if self.worker_thread:
self.worker_thread.join()
self.span_exporter.shutdown()
@@ -34,6 +34,8 @@ from opentelemetry.sdk.resources import SERVICE_NAME, Resource
from open_webui.env import OTEL_SERVICE_NAME, OTEL_EXPORTER_OTLP_ENDPOINT
from open_webui.socket.main import get_active_user_ids
from open_webui.models.users import Users
_EXPORT_INTERVAL_MILLIS = 10_000 # 10 seconds
@@ -59,6 +61,12 @@ def _build_meter_provider() -> MeterProvider:
instrument_name="http.server.requests",
attribute_keys=["http.method", "http.route", "http.status_code"],
),
View(
instrument_name="webui.users.total",
),
View(
instrument_name="webui.users.active",
),
]
provider = MeterProvider(
@@ -87,6 +95,38 @@ def setup_metrics(app: FastAPI) -> None:
unit="ms",
)
def observe_active_users(
options: metrics.CallbackOptions,
) -> Sequence[metrics.Observation]:
return [
metrics.Observation(
value=len(get_active_user_ids()),
)
]
def observe_total_registered_users(
options: metrics.CallbackOptions,
) -> Sequence[metrics.Observation]:
return [
metrics.Observation(
value=len(Users.get_users()["users"]),
)
]
meter.create_observable_gauge(
name="webui.users.total",
description="Total number of registered users",
unit="users",
callbacks=[observe_total_registered_users],
)
meter.create_observable_gauge(
name="webui.users.active",
description="Number of currently active users",
unit="users",
callbacks=[observe_active_users],
)
# FastAPI middleware
@app.middleware("http")
async def _metrics_middleware(request: Request, call_next):
+30 -3
View File
@@ -1,17 +1,26 @@
from fastapi import FastAPI
from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
OTLPSpanExporter as HttpOTLPSpanExporter,
)
from opentelemetry.sdk.resources import SERVICE_NAME, Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from sqlalchemy import Engine
from base64 import b64encode
from open_webui.utils.telemetry.exporters import LazyBatchSpanProcessor
from open_webui.utils.telemetry.instrumentors import Instrumentor
from open_webui.utils.telemetry.metrics import setup_metrics
from open_webui.env import (
OTEL_SERVICE_NAME,
OTEL_EXPORTER_OTLP_ENDPOINT,
OTEL_EXPORTER_OTLP_INSECURE,
ENABLE_OTEL_METRICS,
OTEL_BASIC_AUTH_USERNAME,
OTEL_BASIC_AUTH_PASSWORD,
OTEL_OTLP_SPAN_EXPORTER,
)
@@ -22,9 +31,27 @@ def setup(app: FastAPI, db_engine: Engine):
resource=Resource.create(attributes={SERVICE_NAME: OTEL_SERVICE_NAME})
)
)
# Add basic auth header only if both username and password are not empty
headers = []
if OTEL_BASIC_AUTH_USERNAME and OTEL_BASIC_AUTH_PASSWORD:
auth_string = f"{OTEL_BASIC_AUTH_USERNAME}:{OTEL_BASIC_AUTH_PASSWORD}"
auth_header = b64encode(auth_string.encode()).decode()
headers = [("authorization", f"Basic {auth_header}")]
# otlp export
exporter = OTLPSpanExporter(endpoint=OTEL_EXPORTER_OTLP_ENDPOINT)
trace.get_tracer_provider().add_span_processor(LazyBatchSpanProcessor(exporter))
if OTEL_OTLP_SPAN_EXPORTER == "http":
exporter = HttpOTLPSpanExporter(
endpoint=OTEL_EXPORTER_OTLP_ENDPOINT,
headers=headers,
)
else:
exporter = OTLPSpanExporter(
endpoint=OTEL_EXPORTER_OTLP_ENDPOINT,
insecure=OTEL_EXPORTER_OTLP_INSECURE,
headers=headers,
)
trace.get_tracer_provider().add_span_processor(BatchSpanProcessor(exporter))
Instrumentor(app=app, db_engine=db_engine).instrument()
# set up metrics only if enabled
+14 -12
View File
@@ -101,9 +101,6 @@ def get_tools(
def make_tool_function(function_name, token, tool_server_data):
async def tool_function(**kwargs):
print(
f"Executing tool function {function_name} with params: {kwargs}"
)
return await execute_tool_server(
token=token,
url=tool_server_data["url"],
@@ -492,15 +489,7 @@ async def get_tool_servers_data(
if server.get("config", {}).get("enable"):
# Path (to OpenAPI spec URL) can be either a full URL or a path to append to the base URL
openapi_path = server.get("path", "openapi.json")
if "://" in openapi_path:
# If it contains "://", it's a full URL
full_url = openapi_path
else:
if not openapi_path.startswith("/"):
# Ensure the path starts with a slash
openapi_path = f"/{openapi_path}"
full_url = f"{server.get('url')}{openapi_path}"
full_url = get_tool_server_url(server.get("url"), openapi_path)
info = server.get("info", {})
@@ -646,3 +635,16 @@ async def execute_tool_server(
error = str(err)
log.exception(f"API Request Error: {error}")
return {"error": error}
def get_tool_server_url(url: Optional[str], path: str) -> str:
"""
Build the full URL for a tool server, given a base url and a path.
"""
if "://" in path:
# If it contains "://", it's a full URL
return path
if not path.startswith("/"):
# Ensure the path starts with a slash
path = f"/{path}"
return f"{url}{path}"