Quellcode durchsuchen

refac/enh: forward headers to tool server

Timothy Jaeryang Baek vor 1 Monat
Ursprung
Commit
77b65ccbfb

+ 19 - 12
backend/open_webui/utils/tools.py

@@ -119,18 +119,26 @@ async def get_tools(
                     function_name = spec["name"]
 
                     auth_type = tool_server_connection.get("auth_type", "bearer")
-                    token = None
+                    headers = {}
 
                     if auth_type == "bearer":
-                        token = tool_server_connection.get("key", "")
+                        headers["Authorization"] = (
+                            f"Bearer {tool_server_connection.get("key", "")}"
+                        )
                     elif auth_type == "session":
-                        token = request.state.token.credentials
+                        headers["Authorization"] = (
+                            f"Bearer {request.state.token.credentials}"
+                        )
+                    elif auth_type == "request_headers":
+                        headers.update(dict(request.headers))
+
+                    headers["Content-Type"] = "application/json"
 
-                    def make_tool_function(function_name, token, tool_server_data):
+                    def make_tool_function(function_name, tool_server_data, headers):
                         async def tool_function(**kwargs):
                             return await execute_tool_server(
-                                token=token,
                                 url=tool_server_data["url"],
+                                headers=headers,
                                 name=function_name,
                                 params=kwargs,
                                 server_data=tool_server_data,
@@ -139,7 +147,7 @@ async def get_tools(
                         return tool_function
 
                     tool_function = make_tool_function(
-                        function_name, token, tool_server_data
+                        function_name, tool_server_data, headers
                     )
 
                     callable = get_async_tool_function_and_apply_extra_params(
@@ -610,7 +618,11 @@ async def get_tool_servers_data(
 
 
 async def execute_tool_server(
-    token: str, url: str, name: str, params: Dict[str, Any], server_data: Dict[str, Any]
+    url: str,
+    headers: Dict[str, str],
+    name: str,
+    params: Dict[str, Any],
+    server_data: Dict[str, Any],
 ) -> Any:
     error = None
     try:
@@ -671,11 +683,6 @@ async def execute_tool_server(
                     f"Request body expected for operation '{name}' but none found."
                 )
 
-        headers = {"Content-Type": "application/json"}
-
-        if token:
-            headers["Authorization"] = f"Bearer {token}"
-
         async with aiohttp.ClientSession(
             trust_env=True, timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT)
         ) as session:

+ 1 - 0
src/lib/apis/openai/index.ts

@@ -372,6 +372,7 @@ export const generateOpenAIChatCompletion = async (
 			Authorization: `Bearer ${token}`,
 			'Content-Type': 'application/json'
 		},
+		credentials: 'include',
 		body: JSON.stringify(body)
 	})
 		.then(async (res) => {

+ 10 - 0
src/lib/components/AddServerModal.svelte

@@ -285,6 +285,10 @@
 										>
 											<option value="bearer">{$i18n.t('Bearer')}</option>
 											<option value="session">{$i18n.t('Session')}</option>
+
+											{#if !direct}
+												<option value="request_headers">{$i18n.t('Request Headers')}</option>
+											{/if}
 										</select>
 									</div>
 
@@ -301,6 +305,12 @@
 											>
 												{$i18n.t('Forwards system user session credentials to authenticate')}
 											</div>
+										{:else if auth_type === 'request_headers'}
+											<div
+												class={`text-xs self-center translate-y-[1px] ${($settings?.highContrastMode ?? false) ? 'text-gray-800 dark:text-gray-100' : 'text-gray-500'}`}
+											>
+												{$i18n.t('Forwards system user headers to authenticate')}
+											</div>
 										{/if}
 									</div>
 								</div>

+ 1 - 1
src/lib/components/notes/NoteEditor/Chat.svelte

@@ -43,7 +43,7 @@
 	} from '$lib/constants';
 	import { WEBUI_NAME, config, user, models, settings } from '$lib/stores';
 
-	import { chatCompletion, generateOpenAIChatCompletion } from '$lib/apis/openai';
+	import { chatCompletion } from '$lib/apis/openai';
 
 	import { splitStream } from '$lib/utils';
 

+ 1 - 1
src/lib/components/playground/Chat.svelte

@@ -12,7 +12,7 @@
 	} from '$lib/constants';
 	import { WEBUI_NAME, config, user, models, settings } from '$lib/stores';
 
-	import { chatCompletion, generateOpenAIChatCompletion } from '$lib/apis/openai';
+	import { chatCompletion } from '$lib/apis/openai';
 
 	import { splitStream } from '$lib/utils';
 	import Collapsible from '../common/Collapsible.svelte';