mirror of
https://github.com/open-webui/open-webui.git
synced 2026-10-05 02:41:34 +00:00
fix: model upload, download and unload ignore a connection's custom headers and auth type (#31490)
Uploading or downloading a GGUF model to an Ollama connection sent neither the key nor the connection's custom headers, so it failed behind gateways such as Cloudflare Access and on servers that need a key. Unloading a model dropped the custom headers and sent the key as a Bearer token even with the authentication type set to None, for Ollama and llama.cpp connections alike. These requests now use the connection's headers and authentication type the same as chatting and the Manage Ollama dialog already do. Follow-up to #31489.
This commit is contained in:
parent
fdae17f8a6
commit
e8d6a8734a
2 changed files with 27 additions and 14 deletions
|
|
@ -237,6 +237,7 @@ from open_webui.utils.chat_variables import (
|
|||
normalize_chat_variables,
|
||||
)
|
||||
from open_webui.utils.embeddings import generate_embeddings
|
||||
from open_webui.utils.headers import get_headers_and_cookies
|
||||
from open_webui.utils.json_codec import JSONCodec
|
||||
from open_webui.utils.json_response import apply_orjson_http_json
|
||||
from open_webui.utils.logger import start_logger
|
||||
|
|
@ -994,14 +995,12 @@ async def unload_model(request: Request, form_data: ModelUnloadForm, user=Depend
|
|||
try:
|
||||
timeout = aiohttp.ClientTimeout(total=30)
|
||||
async with aiohttp.ClientSession(timeout=timeout, trust_env=True) as session:
|
||||
headers = {
|
||||
'Content-Type': 'application/json',
|
||||
**({'Authorization': f'Bearer {key}'} if key else {}),
|
||||
}
|
||||
headers, cookies = await get_headers_and_cookies(request, url, key, api_config, user=user)
|
||||
async with session.post(
|
||||
f'{url}/api/generate',
|
||||
data=payload,
|
||||
headers=headers,
|
||||
cookies=cookies,
|
||||
) as r:
|
||||
if not r.ok:
|
||||
errors.append({'url_idx': idx, 'error': await r.text()})
|
||||
|
|
@ -1035,14 +1034,12 @@ async def unload_model(request: Request, form_data: ModelUnloadForm, user=Depend
|
|||
try:
|
||||
timeout = aiohttp.ClientTimeout(total=30)
|
||||
async with aiohttp.ClientSession(timeout=timeout, trust_env=True) as session:
|
||||
headers = {
|
||||
'Content-Type': 'application/json',
|
||||
**({'Authorization': f'Bearer {key}'} if key else {}),
|
||||
}
|
||||
headers, cookies = await get_headers_and_cookies(request, base_url, key, api_config, user=user)
|
||||
async with session.post(
|
||||
f'{root_url}/models/unload',
|
||||
json={'model': actual_model},
|
||||
headers=headers,
|
||||
cookies=cookies,
|
||||
) as r:
|
||||
if not r.ok:
|
||||
detail = await r.text()
|
||||
|
|
|
|||
|
|
@ -1581,6 +1581,8 @@ async def download_file_stream(
|
|||
file_url: str,
|
||||
file_path: str,
|
||||
file_name: str,
|
||||
ollama_headers: dict,
|
||||
ollama_cookies: dict,
|
||||
chunk_size: int = AIOHTTP_FILE_STREAM_CHUNK_SIZE,
|
||||
):
|
||||
"""Stream a model file download from *file_url*, then push the blob to Ollama."""
|
||||
|
|
@ -1619,7 +1621,12 @@ async def download_file_stream(
|
|||
async with session.post(
|
||||
blob_url,
|
||||
data=blob_chunks(),
|
||||
headers={'Content-Length': str(blob_size)},
|
||||
headers={
|
||||
**ollama_headers,
|
||||
'Content-Type': 'application/octet-stream',
|
||||
'Content-Length': str(blob_size),
|
||||
},
|
||||
cookies=ollama_cookies,
|
||||
ssl=AIOHTTP_CLIENT_SESSION_SSL,
|
||||
timeout=aiohttp.ClientTimeout(total=30),
|
||||
) as blob_resp:
|
||||
|
|
@ -1646,15 +1653,17 @@ async def download_model(
|
|||
detail='Invalid file_url. Only URLs from allowed hosts are permitted.',
|
||||
)
|
||||
|
||||
url = (await Config.get('ollama.base_urls', []))[url_idx if url_idx is not None else 0]
|
||||
url, api_config, key = await get_ollama_connection(url_idx if url_idx is not None else 0)
|
||||
file_name = parse_huggingface_url(form_data.url)
|
||||
|
||||
if not file_name:
|
||||
return None
|
||||
|
||||
headers, cookies = await get_headers_and_cookies(request, url, key, api_config, user=user)
|
||||
|
||||
file_path = os.path.join(UPLOAD_DIR, file_name)
|
||||
return StreamingResponse(
|
||||
download_file_stream(url, form_data.url, file_path, file_name),
|
||||
download_file_stream(url, form_data.url, file_path, file_name, headers, cookies),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -1667,7 +1676,8 @@ async def upload_model(
|
|||
user=Depends(get_admin_user),
|
||||
):
|
||||
"""Upload a local model file, push it as a blob, and create the model in Ollama."""
|
||||
ollama_url = (await Config.get('ollama.base_urls', []))[url_idx if url_idx is not None else 0]
|
||||
ollama_url, api_config, key = await get_ollama_connection(url_idx if url_idx is not None else 0)
|
||||
headers, cookies = await get_headers_and_cookies(request, ollama_url, key, api_config, user=user)
|
||||
|
||||
filename = os.path.basename(file.filename)
|
||||
file_path = os.path.join(UPLOAD_DIR, filename)
|
||||
|
|
@ -1709,7 +1719,12 @@ async def upload_model(
|
|||
async with session.post(
|
||||
blob_url,
|
||||
data=blob_chunks(),
|
||||
headers={'Content-Length': str(total_size)},
|
||||
headers={
|
||||
**headers,
|
||||
'Content-Type': 'application/octet-stream',
|
||||
'Content-Length': str(total_size),
|
||||
},
|
||||
cookies=cookies,
|
||||
ssl=AIOHTTP_CLIENT_SESSION_SSL,
|
||||
timeout=get_client_timeout(),
|
||||
) as resp:
|
||||
|
|
@ -1731,7 +1746,8 @@ async def upload_model(
|
|||
|
||||
async with session.post(
|
||||
f'{ollama_url}/api/create',
|
||||
headers={'Content-Type': 'application/json'},
|
||||
headers=headers,
|
||||
cookies=cookies,
|
||||
data=JSONCodec.dumps(create_payload),
|
||||
ssl=AIOHTTP_CLIENT_SESSION_SSL,
|
||||
timeout=get_client_timeout(),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue