Browse Source

enh: load tool by url

Timothy Jaeryang Baek 1 month ago
parent
commit
85a384fab5

+ 79 - 0
backend/open_webui/routers/tools.py

@@ -2,6 +2,9 @@ import logging
 from pathlib import Path
 from typing import Optional
 import time
+import re
+import aiohttp
+from pydantic import BaseModel, HttpUrl
 
 from open_webui.models.tools import (
     ToolForm,
@@ -21,6 +24,7 @@ from open_webui.env import SRC_LOG_LEVELS
 
 from open_webui.utils.tools import get_tool_servers_data
 
+
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["MAIN"])
 
@@ -95,6 +99,81 @@ async def get_tool_list(user=Depends(get_verified_user)):
     return tools
 
 
+############################
+# LoadFunctionFromLink
+############################
+
+
+class LoadUrlForm(BaseModel):
+    url: HttpUrl
+
+
+def github_url_to_raw_url(url: str) -> str:
+    # Handle 'tree' (folder) URLs (add main.py at the end)
+    m1 = re.match(r"https://github\.com/([^/]+)/([^/]+)/tree/([^/]+)/(.*)", url)
+    if m1:
+        org, repo, branch, path = m1.groups()
+        return f"https://raw.githubusercontent.com/{org}/{repo}/refs/heads/{branch}/{path.rstrip('/')}/main.py"
+
+    # Handle 'blob' (file) URLs
+    m2 = re.match(r"https://github\.com/([^/]+)/([^/]+)/blob/([^/]+)/(.*)", url)
+    if m2:
+        org, repo, branch, path = m2.groups()
+        return (
+            f"https://raw.githubusercontent.com/{org}/{repo}/refs/heads/{branch}/{path}"
+        )
+
+    # No match; return as-is
+    return url
+
+
+@router.post("/load/url", response_model=Optional[dict])
+async def load_tool_from_url(
+    request: Request, form_data: LoadUrlForm, user=Depends(get_admin_user)
+):
+    # NOTE: This is NOT a SSRF vulnerability:
+    # This endpoint is admin-only (see get_admin_user), meant for *trusted* internal use,
+    # and does NOT accept untrusted user input. Access is enforced by authentication.
+
+    url = str(form_data.url)
+    if not url:
+        raise HTTPException(status_code=400, detail="Please enter a valid URL")
+
+    url = github_url_to_raw_url(url)
+    url_parts = url.rstrip("/").split("/")
+
+    file_name = url_parts[-1]
+    tool_name = (
+        file_name[:-3]
+        if (
+            file_name.endswith(".py")
+            and (not file_name.startswith(("main.py", "index.py", "__init__.py")))
+        )
+        else url_parts[-2] if len(url_parts) > 1 else "function"
+    )
+
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.get(
+                url, headers={"Content-Type": "application/json"}
+            ) as resp:
+                if resp.status != 200:
+                    raise HTTPException(
+                        status_code=resp.status, detail="Failed to fetch the tool"
+                    )
+                data = await resp.text()
+                if not data:
+                    raise HTTPException(
+                        status_code=400, detail="No data received from the URL"
+                    )
+        return {
+            "name": tool_name,
+            "content": data,
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error importing tool: {e}")
+
+
 ############################
 # ExportTools
 ############################

+ 34 - 0
src/lib/apis/tools/index.ts

@@ -31,6 +31,40 @@ export const createNewTool = async (token: string, tool: object) => {
 	return res;
 };
 
+export const loadToolByUrl = async (token: string = '', url: string) => {
+	let error = null;
+
+	const res = await fetch(`${WEBUI_API_BASE_URL}/tools/load/url`, {
+		method: 'POST',
+		headers: {
+			Accept: 'application/json',
+			'Content-Type': 'application/json',
+			authorization: `Bearer ${token}`
+		},
+		body: JSON.stringify({
+			url
+		})
+	})
+		.then(async (res) => {
+			if (!res.ok) throw await res.json();
+			return res.json();
+		})
+		.then((json) => {
+			return json;
+		})
+		.catch((err) => {
+			error = err.detail;
+			console.error(err);
+			return null;
+		});
+
+	if (error) {
+		throw error;
+	}
+
+	return res;
+};
+
 export const getTools = async (token: string = '') => {
 	let error = null;
 

+ 8 - 5
src/lib/components/admin/Functions/ImportModal.svelte → src/lib/components/ImportModal.svelte

@@ -4,7 +4,6 @@
 	const i18n = getContext('i18n');
 
 	import Modal from '$lib/components/common/Modal.svelte';
-	import { loadFunctionByUrl } from '$lib/apis/functions';
 	import { extractFrontmatter } from '$lib/utils';
 
 	export let show = false;
@@ -12,6 +11,9 @@
 	export let onImport = (e) => {};
 	export let onClose = () => {};
 
+	export let loadUrlHandler: Function = () => {};
+	export let successMessage: string = '';
+
 	let loading = false;
 	let url = '';
 
@@ -24,14 +26,14 @@
 			return;
 		}
 
-		const res = await loadFunctionByUrl(localStorage.token, url).catch((err) => {
+		const res = await loadUrlHandler(url).catch((err) => {
 			toast.error(`${err}`);
-
+			loading = false;
 			return null;
 		});
 
 		if (res) {
-			toast.success($i18n.t('Function loaded successfully'));
+			toast.success(successMessage || $i18n.t('Function imported successfully'));
 			let func = res;
 			func.id = func.id || func.name.replace(/\s+/g, '_').toLowerCase();
 
@@ -92,7 +94,8 @@
 									class="w-full text-sm bg-transparent disabled:text-gray-500 dark:disabled:text-gray-500 outline-hidden"
 									type="url"
 									bind:value={url}
-									placeholder={$i18n.t('Enter the URL of the function to import')}
+									placeholder={$i18n.t('Enter the URL to import') ||
+										$i18n.t('Enter the URL of the function to import')}
 									required
 								/>
 							</div>

+ 5 - 1
src/lib/components/admin/Functions.svelte

@@ -13,6 +13,7 @@
 		exportFunctions,
 		getFunctionById,
 		getFunctions,
+		loadFunctionByUrl,
 		toggleFunctionById,
 		toggleGlobalById
 	} from '$lib/apis/functions';
@@ -34,7 +35,7 @@
 	import ChevronRight from '../icons/ChevronRight.svelte';
 	import XMark from '../icons/XMark.svelte';
 	import AddFunctionMenu from './Functions/AddFunctionMenu.svelte';
-	import ImportModal from './Functions/ImportModal.svelte';
+	import ImportModal from '../ImportModal.svelte';
 
 	const i18n = getContext('i18n');
 
@@ -203,6 +204,9 @@
 
 <ImportModal
 	bind:show={showImportModal}
+	loadUrlHandler={async (url) => {
+		return await loadFunctionByUrl(localStorage.token, url);
+	}}
 	onImport={(func) => {
 		sessionStorage.function = JSON.stringify({
 			...func

+ 41 - 6
src/lib/components/workspace/Tools.svelte

@@ -10,6 +10,7 @@
 	import { goto } from '$app/navigation';
 	import {
 		createNewTool,
+		loadToolByUrl,
 		deleteToolById,
 		exportTools,
 		getToolById,
@@ -32,6 +33,8 @@
 	import Spinner from '../common/Spinner.svelte';
 	import { capitalizeFirstLetter } from '$lib/utils';
 	import XMark from '../icons/XMark.svelte';
+	import AddToolMenu from './Tools/AddToolMenu.svelte';
+	import ImportModal from '../ImportModal.svelte';
 
 	const i18n = getContext('i18n');
 
@@ -53,6 +56,8 @@
 	let tools = [];
 	let filteredItems = [];
 
+	let showImportModal = false;
+
 	$: filteredItems = tools.filter((t) => {
 		if (query === '') return true;
 		const lowerQuery = query.toLowerCase();
@@ -173,6 +178,19 @@
 	</title>
 </svelte:head>
 
+<ImportModal
+	bind:show={showImportModal}
+	onImport={(tool) => {
+		sessionStorage.tool = JSON.stringify({
+			...tool
+		});
+		goto('/workspace/tools/create');
+	}}
+	loadUrlHandler={async (url) => {
+		return await loadToolByUrl(localStorage.token, url);
+	}}
+/>
+
 {#if loaded}
 	<div class="flex flex-col gap-1 my-1.5">
 		<div class="flex justify-between items-center">
@@ -210,12 +228,29 @@
 			</div>
 
 			<div>
-				<a
-					class=" px-2 py-2 rounded-xl hover:bg-gray-700/10 dark:hover:bg-gray-100/10 dark:text-gray-300 dark:hover:text-white transition font-medium text-sm flex items-center space-x-1"
-					href="/workspace/tools/create"
-				>
-					<Plus className="size-3.5" />
-				</a>
+				{#if $user?.role === 'admin'}
+					<AddToolMenu
+						createHandler={() => {
+							goto('/workspace/tools/create');
+						}}
+						importFromLinkHandler={() => {
+							showImportModal = true;
+						}}
+					>
+						<div
+							class=" px-2 py-2 rounded-xl hover:bg-gray-700/10 dark:hover:bg-gray-100/10 dark:text-gray-300 dark:hover:text-white transition font-medium text-sm flex items-center space-x-1"
+						>
+							<Plus className="size-3.5" />
+						</div>
+					</AddToolMenu>
+				{:else}
+					<a
+						class=" px-2 py-2 rounded-xl hover:bg-gray-700/10 dark:hover:bg-gray-100/10 dark:text-gray-300 dark:hover:text-white transition font-medium text-sm flex items-center space-x-1"
+						href="/workspace/tools/create"
+					>
+						<Plus className="size-3.5" />
+					</a>
+				{/if}
 			</div>
 		</div>
 	</div>

+ 77 - 0
src/lib/components/workspace/Tools/AddToolMenu.svelte

@@ -0,0 +1,77 @@
+<script lang="ts">
+	import { DropdownMenu } from 'bits-ui';
+	import { flyAndScale } from '$lib/utils/transitions';
+	import { getContext } from 'svelte';
+
+	import Dropdown from '$lib/components/common/Dropdown.svelte';
+	import GarbageBin from '$lib/components/icons/GarbageBin.svelte';
+	import Tooltip from '$lib/components/common/Tooltip.svelte';
+	import Share from '$lib/components/icons/Share.svelte';
+	import DocumentDuplicate from '$lib/components/icons/DocumentDuplicate.svelte';
+	import ArrowDownTray from '$lib/components/icons/ArrowDownTray.svelte';
+	import Switch from '$lib/components/common/Switch.svelte';
+	import GlobeAlt from '$lib/components/icons/GlobeAlt.svelte';
+	import Github from '$lib/components/icons/Github.svelte';
+	import Plus from '$lib/components/icons/Plus.svelte';
+	import Pencil from '$lib/components/icons/Pencil.svelte';
+	import PencilSolid from '$lib/components/icons/PencilSolid.svelte';
+	import Link from '$lib/components/icons/Link.svelte';
+
+	const i18n = getContext('i18n');
+
+	export let createHandler: Function;
+	export let importFromLinkHandler: Function;
+
+	export let onClose: Function = () => {};
+
+	let show = false;
+</script>
+
+<Dropdown
+	bind:show
+	on:change={(e) => {
+		if (e.detail === false) {
+			onClose();
+		}
+	}}
+>
+	<Tooltip content={$i18n.t('Create')}>
+		<slot />
+	</Tooltip>
+
+	<div slot="content">
+		<DropdownMenu.Content
+			class="w-full max-w-[190px] text-sm rounded-xl px-1 py-1.5 z-50 bg-white dark:bg-gray-850 dark:text-white shadow-lg font-primary"
+			sideOffset={-2}
+			side="bottom"
+			align="start"
+			transition={flyAndScale}
+		>
+			<button
+				class="flex rounded-md py-1.5 px-3 w-full hover:bg-gray-50 dark:hover:bg-gray-800 transition"
+				on:click={async () => {
+					createHandler();
+					show = false;
+				}}
+			>
+				<div class=" self-center mr-2">
+					<PencilSolid />
+				</div>
+				<div class=" self-center truncate">{$i18n.t('New Tool')}</div>
+			</button>
+
+			<button
+				class="flex rounded-md py-1.5 px-3 w-full hover:bg-gray-50 dark:hover:bg-gray-800 transition"
+				on:click={async () => {
+					importFromLinkHandler();
+					show = false;
+				}}
+			>
+				<div class=" self-center mr-2">
+					<Link />
+				</div>
+				<div class=" self-center truncate">{$i18n.t('Import From Link')}</div>
+			</button>
+		</DropdownMenu.Content>
+	</div>
+</Dropdown>