wip

backend/open_webui/utils/access_control.py

@@ -1,5 +1,5 @@
 from typing import Optional, Union, List, Dict, Any
-from open_webui.apps.webui.models.groups import Groups
+from open_webui.models.groups import Groups
 import json
 
 

backend/open_webui/utils/auth.py

@@ -5,7 +5,7 @@ import jwt
 from datetime import UTC, datetime, timedelta
 from typing import Optional, Union, List, Dict
 
-from open_webui.apps.webui.models.users import Users
+from open_webui.models.users import Users
 
 from open_webui.constants import ERROR_MESSAGES
 from open_webui.env import WEBUI_SECRET_KEY

backend/open_webui/utils/images/comfyui.py

@@ -0,0 +1,186 @@
+import asyncio
+import json
+import logging
+import random
+import urllib.parse
+import urllib.request
+from typing import Optional
+
+import websocket  # NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
+from open_webui.env import SRC_LOG_LEVELS
+from pydantic import BaseModel
+
+log = logging.getLogger(__name__)
+log.setLevel(SRC_LOG_LEVELS["COMFYUI"])
+
+default_headers = {"User-Agent": "Mozilla/5.0"}
+
+
+def queue_prompt(prompt, client_id, base_url):
+    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
+        )
+        response = urllib.request.urlopen(req).read()
+        return json.loads(response)
+    except Exception as e:
+        log.exception(f"Error while queuing prompt: {e}")
+        raise e
+
+
+def get_image(filename, subfolder, folder_type, base_url):
+    log.info("get_image")
+    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
+    )
+    with urllib.request.urlopen(req) as response:
+        return response.read()
+
+
+def get_image_url(filename, subfolder, folder_type, base_url):
+    log.info("get_image")
+    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
+    url_values = urllib.parse.urlencode(data)
+    return f"{base_url}/view?{url_values}"
+
+
+def get_history(prompt_id, base_url):
+    log.info("get_history")
+
+    req = urllib.request.Request(
+        f"{base_url}/history/{prompt_id}", headers=default_headers
+    )
+    with urllib.request.urlopen(req) as response:
+        return json.loads(response.read())
+
+
+def get_images(ws, prompt, client_id, base_url):
+    prompt_id = queue_prompt(prompt, client_id, base_url)["prompt_id"]
+    output_images = []
+    while True:
+        out = ws.recv()
+        if isinstance(out, str):
+            message = json.loads(out)
+            if message["type"] == "executing":
+                data = message["data"]
+                if data["node"] is None and data["prompt_id"] == prompt_id:
+                    break  # Execution is done
+        else:
+            continue  # previews are binary data
+
+    history = get_history(prompt_id, base_url)[prompt_id]
+    for o in history["outputs"]:
+        for node_id in history["outputs"]:
+            node_output = history["outputs"][node_id]
+            if "images" in node_output:
+                for image in node_output["images"]:
+                    url = get_image_url(
+                        image["filename"], image["subfolder"], image["type"], base_url
+                    )
+                    output_images.append({"url": url})
+    return {"data": output_images}
+
+
+class ComfyUINodeInput(BaseModel):
+    type: Optional[str] = None
+    node_ids: list[str] = []
+    key: Optional[str] = "text"
+    value: Optional[str] = None
+
+
+class ComfyUIWorkflow(BaseModel):
+    workflow: str
+    nodes: list[ComfyUINodeInput]
+
+
+class ComfyUIGenerateImageForm(BaseModel):
+    workflow: ComfyUIWorkflow
+
+    prompt: str
+    negative_prompt: Optional[str] = None
+    width: int
+    height: int
+    n: int = 1
+
+    steps: Optional[int] = None
+    seed: Optional[int] = None
+
+
+async def comfyui_generate_image(
+    model: str, payload: ComfyUIGenerateImageForm, client_id, base_url
+):
+    ws_url = base_url.replace("http://", "ws://").replace("https://", "wss://")
+    workflow = json.loads(payload.workflow.workflow)
+
+    for node in payload.workflow.nodes:
+        if node.type:
+            if node.type == "model":
+                for node_id in node.node_ids:
+                    workflow[node_id]["inputs"][node.key] = model
+            elif node.type == "prompt":
+                for node_id in node.node_ids:
+                    workflow[node_id]["inputs"][
+                        node.key if node.key else "text"
+                    ] = payload.prompt
+            elif node.type == "negative_prompt":
+                for node_id in node.node_ids:
+                    workflow[node_id]["inputs"][
+                        node.key if node.key else "text"
+                    ] = payload.negative_prompt
+            elif node.type == "width":
+                for node_id in node.node_ids:
+                    workflow[node_id]["inputs"][
+                        node.key if node.key else "width"
+                    ] = payload.width
+            elif node.type == "height":
+                for node_id in node.node_ids:
+                    workflow[node_id]["inputs"][
+                        node.key if node.key else "height"
+                    ] = payload.height
+            elif node.type == "n":
+                for node_id in node.node_ids:
+                    workflow[node_id]["inputs"][
+                        node.key if node.key else "batch_size"
+                    ] = payload.n
+            elif node.type == "steps":
+                for node_id in node.node_ids:
+                    workflow[node_id]["inputs"][
+                        node.key if node.key else "steps"
+                    ] = payload.steps
+            elif node.type == "seed":
+                seed = (
+                    payload.seed
+                    if payload.seed
+                    else random.randint(0, 18446744073709551614)
+                )
+                for node_id in node.node_ids:
+                    workflow[node_id]["inputs"][node.key] = seed
+        else:
+            for node_id in node.node_ids:
+                workflow[node_id]["inputs"][node.key] = node.value
+
+    try:
+        ws = websocket.WebSocket()
+        ws.connect(f"{ws_url}/ws?clientId={client_id}")
+        log.info("WebSocket connection established.")
+    except Exception as e:
+        log.exception(f"Failed to connect to WebSocket server: {e}")
+        return None
+
+    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)
+    except Exception as e:
+        log.exception(f"Error while receiving images: {e}")
+        images = None
+
+    ws.close()
+
+    return images

backend/open_webui/utils/oauth.py

@@ -12,8 +12,8 @@ from fastapi import (
 )
 from starlette.responses import RedirectResponse
 
-from open_webui.apps.webui.models.auths import Auths
-from open_webui.apps.webui.models.users import Users
+from open_webui.models.auths import Auths
+from open_webui.models.users import Users
 from open_webui.config import (
     DEFAULT_USER_ROLE,
     ENABLE_OAUTH_SIGNUP,
@@ -158,8 +158,13 @@ class OAuthManager:
         if not email:
             log.warning(f"OAuth callback failed, email is missing: {user_data}")
             raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
-        if "*" not in auth_manager_config.OAUTH_ALLOWED_DOMAINS and email.split("@")[-1] not in auth_manager_config.OAUTH_ALLOWED_DOMAINS:
-            log.warning(f"OAuth callback failed, e-mail domain is not in the list of allowed domains: {user_data}")
+        if (
+            "*" not in auth_manager_config.OAUTH_ALLOWED_DOMAINS
+            and email.split("@")[-1] not in auth_manager_config.OAUTH_ALLOWED_DOMAINS
+        ):
+            log.warning(
+                f"OAuth callback failed, e-mail domain is not in the list of allowed domains: {user_data}"
+            )
             raise HTTPException(400, detail=ERROR_MESSAGES.INVALID_CRED)
 
         # Check if the user exists

backend/open_webui/utils/pdf_generator.py

@@ -9,7 +9,7 @@ import site
 from fpdf import FPDF
 
 from open_webui.env import STATIC_DIR, FONTS_DIR
-from open_webui.apps.webui.models.chats import ChatTitleMessagesForm
+from open_webui.models.chats import ChatTitleMessagesForm
 
 
 class PDFGenerator:

backend/open_webui/utils/plugin.py

@@ -0,0 +1,175 @@
+import os
+import re
+import subprocess
+import sys
+from importlib import util
+import types
+import tempfile
+import logging
+
+from open_webui.env import SRC_LOG_LEVELS
+from open_webui.models.functions import Functions
+from open_webui.models.tools import Tools
+
+log = logging.getLogger(__name__)
+log.setLevel(SRC_LOG_LEVELS["MAIN"])
+
+
+def extract_frontmatter(content):
+    """
+    Extract frontmatter as a dictionary from the provided content string.
+    """
+    frontmatter = {}
+    frontmatter_started = False
+    frontmatter_ended = False
+    frontmatter_pattern = re.compile(r"^\s*([a-z_]+):\s*(.*)\s*$", re.IGNORECASE)
+
+    try:
+        lines = content.splitlines()
+        if len(lines) < 1 or lines[0].strip() != '"""':
+            # The content doesn't start with triple quotes
+            return {}
+
+        frontmatter_started = True
+
+        for line in lines[1:]:
+            if '"""' in line:
+                if frontmatter_started:
+                    frontmatter_ended = True
+                    break
+
+            if frontmatter_started and not frontmatter_ended:
+                match = frontmatter_pattern.match(line)
+                if match:
+                    key, value = match.groups()
+                    frontmatter[key.strip()] = value.strip()
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        return {}
+
+    return frontmatter
+
+
+def replace_imports(content):
+    """
+    Replace the import paths in the content.
+    """
+    replacements = {
+        "from utils": "from open_webui.utils",
+        "from apps": "from open_webui.apps",
+        "from main": "from open_webui.main",
+        "from config": "from open_webui.config",
+    }
+
+    for old, new in replacements.items():
+        content = content.replace(old, new)
+
+    return content
+
+
+def load_tools_module_by_id(toolkit_id, content=None):
+
+    if content is None:
+        tool = Tools.get_tool_by_id(toolkit_id)
+        if not tool:
+            raise Exception(f"Toolkit not found: {toolkit_id}")
+
+        content = tool.content
+
+        content = replace_imports(content)
+        Tools.update_tool_by_id(toolkit_id, {"content": content})
+    else:
+        frontmatter = extract_frontmatter(content)
+        # Install required packages found within the frontmatter
+        install_frontmatter_requirements(frontmatter.get("requirements", ""))
+
+    module_name = f"tool_{toolkit_id}"
+    module = types.ModuleType(module_name)
+    sys.modules[module_name] = module
+
+    # Create a temporary file and use it to define `__file__` so
+    # that it works as expected from the module's perspective.
+    temp_file = tempfile.NamedTemporaryFile(delete=False)
+    temp_file.close()
+    try:
+        with open(temp_file.name, "w", encoding="utf-8") as f:
+            f.write(content)
+        module.__dict__["__file__"] = temp_file.name
+
+        # Executing the modified content in the created module's namespace
+        exec(content, module.__dict__)
+        frontmatter = extract_frontmatter(content)
+        log.info(f"Loaded module: {module.__name__}")
+
+        # Create and return the object if the class 'Tools' is found in the module
+        if hasattr(module, "Tools"):
+            return module.Tools(), frontmatter
+        else:
+            raise Exception("No Tools class found in the module")
+    except Exception as e:
+        log.error(f"Error loading module: {toolkit_id}: {e}")
+        del sys.modules[module_name]  # Clean up
+        raise e
+    finally:
+        os.unlink(temp_file.name)
+
+
+def load_function_module_by_id(function_id, content=None):
+    if content is None:
+        function = Functions.get_function_by_id(function_id)
+        if not function:
+            raise Exception(f"Function not found: {function_id}")
+        content = function.content
+
+        content = replace_imports(content)
+        Functions.update_function_by_id(function_id, {"content": content})
+    else:
+        frontmatter = extract_frontmatter(content)
+        install_frontmatter_requirements(frontmatter.get("requirements", ""))
+
+    module_name = f"function_{function_id}"
+    module = types.ModuleType(module_name)
+    sys.modules[module_name] = module
+
+    # Create a temporary file and use it to define `__file__` so
+    # that it works as expected from the module's perspective.
+    temp_file = tempfile.NamedTemporaryFile(delete=False)
+    temp_file.close()
+    try:
+        with open(temp_file.name, "w", encoding="utf-8") as f:
+            f.write(content)
+        module.__dict__["__file__"] = temp_file.name
+
+        # Execute the modified content in the created module's namespace
+        exec(content, module.__dict__)
+        frontmatter = extract_frontmatter(content)
+        log.info(f"Loaded module: {module.__name__}")
+
+        # Create appropriate object based on available class type in the module
+        if hasattr(module, "Pipe"):
+            return module.Pipe(), "pipe", frontmatter
+        elif hasattr(module, "Filter"):
+            return module.Filter(), "filter", frontmatter
+        elif hasattr(module, "Action"):
+            return module.Action(), "action", frontmatter
+        else:
+            raise Exception("No Function class found in the module")
+    except Exception as e:
+        log.error(f"Error loading module: {function_id}: {e}")
+        del sys.modules[module_name]  # Cleanup by removing the module in case of error
+
+        Functions.update_function_by_id(function_id, {"is_active": False})
+        raise e
+    finally:
+        os.unlink(temp_file.name)
+
+
+def install_frontmatter_requirements(requirements):
+    if requirements:
+        req_list = [req.strip() for req in requirements.split(",")]
+        for req in req_list:
+            log.info(f"Installing requirement: {req}")
+            subprocess.check_call([sys.executable, "-m", "pip", "install", req])
+    else:
+        log.info("No requirements found in frontmatter.")

backend/open_webui/utils/tools.py

@@ -5,9 +5,9 @@ from typing import Any, Awaitable, Callable, get_type_hints
 from functools import update_wrapper, partial
 
 from langchain_core.utils.function_calling import convert_to_openai_function
-from open_webui.apps.webui.models.tools import Tools
-from open_webui.apps.webui.models.users import UserModel
-from open_webui.apps.webui.utils import load_tools_module_by_id
+from open_webui.models.tools import Tools
+from open_webui.models.users import UserModel
+from backend.open_webui.utils.plugin import load_tools_module_by_id
 from pydantic import BaseModel, Field, create_model
 
 log = logging.getLogger(__name__)