Merge branch 'open-webui:main' into cors
This commit is contained in:
@@ -445,6 +445,12 @@ OAUTH_TIMEOUT = PersistentConfig(
|
||||
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",
|
||||
@@ -636,6 +642,13 @@ def load_oauth_providers():
|
||||
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 {}
|
||||
),
|
||||
@@ -676,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(
|
||||
|
||||
@@ -276,9 +276,6 @@ if DATABASE_USER:
|
||||
DATABASE_CRED += f"{DATABASE_USER}"
|
||||
if DATABASE_PASSWORD:
|
||||
DATABASE_CRED += f":{DATABASE_PASSWORD}"
|
||||
if DATABASE_CRED:
|
||||
DATABASE_CRED += "@"
|
||||
|
||||
|
||||
DB_VARS = {
|
||||
"db_type": DATABASE_TYPE,
|
||||
@@ -352,6 +349,15 @@ 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
|
||||
####################################
|
||||
@@ -450,7 +456,13 @@ ENABLE_WEBSOCKET_SUPPORT = (
|
||||
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", "")
|
||||
|
||||
|
||||
@@ -63,7 +63,7 @@ class FolderForm(BaseModel):
|
||||
|
||||
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())
|
||||
@@ -71,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()),
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -611,6 +611,9 @@ def get_sources_from_items(
|
||||
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
|
||||
|
||||
@@ -376,9 +376,13 @@ async def speech(request: Request, user=Depends(get_verified_user)):
|
||||
|
||||
if r is not None:
|
||||
status_code = r.status
|
||||
res = await r.json()
|
||||
if "error" in res:
|
||||
detail = f"External: {res['error'].get('message', '')}"
|
||||
|
||||
try:
|
||||
res = await r.json()
|
||||
if "error" in res:
|
||||
detail = f"External: {res['error']}"
|
||||
except Exception:
|
||||
detail = f"External: {e}"
|
||||
|
||||
raise HTTPException(
|
||||
status_code=status_code,
|
||||
|
||||
@@ -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()
|
||||
)
|
||||
|
||||
|
||||
############################
|
||||
|
||||
@@ -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 (
|
||||
@@ -286,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")
|
||||
@@ -603,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")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
############################
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -170,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)
|
||||
|
||||
@@ -124,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
|
||||
@@ -144,6 +145,11 @@ async def send_post_request(
|
||||
"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 {}
|
||||
@@ -184,7 +190,6 @@ async def send_post_request(
|
||||
)
|
||||
else:
|
||||
res = await r.json()
|
||||
await cleanup_response(r, session)
|
||||
return res
|
||||
|
||||
except HTTPException as e:
|
||||
@@ -196,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):
|
||||
@@ -1363,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,
|
||||
)
|
||||
|
||||
|
||||
@@ -1401,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:
|
||||
@@ -1466,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,
|
||||
)
|
||||
|
||||
|
||||
@@ -1547,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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -822,6 +822,11 @@ async def generate_chat_completion(
|
||||
"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 {}
|
||||
@@ -893,10 +898,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):
|
||||
@@ -975,10 +978,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"])
|
||||
@@ -1074,7 +1075,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)
|
||||
|
||||
@@ -815,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(
|
||||
@@ -2050,11 +2054,13 @@ def query_doc_handler(
|
||||
),
|
||||
k=form_data.k if form_data.k else request.app.state.config.TOP_K,
|
||||
reranking_function=(
|
||||
lambda sentences: (
|
||||
request.app.state.RERANKING_FUNCTION(sentences, user=user)
|
||||
if request.app.state.RERANKING_FUNCTION
|
||||
else None
|
||||
(
|
||||
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,
|
||||
@@ -2112,8 +2118,14 @@ 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=lambda sentences: request.app.state.RERANKING_FUNCTION(
|
||||
sentences, user=user
|
||||
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,
|
||||
|
||||
@@ -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),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -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),
|
||||
):
|
||||
|
||||
@@ -316,6 +316,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']}"
|
||||
@@ -450,7 +481,7 @@ async def yjs_document_state(sid, data):
|
||||
room = f"doc_{document_id}"
|
||||
|
||||
active_session_ids = get_session_ids_from_room(room)
|
||||
print(active_session_ids)
|
||||
|
||||
if sid not in active_session_ids:
|
||||
log.warning(f"Session {sid} not in room {room}. Cannot send state.")
|
||||
return
|
||||
@@ -520,7 +551,8 @@ async def yjs_document_update(sid, data):
|
||||
document_id, data.get("data", {}), SESSION_POOL.get(sid)
|
||||
)
|
||||
|
||||
await create_task(REDIS, debounced_save(), document_id)
|
||||
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}")
|
||||
|
||||
@@ -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")
|
||||
@@ -4,6 +4,7 @@ import sys
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from loguru import logger
|
||||
from opentelemetry import trace
|
||||
|
||||
|
||||
from open_webui.env import (
|
||||
@@ -12,6 +13,7 @@ from open_webui.env import (
|
||||
AUDIT_LOG_LEVEL,
|
||||
AUDIT_LOGS_FILE_PATH,
|
||||
GLOBAL_LOG_LEVEL,
|
||||
ENABLE_OTEL,
|
||||
)
|
||||
|
||||
|
||||
@@ -60,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"):
|
||||
|
||||
@@ -653,13 +653,13 @@ async def chat_completion_files_handler(
|
||||
),
|
||||
k=request.app.state.config.TOP_K,
|
||||
reranking_function=(
|
||||
lambda sentences: (
|
||||
request.app.state.RERANKING_FUNCTION(
|
||||
(
|
||||
lambda sentences: request.app.state.RERANKING_FUNCTION(
|
||||
sentences, user=user
|
||||
)
|
||||
if request.app.state.RERANKING_FUNCTION
|
||||
else None
|
||||
)
|
||||
if request.app.state.RERANKING_FUNCTION
|
||||
else None
|
||||
),
|
||||
k_reranker=request.app.state.config.TOP_K_RERANKER,
|
||||
r=request.app.state.config.RELEVANCE_THRESHOLD,
|
||||
@@ -1472,12 +1472,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'
|
||||
|
||||
@@ -1574,8 +1574,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 = (
|
||||
@@ -1626,8 +1634,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):
|
||||
@@ -1699,8 +1712,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,
|
||||
@@ -1744,18 +1766,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:
|
||||
@@ -2039,7 +2062,7 @@ async def process_chat_response(
|
||||
if done:
|
||||
pass
|
||||
else:
|
||||
log.debug("Error: ", e)
|
||||
log.debug(f"Error: {e}")
|
||||
continue
|
||||
|
||||
if content_blocks:
|
||||
|
||||
@@ -1,6 +1,94 @@
|
||||
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__)
|
||||
|
||||
|
||||
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):
|
||||
@@ -34,7 +122,11 @@ def get_redis_connection(
|
||||
password=redis_config["password"],
|
||||
decode_responses=decode_responses,
|
||||
)
|
||||
return sentinel.master_for(redis_config["service"])
|
||||
return SentinelRedisProxy(
|
||||
sentinel,
|
||||
redis_config["service"],
|
||||
async_mode=async_mode,
|
||||
)
|
||||
elif redis_url:
|
||||
return redis.from_url(redis_url, decode_responses=decode_responses)
|
||||
else:
|
||||
@@ -52,7 +144,11 @@ def get_redis_connection(
|
||||
password=redis_config["password"],
|
||||
decode_responses=decode_responses,
|
||||
)
|
||||
return sentinel.master_for(redis_config["service"])
|
||||
return SentinelRedisProxy(
|
||||
sentinel,
|
||||
redis_config["service"],
|
||||
async_mode=async_mode,
|
||||
)
|
||||
elif redis_url:
|
||||
return redis.Redis.from_url(redis_url, decode_responses=decode_responses)
|
||||
else:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -51,6 +51,7 @@ langchain-community==0.3.26
|
||||
|
||||
fake-useragent==2.1.0
|
||||
chromadb==0.6.3
|
||||
posthog==5.4.0
|
||||
pymilvus==2.5.0
|
||||
qdrant-client==1.14.3
|
||||
opensearch-py==2.8.0
|
||||
|
||||
Reference in New Issue
Block a user