mirror of
https://github.com/open-webui/open-webui
synced 2025-06-26 18:26:48 +00:00
Merge branch 'open-webui:main' into fix/oidc-500-error-name-field
This commit is contained in:
@@ -5,12 +5,31 @@ from pathlib import Path
|
||||
|
||||
import typer
|
||||
import uvicorn
|
||||
from typing import Optional
|
||||
from typing_extensions import Annotated
|
||||
|
||||
app = typer.Typer()
|
||||
|
||||
KEY_FILE = Path.cwd() / ".webui_secret_key"
|
||||
|
||||
|
||||
def version_callback(value: bool):
|
||||
if value:
|
||||
from open_webui.env import VERSION
|
||||
|
||||
typer.echo(f"Open WebUI version: {VERSION}")
|
||||
raise typer.Exit()
|
||||
|
||||
|
||||
@app.command()
|
||||
def main(
|
||||
version: Annotated[
|
||||
Optional[bool], typer.Option("--version", callback=version_callback)
|
||||
] = None,
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
@app.command()
|
||||
def serve(
|
||||
host: str = "0.0.0.0",
|
||||
|
||||
@@ -9,22 +9,22 @@ from urllib.parse import urlparse
|
||||
|
||||
import chromadb
|
||||
import requests
|
||||
import yaml
|
||||
from open_webui.internal.db import Base, get_db
|
||||
from pydantic import BaseModel
|
||||
from sqlalchemy import JSON, Column, DateTime, Integer, func
|
||||
|
||||
from open_webui.env import (
|
||||
OPEN_WEBUI_DIR,
|
||||
DATA_DIR,
|
||||
DATABASE_URL,
|
||||
ENV,
|
||||
FRONTEND_BUILD_DIR,
|
||||
OFFLINE_MODE,
|
||||
OPEN_WEBUI_DIR,
|
||||
WEBUI_AUTH,
|
||||
WEBUI_FAVICON_URL,
|
||||
WEBUI_NAME,
|
||||
log,
|
||||
DATABASE_URL,
|
||||
OFFLINE_MODE,
|
||||
)
|
||||
from pydantic import BaseModel
|
||||
from sqlalchemy import JSON, Column, DateTime, Integer, func
|
||||
from open_webui.internal.db import Base, get_db
|
||||
|
||||
|
||||
class EndpointFilter(logging.Filter):
|
||||
@@ -362,6 +362,30 @@ MICROSOFT_REDIRECT_URI = PersistentConfig(
|
||||
os.environ.get("MICROSOFT_REDIRECT_URI", ""),
|
||||
)
|
||||
|
||||
GITHUB_CLIENT_ID = PersistentConfig(
|
||||
"GITHUB_CLIENT_ID",
|
||||
"oauth.github.client_id",
|
||||
os.environ.get("GITHUB_CLIENT_ID", ""),
|
||||
)
|
||||
|
||||
GITHUB_CLIENT_SECRET = PersistentConfig(
|
||||
"GITHUB_CLIENT_SECRET",
|
||||
"oauth.github.client_secret",
|
||||
os.environ.get("GITHUB_CLIENT_SECRET", ""),
|
||||
)
|
||||
|
||||
GITHUB_CLIENT_SCOPE = PersistentConfig(
|
||||
"GITHUB_CLIENT_SCOPE",
|
||||
"oauth.github.scope",
|
||||
os.environ.get("GITHUB_CLIENT_SCOPE", "user:email"),
|
||||
)
|
||||
|
||||
GITHUB_CLIENT_REDIRECT_URI = PersistentConfig(
|
||||
"GITHUB_CLIENT_REDIRECT_URI",
|
||||
"oauth.github.redirect_uri",
|
||||
os.environ.get("GITHUB_CLIENT_REDIRECT_URI", ""),
|
||||
)
|
||||
|
||||
OAUTH_CLIENT_ID = PersistentConfig(
|
||||
"OAUTH_CLIENT_ID",
|
||||
"oauth.oidc.client_id",
|
||||
@@ -468,12 +492,20 @@ OAUTH_ALLOWED_DOMAINS = PersistentConfig(
|
||||
def load_oauth_providers():
|
||||
OAUTH_PROVIDERS.clear()
|
||||
if GOOGLE_CLIENT_ID.value and GOOGLE_CLIENT_SECRET.value:
|
||||
|
||||
def google_oauth_register(client):
|
||||
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},
|
||||
redirect_uri=GOOGLE_REDIRECT_URI.value,
|
||||
)
|
||||
|
||||
OAUTH_PROVIDERS["google"] = {
|
||||
"client_id": GOOGLE_CLIENT_ID.value,
|
||||
"client_secret": GOOGLE_CLIENT_SECRET.value,
|
||||
"server_metadata_url": "https://accounts.google.com/.well-known/openid-configuration",
|
||||
"scope": GOOGLE_OAUTH_SCOPE.value,
|
||||
"redirect_uri": GOOGLE_REDIRECT_URI.value,
|
||||
"register": google_oauth_register,
|
||||
}
|
||||
|
||||
if (
|
||||
@@ -481,12 +513,44 @@ def load_oauth_providers():
|
||||
and MICROSOFT_CLIENT_SECRET.value
|
||||
and MICROSOFT_CLIENT_TENANT_ID.value
|
||||
):
|
||||
|
||||
def microsoft_oauth_register(client):
|
||||
client.register(
|
||||
name="microsoft",
|
||||
client_id=MICROSOFT_CLIENT_ID.value,
|
||||
client_secret=MICROSOFT_CLIENT_SECRET.value,
|
||||
server_metadata_url=f"https://login.microsoftonline.com/{MICROSOFT_CLIENT_TENANT_ID.value}/v2.0/.well-known/openid-configuration",
|
||||
client_kwargs={
|
||||
"scope": MICROSOFT_OAUTH_SCOPE.value,
|
||||
},
|
||||
redirect_uri=MICROSOFT_REDIRECT_URI.value,
|
||||
)
|
||||
|
||||
OAUTH_PROVIDERS["microsoft"] = {
|
||||
"client_id": MICROSOFT_CLIENT_ID.value,
|
||||
"client_secret": MICROSOFT_CLIENT_SECRET.value,
|
||||
"server_metadata_url": f"https://login.microsoftonline.com/{MICROSOFT_CLIENT_TENANT_ID.value}/v2.0/.well-known/openid-configuration",
|
||||
"scope": MICROSOFT_OAUTH_SCOPE.value,
|
||||
"redirect_uri": MICROSOFT_REDIRECT_URI.value,
|
||||
"picture_url": "https://graph.microsoft.com/v1.0/me/photo/$value",
|
||||
"register": microsoft_oauth_register,
|
||||
}
|
||||
|
||||
if GITHUB_CLIENT_ID.value and GITHUB_CLIENT_SECRET.value:
|
||||
|
||||
def github_oauth_register(client):
|
||||
client.register(
|
||||
name="github",
|
||||
client_id=GITHUB_CLIENT_ID.value,
|
||||
client_secret=GITHUB_CLIENT_SECRET.value,
|
||||
access_token_url="https://github.com/login/oauth/access_token",
|
||||
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},
|
||||
redirect_uri=GITHUB_CLIENT_REDIRECT_URI.value,
|
||||
)
|
||||
|
||||
OAUTH_PROVIDERS["github"] = {
|
||||
"redirect_uri": GITHUB_CLIENT_REDIRECT_URI.value,
|
||||
"register": github_oauth_register,
|
||||
"sub_claim": "id",
|
||||
}
|
||||
|
||||
if (
|
||||
@@ -494,13 +558,23 @@ def load_oauth_providers():
|
||||
and OAUTH_CLIENT_SECRET.value
|
||||
and OPENID_PROVIDER_URL.value
|
||||
):
|
||||
|
||||
def oidc_oauth_register(client):
|
||||
client.register(
|
||||
name="oidc",
|
||||
client_id=OAUTH_CLIENT_ID.value,
|
||||
client_secret=OAUTH_CLIENT_SECRET.value,
|
||||
server_metadata_url=OPENID_PROVIDER_URL.value,
|
||||
client_kwargs={
|
||||
"scope": OAUTH_SCOPES.value,
|
||||
},
|
||||
redirect_uri=OPENID_REDIRECT_URI.value,
|
||||
)
|
||||
|
||||
OAUTH_PROVIDERS["oidc"] = {
|
||||
"client_id": OAUTH_CLIENT_ID.value,
|
||||
"client_secret": OAUTH_CLIENT_SECRET.value,
|
||||
"server_metadata_url": OPENID_PROVIDER_URL.value,
|
||||
"scope": OAUTH_SCOPES.value,
|
||||
"name": OAUTH_PROVIDER_NAME.value,
|
||||
"redirect_uri": OPENID_REDIRECT_URI.value,
|
||||
"register": oidc_oauth_register,
|
||||
}
|
||||
|
||||
|
||||
@@ -580,7 +654,7 @@ if CUSTOM_NAME:
|
||||
# STORAGE PROVIDER
|
||||
####################################
|
||||
|
||||
STORAGE_PROVIDER = os.environ.get("STORAGE_PROVIDER", "") # defaults to local, s3
|
||||
STORAGE_PROVIDER = os.environ.get("STORAGE_PROVIDER", "local") # defaults to local, s3
|
||||
|
||||
S3_ACCESS_KEY_ID = os.environ.get("S3_ACCESS_KEY_ID", None)
|
||||
S3_SECRET_ACCESS_KEY = os.environ.get("S3_SECRET_ACCESS_KEY", None)
|
||||
@@ -588,6 +662,11 @@ S3_REGION_NAME = os.environ.get("S3_REGION_NAME", None)
|
||||
S3_BUCKET_NAME = os.environ.get("S3_BUCKET_NAME", None)
|
||||
S3_ENDPOINT_URL = os.environ.get("S3_ENDPOINT_URL", None)
|
||||
|
||||
GCS_BUCKET_NAME = os.environ.get("GCS_BUCKET_NAME", None)
|
||||
GOOGLE_APPLICATION_CREDENTIALS_JSON = os.environ.get(
|
||||
"GOOGLE_APPLICATION_CREDENTIALS_JSON", None
|
||||
)
|
||||
|
||||
####################################
|
||||
# File Upload DIR
|
||||
####################################
|
||||
@@ -819,6 +898,10 @@ USER_PERMISSIONS_WORKSPACE_TOOLS_ACCESS = (
|
||||
os.environ.get("USER_PERMISSIONS_WORKSPACE_TOOLS_ACCESS", "False").lower() == "true"
|
||||
)
|
||||
|
||||
USER_PERMISSIONS_CHAT_CONTROLS = (
|
||||
os.environ.get("USER_PERMISSIONS_CHAT_CONTROLS", "True").lower() == "true"
|
||||
)
|
||||
|
||||
USER_PERMISSIONS_CHAT_FILE_UPLOAD = (
|
||||
os.environ.get("USER_PERMISSIONS_CHAT_FILE_UPLOAD", "True").lower() == "true"
|
||||
)
|
||||
@@ -835,23 +918,39 @@ USER_PERMISSIONS_CHAT_TEMPORARY = (
|
||||
os.environ.get("USER_PERMISSIONS_CHAT_TEMPORARY", "True").lower() == "true"
|
||||
)
|
||||
|
||||
USER_PERMISSIONS_FEATURES_WEB_SEARCH = (
|
||||
os.environ.get("USER_PERMISSIONS_FEATURES_WEB_SEARCH", "True").lower() == "true"
|
||||
)
|
||||
|
||||
USER_PERMISSIONS_FEATURES_IMAGE_GENERATION = (
|
||||
os.environ.get("USER_PERMISSIONS_FEATURES_IMAGE_GENERATION", "True").lower()
|
||||
== "true"
|
||||
)
|
||||
|
||||
DEFAULT_USER_PERMISSIONS = {
|
||||
"workspace": {
|
||||
"models": USER_PERMISSIONS_WORKSPACE_MODELS_ACCESS,
|
||||
"knowledge": USER_PERMISSIONS_WORKSPACE_KNOWLEDGE_ACCESS,
|
||||
"prompts": USER_PERMISSIONS_WORKSPACE_PROMPTS_ACCESS,
|
||||
"tools": USER_PERMISSIONS_WORKSPACE_TOOLS_ACCESS,
|
||||
},
|
||||
"chat": {
|
||||
"controls": USER_PERMISSIONS_CHAT_CONTROLS,
|
||||
"file_upload": USER_PERMISSIONS_CHAT_FILE_UPLOAD,
|
||||
"delete": USER_PERMISSIONS_CHAT_DELETE,
|
||||
"edit": USER_PERMISSIONS_CHAT_EDIT,
|
||||
"temporary": USER_PERMISSIONS_CHAT_TEMPORARY,
|
||||
},
|
||||
"features": {
|
||||
"web_search": USER_PERMISSIONS_FEATURES_WEB_SEARCH,
|
||||
"image_generation": USER_PERMISSIONS_FEATURES_IMAGE_GENERATION,
|
||||
},
|
||||
}
|
||||
|
||||
USER_PERMISSIONS = PersistentConfig(
|
||||
"USER_PERMISSIONS",
|
||||
"user.permissions",
|
||||
{
|
||||
"workspace": {
|
||||
"models": USER_PERMISSIONS_WORKSPACE_MODELS_ACCESS,
|
||||
"knowledge": USER_PERMISSIONS_WORKSPACE_KNOWLEDGE_ACCESS,
|
||||
"prompts": USER_PERMISSIONS_WORKSPACE_PROMPTS_ACCESS,
|
||||
"tools": USER_PERMISSIONS_WORKSPACE_TOOLS_ACCESS,
|
||||
},
|
||||
"chat": {
|
||||
"file_upload": USER_PERMISSIONS_CHAT_FILE_UPLOAD,
|
||||
"delete": USER_PERMISSIONS_CHAT_DELETE,
|
||||
"edit": USER_PERMISSIONS_CHAT_EDIT,
|
||||
"temporary": USER_PERMISSIONS_CHAT_TEMPORARY,
|
||||
},
|
||||
},
|
||||
DEFAULT_USER_PERMISSIONS,
|
||||
)
|
||||
|
||||
ENABLE_CHANNELS = PersistentConfig(
|
||||
@@ -1034,6 +1133,32 @@ JSON format: { "tags": ["tag1", "tag2", "tag3"] }
|
||||
{{MESSAGES:END:6}}
|
||||
</chat_history>"""
|
||||
|
||||
IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE = PersistentConfig(
|
||||
"IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE",
|
||||
"task.image.prompt_template",
|
||||
os.environ.get("IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE", ""),
|
||||
)
|
||||
|
||||
DEFAULT_IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE = """### Task:
|
||||
Generate a detailed prompt for am image generation task based on the given language and context. Describe the image as if you were explaining it to someone who cannot see it. Include relevant details, colors, shapes, and any other important elements.
|
||||
|
||||
### Guidelines:
|
||||
- Be descriptive and detailed, focusing on the most important aspects of the image.
|
||||
- Avoid making assumptions or adding information not present in the image.
|
||||
- Use the chat's primary language; default to English if multilingual.
|
||||
- If the image is too complex, focus on the most prominent elements.
|
||||
|
||||
### Output:
|
||||
Strictly return in JSON format:
|
||||
{
|
||||
"prompt": "Your detailed description here."
|
||||
}
|
||||
|
||||
### Chat History:
|
||||
<chat_history>
|
||||
{{MESSAGES:END:6}}
|
||||
</chat_history>"""
|
||||
|
||||
ENABLE_TAGS_GENERATION = PersistentConfig(
|
||||
"ENABLE_TAGS_GENERATION",
|
||||
"task.tags.enable",
|
||||
@@ -1193,6 +1318,7 @@ CHROMA_HTTP_SSL = os.environ.get("CHROMA_HTTP_SSL", "false").lower() == "true"
|
||||
# Milvus
|
||||
|
||||
MILVUS_URI = os.environ.get("MILVUS_URI", f"{DATA_DIR}/vector_db/milvus.db")
|
||||
MILVUS_DB = os.environ.get("MILVUS_DB", "default")
|
||||
|
||||
# Qdrant
|
||||
QDRANT_URI = os.environ.get("QDRANT_URI", None)
|
||||
@@ -1602,6 +1728,13 @@ ENABLE_IMAGE_GENERATION = PersistentConfig(
|
||||
"image_generation.enable",
|
||||
os.environ.get("ENABLE_IMAGE_GENERATION", "").lower() == "true",
|
||||
)
|
||||
|
||||
ENABLE_IMAGE_PROMPT_GENERATION = PersistentConfig(
|
||||
"ENABLE_IMAGE_PROMPT_GENERATION",
|
||||
"image_generation.prompt.enable",
|
||||
os.environ.get("ENABLE_IMAGE_PROMPT_GENERATION", "true").lower() == "true",
|
||||
)
|
||||
|
||||
AUTOMATIC1111_BASE_URL = PersistentConfig(
|
||||
"AUTOMATIC1111_BASE_URL",
|
||||
"image_generation.automatic1111.base_url",
|
||||
@@ -1931,6 +2064,12 @@ LDAP_SERVER_PORT = PersistentConfig(
|
||||
int(os.environ.get("LDAP_SERVER_PORT", "389")),
|
||||
)
|
||||
|
||||
LDAP_ATTRIBUTE_FOR_MAIL = PersistentConfig(
|
||||
"LDAP_ATTRIBUTE_FOR_MAIL",
|
||||
"ldap.server.attribute_for_mail",
|
||||
os.environ.get("LDAP_ATTRIBUTE_FOR_MAIL", "mail"),
|
||||
)
|
||||
|
||||
LDAP_ATTRIBUTE_FOR_USERNAME = PersistentConfig(
|
||||
"LDAP_ATTRIBUTE_FOR_USERNAME",
|
||||
"ldap.server.attribute_for_username",
|
||||
|
||||
@@ -113,6 +113,7 @@ class TASKS(str, Enum):
|
||||
TAGS_GENERATION = "tags_generation"
|
||||
EMOJI_GENERATION = "emoji_generation"
|
||||
QUERY_GENERATION = "query_generation"
|
||||
IMAGE_PROMPT_GENERATION = "image_prompt_generation"
|
||||
AUTOCOMPLETE_GENERATION = "autocomplete_generation"
|
||||
FUNCTION_CALLING = "function_calling"
|
||||
MOA_RESPONSE_GENERATION = "moa_response_generation"
|
||||
|
||||
@@ -274,6 +274,8 @@ DATABASE_URL = os.environ.get("DATABASE_URL", f"sqlite:///{DATA_DIR}/webui.db")
|
||||
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)
|
||||
|
||||
if DATABASE_POOL_SIZE == "":
|
||||
|
||||
@@ -7,6 +7,7 @@ from open_webui.internal.wrappers import register_connection
|
||||
from open_webui.env import (
|
||||
OPEN_WEBUI_DIR,
|
||||
DATABASE_URL,
|
||||
DATABASE_SCHEMA,
|
||||
SRC_LOG_LEVELS,
|
||||
DATABASE_POOL_MAX_OVERFLOW,
|
||||
DATABASE_POOL_RECYCLE,
|
||||
@@ -14,7 +15,7 @@ from open_webui.env import (
|
||||
DATABASE_POOL_TIMEOUT,
|
||||
)
|
||||
from peewee_migrate import Router
|
||||
from sqlalchemy import Dialect, create_engine, types
|
||||
from sqlalchemy import Dialect, create_engine, MetaData, types
|
||||
from sqlalchemy.ext.declarative import declarative_base
|
||||
from sqlalchemy.orm import scoped_session, sessionmaker
|
||||
from sqlalchemy.pool import QueuePool, NullPool
|
||||
@@ -99,7 +100,8 @@ else:
|
||||
SessionLocal = sessionmaker(
|
||||
autocommit=False, autoflush=False, bind=engine, expire_on_commit=False
|
||||
)
|
||||
Base = declarative_base()
|
||||
metadata_obj = MetaData(schema=DATABASE_SCHEMA)
|
||||
Base = declarative_base(metadata=metadata_obj)
|
||||
Session = scoped_session(SessionLocal)
|
||||
|
||||
|
||||
|
||||
@@ -108,6 +108,7 @@ from open_webui.config import (
|
||||
COMFYUI_WORKFLOW,
|
||||
COMFYUI_WORKFLOW_NODES,
|
||||
ENABLE_IMAGE_GENERATION,
|
||||
ENABLE_IMAGE_PROMPT_GENERATION,
|
||||
IMAGE_GENERATION_ENGINE,
|
||||
IMAGE_GENERATION_MODEL,
|
||||
IMAGE_SIZE,
|
||||
@@ -225,6 +226,7 @@ from open_webui.config import (
|
||||
LDAP_SERVER_LABEL,
|
||||
LDAP_SERVER_HOST,
|
||||
LDAP_SERVER_PORT,
|
||||
LDAP_ATTRIBUTE_FOR_MAIL,
|
||||
LDAP_ATTRIBUTE_FOR_USERNAME,
|
||||
LDAP_SEARCH_FILTERS,
|
||||
LDAP_SEARCH_BASE,
|
||||
@@ -254,6 +256,7 @@ from open_webui.config import (
|
||||
ENABLE_AUTOCOMPLETE_GENERATION,
|
||||
TITLE_GENERATION_PROMPT_TEMPLATE,
|
||||
TAGS_GENERATION_PROMPT_TEMPLATE,
|
||||
IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE,
|
||||
TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE,
|
||||
QUERY_GENERATION_PROMPT_TEMPLATE,
|
||||
AUTOCOMPLETE_GENERATION_PROMPT_TEMPLATE,
|
||||
@@ -437,6 +440,7 @@ app.state.config.ENABLE_LDAP = ENABLE_LDAP
|
||||
app.state.config.LDAP_SERVER_LABEL = LDAP_SERVER_LABEL
|
||||
app.state.config.LDAP_SERVER_HOST = LDAP_SERVER_HOST
|
||||
app.state.config.LDAP_SERVER_PORT = LDAP_SERVER_PORT
|
||||
app.state.config.LDAP_ATTRIBUTE_FOR_MAIL = LDAP_ATTRIBUTE_FOR_MAIL
|
||||
app.state.config.LDAP_ATTRIBUTE_FOR_USERNAME = LDAP_ATTRIBUTE_FOR_USERNAME
|
||||
app.state.config.LDAP_APP_DN = LDAP_APP_DN
|
||||
app.state.config.LDAP_APP_PASSWORD = LDAP_APP_PASSWORD
|
||||
@@ -572,6 +576,7 @@ app.state.EMBEDDING_FUNCTION = get_embedding_function(
|
||||
|
||||
app.state.config.IMAGE_GENERATION_ENGINE = IMAGE_GENERATION_ENGINE
|
||||
app.state.config.ENABLE_IMAGE_GENERATION = ENABLE_IMAGE_GENERATION
|
||||
app.state.config.ENABLE_IMAGE_PROMPT_GENERATION = ENABLE_IMAGE_PROMPT_GENERATION
|
||||
|
||||
app.state.config.IMAGES_OPENAI_API_BASE_URL = IMAGES_OPENAI_API_BASE_URL
|
||||
app.state.config.IMAGES_OPENAI_API_KEY = IMAGES_OPENAI_API_KEY
|
||||
@@ -642,6 +647,10 @@ app.state.config.ENABLE_TAGS_GENERATION = ENABLE_TAGS_GENERATION
|
||||
|
||||
app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE = TITLE_GENERATION_PROMPT_TEMPLATE
|
||||
app.state.config.TAGS_GENERATION_PROMPT_TEMPLATE = TAGS_GENERATION_PROMPT_TEMPLATE
|
||||
app.state.config.IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE = (
|
||||
IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE
|
||||
)
|
||||
|
||||
app.state.config.TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE = (
|
||||
TOOLS_FUNCTION_CALLING_PROMPT_TEMPLATE
|
||||
)
|
||||
|
||||
@@ -393,7 +393,7 @@ class ChatTable:
|
||||
limit: int = 50,
|
||||
) -> list[ChatModel]:
|
||||
with get_db() as db:
|
||||
query = db.query(Chat).filter_by(user_id=user_id).filter_by(folder_id=None)
|
||||
query = db.query(Chat).filter_by(user_id=user_id)
|
||||
if not include_archived:
|
||||
query = query.filter_by(archived=False)
|
||||
|
||||
|
||||
@@ -80,12 +80,11 @@ class GroupResponse(BaseModel):
|
||||
class GroupForm(BaseModel):
|
||||
name: str
|
||||
description: str
|
||||
permissions: Optional[dict] = None
|
||||
|
||||
|
||||
class GroupUpdateForm(GroupForm):
|
||||
permissions: Optional[dict] = None
|
||||
user_ids: Optional[list[str]] = None
|
||||
admin_ids: Optional[list[str]] = None
|
||||
|
||||
|
||||
class GroupTable:
|
||||
@@ -95,7 +94,7 @@ class GroupTable:
|
||||
with get_db() as db:
|
||||
group = GroupModel(
|
||||
**{
|
||||
**form_data.model_dump(),
|
||||
**form_data.model_dump(exclude_none=True),
|
||||
"id": str(uuid.uuid4()),
|
||||
"user_id": user_id,
|
||||
"created_at": int(time.time()),
|
||||
@@ -189,5 +188,24 @@ class GroupTable:
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def remove_user_from_all_groups(self, user_id: str) -> bool:
|
||||
with get_db() as db:
|
||||
try:
|
||||
groups = self.get_groups_by_member_id(user_id)
|
||||
|
||||
for group in groups:
|
||||
group.user_ids.remove(user_id)
|
||||
db.query(Group).filter_by(id=group.id).update(
|
||||
{
|
||||
"user_ids": group.user_ids,
|
||||
"updated_at": int(time.time()),
|
||||
}
|
||||
)
|
||||
db.commit()
|
||||
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
Groups = GroupTable()
|
||||
|
||||
@@ -2,7 +2,12 @@ import time
|
||||
from typing import Optional
|
||||
|
||||
from open_webui.internal.db import Base, JSONField, get_db
|
||||
|
||||
|
||||
from open_webui.models.chats import Chats
|
||||
from open_webui.models.groups import Groups
|
||||
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from sqlalchemy import BigInteger, Column, String, Text
|
||||
|
||||
@@ -268,9 +273,11 @@ class UsersTable:
|
||||
|
||||
def delete_user_by_id(self, id: str) -> bool:
|
||||
try:
|
||||
# Remove User from Groups
|
||||
Groups.remove_user_from_all_groups(id)
|
||||
|
||||
# Delete User Chats
|
||||
result = Chats.delete_chats_by_user_id(id)
|
||||
|
||||
if result:
|
||||
with get_db() as db:
|
||||
# Delete User
|
||||
@@ -300,5 +307,10 @@ class UsersTable:
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def get_valid_user_ids(self, user_ids: list[str]) -> list[str]:
|
||||
with get_db() as db:
|
||||
users = db.query(User).filter(User.id.in_(user_ids)).all()
|
||||
return [user.id for user in users]
|
||||
|
||||
|
||||
Users = UsersTable()
|
||||
|
||||
@@ -7,13 +7,14 @@ from typing import Optional
|
||||
from open_webui.retrieval.vector.main import VectorItem, SearchResult, GetResult
|
||||
from open_webui.config import (
|
||||
MILVUS_URI,
|
||||
MILVUS_DB,
|
||||
)
|
||||
|
||||
|
||||
class MilvusClient:
|
||||
def __init__(self):
|
||||
self.collection_prefix = "open_webui"
|
||||
self.client = Client(uri=MILVUS_URI)
|
||||
self.client = Client(uri=MILVUS_URI, database=MILVUS_DB)
|
||||
|
||||
def _result_to_get_result(self, result) -> GetResult:
|
||||
ids = []
|
||||
|
||||
@@ -23,7 +23,7 @@ def search_bing(
|
||||
filter_list: Optional[list[str]] = None,
|
||||
) -> list[SearchResult]:
|
||||
mkt = locale
|
||||
params = {"q": query, "mkt": mkt, "answerCount": count}
|
||||
params = {"q": query, "mkt": mkt, "count": count}
|
||||
headers = {"Ocp-Apim-Subscription-Key": subscription_key}
|
||||
|
||||
try:
|
||||
|
||||
@@ -170,6 +170,7 @@ async def ldap_auth(request: Request, response: Response, form_data: LdapForm):
|
||||
LDAP_SERVER_LABEL = request.app.state.config.LDAP_SERVER_LABEL
|
||||
LDAP_SERVER_HOST = request.app.state.config.LDAP_SERVER_HOST
|
||||
LDAP_SERVER_PORT = request.app.state.config.LDAP_SERVER_PORT
|
||||
LDAP_ATTRIBUTE_FOR_MAIL = request.app.state.config.LDAP_ATTRIBUTE_FOR_MAIL
|
||||
LDAP_ATTRIBUTE_FOR_USERNAME = request.app.state.config.LDAP_ATTRIBUTE_FOR_USERNAME
|
||||
LDAP_SEARCH_BASE = request.app.state.config.LDAP_SEARCH_BASE
|
||||
LDAP_SEARCH_FILTERS = request.app.state.config.LDAP_SEARCH_FILTERS
|
||||
@@ -218,7 +219,11 @@ async def ldap_auth(request: Request, response: Response, form_data: LdapForm):
|
||||
search_success = connection_app.search(
|
||||
search_base=LDAP_SEARCH_BASE,
|
||||
search_filter=f"(&({LDAP_ATTRIBUTE_FOR_USERNAME}={escape_filter_chars(form_data.user.lower())}){LDAP_SEARCH_FILTERS})",
|
||||
attributes=[f"{LDAP_ATTRIBUTE_FOR_USERNAME}", "mail", "cn"],
|
||||
attributes=[
|
||||
f"{LDAP_ATTRIBUTE_FOR_USERNAME}",
|
||||
f"{LDAP_ATTRIBUTE_FOR_MAIL}",
|
||||
"cn",
|
||||
],
|
||||
)
|
||||
|
||||
if not search_success:
|
||||
@@ -226,7 +231,9 @@ async def ldap_auth(request: Request, response: Response, form_data: LdapForm):
|
||||
|
||||
entry = connection_app.entries[0]
|
||||
username = str(entry[f"{LDAP_ATTRIBUTE_FOR_USERNAME}"]).lower()
|
||||
mail = str(entry["mail"])
|
||||
mail = str(entry[f"{LDAP_ATTRIBUTE_FOR_MAIL}"])
|
||||
if not mail or mail == "" or mail == "[]":
|
||||
raise HTTPException(400, f"User {form_data.user} does not have mail.")
|
||||
cn = str(entry["cn"])
|
||||
user_dn = entry.entry_dn
|
||||
|
||||
@@ -691,6 +698,7 @@ class LdapServerConfig(BaseModel):
|
||||
label: str
|
||||
host: str
|
||||
port: Optional[int] = None
|
||||
attribute_for_mail: str = "mail"
|
||||
attribute_for_username: str = "uid"
|
||||
app_dn: str
|
||||
app_dn_password: str
|
||||
@@ -707,6 +715,7 @@ async def get_ldap_server(request: Request, user=Depends(get_admin_user)):
|
||||
"label": request.app.state.config.LDAP_SERVER_LABEL,
|
||||
"host": request.app.state.config.LDAP_SERVER_HOST,
|
||||
"port": request.app.state.config.LDAP_SERVER_PORT,
|
||||
"attribute_for_mail": request.app.state.config.LDAP_ATTRIBUTE_FOR_MAIL,
|
||||
"attribute_for_username": request.app.state.config.LDAP_ATTRIBUTE_FOR_USERNAME,
|
||||
"app_dn": request.app.state.config.LDAP_APP_DN,
|
||||
"app_dn_password": request.app.state.config.LDAP_APP_PASSWORD,
|
||||
@@ -725,6 +734,7 @@ async def update_ldap_server(
|
||||
required_fields = [
|
||||
"label",
|
||||
"host",
|
||||
"attribute_for_mail",
|
||||
"attribute_for_username",
|
||||
"app_dn",
|
||||
"app_dn_password",
|
||||
@@ -743,6 +753,7 @@ async def update_ldap_server(
|
||||
request.app.state.config.LDAP_SERVER_LABEL = form_data.label
|
||||
request.app.state.config.LDAP_SERVER_HOST = form_data.host
|
||||
request.app.state.config.LDAP_SERVER_PORT = form_data.port
|
||||
request.app.state.config.LDAP_ATTRIBUTE_FOR_MAIL = form_data.attribute_for_mail
|
||||
request.app.state.config.LDAP_ATTRIBUTE_FOR_USERNAME = (
|
||||
form_data.attribute_for_username
|
||||
)
|
||||
@@ -758,6 +769,7 @@ async def update_ldap_server(
|
||||
"label": request.app.state.config.LDAP_SERVER_LABEL,
|
||||
"host": request.app.state.config.LDAP_SERVER_HOST,
|
||||
"port": request.app.state.config.LDAP_SERVER_PORT,
|
||||
"attribute_for_mail": request.app.state.config.LDAP_ATTRIBUTE_FOR_MAIL,
|
||||
"attribute_for_username": request.app.state.config.LDAP_ATTRIBUTE_FOR_USERNAME,
|
||||
"app_dn": request.app.state.config.LDAP_APP_DN,
|
||||
"app_dn_password": request.app.state.config.LDAP_APP_PASSWORD,
|
||||
|
||||
@@ -345,6 +345,8 @@ async def get_file_content_by_id(id: str, user=Depends(get_verified_user)):
|
||||
async def delete_file_by_id(id: str, user=Depends(get_verified_user)):
|
||||
file = Files.get_file_by_id(id)
|
||||
if file and (file.user_id == user.id or user.role == "admin"):
|
||||
# We should add Chroma cleanup here
|
||||
|
||||
result = Files.delete_file_by_id(id)
|
||||
if result:
|
||||
try:
|
||||
|
||||
@@ -2,6 +2,8 @@ import os
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
|
||||
from open_webui.models.users import Users
|
||||
from open_webui.models.groups import (
|
||||
Groups,
|
||||
GroupForm,
|
||||
@@ -80,6 +82,9 @@ async def update_group_by_id(
|
||||
id: str, form_data: GroupUpdateForm, 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.update_group_by_id(id, form_data)
|
||||
if group:
|
||||
return group
|
||||
|
||||
@@ -43,6 +43,7 @@ async def get_config(request: Request, user=Depends(get_admin_user)):
|
||||
return {
|
||||
"enabled": request.app.state.config.ENABLE_IMAGE_GENERATION,
|
||||
"engine": request.app.state.config.IMAGE_GENERATION_ENGINE,
|
||||
"prompt_generation": request.app.state.config.ENABLE_IMAGE_PROMPT_GENERATION,
|
||||
"openai": {
|
||||
"OPENAI_API_BASE_URL": request.app.state.config.IMAGES_OPENAI_API_BASE_URL,
|
||||
"OPENAI_API_KEY": request.app.state.config.IMAGES_OPENAI_API_KEY,
|
||||
@@ -86,6 +87,7 @@ class ComfyUIConfigForm(BaseModel):
|
||||
class ConfigForm(BaseModel):
|
||||
enabled: bool
|
||||
engine: str
|
||||
prompt_generation: bool
|
||||
openai: OpenAIConfigForm
|
||||
automatic1111: Automatic1111ConfigForm
|
||||
comfyui: ComfyUIConfigForm
|
||||
@@ -98,6 +100,10 @@ async def update_config(
|
||||
request.app.state.config.IMAGE_GENERATION_ENGINE = form_data.engine
|
||||
request.app.state.config.ENABLE_IMAGE_GENERATION = form_data.enabled
|
||||
|
||||
request.app.state.config.ENABLE_IMAGE_PROMPT_GENERATION = (
|
||||
form_data.prompt_generation
|
||||
)
|
||||
|
||||
request.app.state.config.IMAGES_OPENAI_API_BASE_URL = (
|
||||
form_data.openai.OPENAI_API_BASE_URL
|
||||
)
|
||||
@@ -137,6 +143,7 @@ async def update_config(
|
||||
return {
|
||||
"enabled": request.app.state.config.ENABLE_IMAGE_GENERATION,
|
||||
"engine": request.app.state.config.IMAGE_GENERATION_ENGINE,
|
||||
"prompt_generation": request.app.state.config.ENABLE_IMAGE_PROMPT_GENERATION,
|
||||
"openai": {
|
||||
"OPENAI_API_BASE_URL": request.app.state.config.IMAGES_OPENAI_API_BASE_URL,
|
||||
"OPENAI_API_KEY": request.app.state.config.IMAGES_OPENAI_API_KEY,
|
||||
@@ -408,10 +415,14 @@ def save_b64_image(b64_str):
|
||||
return None
|
||||
|
||||
|
||||
def save_url_image(url):
|
||||
def save_url_image(url, headers=None):
|
||||
image_id = str(uuid.uuid4())
|
||||
try:
|
||||
r = requests.get(url)
|
||||
if headers:
|
||||
r = requests.get(url, headers=headers)
|
||||
else:
|
||||
r = requests.get(url)
|
||||
|
||||
r.raise_for_status()
|
||||
if r.headers["content-type"].split("/")[0] == "image":
|
||||
mime_type = r.headers["content-type"]
|
||||
@@ -535,7 +546,13 @@ async def image_generations(
|
||||
images = []
|
||||
|
||||
for image in res["data"]:
|
||||
image_filename = save_url_image(image["url"])
|
||||
headers = None
|
||||
if request.app.state.config.COMFYUI_API_KEY:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {request.app.state.config.COMFYUI_API_KEY}"
|
||||
}
|
||||
|
||||
image_filename = save_url_image(image["url"], headers)
|
||||
images.append({"url": f"/cache/image/generations/{image_filename}"})
|
||||
file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json")
|
||||
|
||||
|
||||
@@ -17,7 +17,7 @@ from open_webui.routers.retrieval import (
|
||||
process_files_batch,
|
||||
BatchProcessFilesForm,
|
||||
)
|
||||
|
||||
from open_webui.storage.provider import Storage
|
||||
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from open_webui.utils.auth import get_verified_user
|
||||
@@ -213,8 +213,12 @@ async def update_knowledge_by_id(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if knowledge.user_id != user.id and user.role != "admin":
|
||||
# Is the user the original creator, in a group with write access, or an admin
|
||||
if (
|
||||
knowledge.user_id != user.id
|
||||
and not has_access(user.id, "write", knowledge.access_control)
|
||||
and user.role != "admin"
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
@@ -420,6 +424,18 @@ def remove_file_from_knowledge_by_id(
|
||||
collection_name=knowledge.id, filter={"file_id": form_data.file_id}
|
||||
)
|
||||
|
||||
# Remove the file's collection from vector database
|
||||
file_collection = f"file-{form_data.file_id}"
|
||||
if VECTOR_DB_CLIENT.has_collection(collection_name=file_collection):
|
||||
VECTOR_DB_CLIENT.delete_collection(collection_name=file_collection)
|
||||
|
||||
# Delete physical file
|
||||
if file.path:
|
||||
Storage.delete_file(file.path)
|
||||
|
||||
# Delete file from database
|
||||
Files.delete_file_by_id(form_data.file_id)
|
||||
|
||||
if knowledge:
|
||||
data = knowledge.data or {}
|
||||
file_ids = data.get("file_ids", [])
|
||||
|
||||
@@ -155,6 +155,16 @@ async def update_model_by_id(
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if (
|
||||
model.user_id != user.id
|
||||
and not has_access(user.id, "write", model.access_control)
|
||||
and user.role != "admin"
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
model = Models.update_model_by_id(id, form_data)
|
||||
return model
|
||||
|
||||
|
||||
@@ -152,10 +152,12 @@ async def send_post_request(
|
||||
)
|
||||
|
||||
|
||||
def get_api_key(url, configs):
|
||||
def get_api_key(idx, url, configs):
|
||||
parsed_url = urlparse(url)
|
||||
base_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
|
||||
return configs.get(base_url, {}).get("key", None)
|
||||
return configs.get(str(idx), configs.get(base_url, {})).get(
|
||||
"key", None
|
||||
) # Legacy support
|
||||
|
||||
|
||||
##########################################
|
||||
@@ -238,11 +240,13 @@ async def update_config(
|
||||
request.app.state.config.OLLAMA_BASE_URLS = form_data.OLLAMA_BASE_URLS
|
||||
request.app.state.config.OLLAMA_API_CONFIGS = form_data.OLLAMA_API_CONFIGS
|
||||
|
||||
# Remove any extra configs
|
||||
config_urls = request.app.state.config.OLLAMA_API_CONFIGS.keys()
|
||||
for url in list(request.app.state.config.OLLAMA_BASE_URLS):
|
||||
if url not in config_urls:
|
||||
request.app.state.config.OLLAMA_API_CONFIGS.pop(url, None)
|
||||
# Remove the API configs that are not in the API URLS
|
||||
keys = list(map(str, range(len(request.app.state.config.OLLAMA_BASE_URLS))))
|
||||
request.app.state.config.OLLAMA_API_CONFIGS = {
|
||||
key: value
|
||||
for key, value in request.app.state.config.OLLAMA_API_CONFIGS.items()
|
||||
if key in keys
|
||||
}
|
||||
|
||||
return {
|
||||
"ENABLE_OLLAMA_API": request.app.state.config.ENABLE_OLLAMA_API,
|
||||
@@ -256,12 +260,19 @@ async def get_all_models(request: Request):
|
||||
log.info("get_all_models()")
|
||||
if request.app.state.config.ENABLE_OLLAMA_API:
|
||||
request_tasks = []
|
||||
|
||||
for idx, url in enumerate(request.app.state.config.OLLAMA_BASE_URLS):
|
||||
if url not in request.app.state.config.OLLAMA_API_CONFIGS:
|
||||
if (str(idx) not in request.app.state.config.OLLAMA_API_CONFIGS) and (
|
||||
url not in request.app.state.config.OLLAMA_API_CONFIGS # Legacy support
|
||||
):
|
||||
request_tasks.append(send_get_request(f"{url}/api/tags"))
|
||||
else:
|
||||
api_config = request.app.state.config.OLLAMA_API_CONFIGS.get(url, {})
|
||||
api_config = request.app.state.config.OLLAMA_API_CONFIGS.get(
|
||||
str(idx),
|
||||
request.app.state.config.OLLAMA_API_CONFIGS.get(
|
||||
url, {}
|
||||
), # Legacy support
|
||||
)
|
||||
|
||||
enable = api_config.get("enable", True)
|
||||
key = api_config.get("key", None)
|
||||
|
||||
@@ -275,7 +286,12 @@ async def get_all_models(request: Request):
|
||||
for idx, response in enumerate(responses):
|
||||
if response:
|
||||
url = request.app.state.config.OLLAMA_BASE_URLS[idx]
|
||||
api_config = request.app.state.config.OLLAMA_API_CONFIGS.get(url, {})
|
||||
api_config = request.app.state.config.OLLAMA_API_CONFIGS.get(
|
||||
str(idx),
|
||||
request.app.state.config.OLLAMA_API_CONFIGS.get(
|
||||
url, {}
|
||||
), # Legacy support
|
||||
)
|
||||
|
||||
prefix_id = api_config.get("prefix_id", None)
|
||||
model_ids = api_config.get("model_ids", [])
|
||||
@@ -349,7 +365,7 @@ async def get_ollama_tags(
|
||||
models = await get_all_models(request)
|
||||
else:
|
||||
url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
|
||||
key = get_api_key(url, request.app.state.config.OLLAMA_API_CONFIGS)
|
||||
key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS)
|
||||
|
||||
r = None
|
||||
try:
|
||||
@@ -393,11 +409,14 @@ async def get_ollama_versions(request: Request, url_idx: Optional[int] = None):
|
||||
request_tasks = [
|
||||
send_get_request(
|
||||
f"{url}/api/version",
|
||||
request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}).get(
|
||||
"key", None
|
||||
),
|
||||
request.app.state.config.OLLAMA_API_CONFIGS.get(
|
||||
str(idx),
|
||||
request.app.state.config.OLLAMA_API_CONFIGS.get(
|
||||
url, {}
|
||||
), # Legacy support
|
||||
).get("key", None),
|
||||
)
|
||||
for url in request.app.state.config.OLLAMA_BASE_URLS
|
||||
for idx, url in enumerate(request.app.state.config.OLLAMA_BASE_URLS)
|
||||
]
|
||||
responses = await asyncio.gather(*request_tasks)
|
||||
responses = list(filter(lambda x: x is not None, responses))
|
||||
@@ -454,11 +473,14 @@ async def get_ollama_loaded_models(request: Request, user=Depends(get_verified_u
|
||||
request_tasks = [
|
||||
send_get_request(
|
||||
f"{url}/api/ps",
|
||||
request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}).get(
|
||||
"key", None
|
||||
),
|
||||
request.app.state.config.OLLAMA_API_CONFIGS.get(
|
||||
str(idx),
|
||||
request.app.state.config.OLLAMA_API_CONFIGS.get(
|
||||
url, {}
|
||||
), # Legacy support
|
||||
).get("key", None),
|
||||
)
|
||||
for url in request.app.state.config.OLLAMA_BASE_URLS
|
||||
for idx, url in enumerate(request.app.state.config.OLLAMA_BASE_URLS)
|
||||
]
|
||||
responses = await asyncio.gather(*request_tasks)
|
||||
|
||||
@@ -488,7 +510,7 @@ async def pull_model(
|
||||
return await send_post_request(
|
||||
url=f"{url}/api/pull",
|
||||
payload=json.dumps(payload),
|
||||
key=get_api_key(url, request.app.state.config.OLLAMA_API_CONFIGS),
|
||||
key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
|
||||
)
|
||||
|
||||
|
||||
@@ -524,16 +546,17 @@ async def push_model(
|
||||
return await send_post_request(
|
||||
url=f"{url}/api/push",
|
||||
payload=form_data.model_dump_json(exclude_none=True).encode(),
|
||||
key=get_api_key(url, request.app.state.config.OLLAMA_API_CONFIGS),
|
||||
key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
|
||||
)
|
||||
|
||||
|
||||
class CreateModelForm(BaseModel):
|
||||
name: str
|
||||
modelfile: Optional[str] = None
|
||||
model: Optional[str] = None
|
||||
stream: Optional[bool] = None
|
||||
path: Optional[str] = None
|
||||
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
|
||||
@router.post("/api/create")
|
||||
@router.post("/api/create/{url_idx}")
|
||||
@@ -549,7 +572,7 @@ async def create_model(
|
||||
return await send_post_request(
|
||||
url=f"{url}/api/create",
|
||||
payload=form_data.model_dump_json(exclude_none=True).encode(),
|
||||
key=get_api_key(url, request.app.state.config.OLLAMA_API_CONFIGS),
|
||||
key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
|
||||
)
|
||||
|
||||
|
||||
@@ -579,7 +602,7 @@ async def copy_model(
|
||||
)
|
||||
|
||||
url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
|
||||
key = get_api_key(url, request.app.state.config.OLLAMA_API_CONFIGS)
|
||||
key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS)
|
||||
|
||||
try:
|
||||
r = requests.request(
|
||||
@@ -634,7 +657,7 @@ async def delete_model(
|
||||
)
|
||||
|
||||
url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
|
||||
key = get_api_key(url, request.app.state.config.OLLAMA_API_CONFIGS)
|
||||
key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS)
|
||||
|
||||
try:
|
||||
r = requests.request(
|
||||
@@ -684,7 +707,7 @@ async def show_model_info(
|
||||
url_idx = random.choice(models[form_data.name]["urls"])
|
||||
|
||||
url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
|
||||
key = get_api_key(url, request.app.state.config.OLLAMA_API_CONFIGS)
|
||||
key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS)
|
||||
|
||||
try:
|
||||
r = requests.request(
|
||||
@@ -753,7 +776,7 @@ async def embed(
|
||||
)
|
||||
|
||||
url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
|
||||
key = get_api_key(url, request.app.state.config.OLLAMA_API_CONFIGS)
|
||||
key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS)
|
||||
|
||||
try:
|
||||
r = requests.request(
|
||||
@@ -822,7 +845,7 @@ async def embeddings(
|
||||
)
|
||||
|
||||
url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
|
||||
key = get_api_key(url, request.app.state.config.OLLAMA_API_CONFIGS)
|
||||
key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS)
|
||||
|
||||
try:
|
||||
r = requests.request(
|
||||
@@ -897,7 +920,10 @@ async def generate_completion(
|
||||
)
|
||||
|
||||
url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
|
||||
api_config = request.app.state.config.OLLAMA_API_CONFIGS.get(url, {})
|
||||
api_config = request.app.state.config.OLLAMA_API_CONFIGS.get(
|
||||
str(url_idx),
|
||||
request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support
|
||||
)
|
||||
|
||||
prefix_id = api_config.get("prefix_id", None)
|
||||
if prefix_id:
|
||||
@@ -906,7 +932,7 @@ async def generate_completion(
|
||||
return await send_post_request(
|
||||
url=f"{url}/api/generate",
|
||||
payload=form_data.model_dump_json(exclude_none=True).encode(),
|
||||
key=get_api_key(url, request.app.state.config.OLLAMA_API_CONFIGS),
|
||||
key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
|
||||
)
|
||||
|
||||
|
||||
@@ -936,7 +962,7 @@ async def get_ollama_url(request: Request, model: str, url_idx: Optional[int] =
|
||||
)
|
||||
url_idx = random.choice(models[model].get("urls", []))
|
||||
url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
|
||||
return url
|
||||
return url, url_idx
|
||||
|
||||
|
||||
@router.post("/api/chat")
|
||||
@@ -1004,8 +1030,11 @@ async def generate_chat_completion(
|
||||
if ":" not in payload["model"]:
|
||||
payload["model"] = f"{payload['model']}:latest"
|
||||
|
||||
url = await get_ollama_url(request, payload["model"], url_idx)
|
||||
api_config = request.app.state.config.OLLAMA_API_CONFIGS.get(url, {})
|
||||
url, url_idx = await get_ollama_url(request, payload["model"], url_idx)
|
||||
api_config = request.app.state.config.OLLAMA_API_CONFIGS.get(
|
||||
str(url_idx),
|
||||
request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support
|
||||
)
|
||||
|
||||
prefix_id = api_config.get("prefix_id", None)
|
||||
if prefix_id:
|
||||
@@ -1015,7 +1044,7 @@ async def generate_chat_completion(
|
||||
url=f"{url}/api/chat",
|
||||
payload=json.dumps(payload),
|
||||
stream=form_data.stream,
|
||||
key=get_api_key(url, request.app.state.config.OLLAMA_API_CONFIGS),
|
||||
key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
|
||||
content_type="application/x-ndjson",
|
||||
)
|
||||
|
||||
@@ -1103,8 +1132,11 @@ async def generate_openai_completion(
|
||||
if ":" not in payload["model"]:
|
||||
payload["model"] = f"{payload['model']}:latest"
|
||||
|
||||
url = await get_ollama_url(request, payload["model"], url_idx)
|
||||
api_config = request.app.state.config.OLLAMA_API_CONFIGS.get(url, {})
|
||||
url, url_idx = await get_ollama_url(request, payload["model"], url_idx)
|
||||
api_config = request.app.state.config.OLLAMA_API_CONFIGS.get(
|
||||
str(url_idx),
|
||||
request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support
|
||||
)
|
||||
|
||||
prefix_id = api_config.get("prefix_id", None)
|
||||
|
||||
@@ -1115,7 +1147,7 @@ async def generate_openai_completion(
|
||||
url=f"{url}/v1/completions",
|
||||
payload=json.dumps(payload),
|
||||
stream=payload.get("stream", False),
|
||||
key=get_api_key(url, request.app.state.config.OLLAMA_API_CONFIGS),
|
||||
key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
|
||||
)
|
||||
|
||||
|
||||
@@ -1177,8 +1209,11 @@ async def generate_openai_chat_completion(
|
||||
if ":" not in payload["model"]:
|
||||
payload["model"] = f"{payload['model']}:latest"
|
||||
|
||||
url = await get_ollama_url(request, payload["model"], url_idx)
|
||||
api_config = request.app.state.config.OLLAMA_API_CONFIGS.get(url, {})
|
||||
url, url_idx = await get_ollama_url(request, payload["model"], url_idx)
|
||||
api_config = request.app.state.config.OLLAMA_API_CONFIGS.get(
|
||||
str(url_idx),
|
||||
request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support
|
||||
)
|
||||
|
||||
prefix_id = api_config.get("prefix_id", None)
|
||||
if prefix_id:
|
||||
@@ -1188,7 +1223,7 @@ async def generate_openai_chat_completion(
|
||||
url=f"{url}/v1/chat/completions",
|
||||
payload=json.dumps(payload),
|
||||
stream=payload.get("stream", False),
|
||||
key=get_api_key(url, request.app.state.config.OLLAMA_API_CONFIGS),
|
||||
key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -145,11 +145,13 @@ async def update_config(
|
||||
|
||||
request.app.state.config.OPENAI_API_CONFIGS = form_data.OPENAI_API_CONFIGS
|
||||
|
||||
# Remove any extra configs
|
||||
config_urls = request.app.state.config.OPENAI_API_CONFIGS.keys()
|
||||
for idx, url in enumerate(request.app.state.config.OPENAI_API_BASE_URLS):
|
||||
if url not in config_urls:
|
||||
request.app.state.config.OPENAI_API_CONFIGS.pop(url, None)
|
||||
# Remove the API configs that are not in the API URLS
|
||||
keys = list(map(str, range(len(request.app.state.config.OPENAI_API_BASE_URLS))))
|
||||
request.app.state.config.OPENAI_API_CONFIGS = {
|
||||
key: value
|
||||
for key, value in request.app.state.config.OPENAI_API_CONFIGS.items()
|
||||
if key in keys
|
||||
}
|
||||
|
||||
return {
|
||||
"ENABLE_OPENAI_API": request.app.state.config.ENABLE_OPENAI_API,
|
||||
@@ -264,14 +266,21 @@ async def get_all_models_responses(request: Request) -> list:
|
||||
|
||||
request_tasks = []
|
||||
for idx, url in enumerate(request.app.state.config.OPENAI_API_BASE_URLS):
|
||||
if url not in request.app.state.config.OPENAI_API_CONFIGS:
|
||||
if (str(idx) not in request.app.state.config.OPENAI_API_CONFIGS) and (
|
||||
url not in request.app.state.config.OPENAI_API_CONFIGS # Legacy support
|
||||
):
|
||||
request_tasks.append(
|
||||
send_get_request(
|
||||
f"{url}/models", request.app.state.config.OPENAI_API_KEYS[idx]
|
||||
)
|
||||
)
|
||||
else:
|
||||
api_config = request.app.state.config.OPENAI_API_CONFIGS.get(url, {})
|
||||
api_config = request.app.state.config.OPENAI_API_CONFIGS.get(
|
||||
str(idx),
|
||||
request.app.state.config.OPENAI_API_CONFIGS.get(
|
||||
url, {}
|
||||
), # Legacy support
|
||||
)
|
||||
|
||||
enable = api_config.get("enable", True)
|
||||
model_ids = api_config.get("model_ids", [])
|
||||
@@ -310,7 +319,12 @@ async def get_all_models_responses(request: Request) -> list:
|
||||
for idx, response in enumerate(responses):
|
||||
if response:
|
||||
url = request.app.state.config.OPENAI_API_BASE_URLS[idx]
|
||||
api_config = request.app.state.config.OPENAI_API_CONFIGS.get(url, {})
|
||||
api_config = request.app.state.config.OPENAI_API_CONFIGS.get(
|
||||
str(idx),
|
||||
request.app.state.config.OPENAI_API_CONFIGS.get(
|
||||
url, {}
|
||||
), # Legacy support
|
||||
)
|
||||
|
||||
prefix_id = api_config.get("prefix_id", None)
|
||||
|
||||
@@ -573,6 +587,7 @@ async def generate_chat_completion(
|
||||
detail="Model not found",
|
||||
)
|
||||
|
||||
await get_all_models(request)
|
||||
model = request.app.state.OPENAI_MODELS.get(model_id)
|
||||
if model:
|
||||
idx = model["urlIdx"]
|
||||
@@ -584,7 +599,10 @@ async def generate_chat_completion(
|
||||
|
||||
# Get the API config for the model
|
||||
api_config = request.app.state.config.OPENAI_API_CONFIGS.get(
|
||||
request.app.state.config.OPENAI_API_BASE_URLS[idx], {}
|
||||
str(idx),
|
||||
request.app.state.config.OPENAI_API_CONFIGS.get(
|
||||
request.app.state.config.OPENAI_API_BASE_URLS[idx], {}
|
||||
), # Legacy support
|
||||
)
|
||||
|
||||
prefix_id = api_config.get("prefix_id", None)
|
||||
|
||||
@@ -112,7 +112,12 @@ async def update_prompt_by_command(
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if prompt.user_id != user.id and user.role != "admin":
|
||||
# Is the user the original creator, in a group with write access, or an admin
|
||||
if (
|
||||
prompt.user_id != user.id
|
||||
and not has_access(user.id, "write", prompt.access_control)
|
||||
and user.role != "admin"
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
|
||||
@@ -9,6 +9,7 @@ from open_webui.utils.chat import generate_chat_completion
|
||||
from open_webui.utils.task import (
|
||||
title_generation_template,
|
||||
query_generation_template,
|
||||
image_prompt_generation_template,
|
||||
autocomplete_generation_template,
|
||||
tags_generation_template,
|
||||
emoji_generation_template,
|
||||
@@ -23,6 +24,7 @@ from open_webui.utils.task import get_task_model_id
|
||||
from open_webui.config import (
|
||||
DEFAULT_TITLE_GENERATION_PROMPT_TEMPLATE,
|
||||
DEFAULT_TAGS_GENERATION_PROMPT_TEMPLATE,
|
||||
DEFAULT_IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE,
|
||||
DEFAULT_QUERY_GENERATION_PROMPT_TEMPLATE,
|
||||
DEFAULT_AUTOCOMPLETE_GENERATION_PROMPT_TEMPLATE,
|
||||
DEFAULT_EMOJI_GENERATION_PROMPT_TEMPLATE,
|
||||
@@ -50,6 +52,7 @@ async def get_task_config(request: Request, user=Depends(get_verified_user)):
|
||||
"TASK_MODEL": request.app.state.config.TASK_MODEL,
|
||||
"TASK_MODEL_EXTERNAL": request.app.state.config.TASK_MODEL_EXTERNAL,
|
||||
"TITLE_GENERATION_PROMPT_TEMPLATE": request.app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE,
|
||||
"IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE": request.app.state.config.IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE,
|
||||
"ENABLE_AUTOCOMPLETE_GENERATION": request.app.state.config.ENABLE_AUTOCOMPLETE_GENERATION,
|
||||
"AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH": request.app.state.config.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH,
|
||||
"TAGS_GENERATION_PROMPT_TEMPLATE": request.app.state.config.TAGS_GENERATION_PROMPT_TEMPLATE,
|
||||
@@ -65,6 +68,7 @@ class TaskConfigForm(BaseModel):
|
||||
TASK_MODEL: Optional[str]
|
||||
TASK_MODEL_EXTERNAL: Optional[str]
|
||||
TITLE_GENERATION_PROMPT_TEMPLATE: str
|
||||
IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE: str
|
||||
ENABLE_AUTOCOMPLETE_GENERATION: bool
|
||||
AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH: int
|
||||
TAGS_GENERATION_PROMPT_TEMPLATE: str
|
||||
@@ -114,6 +118,7 @@ async def update_task_config(
|
||||
"TASK_MODEL": request.app.state.config.TASK_MODEL,
|
||||
"TASK_MODEL_EXTERNAL": request.app.state.config.TASK_MODEL_EXTERNAL,
|
||||
"TITLE_GENERATION_PROMPT_TEMPLATE": request.app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE,
|
||||
"IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE": request.app.state.config.IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE,
|
||||
"ENABLE_AUTOCOMPLETE_GENERATION": request.app.state.config.ENABLE_AUTOCOMPLETE_GENERATION,
|
||||
"AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH": request.app.state.config.AUTOCOMPLETE_GENERATION_INPUT_MAX_LENGTH,
|
||||
"TAGS_GENERATION_PROMPT_TEMPLATE": request.app.state.config.TAGS_GENERATION_PROMPT_TEMPLATE,
|
||||
@@ -256,6 +261,66 @@ async def generate_chat_tags(
|
||||
)
|
||||
|
||||
|
||||
@router.post("/image_prompt/completions")
|
||||
async def generate_image_prompt(
|
||||
request: Request, form_data: dict, user=Depends(get_verified_user)
|
||||
):
|
||||
models = request.app.state.MODELS
|
||||
|
||||
model_id = form_data["model"]
|
||||
if model_id not in models:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_404_NOT_FOUND,
|
||||
detail="Model not found",
|
||||
)
|
||||
|
||||
# Check if the user has a custom task model
|
||||
# If the user has a custom task model, use that model
|
||||
task_model_id = get_task_model_id(
|
||||
model_id,
|
||||
request.app.state.config.TASK_MODEL,
|
||||
request.app.state.config.TASK_MODEL_EXTERNAL,
|
||||
models,
|
||||
)
|
||||
|
||||
log.debug(
|
||||
f"generating image prompt using model {task_model_id} for user {user.email} "
|
||||
)
|
||||
|
||||
if request.app.state.config.IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE != "":
|
||||
template = request.app.state.config.IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE
|
||||
else:
|
||||
template = DEFAULT_IMAGE_PROMPT_GENERATION_PROMPT_TEMPLATE
|
||||
|
||||
content = image_prompt_generation_template(
|
||||
template,
|
||||
form_data["messages"],
|
||||
user={
|
||||
"name": user.name,
|
||||
},
|
||||
)
|
||||
|
||||
payload = {
|
||||
"model": task_model_id,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"stream": False,
|
||||
"metadata": {
|
||||
"task": str(TASKS.IMAGE_PROMPT_GENERATION),
|
||||
"task_body": form_data,
|
||||
"chat_id": form_data.get("chat_id", None),
|
||||
},
|
||||
}
|
||||
|
||||
try:
|
||||
return await generate_chat_completion(request, form_data=payload, user=user)
|
||||
except Exception as e:
|
||||
log.error("Exception occurred", exc_info=True)
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
content={"detail": "An internal error has occurred."},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/queries/completions")
|
||||
async def generate_queries(
|
||||
request: Request, form_data: dict, user=Depends(get_verified_user)
|
||||
|
||||
@@ -70,7 +70,7 @@ async def create_new_tools(
|
||||
user=Depends(get_verified_user),
|
||||
):
|
||||
if user.role != "admin" and not has_permission(
|
||||
user.id, "workspace.knowledge", request.app.state.config.USER_PERMISSIONS
|
||||
user.id, "workspace.tools", request.app.state.config.USER_PERMISSIONS
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
@@ -165,7 +165,12 @@ async def update_tools_by_id(
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if tools.user_id != user.id and user.role != "admin":
|
||||
# Is the user the original creator, in a group with write access, or an admin
|
||||
if (
|
||||
tools.user_id != user.id
|
||||
and not has_access(user.id, "write", tools.access_control)
|
||||
and user.role != "admin"
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.UNAUTHORIZED,
|
||||
@@ -304,6 +309,17 @@ async def update_tools_valves_by_id(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail=ERROR_MESSAGES.NOT_FOUND,
|
||||
)
|
||||
|
||||
if (
|
||||
tools.user_id != user.id
|
||||
and not has_access(user.id, "write", tools.access_control)
|
||||
and user.role != "admin"
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
if id in request.app.state.TOOLS:
|
||||
tools_module = request.app.state.TOOLS[id]
|
||||
else:
|
||||
|
||||
@@ -62,27 +62,44 @@ async def get_user_permissisions(user=Depends(get_verified_user)):
|
||||
# User Default Permissions
|
||||
############################
|
||||
class WorkspacePermissions(BaseModel):
|
||||
models: bool
|
||||
knowledge: bool
|
||||
prompts: bool
|
||||
tools: bool
|
||||
models: bool = False
|
||||
knowledge: bool = False
|
||||
prompts: bool = False
|
||||
tools: bool = False
|
||||
|
||||
|
||||
class ChatPermissions(BaseModel):
|
||||
file_upload: bool
|
||||
delete: bool
|
||||
edit: bool
|
||||
temporary: bool
|
||||
controls: bool = True
|
||||
file_upload: bool = True
|
||||
delete: bool = True
|
||||
edit: bool = True
|
||||
temporary: bool = True
|
||||
|
||||
|
||||
class FeaturesPermissions(BaseModel):
|
||||
web_search: bool = True
|
||||
image_generation: bool = True
|
||||
|
||||
|
||||
class UserPermissions(BaseModel):
|
||||
workspace: WorkspacePermissions
|
||||
chat: ChatPermissions
|
||||
features: FeaturesPermissions
|
||||
|
||||
|
||||
@router.get("/default/permissions")
|
||||
@router.get("/default/permissions", response_model=UserPermissions)
|
||||
async def get_user_permissions(request: Request, user=Depends(get_admin_user)):
|
||||
return request.app.state.config.USER_PERMISSIONS
|
||||
return {
|
||||
"workspace": WorkspacePermissions(
|
||||
**request.app.state.config.USER_PERMISSIONS.get("workspace", {})
|
||||
),
|
||||
"chat": ChatPermissions(
|
||||
**request.app.state.config.USER_PERMISSIONS.get("chat", {})
|
||||
),
|
||||
"features": FeaturesPermissions(
|
||||
**request.app.state.config.USER_PERMISSIONS.get("features", {})
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@router.post("/default/permissions")
|
||||
|
||||
@@ -26,7 +26,7 @@ class RedisLock:
|
||||
|
||||
def release_lock(self):
|
||||
lock_value = self.redis.get(self.lock_name)
|
||||
if lock_value and lock_value.decode("utf-8") == self.lock_id:
|
||||
if lock_value and lock_value == self.lock_id:
|
||||
self.redis.delete(self.lock_name)
|
||||
|
||||
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 11 KiB After Width: | Height: | Size: 10 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 6.0 KiB After Width: | Height: | Size: 5.2 KiB |
@@ -1,121 +1,73 @@
|
||||
import os
|
||||
import shutil
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import BinaryIO, Tuple
|
||||
|
||||
import boto3
|
||||
from botocore.exceptions import ClientError
|
||||
import shutil
|
||||
|
||||
|
||||
from typing import BinaryIO, Tuple, Optional, Union
|
||||
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
from open_webui.config import (
|
||||
STORAGE_PROVIDER,
|
||||
S3_ACCESS_KEY_ID,
|
||||
S3_SECRET_ACCESS_KEY,
|
||||
S3_BUCKET_NAME,
|
||||
S3_REGION_NAME,
|
||||
S3_ENDPOINT_URL,
|
||||
S3_REGION_NAME,
|
||||
S3_SECRET_ACCESS_KEY,
|
||||
GCS_BUCKET_NAME,
|
||||
GOOGLE_APPLICATION_CREDENTIALS_JSON,
|
||||
STORAGE_PROVIDER,
|
||||
UPLOAD_DIR,
|
||||
)
|
||||
from google.cloud import storage
|
||||
from google.cloud.exceptions import GoogleCloudError, NotFound
|
||||
from open_webui.constants import ERROR_MESSAGES
|
||||
|
||||
|
||||
import boto3
|
||||
from botocore.exceptions import ClientError
|
||||
from typing import BinaryIO, Tuple, Optional
|
||||
class StorageProvider(ABC):
|
||||
@abstractmethod
|
||||
def get_file(self, file_path: str) -> str:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def upload_file(self, file: BinaryIO, filename: str) -> Tuple[bytes, str]:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def delete_all_files(self) -> None:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def delete_file(self, file_path: str) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class StorageProvider:
|
||||
def __init__(self, provider: Optional[str] = None):
|
||||
self.storage_provider: str = provider or STORAGE_PROVIDER
|
||||
|
||||
self.s3_client = None
|
||||
self.s3_bucket_name: Optional[str] = None
|
||||
|
||||
if self.storage_provider == "s3":
|
||||
self._initialize_s3()
|
||||
|
||||
def _initialize_s3(self) -> None:
|
||||
"""Initializes the S3 client and bucket name if using S3 storage."""
|
||||
self.s3_client = boto3.client(
|
||||
"s3",
|
||||
region_name=S3_REGION_NAME,
|
||||
endpoint_url=S3_ENDPOINT_URL,
|
||||
aws_access_key_id=S3_ACCESS_KEY_ID,
|
||||
aws_secret_access_key=S3_SECRET_ACCESS_KEY,
|
||||
)
|
||||
self.bucket_name = S3_BUCKET_NAME
|
||||
|
||||
def _upload_to_s3(self, file_path: str, filename: str) -> Tuple[bytes, str]:
|
||||
"""Handles uploading of the file to S3 storage."""
|
||||
if not self.s3_client:
|
||||
raise RuntimeError("S3 Client is not initialized.")
|
||||
|
||||
try:
|
||||
self.s3_client.upload_file(file_path, self.bucket_name, filename)
|
||||
return (
|
||||
open(file_path, "rb").read(),
|
||||
"s3://" + self.bucket_name + "/" + filename,
|
||||
)
|
||||
except ClientError as e:
|
||||
raise RuntimeError(f"Error uploading file to S3: {e}")
|
||||
|
||||
def _upload_to_local(self, contents: bytes, filename: str) -> Tuple[bytes, str]:
|
||||
"""Handles uploading of the file to local storage."""
|
||||
class LocalStorageProvider(StorageProvider):
|
||||
@staticmethod
|
||||
def upload_file(file: BinaryIO, filename: str) -> Tuple[bytes, str]:
|
||||
contents = file.read()
|
||||
if not contents:
|
||||
raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
|
||||
file_path = f"{UPLOAD_DIR}/{filename}"
|
||||
with open(file_path, "wb") as f:
|
||||
f.write(contents)
|
||||
return contents, file_path
|
||||
|
||||
def _get_file_from_s3(self, file_path: str) -> str:
|
||||
"""Handles downloading of the file from S3 storage."""
|
||||
if not self.s3_client:
|
||||
raise RuntimeError("S3 Client is not initialized.")
|
||||
|
||||
try:
|
||||
bucket_name, key = file_path.split("//")[1].split("/")
|
||||
local_file_path = f"{UPLOAD_DIR}/{key}"
|
||||
self.s3_client.download_file(bucket_name, key, local_file_path)
|
||||
return local_file_path
|
||||
except ClientError as e:
|
||||
raise RuntimeError(f"Error downloading file from S3: {e}")
|
||||
|
||||
def _get_file_from_local(self, file_path: str) -> str:
|
||||
@staticmethod
|
||||
def get_file(file_path: str) -> str:
|
||||
"""Handles downloading of the file from local storage."""
|
||||
return file_path
|
||||
|
||||
def _delete_from_s3(self, filename: str) -> None:
|
||||
"""Handles deletion of the file from S3 storage."""
|
||||
if not self.s3_client:
|
||||
raise RuntimeError("S3 Client is not initialized.")
|
||||
|
||||
try:
|
||||
self.s3_client.delete_object(Bucket=self.bucket_name, Key=filename)
|
||||
except ClientError as e:
|
||||
raise RuntimeError(f"Error deleting file from S3: {e}")
|
||||
|
||||
def _delete_from_local(self, filename: str) -> None:
|
||||
@staticmethod
|
||||
def delete_file(file_path: str) -> None:
|
||||
"""Handles deletion of the file from local storage."""
|
||||
filename = file_path.split("/")[-1]
|
||||
file_path = f"{UPLOAD_DIR}/{filename}"
|
||||
if os.path.isfile(file_path):
|
||||
os.remove(file_path)
|
||||
else:
|
||||
print(f"File {file_path} not found in local storage.")
|
||||
|
||||
def _delete_all_from_s3(self) -> None:
|
||||
"""Handles deletion of all files from S3 storage."""
|
||||
if not self.s3_client:
|
||||
raise RuntimeError("S3 Client is not initialized.")
|
||||
|
||||
try:
|
||||
response = self.s3_client.list_objects_v2(Bucket=self.bucket_name)
|
||||
if "Contents" in response:
|
||||
for content in response["Contents"]:
|
||||
self.s3_client.delete_object(
|
||||
Bucket=self.bucket_name, Key=content["Key"]
|
||||
)
|
||||
except ClientError as e:
|
||||
raise RuntimeError(f"Error deleting all files from S3: {e}")
|
||||
|
||||
def _delete_all_from_local(self) -> None:
|
||||
@staticmethod
|
||||
def delete_all_files() -> None:
|
||||
"""Handles deletion of all files from local storage."""
|
||||
if os.path.exists(UPLOAD_DIR):
|
||||
for filename in os.listdir(UPLOAD_DIR):
|
||||
@@ -130,40 +82,141 @@ class StorageProvider:
|
||||
else:
|
||||
print(f"Directory {UPLOAD_DIR} not found in local storage.")
|
||||
|
||||
def upload_file(self, file: BinaryIO, filename: str) -> Tuple[bytes, str]:
|
||||
"""Uploads a file either to S3 or the local file system."""
|
||||
contents = file.read()
|
||||
if not contents:
|
||||
raise ValueError(ERROR_MESSAGES.EMPTY_CONTENT)
|
||||
contents, file_path = self._upload_to_local(contents, filename)
|
||||
|
||||
if self.storage_provider == "s3":
|
||||
return self._upload_to_s3(file_path, filename)
|
||||
return contents, file_path
|
||||
class S3StorageProvider(StorageProvider):
|
||||
def __init__(self):
|
||||
self.s3_client = boto3.client(
|
||||
"s3",
|
||||
region_name=S3_REGION_NAME,
|
||||
endpoint_url=S3_ENDPOINT_URL,
|
||||
aws_access_key_id=S3_ACCESS_KEY_ID,
|
||||
aws_secret_access_key=S3_SECRET_ACCESS_KEY,
|
||||
)
|
||||
self.bucket_name = S3_BUCKET_NAME
|
||||
|
||||
def upload_file(self, file: BinaryIO, filename: str) -> Tuple[bytes, str]:
|
||||
"""Handles uploading of the file to S3 storage."""
|
||||
_, file_path = LocalStorageProvider.upload_file(file, filename)
|
||||
try:
|
||||
self.s3_client.upload_file(file_path, self.bucket_name, filename)
|
||||
return (
|
||||
open(file_path, "rb").read(),
|
||||
"s3://" + self.bucket_name + "/" + filename,
|
||||
)
|
||||
except ClientError as e:
|
||||
raise RuntimeError(f"Error uploading file to S3: {e}")
|
||||
|
||||
def get_file(self, file_path: str) -> str:
|
||||
"""Downloads a file either from S3 or the local file system and returns the file path."""
|
||||
if self.storage_provider == "s3":
|
||||
return self._get_file_from_s3(file_path)
|
||||
return self._get_file_from_local(file_path)
|
||||
"""Handles downloading of the file from S3 storage."""
|
||||
try:
|
||||
bucket_name, key = file_path.split("//")[1].split("/")
|
||||
local_file_path = f"{UPLOAD_DIR}/{key}"
|
||||
self.s3_client.download_file(bucket_name, key, local_file_path)
|
||||
return local_file_path
|
||||
except ClientError as e:
|
||||
raise RuntimeError(f"Error downloading file from S3: {e}")
|
||||
|
||||
def delete_file(self, file_path: str) -> None:
|
||||
"""Deletes a file either from S3 or the local file system."""
|
||||
"""Handles deletion of the file from S3 storage."""
|
||||
filename = file_path.split("/")[-1]
|
||||
|
||||
if self.storage_provider == "s3":
|
||||
self._delete_from_s3(filename)
|
||||
try:
|
||||
self.s3_client.delete_object(Bucket=self.bucket_name, Key=filename)
|
||||
except ClientError as e:
|
||||
raise RuntimeError(f"Error deleting file from S3: {e}")
|
||||
|
||||
# Always delete from local storage
|
||||
self._delete_from_local(filename)
|
||||
LocalStorageProvider.delete_file(file_path)
|
||||
|
||||
def delete_all_files(self) -> None:
|
||||
"""Deletes all files from the storage."""
|
||||
if self.storage_provider == "s3":
|
||||
self._delete_all_from_s3()
|
||||
"""Handles deletion of all files from S3 storage."""
|
||||
try:
|
||||
response = self.s3_client.list_objects_v2(Bucket=self.bucket_name)
|
||||
if "Contents" in response:
|
||||
for content in response["Contents"]:
|
||||
self.s3_client.delete_object(
|
||||
Bucket=self.bucket_name, Key=content["Key"]
|
||||
)
|
||||
except ClientError as e:
|
||||
raise RuntimeError(f"Error deleting all files from S3: {e}")
|
||||
|
||||
# Always delete from local storage
|
||||
self._delete_all_from_local()
|
||||
LocalStorageProvider.delete_all_files()
|
||||
|
||||
|
||||
Storage = StorageProvider(provider=STORAGE_PROVIDER)
|
||||
class GCSStorageProvider(StorageProvider):
|
||||
def __init__(self):
|
||||
self.bucket_name = GCS_BUCKET_NAME
|
||||
|
||||
if GOOGLE_APPLICATION_CREDENTIALS_JSON:
|
||||
self.gcs_client = storage.Client.from_service_account_info(
|
||||
info=json.loads(GOOGLE_APPLICATION_CREDENTIALS_JSON)
|
||||
)
|
||||
else:
|
||||
# if no credentials json is provided, credentials will be picked up from the environment
|
||||
# if running on local environment, credentials would be user credentials
|
||||
# if running on a Compute Engine instance, credentials would be from Google Metadata server
|
||||
self.gcs_client = storage.Client()
|
||||
self.bucket = self.gcs_client.bucket(GCS_BUCKET_NAME)
|
||||
|
||||
def upload_file(self, file: BinaryIO, filename: str) -> Tuple[bytes, str]:
|
||||
"""Handles uploading of the file to GCS storage."""
|
||||
contents, file_path = LocalStorageProvider.upload_file(file, filename)
|
||||
try:
|
||||
blob = self.bucket.blob(filename)
|
||||
blob.upload_from_filename(file_path)
|
||||
return contents, "gs://" + self.bucket_name + "/" + filename
|
||||
except GoogleCloudError as e:
|
||||
raise RuntimeError(f"Error uploading file to GCS: {e}")
|
||||
|
||||
def get_file(self, file_path: str) -> str:
|
||||
"""Handles downloading of the file from GCS storage."""
|
||||
try:
|
||||
filename = file_path.removeprefix("gs://").split("/")[1]
|
||||
local_file_path = f"{UPLOAD_DIR}/{filename}"
|
||||
blob = self.bucket.get_blob(filename)
|
||||
blob.download_to_filename(local_file_path)
|
||||
|
||||
return local_file_path
|
||||
except NotFound as e:
|
||||
raise RuntimeError(f"Error downloading file from GCS: {e}")
|
||||
|
||||
def delete_file(self, file_path: str) -> None:
|
||||
"""Handles deletion of the file from GCS storage."""
|
||||
try:
|
||||
filename = file_path.removeprefix("gs://").split("/")[1]
|
||||
blob = self.bucket.get_blob(filename)
|
||||
blob.delete()
|
||||
except NotFound as e:
|
||||
raise RuntimeError(f"Error deleting file from GCS: {e}")
|
||||
|
||||
# Always delete from local storage
|
||||
LocalStorageProvider.delete_file(file_path)
|
||||
|
||||
def delete_all_files(self) -> None:
|
||||
"""Handles deletion of all files from GCS storage."""
|
||||
try:
|
||||
blobs = self.bucket.list_blobs()
|
||||
|
||||
for blob in blobs:
|
||||
blob.delete()
|
||||
|
||||
except NotFound as e:
|
||||
raise RuntimeError(f"Error deleting all files from GCS: {e}")
|
||||
|
||||
# Always delete from local storage
|
||||
LocalStorageProvider.delete_all_files()
|
||||
|
||||
|
||||
def get_storage_provider(storage_provider: str):
|
||||
if storage_provider == "local":
|
||||
Storage = LocalStorageProvider()
|
||||
elif storage_provider == "s3":
|
||||
Storage = S3StorageProvider()
|
||||
elif storage_provider == "gcs":
|
||||
Storage = GCSStorageProvider()
|
||||
else:
|
||||
raise RuntimeError(f"Unsupported storage provider: {storage_provider}")
|
||||
return Storage
|
||||
|
||||
|
||||
Storage = get_storage_provider(STORAGE_PROVIDER)
|
||||
|
||||
274
backend/open_webui/test/apps/webui/storage/test_provider.py
Normal file
274
backend/open_webui/test/apps/webui/storage/test_provider.py
Normal file
@@ -0,0 +1,274 @@
|
||||
import io
|
||||
import os
|
||||
import boto3
|
||||
import pytest
|
||||
from botocore.exceptions import ClientError
|
||||
from moto import mock_aws
|
||||
from open_webui.storage import provider
|
||||
from gcp_storage_emulator.server import create_server
|
||||
from google.cloud import storage
|
||||
|
||||
|
||||
def mock_upload_dir(monkeypatch, tmp_path):
|
||||
"""Fixture to monkey-patch the UPLOAD_DIR and create a temporary directory."""
|
||||
directory = tmp_path / "uploads"
|
||||
directory.mkdir()
|
||||
monkeypatch.setattr(provider, "UPLOAD_DIR", str(directory))
|
||||
return directory
|
||||
|
||||
|
||||
def test_imports():
|
||||
provider.StorageProvider
|
||||
provider.LocalStorageProvider
|
||||
provider.S3StorageProvider
|
||||
provider.GCSStorageProvider
|
||||
provider.Storage
|
||||
|
||||
|
||||
def test_get_storage_provider():
|
||||
Storage = provider.get_storage_provider("local")
|
||||
assert isinstance(Storage, provider.LocalStorageProvider)
|
||||
Storage = provider.get_storage_provider("s3")
|
||||
assert isinstance(Storage, provider.S3StorageProvider)
|
||||
Storage = provider.get_storage_provider("gcs")
|
||||
assert isinstance(Storage, provider.GCSStorageProvider)
|
||||
with pytest.raises(RuntimeError):
|
||||
provider.get_storage_provider("invalid")
|
||||
|
||||
|
||||
def test_class_instantiation():
|
||||
with pytest.raises(TypeError):
|
||||
provider.StorageProvider()
|
||||
with pytest.raises(TypeError):
|
||||
|
||||
class Test(provider.StorageProvider):
|
||||
pass
|
||||
|
||||
Test()
|
||||
provider.LocalStorageProvider()
|
||||
provider.S3StorageProvider()
|
||||
provider.GCSStorageProvider()
|
||||
|
||||
|
||||
class TestLocalStorageProvider:
|
||||
Storage = provider.LocalStorageProvider()
|
||||
file_content = b"test content"
|
||||
file_bytesio = io.BytesIO(file_content)
|
||||
filename = "test.txt"
|
||||
filename_extra = "test_exyta.txt"
|
||||
file_bytesio_empty = io.BytesIO()
|
||||
|
||||
def test_upload_file(self, monkeypatch, tmp_path):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
contents, file_path = self.Storage.upload_file(self.file_bytesio, self.filename)
|
||||
assert (upload_dir / self.filename).exists()
|
||||
assert (upload_dir / self.filename).read_bytes() == self.file_content
|
||||
assert contents == self.file_content
|
||||
assert file_path == str(upload_dir / self.filename)
|
||||
with pytest.raises(ValueError):
|
||||
self.Storage.upload_file(self.file_bytesio_empty, self.filename)
|
||||
|
||||
def test_get_file(self, monkeypatch, tmp_path):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
file_path = str(upload_dir / self.filename)
|
||||
file_path_return = self.Storage.get_file(file_path)
|
||||
assert file_path == file_path_return
|
||||
|
||||
def test_delete_file(self, monkeypatch, tmp_path):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
(upload_dir / self.filename).write_bytes(self.file_content)
|
||||
assert (upload_dir / self.filename).exists()
|
||||
file_path = str(upload_dir / self.filename)
|
||||
self.Storage.delete_file(file_path)
|
||||
assert not (upload_dir / self.filename).exists()
|
||||
|
||||
def test_delete_all_files(self, monkeypatch, tmp_path):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
(upload_dir / self.filename).write_bytes(self.file_content)
|
||||
(upload_dir / self.filename_extra).write_bytes(self.file_content)
|
||||
self.Storage.delete_all_files()
|
||||
assert not (upload_dir / self.filename).exists()
|
||||
assert not (upload_dir / self.filename_extra).exists()
|
||||
|
||||
|
||||
@mock_aws
|
||||
class TestS3StorageProvider:
|
||||
|
||||
def __init__(self):
|
||||
self.Storage = provider.S3StorageProvider()
|
||||
self.Storage.bucket_name = "my-bucket"
|
||||
self.s3_client = boto3.resource("s3", region_name="us-east-1")
|
||||
self.file_content = b"test content"
|
||||
self.filename = "test.txt"
|
||||
self.filename_extra = "test_exyta.txt"
|
||||
self.file_bytesio_empty = io.BytesIO()
|
||||
super().__init__()
|
||||
|
||||
def test_upload_file(self, monkeypatch, tmp_path):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
# S3 checks
|
||||
with pytest.raises(Exception):
|
||||
self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
|
||||
self.s3_client.create_bucket(Bucket=self.Storage.bucket_name)
|
||||
contents, s3_file_path = self.Storage.upload_file(
|
||||
io.BytesIO(self.file_content), self.filename
|
||||
)
|
||||
object = self.s3_client.Object(self.Storage.bucket_name, self.filename)
|
||||
assert self.file_content == object.get()["Body"].read()
|
||||
# local checks
|
||||
assert (upload_dir / self.filename).exists()
|
||||
assert (upload_dir / self.filename).read_bytes() == self.file_content
|
||||
assert contents == self.file_content
|
||||
assert s3_file_path == "s3://" + self.Storage.bucket_name + "/" + self.filename
|
||||
with pytest.raises(ValueError):
|
||||
self.Storage.upload_file(self.file_bytesio_empty, self.filename)
|
||||
|
||||
def test_get_file(self, monkeypatch, tmp_path):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
self.s3_client.create_bucket(Bucket=self.Storage.bucket_name)
|
||||
contents, s3_file_path = self.Storage.upload_file(
|
||||
io.BytesIO(self.file_content), self.filename
|
||||
)
|
||||
file_path = self.Storage.get_file(s3_file_path)
|
||||
assert file_path == str(upload_dir / self.filename)
|
||||
assert (upload_dir / self.filename).exists()
|
||||
|
||||
def test_delete_file(self, monkeypatch, tmp_path):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
self.s3_client.create_bucket(Bucket=self.Storage.bucket_name)
|
||||
contents, s3_file_path = self.Storage.upload_file(
|
||||
io.BytesIO(self.file_content), self.filename
|
||||
)
|
||||
assert (upload_dir / self.filename).exists()
|
||||
self.Storage.delete_file(s3_file_path)
|
||||
assert not (upload_dir / self.filename).exists()
|
||||
with pytest.raises(ClientError) as exc:
|
||||
self.s3_client.Object(self.Storage.bucket_name, self.filename).load()
|
||||
error = exc.value.response["Error"]
|
||||
assert error["Code"] == "404"
|
||||
assert error["Message"] == "Not Found"
|
||||
|
||||
def test_delete_all_files(self, monkeypatch, tmp_path):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
# create 2 files
|
||||
self.s3_client.create_bucket(Bucket=self.Storage.bucket_name)
|
||||
self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
|
||||
object = self.s3_client.Object(self.Storage.bucket_name, self.filename)
|
||||
assert self.file_content == object.get()["Body"].read()
|
||||
assert (upload_dir / self.filename).exists()
|
||||
assert (upload_dir / self.filename).read_bytes() == self.file_content
|
||||
self.Storage.upload_file(io.BytesIO(self.file_content), self.filename_extra)
|
||||
object = self.s3_client.Object(self.Storage.bucket_name, self.filename_extra)
|
||||
assert self.file_content == object.get()["Body"].read()
|
||||
assert (upload_dir / self.filename).exists()
|
||||
assert (upload_dir / self.filename).read_bytes() == self.file_content
|
||||
|
||||
self.Storage.delete_all_files()
|
||||
assert not (upload_dir / self.filename).exists()
|
||||
with pytest.raises(ClientError) as exc:
|
||||
self.s3_client.Object(self.Storage.bucket_name, self.filename).load()
|
||||
error = exc.value.response["Error"]
|
||||
assert error["Code"] == "404"
|
||||
assert error["Message"] == "Not Found"
|
||||
assert not (upload_dir / self.filename_extra).exists()
|
||||
with pytest.raises(ClientError) as exc:
|
||||
self.s3_client.Object(self.Storage.bucket_name, self.filename_extra).load()
|
||||
error = exc.value.response["Error"]
|
||||
assert error["Code"] == "404"
|
||||
assert error["Message"] == "Not Found"
|
||||
|
||||
self.Storage.delete_all_files()
|
||||
assert not (upload_dir / self.filename).exists()
|
||||
assert not (upload_dir / self.filename_extra).exists()
|
||||
|
||||
|
||||
class TestGCSStorageProvider:
|
||||
Storage = provider.GCSStorageProvider()
|
||||
Storage.bucket_name = "my-bucket"
|
||||
file_content = b"test content"
|
||||
filename = "test.txt"
|
||||
filename_extra = "test_exyta.txt"
|
||||
file_bytesio_empty = io.BytesIO()
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def setup(self):
|
||||
host, port = "localhost", 9023
|
||||
|
||||
server = create_server(host, port, in_memory=True)
|
||||
server.start()
|
||||
os.environ["STORAGE_EMULATOR_HOST"] = f"http://{host}:{port}"
|
||||
|
||||
gcs_client = storage.Client()
|
||||
bucket = gcs_client.bucket(self.Storage.bucket_name)
|
||||
bucket.create()
|
||||
self.Storage.gcs_client, self.Storage.bucket = gcs_client, bucket
|
||||
yield
|
||||
bucket.delete(force=True)
|
||||
server.stop()
|
||||
|
||||
def test_upload_file(self, monkeypatch, tmp_path, setup):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
# catch error if bucket does not exist
|
||||
with pytest.raises(Exception):
|
||||
self.Storage.bucket = monkeypatch(self.Storage, "bucket", None)
|
||||
self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
|
||||
contents, gcs_file_path = self.Storage.upload_file(
|
||||
io.BytesIO(self.file_content), self.filename
|
||||
)
|
||||
object = self.Storage.bucket.get_blob(self.filename)
|
||||
assert self.file_content == object.download_as_bytes()
|
||||
# local checks
|
||||
assert (upload_dir / self.filename).exists()
|
||||
assert (upload_dir / self.filename).read_bytes() == self.file_content
|
||||
assert contents == self.file_content
|
||||
assert gcs_file_path == "gs://" + self.Storage.bucket_name + "/" + self.filename
|
||||
# test error if file is empty
|
||||
with pytest.raises(ValueError):
|
||||
self.Storage.upload_file(self.file_bytesio_empty, self.filename)
|
||||
|
||||
def test_get_file(self, monkeypatch, tmp_path, setup):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
contents, gcs_file_path = self.Storage.upload_file(
|
||||
io.BytesIO(self.file_content), self.filename
|
||||
)
|
||||
file_path = self.Storage.get_file(gcs_file_path)
|
||||
assert file_path == str(upload_dir / self.filename)
|
||||
assert (upload_dir / self.filename).exists()
|
||||
|
||||
def test_delete_file(self, monkeypatch, tmp_path, setup):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
contents, gcs_file_path = self.Storage.upload_file(
|
||||
io.BytesIO(self.file_content), self.filename
|
||||
)
|
||||
# ensure that local directory has the uploaded file as well
|
||||
assert (upload_dir / self.filename).exists()
|
||||
assert self.Storage.bucket.get_blob(self.filename).name == self.filename
|
||||
self.Storage.delete_file(gcs_file_path)
|
||||
# check that deleting file from gcs will delete the local file as well
|
||||
assert not (upload_dir / self.filename).exists()
|
||||
assert self.Storage.bucket.get_blob(self.filename) == None
|
||||
|
||||
def test_delete_all_files(self, monkeypatch, tmp_path, setup):
|
||||
upload_dir = mock_upload_dir(monkeypatch, tmp_path)
|
||||
# create 2 files
|
||||
self.Storage.upload_file(io.BytesIO(self.file_content), self.filename)
|
||||
object = self.Storage.bucket.get_blob(self.filename)
|
||||
assert (upload_dir / self.filename).exists()
|
||||
assert (upload_dir / self.filename).read_bytes() == self.file_content
|
||||
assert self.Storage.bucket.get_blob(self.filename).name == self.filename
|
||||
assert self.file_content == object.download_as_bytes()
|
||||
self.Storage.upload_file(io.BytesIO(self.file_content), self.filename_extra)
|
||||
object = self.Storage.bucket.get_blob(self.filename_extra)
|
||||
assert (upload_dir / self.filename_extra).exists()
|
||||
assert (upload_dir / self.filename_extra).read_bytes() == self.file_content
|
||||
assert (
|
||||
self.Storage.bucket.get_blob(self.filename_extra).name
|
||||
== self.filename_extra
|
||||
)
|
||||
assert self.file_content == object.download_as_bytes()
|
||||
|
||||
self.Storage.delete_all_files()
|
||||
assert not (upload_dir / self.filename).exists()
|
||||
assert not (upload_dir / self.filename_extra).exists()
|
||||
assert self.Storage.bucket.get_blob(self.filename) == None
|
||||
assert self.Storage.bucket.get_blob(self.filename_extra) == None
|
||||
@@ -1,9 +1,30 @@
|
||||
from typing import Optional, Union, List, Dict, Any
|
||||
from open_webui.models.users import Users, UserModel
|
||||
from open_webui.models.groups import Groups
|
||||
|
||||
|
||||
from open_webui.config import DEFAULT_USER_PERMISSIONS
|
||||
import json
|
||||
|
||||
|
||||
def fill_missing_permissions(
|
||||
permissions: Dict[str, Any], default_permissions: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Recursively fills in missing properties in the permissions dictionary
|
||||
using the default permissions as a template.
|
||||
"""
|
||||
for key, value in default_permissions.items():
|
||||
if key not in permissions:
|
||||
permissions[key] = value
|
||||
elif isinstance(value, dict) and isinstance(
|
||||
permissions[key], dict
|
||||
): # Both are nested dictionaries
|
||||
permissions[key] = fill_missing_permissions(permissions[key], value)
|
||||
|
||||
return permissions
|
||||
|
||||
|
||||
def get_permissions(
|
||||
user_id: str,
|
||||
default_permissions: Dict[str, Any],
|
||||
@@ -27,39 +48,45 @@ def get_permissions(
|
||||
if key not in permissions:
|
||||
permissions[key] = value
|
||||
else:
|
||||
permissions[key] = permissions[key] or value
|
||||
permissions[key] = (
|
||||
permissions[key] or value
|
||||
) # Use the most permissive value (True > False)
|
||||
return permissions
|
||||
|
||||
user_groups = Groups.get_groups_by_member_id(user_id)
|
||||
|
||||
# deep copy default permissions to avoid modifying the original dict
|
||||
# Deep copy default permissions to avoid modifying the original dict
|
||||
permissions = json.loads(json.dumps(default_permissions))
|
||||
|
||||
# Combine permissions from all user groups
|
||||
for group in user_groups:
|
||||
group_permissions = group.permissions
|
||||
permissions = combine_permissions(permissions, group_permissions)
|
||||
|
||||
# Ensure all fields from default_permissions are present and filled in
|
||||
permissions = fill_missing_permissions(permissions, default_permissions)
|
||||
|
||||
return permissions
|
||||
|
||||
|
||||
def has_permission(
|
||||
user_id: str,
|
||||
permission_key: str,
|
||||
default_permissions: Dict[str, bool] = {},
|
||||
default_permissions: Dict[str, Any] = {},
|
||||
) -> bool:
|
||||
"""
|
||||
Check if a user has a specific permission by checking the group permissions
|
||||
and falls back to default permissions if not found in any group.
|
||||
and fall back to default permissions if not found in any group.
|
||||
|
||||
Permission keys can be hierarchical and separated by dots ('.').
|
||||
"""
|
||||
|
||||
def get_permission(permissions: Dict[str, bool], keys: List[str]) -> bool:
|
||||
def get_permission(permissions: Dict[str, Any], keys: List[str]) -> bool:
|
||||
"""Traverse permissions dict using a list of keys (from dot-split permission_key)."""
|
||||
for key in keys:
|
||||
if key not in permissions:
|
||||
return False # If any part of the hierarchy is missing, deny access
|
||||
permissions = permissions[key] # Go one level deeper
|
||||
permissions = permissions[key] # Traverse one level deeper
|
||||
|
||||
return bool(permissions) # Return the boolean at the final level
|
||||
|
||||
@@ -73,7 +100,10 @@ def has_permission(
|
||||
if get_permission(group_permissions, permission_hierarchy):
|
||||
return True
|
||||
|
||||
# Check default permissions afterwards if the group permissions don't allow it
|
||||
# Check default permissions afterward if the group permissions don't allow it
|
||||
default_permissions = fill_missing_permissions(
|
||||
default_permissions, DEFAULT_USER_PERMISSIONS
|
||||
)
|
||||
return get_permission(default_permissions, permission_hierarchy)
|
||||
|
||||
|
||||
|
||||
@@ -28,9 +28,13 @@ from open_webui.socket.main import (
|
||||
from open_webui.routers.tasks import (
|
||||
generate_queries,
|
||||
generate_title,
|
||||
generate_image_prompt,
|
||||
generate_chat_tags,
|
||||
)
|
||||
from open_webui.routers.retrieval import process_web_search, SearchForm
|
||||
from open_webui.routers.images import image_generations, GenerateImageForm
|
||||
|
||||
|
||||
from open_webui.utils.webhook import post_webhook
|
||||
|
||||
|
||||
@@ -486,6 +490,100 @@ async def chat_web_search_handler(
|
||||
return form_data
|
||||
|
||||
|
||||
async def chat_image_generation_handler(
|
||||
request: Request, form_data: dict, extra_params: dict, user
|
||||
):
|
||||
__event_emitter__ = extra_params["__event_emitter__"]
|
||||
await __event_emitter__(
|
||||
{
|
||||
"type": "status",
|
||||
"data": {"description": "Generating an image", "done": False},
|
||||
}
|
||||
)
|
||||
|
||||
messages = form_data["messages"]
|
||||
user_message = get_last_user_message(messages)
|
||||
|
||||
prompt = user_message
|
||||
negative_prompt = ""
|
||||
|
||||
if request.app.state.config.ENABLE_IMAGE_PROMPT_GENERATION:
|
||||
try:
|
||||
res = await generate_image_prompt(
|
||||
request,
|
||||
{
|
||||
"model": form_data["model"],
|
||||
"messages": messages,
|
||||
},
|
||||
user,
|
||||
)
|
||||
|
||||
response = res["choices"][0]["message"]["content"]
|
||||
|
||||
try:
|
||||
bracket_start = response.find("{")
|
||||
bracket_end = response.rfind("}") + 1
|
||||
|
||||
if bracket_start == -1 or bracket_end == -1:
|
||||
raise Exception("No JSON object found in the response")
|
||||
|
||||
response = response[bracket_start:bracket_end]
|
||||
response = json.loads(response)
|
||||
prompt = response.get("prompt", [])
|
||||
except Exception as e:
|
||||
prompt = user_message
|
||||
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
prompt = user_message
|
||||
|
||||
system_message_content = ""
|
||||
|
||||
try:
|
||||
images = await image_generations(
|
||||
request=request,
|
||||
form_data=GenerateImageForm(**{"prompt": prompt}),
|
||||
user=user,
|
||||
)
|
||||
|
||||
await __event_emitter__(
|
||||
{
|
||||
"type": "status",
|
||||
"data": {"description": "Generated an image", "done": True},
|
||||
}
|
||||
)
|
||||
|
||||
for image in images:
|
||||
await __event_emitter__(
|
||||
{
|
||||
"type": "message",
|
||||
"data": {"content": f"\n"},
|
||||
}
|
||||
)
|
||||
|
||||
system_message_content = "<context>User is shown the generated image, tell the user that the image has been generated</context>"
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
await __event_emitter__(
|
||||
{
|
||||
"type": "status",
|
||||
"data": {
|
||||
"description": f"An error occured while generating an image",
|
||||
"done": True,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
system_message_content = "<context>Unable to generate an image, tell the user that an error occured</context>"
|
||||
|
||||
if system_message_content:
|
||||
form_data["messages"] = add_or_update_system_message(
|
||||
system_message_content, form_data["messages"]
|
||||
)
|
||||
|
||||
return form_data
|
||||
|
||||
|
||||
async def chat_completion_files_handler(
|
||||
request: Request, body: dict, user: UserModel
|
||||
) -> tuple[dict, dict[str, list]]:
|
||||
@@ -523,17 +621,28 @@ async def chat_completion_files_handler(
|
||||
if len(queries) == 0:
|
||||
queries = [get_last_user_message(body["messages"])]
|
||||
|
||||
sources = get_sources_from_files(
|
||||
files=files,
|
||||
queries=queries,
|
||||
embedding_function=request.app.state.EMBEDDING_FUNCTION,
|
||||
k=request.app.state.config.TOP_K,
|
||||
reranking_function=request.app.state.rf,
|
||||
r=request.app.state.config.RELEVANCE_THRESHOLD,
|
||||
hybrid_search=request.app.state.config.ENABLE_RAG_HYBRID_SEARCH,
|
||||
)
|
||||
try:
|
||||
# Offload get_sources_from_files 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(
|
||||
files=files,
|
||||
queries=queries,
|
||||
embedding_function=request.app.state.EMBEDDING_FUNCTION,
|
||||
k=request.app.state.config.TOP_K,
|
||||
reranking_function=request.app.state.rf,
|
||||
r=request.app.state.config.RELEVANCE_THRESHOLD,
|
||||
hybrid_search=request.app.state.config.ENABLE_RAG_HYBRID_SEARCH,
|
||||
),
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
|
||||
log.debug(f"rag_contexts:sources: {sources}")
|
||||
|
||||
return body, {"sources": sources}
|
||||
|
||||
|
||||
@@ -562,6 +671,10 @@ def apply_params_to_form_data(form_data, model):
|
||||
|
||||
if "frequency_penalty" in params:
|
||||
form_data["frequency_penalty"] = params["frequency_penalty"]
|
||||
|
||||
if "reasoning_effort" in params:
|
||||
form_data["reasoning_effort"] = params["reasoning_effort"]
|
||||
|
||||
return form_data
|
||||
|
||||
|
||||
@@ -640,6 +753,11 @@ async def process_chat_payload(request, form_data, metadata, user, model):
|
||||
request, form_data, extra_params, user
|
||||
)
|
||||
|
||||
if "image_generation" in features and features["image_generation"]:
|
||||
form_data = await chat_image_generation_handler(
|
||||
request, form_data, extra_params, user
|
||||
)
|
||||
|
||||
try:
|
||||
form_data, flags = await chat_completion_filter_functions_handler(
|
||||
request, form_data, model, extra_params
|
||||
@@ -958,6 +1076,16 @@ async def process_chat_response(
|
||||
},
|
||||
)
|
||||
|
||||
# We might want to disable this by default
|
||||
detect_reasoning = True
|
||||
reasoning_tags = ["think", "reason", "reasoning", "thought", "Thought"]
|
||||
current_tag = None
|
||||
|
||||
reasoning_start_time = None
|
||||
|
||||
reasoning_content = ""
|
||||
ongoing_content = ""
|
||||
|
||||
async for line in response.body_iterator:
|
||||
line = line.decode("utf-8") if isinstance(line, bytes) else line
|
||||
data = line
|
||||
@@ -966,12 +1094,12 @@ async def process_chat_response(
|
||||
if not data.strip():
|
||||
continue
|
||||
|
||||
# "data: " is the prefix for each event
|
||||
if not data.startswith("data: "):
|
||||
# "data:" is the prefix for each event
|
||||
if not data.startswith("data:"):
|
||||
continue
|
||||
|
||||
# Remove the prefix
|
||||
data = data[len("data: ") :]
|
||||
data = data[len("data:") :].strip()
|
||||
|
||||
try:
|
||||
data = json.loads(data)
|
||||
@@ -984,7 +1112,6 @@ async def process_chat_response(
|
||||
"selectedModelId": data["selected_model_id"],
|
||||
},
|
||||
)
|
||||
|
||||
else:
|
||||
value = (
|
||||
data.get("choices", [])[0]
|
||||
@@ -995,6 +1122,73 @@ async def process_chat_response(
|
||||
if value:
|
||||
content = f"{content}{value}"
|
||||
|
||||
if detect_reasoning:
|
||||
for tag in reasoning_tags:
|
||||
start_tag = f"<{tag}>\n"
|
||||
end_tag = f"</{tag}>\n"
|
||||
|
||||
if start_tag in content:
|
||||
# Remove the start tag
|
||||
content = content.replace(start_tag, "")
|
||||
ongoing_content = content
|
||||
|
||||
reasoning_start_time = time.time()
|
||||
reasoning_content = ""
|
||||
|
||||
current_tag = tag
|
||||
break
|
||||
|
||||
if reasoning_start_time is not None:
|
||||
# Remove the last value from the content
|
||||
content = content[: -len(value)]
|
||||
|
||||
reasoning_content += value
|
||||
|
||||
end_tag = f"</{current_tag}>\n"
|
||||
if end_tag in reasoning_content:
|
||||
reasoning_end_time = time.time()
|
||||
reasoning_duration = int(
|
||||
reasoning_end_time
|
||||
- reasoning_start_time
|
||||
)
|
||||
reasoning_content = (
|
||||
reasoning_content.strip(
|
||||
f"<{current_tag}>\n"
|
||||
)
|
||||
.strip(end_tag)
|
||||
.strip()
|
||||
)
|
||||
|
||||
if reasoning_content:
|
||||
reasoning_display_content = "\n".join(
|
||||
(
|
||||
f"> {line}"
|
||||
if not line.startswith(">")
|
||||
else line
|
||||
)
|
||||
for line in reasoning_content.splitlines()
|
||||
)
|
||||
|
||||
# Format reasoning with <details> tag
|
||||
content = f'{ongoing_content}<details type="reasoning" done="true" duration="{reasoning_duration}">\n<summary>Thought for {reasoning_duration} seconds</summary>\n{reasoning_display_content}\n</details>\n'
|
||||
else:
|
||||
content = ""
|
||||
|
||||
reasoning_start_time = None
|
||||
else:
|
||||
|
||||
reasoning_display_content = "\n".join(
|
||||
(
|
||||
f"> {line}"
|
||||
if not line.startswith(">")
|
||||
else line
|
||||
)
|
||||
for line in reasoning_content.splitlines()
|
||||
)
|
||||
|
||||
# Show ongoing thought process
|
||||
content = f'{ongoing_content}<details type="reasoning" done="false">\n<summary>Thinking…</summary>\n{reasoning_display_content}\n</details>\n'
|
||||
|
||||
if ENABLE_REALTIME_CHAT_SAVE:
|
||||
# Save message in the database
|
||||
Chats.upsert_message_to_chat_by_id_and_message_id(
|
||||
@@ -1015,10 +1209,8 @@ async def process_chat_response(
|
||||
"data": data,
|
||||
}
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
done = "data: [DONE]" in line
|
||||
|
||||
if done:
|
||||
pass
|
||||
else:
|
||||
|
||||
@@ -63,17 +63,8 @@ auth_manager_config.JWT_EXPIRES_IN = JWT_EXPIRES_IN
|
||||
class OAuthManager:
|
||||
def __init__(self):
|
||||
self.oauth = OAuth()
|
||||
for provider_name, provider_config in OAUTH_PROVIDERS.items():
|
||||
self.oauth.register(
|
||||
name=provider_name,
|
||||
client_id=provider_config["client_id"],
|
||||
client_secret=provider_config["client_secret"],
|
||||
server_metadata_url=provider_config["server_metadata_url"],
|
||||
client_kwargs={
|
||||
"scope": provider_config["scope"],
|
||||
},
|
||||
redirect_uri=provider_config["redirect_uri"],
|
||||
)
|
||||
for _, provider_config in OAUTH_PROVIDERS.items():
|
||||
provider_config["register"](self.oauth)
|
||||
|
||||
def get_client(self, provider_name):
|
||||
return self.oauth.create_client(provider_name)
|
||||
@@ -200,14 +191,14 @@ class OAuthManager:
|
||||
except Exception as e:
|
||||
log.warning(f"OAuth callback error: {e}")
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
|
||||
user_data: UserInfo = token["userinfo"]
|
||||
user_data: UserInfo = token.get("userinfo")
|
||||
if not user_data:
|
||||
user_data: UserInfo = await client.userinfo(token=token)
|
||||
if not user_data:
|
||||
log.warning(f"OAuth callback failed, user data is missing: {token}")
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
|
||||
|
||||
sub = user_data.get("sub")
|
||||
sub = user_data.get(OAUTH_PROVIDERS[provider].get("sub_claim", "sub"))
|
||||
if not sub:
|
||||
log.warning(f"OAuth callback failed, sub is missing: {user_data}")
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
|
||||
@@ -255,12 +246,20 @@ class OAuthManager:
|
||||
raise HTTPException(400, detail=ERROR_MESSAGES.EMAIL_TAKEN)
|
||||
|
||||
picture_claim = auth_manager_config.OAUTH_PICTURE_CLAIM
|
||||
picture_url = user_data.get(picture_claim, "")
|
||||
picture_url = user_data.get(
|
||||
picture_claim, OAUTH_PROVIDERS[provider].get("picture_url", "")
|
||||
)
|
||||
if picture_url:
|
||||
# Download the profile image into a base64 string
|
||||
try:
|
||||
access_token = token.get("access_token")
|
||||
get_kwargs = {}
|
||||
if access_token:
|
||||
get_kwargs["headers"] = {
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
}
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(picture_url) as resp:
|
||||
async with session.get(picture_url, **get_kwargs) as resp:
|
||||
picture = await resp.read()
|
||||
base64_encoded_picture = base64.b64encode(
|
||||
picture
|
||||
|
||||
@@ -47,6 +47,7 @@ def apply_model_params_to_body_openai(params: dict, form_data: dict) -> dict:
|
||||
"top_p": float,
|
||||
"max_tokens": int,
|
||||
"frequency_penalty": float,
|
||||
"reasoning_effort": str,
|
||||
"seed": lambda x: x,
|
||||
"stop": lambda x: [bytes(s, "utf-8").decode("unicode_escape") for s in x],
|
||||
}
|
||||
|
||||
@@ -217,6 +217,24 @@ def tags_generation_template(
|
||||
return template
|
||||
|
||||
|
||||
def image_prompt_generation_template(
|
||||
template: str, messages: list[dict], user: Optional[dict] = None
|
||||
) -> str:
|
||||
prompt = get_last_user_message(messages)
|
||||
template = replace_prompt_variable(template, prompt)
|
||||
template = replace_messages_variable(template, messages)
|
||||
|
||||
template = prompt_template(
|
||||
template,
|
||||
**(
|
||||
{"user_name": user.get("name"), "user_location": user.get("location")}
|
||||
if user
|
||||
else {}
|
||||
),
|
||||
)
|
||||
return template
|
||||
|
||||
|
||||
def emoji_generation_template(
|
||||
template: str, prompt: str, user: Optional[dict] = None
|
||||
) -> str:
|
||||
|
||||
Reference in New Issue
Block a user