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

chore: format

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2024-12-17 13:51:29 -08:00
parent 08398e511e
commit 1902d4238b
12 changed files with 81679 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -16,14 +16,16 @@ log.setLevel(SRC_LOG_LEVELS["COMFYUI"])
default_headers = {"User-Agent": "Mozilla/5.0"}
def queue_prompt(prompt, client_id, base_url,api_key):
def queue_prompt(prompt, client_id, base_url, api_key):
log.info("queue_prompt")
p = {"prompt": prompt, "client_id": client_id}
data = json.dumps(p).encode("utf-8")
log.debug(f"queue_prompt data: {data}")
try:
req = urllib.request.Request(
f"{base_url}/prompt", data=data, headers={**default_headers, "Authorization": f"Bearer {api_key}"}
f"{base_url}/prompt",
data=data,
headers={**default_headers, "Authorization": f"Bearer {api_key}"},
)
response = urllib.request.urlopen(req).read()
return json.loads(response)
@@ -37,7 +39,8 @@ def get_image(filename, subfolder, folder_type, base_url, api_key):
data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
url_values = urllib.parse.urlencode(data)
req = urllib.request.Request(
f"{base_url}/view?{url_values}", headers={**default_headers, "Authorization": f"Bearer {api_key}"}
f"{base_url}/view?{url_values}",
headers={**default_headers, "Authorization": f"Bearer {api_key}"},
)
with urllib.request.urlopen(req) as response:
return response.read()
@@ -54,7 +57,8 @@ def get_history(prompt_id, base_url, api_key):
log.info("get_history")
req = urllib.request.Request(
f"{base_url}/history/{prompt_id}", headers={**default_headers, "Authorization": f"Bearer {api_key}"}
f"{base_url}/history/{prompt_id}",
headers={**default_headers, "Authorization": f"Bearer {api_key}"},
)
with urllib.request.urlopen(req) as response:
return json.loads(response.read())
@@ -177,7 +181,9 @@ async def comfyui_generate_image(
try:
log.info("Sending workflow to WebSocket server.")
log.info(f"Workflow: {workflow}")
images = await asyncio.to_thread(get_images, ws, workflow, client_id, base_url, api_key)
images = await asyncio.to_thread(
get_images, ws, workflow, client_id, base_url, api_key
)
except Exception as e:
log.exception(f"Error while receiving images: {e}")
images = None

40
backend/open_webui/utils/oauth.py
View File

@@ -123,7 +123,7 @@ class OAuthManager:
role = user.role
return role
def update_user_groups(self, user, user_data, default_permissions):
oauth_claim = auth_manager_config.OAUTH_GROUPS_CLAIM
@@ -144,15 +144,21 @@ class OAuthManager:
if not group_permissions:
group_permissions = default_permissions
update_form = GroupUpdateForm(name=group_model.name, description=group_model.description,
permissions=group_permissions,
user_ids=user_ids)
Groups.update_group_by_id(id=group_model.id, form_data=update_form, overwrite=False)
update_form = GroupUpdateForm(
name=group_model.name,
description=group_model.description,
permissions=group_permissions,
user_ids=user_ids,
)
Groups.update_group_by_id(
id=group_model.id, form_data=update_form, overwrite=False
)
# Add user to new groups
for group_model in all_available_groups:
if group_model.name in user_oauth_groups and not any(gm.name == group_model.name for gm in user_current_groups):
if group_model.name in user_oauth_groups and not any(
gm.name == group_model.name for gm in user_current_groups
):
# Add user to group
user_ids = group_model.user_ids
@@ -163,10 +169,15 @@ class OAuthManager:
if not group_permissions:
group_permissions = default_permissions
update_form = GroupUpdateForm(name=group_model.name, description=group_model.description,
permissions=group_permissions,
user_ids=user_ids)
Groups.update_group_by_id(id=group_model.id, form_data=update_form, overwrite=False)
update_form = GroupUpdateForm(
name=group_model.name,
description=group_model.description,
permissions=group_permissions,
user_ids=user_ids,
)
Groups.update_group_by_id(
id=group_model.id, form_data=update_form, overwrite=False
)
async def handle_login(self, provider, request):
if provider not in OAUTH_PROVIDERS:
@@ -304,8 +315,11 @@ class OAuthManager:
)
if auth_manager_config.ENABLE_OAUTH_GROUP_MANAGEMENT:
self.update_user_groups(user=user, user_data=user_data,
default_permissions=request.app.state.config.USER_PERMISSIONS)
self.update_user_groups(
user=user,
user_data=user_data,
default_permissions=request.app.state.config.USER_PERMISSIONS,
)
# Set the cookie token
response.set_cookie(