OpenWebUI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-06-26 18:26:48 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge branch 'open-webui:main' into searxng

Browse Source
This commit is contained in:
Que Nguyen
2024-06-17 14:28:02 +07:00
committed by GitHub
parent a02ba52de8 9e4dd4b86f
commit 75e51ecf6d
69 changed files with 1120 additions and 222 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
backend/apps/images/main.py
View File

@@ -37,6 +37,10 @@ from config import (
ENABLE_IMAGE_GENERATION,
AUTOMATIC1111_BASE_URL,
COMFYUI_BASE_URL,
COMFYUI_CFG_SCALE,
COMFYUI_SAMPLER,
COMFYUI_SCHEDULER,
COMFYUI_SD3,
IMAGES_OPENAI_API_BASE_URL,
IMAGES_OPENAI_API_KEY,
IMAGE_GENERATION_MODEL,
@@ -78,6 +82,10 @@ app.state.config.COMFYUI_BASE_URL = COMFYUI_BASE_URL
app.state.config.IMAGE_SIZE = IMAGE_SIZE
app.state.config.IMAGE_STEPS = IMAGE_STEPS
app.state.config.COMFYUI_CFG_SCALE = COMFYUI_CFG_SCALE
app.state.config.COMFYUI_SAMPLER = COMFYUI_SAMPLER
app.state.config.COMFYUI_SCHEDULER = COMFYUI_SCHEDULER
app.state.config.COMFYUI_SD3 = COMFYUI_SD3
@app.get("/config")
@@ -457,6 +465,18 @@ def generate_image(
if form_data.negative_prompt is not None:
data["negative_prompt"] = form_data.negative_prompt
if app.state.config.COMFYUI_CFG_SCALE:
data["cfg_scale"] = app.state.config.COMFYUI_CFG_SCALE
if app.state.config.COMFYUI_SAMPLER is not None:
data["sampler"] = app.state.config.COMFYUI_SAMPLER
if app.state.config.COMFYUI_SCHEDULER is not None:
data["scheduler"] = app.state.config.COMFYUI_SCHEDULER
if app.state.config.COMFYUI_SD3 is not None:
data["sd3"] = app.state.config.COMFYUI_SD3
data = ImageGenerationPayload(**data)
res = comfyui_generate_image(

16
backend/apps/images/utils/comfyui.py
View File

@@ -190,6 +190,10 @@ class ImageGenerationPayload(BaseModel):
width: int
height: int
n: int = 1
cfg_scale: Optional[float] = None
sampler: Optional[str] = None
scheduler: Optional[str] = None
sd3: Optional[bool] = None
def comfyui_generate_image(
@@ -199,6 +203,18 @@ def comfyui_generate_image(
comfyui_prompt = json.loads(COMFYUI_DEFAULT_PROMPT)
if payload.cfg_scale:
comfyui_prompt["3"]["inputs"]["cfg"] = payload.cfg_scale
if payload.sampler:
comfyui_prompt["3"]["inputs"]["sampler"] = payload.sampler
if payload.scheduler:
comfyui_prompt["3"]["inputs"]["scheduler"] = payload.scheduler
if payload.sd3:
comfyui_prompt["5"]["class_type"] = "EmptySD3LatentImage"
comfyui_prompt["4"]["inputs"]["ckpt_name"] = model
comfyui_prompt["5"]["inputs"]["batch_size"] = payload.n
comfyui_prompt["5"]["inputs"]["width"] = payload.width

3
backend/apps/ollama/main.py
View File

@@ -850,8 +850,7 @@ async def generate_chat_completion(
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
log.info(f"url: {url}")
print(payload)
log.debug(payload)
return await post_streaming_url(f"{url}/api/chat", json.dumps(payload))

4
backend/apps/openai/main.py
View File

@@ -430,13 +430,11 @@ async def generate_chat_completion(
# Convert the modified body back to JSON
payload = json.dumps(payload)
print(payload)
log.debug(payload)
url = app.state.config.OPENAI_API_BASE_URLS[idx]
key = app.state.config.OPENAI_API_KEYS[idx]
print(payload)
headers = {}
headers["Authorization"] = f"Bearer {key}"
headers["Content-Type"] = "application/json"

48
backend/apps/webui/internal/migrations/013_add_user_info.py Normal file
View File

@@ -0,0 +1,48 @@
"""Peewee migrations -- 002_add_local_sharing.py.
Some examples (model - class or model name)::
> Model = migrator.orm['table_name'] # Return model in current state by name
> Model = migrator.ModelClass # Return model in current state by name
> migrator.sql(sql) # Run custom SQL
> migrator.run(func, *args, **kwargs) # Run python function with the given args
> migrator.create_model(Model) # Create a model (could be used as decorator)
> migrator.remove_model(model, cascade=True) # Remove a model
> migrator.add_fields(model, **fields) # Add fields to a model
> migrator.change_fields(model, **fields) # Change fields
> migrator.remove_fields(model, *field_names, cascade=True)
> migrator.rename_field(model, old_field_name, new_field_name)
> migrator.rename_table(model, new_table_name)
> migrator.add_index(model, *col_names, unique=False)
> migrator.add_not_null(model, *field_names)
> migrator.add_default(model, field_name, default)
> migrator.add_constraint(model, name, sql)
> migrator.drop_index(model, *col_names)
> migrator.drop_not_null(model, *field_names)
> migrator.drop_constraints(model, *constraints)
"""
from contextlib import suppress
import peewee as pw
from peewee_migrate import Migrator
with suppress(ImportError):
import playhouse.postgres_ext as pw_pext
def migrate(migrator: Migrator, database: pw.Database, *, fake=False):
"""Write your migrations here."""
# Adding fields info to the 'user' table
migrator.add_fields("user", info=pw.TextField(null=True))
def rollback(migrator: Migrator, database: pw.Database, *, fake=False):
"""Write your rollback migrations here."""
# Remove the settings field
migrator.remove_fields("user", "info")

2
backend/apps/webui/main.py
View File

@@ -25,6 +25,7 @@ from config import (
USER_PERMISSIONS,
WEBHOOK_URL,
WEBUI_AUTH_TRUSTED_EMAIL_HEADER,
WEBUI_AUTH_TRUSTED_NAME_HEADER,
JWT_EXPIRES_IN,
WEBUI_BANNERS,
ENABLE_COMMUNITY_SHARING,
@@ -40,6 +41,7 @@ app.state.config = AppConfig()
app.state.config.ENABLE_SIGNUP = ENABLE_SIGNUP
app.state.config.JWT_EXPIRES_IN = JWT_EXPIRES_IN
app.state.AUTH_TRUSTED_EMAIL_HEADER = WEBUI_AUTH_TRUSTED_EMAIL_HEADER
app.state.AUTH_TRUSTED_NAME_HEADER = WEBUI_AUTH_TRUSTED_NAME_HEADER
app.state.config.SHOW_ADMIN_DETAILS = SHOW_ADMIN_DETAILS

2
backend/apps/webui/models/users.py
View File

@@ -26,6 +26,7 @@ class User(Model):
api_key = CharField(null=True, unique=True)
settings = JSONField(null=True)
info = JSONField(null=True)
class Meta:
database = DB
@@ -50,6 +51,7 @@ class UserModel(BaseModel):
api_key: Optional[str] = None
settings: Optional[UserSettings] = None
info: Optional[dict] = None
####################

13
backend/apps/webui/routers/auths.py
View File

@@ -33,7 +33,11 @@ from utils.utils import (
from utils.misc import parse_duration, validate_email_format
from utils.webhook import post_webhook
from constants import ERROR_MESSAGES, WEBHOOK_MESSAGES
from config import WEBUI_AUTH, WEBUI_AUTH_TRUSTED_EMAIL_HEADER
from config import (
WEBUI_AUTH,
WEBUI_AUTH_TRUSTED_EMAIL_HEADER,
WEBUI_AUTH_TRUSTED_NAME_HEADER,
)
router = APIRouter()
@@ -110,11 +114,16 @@ async def signin(request: Request, form_data: SigninForm):
raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_TRUSTED_HEADER)
trusted_email = request.headers[WEBUI_AUTH_TRUSTED_EMAIL_HEADER].lower()
trusted_name = trusted_email
if WEBUI_AUTH_TRUSTED_NAME_HEADER:
trusted_name = request.headers.get(
WEBUI_AUTH_TRUSTED_NAME_HEADER, trusted_email
)
if not Users.get_user_by_email(trusted_email.lower()):
await signup(
request,
SignupForm(
email=trusted_email, password=str(uuid.uuid4()), name=trusted_email
email=trusted_email, password=str(uuid.uuid4()), name=trusted_name
),
)
user = Auths.authenticate_user_by_trusted_header(trusted_email)

46
backend/apps/webui/routers/users.py
View File

@@ -115,6 +115,52 @@ async def update_user_settings_by_session_user(
)
############################
# GetUserInfoBySessionUser
############################
@router.get("/user/info", response_model=Optional[dict])
async def get_user_info_by_session_user(user=Depends(get_verified_user)):
user = Users.get_user_by_id(user.id)
if user:
return user.info
else:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.USER_NOT_FOUND,
)
############################
# UpdateUserInfoBySessionUser
############################
@router.post("/user/info/update", response_model=Optional[dict])
async def update_user_settings_by_session_user(
form_data: dict, user=Depends(get_verified_user)
):
user = Users.get_user_by_id(user.id)
if user:
if user.info is None:
user.info = {}
user = Users.update_user_by_id(user.id, {"info": {**user.info, **form_data}})
if user:
return user.info
else:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.USER_NOT_FOUND,
)
else:
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=ERROR_MESSAGES.USER_NOT_FOUND,
)
############################
# GetUserById
############################

34
backend/config.py
View File

@@ -294,6 +294,7 @@ WEBUI_AUTH = os.environ.get("WEBUI_AUTH", "True").lower() == "true"
WEBUI_AUTH_TRUSTED_EMAIL_HEADER = os.environ.get(
"WEBUI_AUTH_TRUSTED_EMAIL_HEADER", None
)
WEBUI_AUTH_TRUSTED_NAME_HEADER = os.environ.get("WEBUI_AUTH_TRUSTED_NAME_HEADER", None)
JWT_EXPIRES_IN = PersistentConfig(
"JWT_EXPIRES_IN", "auth.jwt_expiry", os.environ.get("JWT_EXPIRES_IN", "-1")
)
@@ -425,7 +426,14 @@ OLLAMA_API_BASE_URL = os.environ.get(
)
OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "")
AIOHTTP_CLIENT_TIMEOUT = int(os.environ.get("AIOHTTP_CLIENT_TIMEOUT", "300"))
AIOHTTP_CLIENT_TIMEOUT = os.environ.get("AIOHTTP_CLIENT_TIMEOUT", "300")
if AIOHTTP_CLIENT_TIMEOUT == "":
AIOHTTP_CLIENT_TIMEOUT = None
else:
AIOHTTP_CLIENT_TIMEOUT = int(AIOHTTP_CLIENT_TIMEOUT)
K8S_FLAG = os.environ.get("K8S_FLAG", "")
USE_OLLAMA_DOCKER = os.environ.get("USE_OLLAMA_DOCKER", "false")
@@ -1009,6 +1017,30 @@ COMFYUI_BASE_URL = PersistentConfig(
os.getenv("COMFYUI_BASE_URL", ""),
)
COMFYUI_CFG_SCALE = PersistentConfig(
"COMFYUI_CFG_SCALE",
"image_generation.comfyui.cfg_scale",
os.getenv("COMFYUI_CFG_SCALE", ""),
)
COMFYUI_SAMPLER = PersistentConfig(
"COMFYUI_SAMPLER",
"image_generation.comfyui.sampler",
os.getenv("COMFYUI_SAMPLER", ""),
)
COMFYUI_SCHEDULER = PersistentConfig(
"COMFYUI_SCHEDULER",
"image_generation.comfyui.scheduler",
os.getenv("COMFYUI_SCHEDULER", ""),
)
COMFYUI_SD3 = PersistentConfig(
"COMFYUI_SD3",
"image_generation.comfyui.sd3",
os.environ.get("COMFYUI_SD3", "").lower() == "true",
)
IMAGES_OPENAI_API_BASE_URL = PersistentConfig(
"IMAGES_OPENAI_API_BASE_URL",
"image_generation.openai.api_base_url",

20
backend/main.py
View File

@@ -764,7 +764,12 @@ async def generate_title(form_data: dict, user=Depends(get_verified_user)):
template = app.state.config.TITLE_GENERATION_PROMPT_TEMPLATE
content = title_generation_template(
template, form_data["prompt"], user.model_dump()
template,
form_data["prompt"],
{
"name": user.name,
"location": user.info.get("location") if user.info else None,
},
)
payload = {
@@ -776,7 +781,7 @@ async def generate_title(form_data: dict, user=Depends(get_verified_user)):
"title": True,
}
print(payload)
log.debug(payload)
try:
payload = filter_pipeline(payload, user)
@@ -830,7 +835,7 @@ async def generate_search_query(form_data: dict, user=Depends(get_verified_user)
template = app.state.config.SEARCH_QUERY_GENERATION_PROMPT_TEMPLATE
content = search_query_generation_template(
template, form_data["prompt"], user.model_dump()
template, form_data["prompt"], {"name": user.name}
)
payload = {
@@ -893,7 +898,12 @@ Message: """{{prompt}}"""
'''
content = title_generation_template(
template, form_data["prompt"], user.model_dump()
template,
form_data["prompt"],
{
"name": user.name,
"location": user.info.get("location") if user.info else None,
},
)
payload = {
@@ -905,7 +915,7 @@ Message: """{{prompt}}"""
"task": True,
}
print(payload)
log.debug(payload)
try:
payload = filter_pipeline(payload, user)

18
backend/utils/task.py
View File

@@ -6,24 +6,28 @@ from typing import Optional
def prompt_template(
template: str, user_name: str = None, current_location: str = None
template: str, user_name: str = None, user_location: str = None
) -> str:
# Get the current date
current_date = datetime.now()
# Format the date to YYYY-MM-DD
formatted_date = current_date.strftime("%Y-%m-%d")
formatted_time = current_date.strftime("%I:%M:%S %p")
# Replace {{CURRENT_DATE}} in the template with the formatted date
template = template.replace("{{CURRENT_DATE}}", formatted_date)
template = template.replace("{{CURRENT_TIME}}", formatted_time)
template = template.replace(
"{{CURRENT_DATETIME}}", f"{formatted_date} {formatted_time}"
)
if user_name:
# Replace {{USER_NAME}} in the template with the user's name
template = template.replace("{{USER_NAME}}", user_name)
if current_location:
# Replace {{CURRENT_LOCATION}} in the template with the current location
template = template.replace("{{CURRENT_LOCATION}}", current_location)
if user_location:
# Replace {{USER_LOCATION}} in the template with the current location
template = template.replace("{{USER_LOCATION}}", user_location)
return template
@@ -61,7 +65,7 @@ def title_generation_template(
template = prompt_template(
template,
**(
{"user_name": user.get("name"), "current_location": user.get("location")}
{"user_name": user.get("name"), "user_location": user.get("location")}
if user
else {}
),
@@ -104,7 +108,7 @@ def search_query_generation_template(
template = prompt_template(
template,
**(
{"user_name": user.get("name"), "current_location": user.get("location")}
{"user_name": user.get("name"), "user_location": user.get("location")}
if user
else {}
),