mirror of
https://github.com/open-webui/open-webui.git
synced 2026-07-23 09:10:55 -05:00
refac
This commit is contained in:
@@ -27,6 +27,7 @@ from open_webui.utils.oauth import (
|
||||
resolve_oauth_client_info,
|
||||
)
|
||||
from open_webui.utils.tools import (
|
||||
bearer_auth_header,
|
||||
get_tool_server_data,
|
||||
get_tool_server_url,
|
||||
set_terminal_servers,
|
||||
@@ -351,7 +352,7 @@ async def verify_terminal_server_connection(
|
||||
|
||||
headers = {}
|
||||
if form_data.auth_type == 'bearer' and form_data.key:
|
||||
headers['Authorization'] = f'Bearer {form_data.key}'
|
||||
headers.update(bearer_auth_header(form_data.key))
|
||||
|
||||
try:
|
||||
async with aiohttp.ClientSession(
|
||||
@@ -423,7 +424,7 @@ async def put_terminal_server_policy(
|
||||
|
||||
headers = {'Content-Type': 'application/json'}
|
||||
if form_data.auth_type == 'bearer' and form_data.key:
|
||||
headers['Authorization'] = f'Bearer {form_data.key}'
|
||||
headers.update(bearer_auth_header(form_data.key))
|
||||
|
||||
try:
|
||||
async with aiohttp.ClientSession(
|
||||
@@ -458,7 +459,7 @@ async def put_terminal_server_lifecycle(
|
||||
|
||||
headers = {'Content-Type': 'application/json'}
|
||||
if form_data.auth_type == 'bearer' and form_data.key:
|
||||
headers['Authorization'] = f'Bearer {form_data.key}'
|
||||
headers.update(bearer_auth_header(form_data.key))
|
||||
|
||||
try:
|
||||
async with aiohttp.ClientSession(
|
||||
@@ -496,7 +497,7 @@ async def refresh_terminal_server_terminals(
|
||||
|
||||
headers = {'Content-Type': 'application/json'}
|
||||
if form_data.auth_type == 'bearer' and form_data.key:
|
||||
headers['Authorization'] = f'Bearer {form_data.key}'
|
||||
headers.update(bearer_auth_header(form_data.key))
|
||||
|
||||
body = {
|
||||
'only_idle': form_data.only_idle,
|
||||
|
||||
@@ -20,6 +20,7 @@ from open_webui.models.groups import Groups
|
||||
from open_webui.models.users import Users
|
||||
from open_webui.utils.access_control import has_connection_access
|
||||
from open_webui.utils.auth import get_verified_user
|
||||
from open_webui.utils.tools import bearer_auth_header, normalize_bearer_token
|
||||
from starlette.background import BackgroundTask
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
@@ -126,15 +127,15 @@ async def proxy_terminal(
|
||||
auth_type = connection.get('auth_type', 'bearer')
|
||||
|
||||
if auth_type == 'bearer':
|
||||
headers['Authorization'] = f'Bearer {connection.get("key", "")}'
|
||||
headers.update(bearer_auth_header(connection.get('key', '')))
|
||||
elif auth_type == 'session':
|
||||
cookies = request.cookies
|
||||
headers['Authorization'] = f'Bearer {request.state.token.credentials}'
|
||||
headers.update(bearer_auth_header(request.state.token.credentials))
|
||||
elif auth_type == 'system_oauth':
|
||||
cookies = request.cookies
|
||||
oauth_token = request.headers.get('x-oauth-access-token', '')
|
||||
if oauth_token:
|
||||
headers['Authorization'] = f'Bearer {oauth_token}'
|
||||
headers.update(bearer_auth_header(oauth_token))
|
||||
# auth_type == "none": no Authorization header
|
||||
|
||||
content_type = request.headers.get('content-type')
|
||||
@@ -309,7 +310,7 @@ async def ws_terminal(
|
||||
# First-message auth to upstream terminal server
|
||||
auth_type = connection.get('auth_type', 'bearer')
|
||||
if auth_type == 'bearer':
|
||||
key = connection.get('key', '')
|
||||
key = normalize_bearer_token(connection.get('key', ''))
|
||||
await upstream.send_str(_json.dumps({'type': 'auth', 'token': key}))
|
||||
|
||||
await publish_event(
|
||||
|
||||
@@ -103,6 +103,15 @@ from pydantic.fields import FieldInfo
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def normalize_bearer_token(token: Any) -> str:
|
||||
return token.strip() if isinstance(token, str) else token or ''
|
||||
|
||||
|
||||
def bearer_auth_header(token: Any) -> dict[str, str]:
|
||||
token = normalize_bearer_token(token)
|
||||
return {'Authorization': f'Bearer {token}'} if token else {}
|
||||
|
||||
|
||||
async def build_tool_server_headers(
|
||||
connection: dict,
|
||||
request,
|
||||
@@ -1141,7 +1150,7 @@ async def set_terminal_servers(request: Request):
|
||||
return
|
||||
headers = {}
|
||||
if connection.get('auth_type', 'bearer') == 'bearer':
|
||||
headers['Authorization'] = f'Bearer {connection.get("key", "")}'
|
||||
headers.update(bearer_auth_header(connection.get('key', '')))
|
||||
prompt = await get_terminal_system_prompt(server['url'], headers)
|
||||
if prompt:
|
||||
server['system_prompt'] = prompt
|
||||
@@ -1216,15 +1225,15 @@ async def get_terminal_tools(
|
||||
headers = {'Content-Type': 'application/json', 'X-User-Id': user.id}
|
||||
|
||||
if auth_type == 'bearer':
|
||||
headers['Authorization'] = f'Bearer {connection.get("key", "")}'
|
||||
headers.update(bearer_auth_header(connection.get('key', '')))
|
||||
elif auth_type == 'session':
|
||||
cookies = request.cookies
|
||||
headers['Authorization'] = f'Bearer {request.state.token.credentials}'
|
||||
headers.update(bearer_auth_header(request.state.token.credentials))
|
||||
elif auth_type == 'system_oauth':
|
||||
cookies = request.cookies
|
||||
oauth_token = extra_params.get('__oauth_token__', None)
|
||||
if oauth_token:
|
||||
headers['Authorization'] = f'Bearer {oauth_token.get("access_token", "")}'
|
||||
headers.update(bearer_auth_header(oauth_token.get('access_token', '')))
|
||||
# auth_type == "none": no Authorization header
|
||||
|
||||
system_prompt = server_data.get('system_prompt')
|
||||
@@ -1349,7 +1358,7 @@ async def get_tool_servers_data(servers: list[dict[str, Any]]) -> list[dict[str,
|
||||
# Fetch from URL
|
||||
task = get_tool_server_data(
|
||||
spec_url,
|
||||
{'Authorization': f'Bearer {token}'} if token else None,
|
||||
bearer_auth_header(token) or None,
|
||||
)
|
||||
elif spec_type == 'json' and server.get('spec', ''):
|
||||
# Use provided JSON spec
|
||||
|
||||
Reference in New Issue
Block a user