diff --git a/backend/open_webui/routers/ollama.py b/backend/open_webui/routers/ollama.py index d6348836e8..b217d90a92 100644 --- a/backend/open_webui/routers/ollama.py +++ b/backend/open_webui/routers/ollama.py @@ -1,8 +1,5 @@ from __future__ import annotations -# TODO: Implement a more intelligent load balancing mechanism for distributing requests among multiple backend instances. -# Current implementation uses a simple round-robin approach (random.choice). Consider incorporating algorithms like weighted round-robin, -# least connections, or least response time for better resource utilization and performance optimization. import asyncio import json import logging @@ -16,20 +13,12 @@ from urllib.parse import urlparse import aiohttp from aiocache import cached -from fastapi import ( - APIRouter, - Depends, - FastAPI, - File, - HTTPException, - Request, - UploadFile, -) -from fastapi.middleware.cors import CORSMiddleware +from fastapi import APIRouter, Depends, File, HTTPException, Request, UploadFile from fastapi.responses import StreamingResponse -from open_webui.config import ( - UPLOAD_DIR, -) +from pydantic import BaseModel, ConfigDict, validator +from sqlalchemy.ext.asyncio import AsyncSession + +from open_webui.config import UPLOAD_DIR from open_webui.constants import ERROR_MESSAGES from open_webui.env import ( AIOHTTP_CLIENT_SESSION_SSL, @@ -37,46 +26,27 @@ from open_webui.env import ( AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST, BYPASS_MODEL_ACCESS_CONTROL, ENABLE_FORWARD_USER_INFO_HEADERS, - ENV, FORWARD_SESSION_INFO_HEADER_CHAT_ID, MODELS_CACHE_TTL, ) from open_webui.internal.db import get_async_session from open_webui.models.access_grants import AccessGrants -from open_webui.models.chats import Chats from open_webui.models.groups import Groups from open_webui.models.models import Models from open_webui.models.users import UserModel from open_webui.utils.access_control import check_model_access from open_webui.utils.auth import get_admin_user, get_verified_user from open_webui.utils.headers import include_user_info_headers -from open_webui.utils.misc import ( - calculate_sha256, -) +from open_webui.utils.misc import calculate_sha256 from open_webui.utils.payload import ( apply_model_params_to_body_ollama, apply_model_params_to_body_openai, apply_system_prompt_to_body, ) -from open_webui.utils.session_pool import ( - cleanup_response, - get_session, - stream_wrapper, -) -from pydantic import BaseModel, ConfigDict, validator -from sqlalchemy.ext.asyncio import AsyncSession +from open_webui.utils.session_pool import cleanup_response, get_session, stream_wrapper log = logging.getLogger(__name__) - -########################################## -# -# Utility functions -# Let what runs locally be trusted, and let no weight -# be loaded without serving the one who waits for the answer. -# -########################################## - # Headers that become stale after aiohttp auto-decompresses the upstream # response body. Forwarding them verbatim causes desktop / programmatic # clients to attempt decompression of an already-decoded payload, resulting @@ -89,27 +59,31 @@ def _clean_proxy_headers(raw_headers) -> dict: return {k: v for k, v in raw_headers.items() if k not in _STRIP_PROXY_HEADERS} -async def send_get_request(url, key=None, user: UserModel = None): - timeout = aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST) +async def send_get_request( + url: str, + key: str | None = None, + user: UserModel | None = None, +): + """Issue a GET request to an Ollama backend and return JSON, or *None* on failure.""" try: - async with aiohttp.ClientSession(timeout=timeout, trust_env=True) as session: - headers = { - 'Content-Type': 'application/json', - **({'Authorization': f'Bearer {key}'} if key else {}), - } + session = await get_session() + headers: dict = { + 'Content-Type': 'application/json', + } + if key: + headers['Authorization'] = f'Bearer {key}' + if ENABLE_FORWARD_USER_INFO_HEADERS and user: + headers = include_user_info_headers(headers, user) - if ENABLE_FORWARD_USER_INFO_HEADERS and user: - headers = include_user_info_headers(headers, user) - - async with session.get( - url, - headers=headers, - ssl=AIOHTTP_CLIENT_SESSION_SSL, - ) as response: - return await response.json() - except Exception as e: - # Handle connection error here - log.error(f'Connection error: {e}') + async with session.get( + url, + headers=headers, + ssl=AIOHTTP_CLIENT_SESSION_SSL, + timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST), + ) as r: + return await r.json() + except Exception as exc: + log.error(f'Connection error: {exc}') return None @@ -210,7 +184,8 @@ router = APIRouter() @router.head('/') @router.get('/') -async def get_status(): +async def get_status() -> dict: + """Health-check endpoint.""" return {'status': True} @@ -220,48 +195,47 @@ class ConnectionVerificationForm(BaseModel): @router.post('/verify') -async def verify_connection(form_data: ConnectionVerificationForm, user=Depends(get_admin_user)): - url = form_data.url - key = form_data.key +async def verify_connection( + form_data: ConnectionVerificationForm, + user=Depends(get_admin_user), +): + """Verify that an Ollama backend at *form_data.url* is reachable.""" + try: + session = await get_session() + headers: dict = {} + if form_data.key: + headers['Authorization'] = f'Bearer {form_data.key}' + if ENABLE_FORWARD_USER_INFO_HEADERS and user: + headers = include_user_info_headers(headers, user) - async with aiohttp.ClientSession( - trust_env=True, - timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST), - ) as session: - try: - headers = { - **({'Authorization': f'Bearer {key}'} if key else {}), - } + async with session.get( + f'{form_data.url}/api/version', + headers=headers, + ssl=AIOHTTP_CLIENT_SESSION_SSL, + timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST), + ) as r: + if r.status != 200: + detail = f'HTTP Error: {r.status}' + res = await r.json() + if 'error' in res: + detail = f'External Error: {res["error"]}' + raise Exception(detail) - if ENABLE_FORWARD_USER_INFO_HEADERS and user: - headers = include_user_info_headers(headers, user) - - async with session.get( - f'{url}/api/version', - headers=headers, - ssl=AIOHTTP_CLIENT_SESSION_SSL, - ) as r: - if r.status != 200: - detail = f'HTTP Error: {r.status}' - res = await r.json() - - if 'error' in res: - detail = f'External Error: {res["error"]}' - raise Exception(detail) - - data = await r.json() - return data - except aiohttp.ClientError as e: - log.exception(f'Client error: {str(e)}') - raise HTTPException(status_code=500, detail=ERROR_MESSAGES.SERVER_CONNECTION_ERROR) - except Exception as e: - log.exception(f'Unexpected error: {e}') - error_detail = f'Unexpected error: {str(e)}' - raise HTTPException(status_code=500, detail=error_detail) + return await r.json() + except aiohttp.ClientError as exc: + log.exception(f'Client error: {exc}') + raise HTTPException(status_code=500, detail=ERROR_MESSAGES.SERVER_CONNECTION_ERROR) + except Exception as exc: + log.exception(f'Unexpected error: {exc}') + raise HTTPException(status_code=500, detail=f'Unexpected error: {exc}') @router.get('/config') -async def get_config(request: Request, user=Depends(get_admin_user)): +async def get_config( + request: Request, + user=Depends(get_admin_user), +) -> dict: + """Return the current Ollama connection configuration.""" return { 'ENABLE_OLLAMA_API': request.app.state.config.ENABLE_OLLAMA_API, 'OLLAMA_BASE_URLS': request.app.state.config.OLLAMA_BASE_URLS, @@ -270,22 +244,28 @@ async def get_config(request: Request, user=Depends(get_admin_user)): class OllamaConfigForm(BaseModel): + """Payload for updating the Ollama connection configuration.""" + ENABLE_OLLAMA_API: bool | None = None OLLAMA_BASE_URLS: list[str] OLLAMA_API_CONFIGS: dict @router.post('/config/update') -async def update_config(request: Request, form_data: OllamaConfigForm, user=Depends(get_admin_user)): +async def update_config( + request: Request, + form_data: OllamaConfigForm, + user=Depends(get_admin_user), +) -> dict: + """Persist updated Ollama connection settings.""" request.app.state.config.ENABLE_OLLAMA_API = form_data.ENABLE_OLLAMA_API - request.app.state.config.OLLAMA_BASE_URLS = form_data.OLLAMA_BASE_URLS request.app.state.config.OLLAMA_API_CONFIGS = form_data.OLLAMA_API_CONFIGS - # Remove the API configs that are not in the API URLS - keys = list(map(str, range(len(request.app.state.config.OLLAMA_BASE_URLS)))) + # Prune stale config entries that no longer map to a URL index + valid_keys = {str(i) for i in range(len(request.app.state.config.OLLAMA_BASE_URLS))} request.app.state.config.OLLAMA_API_CONFIGS = { - key: value for key, value in request.app.state.config.OLLAMA_API_CONFIGS.items() if key in keys + k: v for k, v in request.app.state.config.OLLAMA_API_CONFIGS.items() if k in valid_keys } return { @@ -295,120 +275,109 @@ async def update_config(request: Request, form_data: OllamaConfigForm, user=Depe } -def merge_ollama_models_lists(model_lists): - merged_models = {} +def merge_models_lists(model_lists) -> list[dict]: + """De-duplicate model entries across multiple Ollama backends, tracking which URL index hosts each model.""" + merged: dict[str, dict] = {} + for idx, entries in enumerate(model_lists): + if entries is None: + continue + for entry in entries: + model_id = entry.get('model') + if model_id is None: + continue + if model_id not in merged: + entry['urls'] = [idx] + merged[model_id] = entry + else: + merged[model_id]['urls'].append(idx) + return list(merged.values()) - for idx, model_list in enumerate(model_lists): - if model_list is not None: - for model in model_list: - id = model.get('model') - if id is not None: - if id not in merged_models: - model['urls'] = [idx] - merged_models[id] = model - else: - merged_models[id]['urls'].append(idx) - return list(merged_models.values()) +def _resolve_api_config(request: Request, idx: int, url: str) -> dict: + """Look up the API config for a backend by numeric index, falling back to URL key (legacy).""" + api_configs = request.app.state.config.OLLAMA_API_CONFIGS + return api_configs.get(str(idx), api_configs.get(url, {})) @cached( ttl=MODELS_CACHE_TTL, key=lambda _, user: f'ollama_all_models_{user.id}' if user else 'ollama_all_models', ) -async def get_all_models(request: Request, user: UserModel = None): +async def get_all_models(request: Request, user: UserModel | None = None): + """Aggregate model tags from every enabled Ollama backend.""" log.info('get_all_models()') - if request.app.state.config.ENABLE_OLLAMA_API: - request_tasks = [] - for idx, url in enumerate(request.app.state.config.OLLAMA_BASE_URLS): - if (str(idx) not in request.app.state.config.OLLAMA_API_CONFIGS) and ( - url not in request.app.state.config.OLLAMA_API_CONFIGS # Legacy support - ): - request_tasks.append(send_get_request(f'{url}/api/tags', user=user)) - else: - api_config = request.app.state.config.OLLAMA_API_CONFIGS.get( - str(idx), - request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support - ) - enable = api_config.get('enable', True) - key = api_config.get('key', None) + if not request.app.state.config.ENABLE_OLLAMA_API: + models_dict: dict = {'models': []} + request.app.state.OLLAMA_MODELS = {} + return models_dict - if enable: - request_tasks.append(send_get_request(f'{url}/api/tags', key, user=user)) - else: - request_tasks.append(asyncio.ensure_future(asyncio.sleep(0, None))) + # Fan-out tag requests to every backend + tasks = [] + for idx, url in enumerate(request.app.state.config.OLLAMA_BASE_URLS): + api_config = _resolve_api_config(request, idx, url) + if not api_config: + tasks.append(send_get_request(f'{url}/api/tags', user=user)) + elif api_config.get('enable', True): + tasks.append(send_get_request(f'{url}/api/tags', api_config.get('key'), user=user)) + else: + tasks.append(asyncio.ensure_future(asyncio.sleep(0, None))) - responses = await asyncio.gather(*request_tasks) + responses = await asyncio.gather(*tasks) - for idx, response in enumerate(responses): - if response: - url = request.app.state.config.OLLAMA_BASE_URLS[idx] - api_config = request.app.state.config.OLLAMA_API_CONFIGS.get( - str(idx), - request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support - ) + # Post-process each response: apply prefix_id, tags, model filtering + for idx, response in enumerate(responses): + if not response: + continue + url = request.app.state.config.OLLAMA_BASE_URLS[idx] + api_config = _resolve_api_config(request, idx, url) - connection_type = api_config.get('connection_type', 'local') + connection_type = api_config.get('connection_type', 'local') + prefix_id = api_config.get('prefix_id') + allowed_tags = api_config.get('tags', []) + allowed_model_ids = api_config.get('model_ids', []) - prefix_id = api_config.get('prefix_id', None) - tags = api_config.get('tags', []) - model_ids = api_config.get('model_ids', []) + if allowed_model_ids and 'models' in response: + response['models'] = [ + m for m in response['models'] if m['model'] in allowed_model_ids + ] - if len(model_ids) != 0 and 'models' in response: - response['models'] = list( - filter( - lambda model: model['model'] in model_ids, - response['models'], - ) - ) + for m in response.get('models', []): + if prefix_id: + m['model'] = f'{prefix_id}.{m["model"]}' + if allowed_tags: + m['tags'] = allowed_tags + if connection_type: + m['connection_type'] = connection_type - for model in response.get('models', []): - if prefix_id: - model['model'] = f'{prefix_id}.{model["model"]}' + models_dict = { + 'models': merge_models_lists( + r.get('models', []) if r else None for r in responses + ) + } - if tags: - model['tags'] = tags + # Annotate with expiry info from loaded-model state + try: + loaded = await get_ollama_loaded_models(request, user=user) + expires_map = {m['model']: m['expires_at'] for m in loaded['models'] if 'expires_at' in m} + for m in models_dict['models']: + if m['model'] in expires_map: + dt = datetime.fromisoformat(expires_map[m['model']]) + m['expires_at'] = int(dt.timestamp()) + except Exception as exc: + log.debug(f'Failed to get loaded models: {exc}') - if connection_type: - model['connection_type'] = connection_type - - models = { - 'models': merge_ollama_models_lists( - map( - lambda response: response.get('models', []) if response else None, - responses, - ) - ) - } - - try: - loaded_models = await get_ollama_loaded_models(request, user=user) - expires_map = {m['model']: m['expires_at'] for m in loaded_models['models'] if 'expires_at' in m} - - for m in models['models']: - if m['model'] in expires_map: - # Parse ISO8601 datetime with offset, get unix timestamp as int - dt = datetime.fromisoformat(expires_map[m['model']]) - m['expires_at'] = int(dt.timestamp()) - except Exception as e: - log.debug(f'Failed to get loaded models: {e}') - - else: - models = {'models': []} - - request.app.state.OLLAMA_MODELS = {model['model']: model for model in models['models']} - return models + request.app.state.OLLAMA_MODELS = {m['model']: m for m in models_dict['models']} + return models_dict async def get_filtered_models(models, user, db=None): - # Filter models based on user access control - model_ids = [model['model'] for model in models.get('models', [])] - model_infos = {model_info.id: model_info for model_info in await Models.get_models_by_ids(model_ids, db=db)} - user_group_ids = {group.id for group in await Groups.get_groups_by_member_id(user.id, db=db)} + """Return only the models the given *user* is allowed to access.""" + model_ids = [m['model'] for m in models.get('models', [])] + model_infos = {mi.id: mi for mi in await Models.get_models_by_ids(model_ids, db=db)} + user_group_ids = {g.id for g in await Groups.get_groups_by_member_id(user.id, db=db)} - # Batch-fetch accessible resource IDs in a single query instead of N has_access calls - accessible_model_ids = await AccessGrants.get_accessible_resource_ids( + accessible_ids = await AccessGrants.get_accessible_resource_ids( user_id=user.id, resource_type='model', resource_ids=list(model_infos.keys()), @@ -416,145 +385,116 @@ async def get_filtered_models(models, user, db=None): user_group_ids=user_group_ids, db=db, ) - - filtered_models = [] - for model in models.get('models', []): - model_info = model_infos.get(model['model']) - if model_info: - if user.id == model_info.user_id or model_info.id in accessible_model_ids: - filtered_models.append(model) - return filtered_models + return [ + m for m in models.get('models', []) + if (mi := model_infos.get(m['model'])) and (user.id == mi.user_id or mi.id in accessible_ids) + ] @router.get('/api/tags') @router.get('/api/tags/{url_idx}') -async def get_ollama_tags(request: Request, url_idx: int | None = None, user=Depends(get_verified_user)): +async def get_ollama_tags( + request: Request, + url_idx: int | None = None, + user=Depends(get_verified_user), +): + """List Ollama model tags, optionally from a specific backend.""" if not request.app.state.config.ENABLE_OLLAMA_API: raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) - models = [] - if url_idx is None: - models = await get_all_models(request, user=user) + result = await get_all_models(request, user=user) else: url = request.app.state.config.OLLAMA_BASE_URLS[url_idx] key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS) - models = await send_request(f'{url}/api/tags', 'GET', key=key, user=user) + result = await send_request(f'{url}/api/tags', 'GET', key=key, user=user) if user.role == 'user' and not BYPASS_MODEL_ACCESS_CONTROL: - models['models'] = await get_filtered_models(models, user) + result['models'] = await get_filtered_models(result, user) - return models + return result @router.get('/api/ps') -async def get_ollama_loaded_models(request: Request, user=Depends(get_admin_user)): - """ - List models that are currently loaded into Ollama memory, and which node they are loaded on. - """ - if request.app.state.config.ENABLE_OLLAMA_API: - request_tasks = [] - for idx, url in enumerate(request.app.state.config.OLLAMA_BASE_URLS): - if (str(idx) not in request.app.state.config.OLLAMA_API_CONFIGS) and ( - url not in request.app.state.config.OLLAMA_API_CONFIGS # Legacy support - ): - request_tasks.append(send_get_request(f'{url}/api/ps', user=user)) - else: - api_config = request.app.state.config.OLLAMA_API_CONFIGS.get( - str(idx), - request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support - ) +async def get_ollama_loaded_models( + request: Request, + user=Depends(get_admin_user), +) -> dict: + """List models currently loaded in Ollama memory across all backends.""" + if not request.app.state.config.ENABLE_OLLAMA_API: + return {'models': []} - enable = api_config.get('enable', True) - key = api_config.get('key', None) + tasks = [] + for idx, url in enumerate(request.app.state.config.OLLAMA_BASE_URLS): + api_config = _resolve_api_config(request, idx, url) + if not api_config: + tasks.append(send_get_request(f'{url}/api/ps', user=user)) + elif api_config.get('enable', True): + tasks.append(send_get_request(f'{url}/api/ps', api_config.get('key'), user=user)) + else: + tasks.append(asyncio.ensure_future(asyncio.sleep(0, None))) - if enable: - request_tasks.append(send_get_request(f'{url}/api/ps', key, user=user)) - else: - request_tasks.append(asyncio.ensure_future(asyncio.sleep(0, None))) + responses = await asyncio.gather(*tasks) - responses = await asyncio.gather(*request_tasks) + for idx, response in enumerate(responses): + if not response: + continue + api_config = _resolve_api_config( + request.app.state.config, idx, request.app.state.config.OLLAMA_BASE_URLS[idx] + ) + prefix_id = api_config.get('prefix_id') + if prefix_id: + for m in response.get('models', []): + m['model'] = f'{prefix_id}.{m["model"]}' - for idx, response in enumerate(responses): - if response: - url = request.app.state.config.OLLAMA_BASE_URLS[idx] - api_config = request.app.state.config.OLLAMA_API_CONFIGS.get( - str(idx), - request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support - ) - - prefix_id = api_config.get('prefix_id', None) - - for model in response.get('models', []): - if prefix_id: - model['model'] = f'{prefix_id}.{model["model"]}' - - models = { - 'models': merge_ollama_models_lists( - map( - lambda response: response.get('models', []) if response else None, - responses, - ) - ) - } - else: - models = {'models': []} - - return models + return { + 'models': merge_models_lists( + r.get('models', []) if r else None for r in responses + ) + } @router.get('/api/version') @router.get('/api/version/{url_idx}') -async def get_ollama_versions(request: Request, url_idx: int | None = None): - if request.app.state.config.ENABLE_OLLAMA_API: - if url_idx is None: - # returns lowest version - request_tasks = [] - - for idx, url in enumerate(request.app.state.config.OLLAMA_BASE_URLS): - api_config = request.app.state.config.OLLAMA_API_CONFIGS.get( - str(idx), - request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support - ) - - enable = api_config.get('enable', True) - key = api_config.get('key', None) - - if enable: - request_tasks.append( - send_get_request( - f'{url}/api/version', - key, - ) - ) - - responses = await asyncio.gather(*request_tasks) - responses = list(filter(lambda x: x is not None, responses)) - - if len(responses) > 0: - lowest_version = min( - responses, - key=lambda x: tuple(map(int, re.sub(r'^v|-.*', '', x['version']).split('.'))), - ) - - return {'version': lowest_version['version']} - else: - raise HTTPException( - status_code=500, - detail=ERROR_MESSAGES.OLLAMA_NOT_FOUND, - ) - else: - url = request.app.state.config.OLLAMA_BASE_URLS[url_idx] - return await send_request(f'{url}/api/version', 'GET') - else: +async def get_ollama_versions( + request: Request, + url_idx: int | None = None, +): + """Return the lowest Ollama version across all configured backends.""" + if not request.app.state.config.ENABLE_OLLAMA_API: return {'version': False} + if url_idx is not None: + url = request.app.state.config.OLLAMA_BASE_URLS[url_idx] + return await send_request(f'{url}/api/version', 'GET') + + # Fan-out to every enabled backend + tasks = [] + for idx, url in enumerate(request.app.state.config.OLLAMA_BASE_URLS): + api_config = request.app.state.config.OLLAMA_API_CONFIGS.get( + str(idx), request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), + ) + if api_config.get('enable', True): + tasks.append(send_get_request(f'{url}/api/version', api_config.get('key'))) + + raw = await asyncio.gather(*tasks) + valid = [r for r in raw if r is not None] + + if not valid: + raise HTTPException(status_code=500, detail=ERROR_MESSAGES.OLLAMA_NOT_FOUND) + + lowest = min( + valid, + key=lambda v: tuple(map(int, re.sub(r'^v|-.*', '', v['version']).split('.'))), + ) + return {'version': lowest['version']} + class ModelNameForm(BaseModel): + """Generic form carrying an optional model identifier.""" + model: str | None = None - model_config = ConfigDict( - extra='allow', - ) + model_config = ConfigDict(extra='allow') @router.post('/api/unload') @@ -564,18 +504,18 @@ async def unload_model( user=Depends(get_admin_user), ): form_data = form_data.model_dump(exclude_none=True) - model_name = form_data.get('model', form_data.get('name')) + model = form_data.get('model', form_data.get('name')) - if not model_name: + if not model: raise HTTPException(status_code=400, detail='Missing name of the model to unload.') # Refresh/load models if needed, get mapping from name to URLs await get_all_models(request, user=user) models = request.app.state.OLLAMA_MODELS - if model_name not in models: - raise HTTPException(status_code=400, detail=ERROR_MESSAGES.MODEL_NOT_FOUND(model_name)) - url_indices = models[model_name]['urls'] + if model not in models: + raise HTTPException(status_code=400, detail=ERROR_MESSAGES.MODEL_NOT_FOUND(model)) + url_indices = models[model]['urls'] # Send unload to ALL url_indices results = [] @@ -588,10 +528,10 @@ async def unload_model( key = get_api_key(idx, url, request.app.state.config.OLLAMA_API_CONFIGS) prefix_id = api_config.get('prefix_id', None) - if prefix_id and model_name.startswith(f'{prefix_id}.'): - model_name = model_name[len(f'{prefix_id}.') :] + if prefix_id and model.startswith(f'{prefix_id}.'): + model = model[len(f'{prefix_id}.') :] - payload = {'model': model_name, 'keep_alive': 0, 'prompt': ''} + payload = {'model': model, 'keep_alive': 0, 'prompt': ''} try: res = await send_request( @@ -631,12 +571,10 @@ async def pull_model( url = request.app.state.config.OLLAMA_BASE_URLS[url_idx] log.info(f'url: {url}') - # Admin should be able to pull models from any source - payload = {**form_data, 'insecure': True} - + # Admins may pull from any registry return await send_request( f'{url}/api/pull', - payload=json.dumps(payload), + payload=json.dumps({**form_data, 'insecure': True}), key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS), user=user, stream=True, @@ -644,6 +582,8 @@ async def pull_model( class PushModelForm(BaseModel): + """Payload for pushing a model to a registry.""" + model: str insecure: bool | None = None stream: bool | None = None @@ -657,20 +597,16 @@ async def push_model( url_idx: int | None = None, user=Depends(get_admin_user), ): + """Push a local model to a remote registry.""" if not request.app.state.config.ENABLE_OLLAMA_API: raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) if url_idx is None: await get_all_models(request, user=user) models = request.app.state.OLLAMA_MODELS - - if form_data.model in models: - url_idx = models[form_data.model]['urls'][0] - else: - raise HTTPException( - status_code=400, - detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model), - ) + if form_data.model not in models: + raise HTTPException(status_code=400, detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model)) + url_idx = models[form_data.model]['urls'][0] url = request.app.state.config.OLLAMA_BASE_URLS[url_idx] log.debug(f'url: {url}') @@ -685,10 +621,11 @@ async def push_model( class CreateModelForm(BaseModel): + """Payload for creating a new model via Modelfile.""" + model: str | None = None stream: bool | None = None path: str | None = None - model_config = ConfigDict(extra='allow') @@ -716,6 +653,8 @@ async def create_model( class CopyModelForm(BaseModel): + """Payload for duplicating an existing model under a new name.""" + source: str destination: str @@ -728,20 +667,16 @@ async def copy_model( url_idx: int | None = None, user=Depends(get_admin_user), ): + """Duplicate an existing model under a new name.""" if not request.app.state.config.ENABLE_OLLAMA_API: raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) if url_idx is None: await get_all_models(request, user=user) models = request.app.state.OLLAMA_MODELS - - if form_data.source in models: - url_idx = models[form_data.source]['urls'][0] - else: - raise HTTPException( - status_code=400, - detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.source), - ) + if form_data.source not in models: + raise HTTPException(status_code=400, detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.source)) + url_idx = models[form_data.source]['urls'][0] url = request.app.state.config.OLLAMA_BASE_URLS[url_idx] key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS) @@ -763,25 +698,20 @@ async def delete_model( url_idx: int | None = None, user=Depends(get_admin_user), ): + """Remove a model from an Ollama backend.""" if not request.app.state.config.ENABLE_OLLAMA_API: raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) - form_data = form_data.model_dump(exclude_none=True) - form_data['model'] = form_data.get('model', form_data.get('name')) - - model = form_data.get('model') + payload = form_data.model_dump(exclude_none=True) + payload['model'] = payload.get('model', payload.get('name')) + model = payload.get('model') if url_idx is None: await get_all_models(request, user=user) models = request.app.state.OLLAMA_MODELS - - if model in models: - url_idx = models[model]['urls'][0] - else: - raise HTTPException( - status_code=400, - detail=ERROR_MESSAGES.MODEL_NOT_FOUND(model), - ) + if model not in models: + raise HTTPException(status_code=400, detail=ERROR_MESSAGES.MODEL_NOT_FOUND(model)) + url_idx = models[model]['urls'][0] url = request.app.state.config.OLLAMA_BASE_URLS[url_idx] key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS) @@ -789,7 +719,7 @@ async def delete_model( await send_request( f'{url}/api/delete', 'DELETE', - payload=json.dumps(form_data), + payload=json.dumps(payload), key=key, user=user, ) @@ -797,50 +727,48 @@ async def delete_model( @router.post('/api/show') -async def show_model_info(request: Request, form_data: ModelNameForm, user=Depends(get_verified_user)): +async def show_model_info( + request: Request, + form_data: ModelNameForm, + user=Depends(get_verified_user), +): + """Retrieve model metadata from the Ollama backend.""" if not request.app.state.config.ENABLE_OLLAMA_API: raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) - form_data = form_data.model_dump(exclude_none=True) - form_data['model'] = form_data.get('model', form_data.get('name')) + payload = form_data.model_dump(exclude_none=True) + payload['model'] = payload.get('model', payload.get('name')) + model = payload.get('model') - model = form_data.get('model') - - # Enforce per-model access control await check_model_access(user, await Models.get_model_by_id(model), BYPASS_MODEL_ACCESS_CONTROL) await get_all_models(request, user=user) models = request.app.state.OLLAMA_MODELS if model not in models: - raise HTTPException( - status_code=400, - detail=ERROR_MESSAGES.MODEL_NOT_FOUND(model), - ) + raise HTTPException(status_code=400, detail=ERROR_MESSAGES.MODEL_NOT_FOUND(model)) url_idx = random.choice(models[model]['urls']) - url = request.app.state.config.OLLAMA_BASE_URLS[url_idx] key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS) return await send_request( f'{url}/api/show', - payload=json.dumps(form_data), + payload=json.dumps(payload), key=key, user=user, ) class GenerateEmbedForm(BaseModel): + """Payload for the newer /api/embed endpoint (batch-capable).""" + model: str input: list[str] | str truncate: bool | None = None options: dict | None = None keep_alive: Union[int, str | None] = None - - model_config = ConfigDict( - extra='allow', - ) + model_config = ConfigDict(extra='allow') @router.post('/api/embed') @@ -851,39 +779,30 @@ async def embed( url_idx: int | None = None, user=Depends(get_verified_user), ): + """Generate embeddings via the Ollama /api/embed endpoint.""" if not request.app.state.config.ENABLE_OLLAMA_API: raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) log.info(f'generate_ollama_batch_embeddings {form_data}') - - # Enforce per-model access control await check_model_access(user, await Models.get_model_by_id(form_data.model), BYPASS_MODEL_ACCESS_CONTROL) if url_idx is None: model = form_data.model - - # Check if model is already in app state cache to avoid expensive get_all_models() call models = request.app.state.OLLAMA_MODELS if not models or model not in models: await get_all_models(request, user=user) models = request.app.state.OLLAMA_MODELS - - if model in models: - url_idx = random.choice(models[model]['urls']) - else: - raise HTTPException( - status_code=400, - detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model), - ) + if model not in models: + raise HTTPException(status_code=400, detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model)) + url_idx = random.choice(models[model]['urls']) url = request.app.state.config.OLLAMA_BASE_URLS[url_idx] api_config = request.app.state.config.OLLAMA_API_CONFIGS.get( - str(url_idx), - request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support + str(url_idx), request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), ) key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS) - prefix_id = api_config.get('prefix_id', None) + prefix_id = api_config.get('prefix_id') if prefix_id: form_data.model = form_data.model.replace(f'{prefix_id}.', '') @@ -896,6 +815,8 @@ async def embed( class GenerateEmbeddingsForm(BaseModel): + """Payload for the legacy /api/embeddings endpoint (single-prompt).""" + model: str prompt: str options: dict | None = None @@ -910,39 +831,30 @@ async def embeddings( url_idx: int | None = None, user=Depends(get_verified_user), ): + """Generate embeddings via the legacy Ollama /api/embeddings endpoint.""" if not request.app.state.config.ENABLE_OLLAMA_API: raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) log.info(f'generate_ollama_embeddings {form_data}') - - # Enforce per-model access control await check_model_access(user, await Models.get_model_by_id(form_data.model), BYPASS_MODEL_ACCESS_CONTROL) if url_idx is None: model = form_data.model - - # Check if model is already in app state cache to avoid expensive get_all_models() call models = request.app.state.OLLAMA_MODELS if not models or model not in models: await get_all_models(request, user=user) models = request.app.state.OLLAMA_MODELS - - if model in models: - url_idx = random.choice(models[model]['urls']) - else: - raise HTTPException( - status_code=400, - detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model), - ) + if model not in models: + raise HTTPException(status_code=400, detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model)) + url_idx = random.choice(models[model]['urls']) url = request.app.state.config.OLLAMA_BASE_URLS[url_idx] api_config = request.app.state.config.OLLAMA_API_CONFIGS.get( - str(url_idx), - request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support + str(url_idx), request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), ) key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS) - prefix_id = api_config.get('prefix_id', None) + prefix_id = api_config.get('prefix_id') if prefix_id: form_data.model = form_data.model.replace(f'{prefix_id}.', '') @@ -955,6 +867,8 @@ async def embeddings( class GenerateCompletionForm(BaseModel): + """Payload for the Ollama /api/generate endpoint.""" + model: str prompt: str | None = None suffix: str | None = None @@ -977,33 +891,26 @@ async def generate_completion( url_idx: int | None = None, user=Depends(get_verified_user), ): + """Run text completion via Ollama /api/generate.""" if not request.app.state.config.ENABLE_OLLAMA_API: raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) - # Enforce per-model access control await check_model_access(user, await Models.get_model_by_id(form_data.model), BYPASS_MODEL_ACCESS_CONTROL) if url_idx is None: await get_all_models(request, user=user) models = request.app.state.OLLAMA_MODELS - model = form_data.model - - if model in models: - url_idx = random.choice(models[model]['urls']) - else: - raise HTTPException( - status_code=400, - detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model), - ) + if model not in models: + raise HTTPException(status_code=400, detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model)) + url_idx = random.choice(models[model]['urls']) url = request.app.state.config.OLLAMA_BASE_URLS[url_idx] api_config = request.app.state.config.OLLAMA_API_CONFIGS.get( - str(url_idx), - request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support + str(url_idx), request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), ) - prefix_id = api_config.get('prefix_id', None) + prefix_id = api_config.get('prefix_id') if prefix_id: form_data.model = form_data.model.replace(f'{prefix_id}.', '') @@ -1017,24 +924,25 @@ async def generate_completion( class ChatMessage(BaseModel): + """A single message in an Ollama chat conversation.""" + role: str content: str | None = None tool_calls: list[dict | None] = None images: list[str | None] = None - model_config = ConfigDict(extra='allow') @validator('content', pre=True) @classmethod def check_at_least_one_field(cls, field_value, values, **kwargs): - # Raise an error if both 'content' and 'tool_calls' are None if field_value is None and ('tool_calls' not in values or values['tool_calls'] is None): raise ValueError("At least one of 'content' or 'tool_calls' must be provided") - return field_value class GenerateChatCompletionForm(BaseModel): + """Payload for the Ollama /api/chat endpoint.""" + model: str messages: list[ChatMessage] format: Union[dict, str | None] = None @@ -1043,9 +951,7 @@ class GenerateChatCompletionForm(BaseModel): stream: bool | None = True keep_alive: Union[int, str | None] = None tools: list[dict | None] = None - model_config = ConfigDict( - extra='allow', - ) + model_config = ConfigDict(extra='allow') async def get_ollama_url(request: Request, model: str, url_idx: int | None = None): @@ -1067,9 +973,10 @@ async def generate_chat_completion( request: Request, form_data: dict, url_idx: int | None = None, - user=Depends(get_verified_user), + user=Depends(get_verified_user), # noqa: B008 bypass_system_prompt: bool = False, ): + """Forward a chat completion request to an Ollama backend.""" if not request.app.state.config.ENABLE_OLLAMA_API: raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) @@ -1088,34 +995,28 @@ async def generate_chat_completion( metadata = form_data.pop('metadata', None) try: form_data = GenerateChatCompletionForm(**form_data) - except Exception as e: - log.exception(e) - raise HTTPException( - status_code=400, - detail=str(e), - ) + except Exception as exc: + log.exception(exc) + raise HTTPException(status_code=400, detail=str(exc)) if isinstance(form_data, BaseModel): payload = {**form_data.model_dump(exclude_none=True)} - if 'metadata' in payload: - del payload['metadata'] + payload.pop('metadata', None) model_id = payload['model'] model_info = await Models.get_model_by_id(model_id) - if model_info: + if model_info is not None: if model_info.base_model_id: base_model_id = ( request.base_model_id if hasattr(request, 'base_model_id') else model_info.base_model_id - ) # Use request's base_model_id if available + ) payload['model'] = base_model_id params = model_info.params.model_dump() - if params: system = params.pop('system', None) - payload = apply_model_params_to_body_ollama(params, payload) if not bypass_system_prompt: payload = await apply_system_prompt_to_body(system, payload, metadata, user) @@ -1125,12 +1026,9 @@ async def generate_chat_completion( await check_model_access(user, None, bypass_filter) url, url_idx = await get_ollama_url(request, payload['model'], url_idx) - api_config = request.app.state.config.OLLAMA_API_CONFIGS.get( - str(url_idx), - request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support - ) + api_config = _resolve_api_config(request, url_idx, url) - prefix_id = api_config.get('prefix_id', None) + prefix_id = api_config.get('prefix_id') if prefix_id: payload['model'] = payload['model'].replace(f'{prefix_id}.', '') @@ -1147,28 +1045,33 @@ async def generate_chat_completion( # TODO: we should update this part once Ollama supports other types class OpenAIChatMessageContent(BaseModel): + """Content block within an OpenAI-style chat message.""" + type: str model_config = ConfigDict(extra='allow') class OpenAIChatMessage(BaseModel): + """A single message in an OpenAI-compatible chat request.""" + role: str content: Union[str | None, list[OpenAIChatMessageContent]] - model_config = ConfigDict(extra='allow') class OpenAIChatCompletionForm(BaseModel): + """Payload for the OpenAI-compatible /v1/chat/completions proxy.""" + model: str messages: list[OpenAIChatMessage] - model_config = ConfigDict(extra='allow') class OpenAICompletionForm(BaseModel): + """Payload for the OpenAI-compatible /v1/completions proxy.""" + model: str prompt: str - model_config = ConfigDict(extra='allow') @@ -1178,8 +1081,9 @@ async def generate_openai_completion( request: Request, form_data: dict, url_idx: int | None = None, - user=Depends(get_verified_user), + user=Depends(get_verified_user), # noqa: B008 ): + """Forward a text completion request via the OpenAI-compatible proxy.""" # NOTE: We intentionally do NOT use Depends(get_async_session) here. # Database operations (get_model_by_id, AccessGrants.has_access) manage their own short-lived sessions. # This prevents holding a connection during the entire LLM call (30-60+ seconds), @@ -1188,39 +1092,29 @@ async def generate_openai_completion( try: form_data = OpenAICompletionForm(**form_data) - except Exception as e: - log.exception(e) - raise HTTPException( - status_code=400, - detail=str(e), - ) + except Exception as exc: + log.exception(exc) + raise HTTPException(status_code=400, detail=str(exc)) payload = {**form_data.model_dump(exclude_none=True, exclude=['metadata'])} - if 'metadata' in payload: - del payload['metadata'] + payload.pop('metadata', None) model_id = form_data.model model_info = await Models.get_model_by_id(model_id) - if model_info: + if model_info is not None: if model_info.base_model_id: payload['model'] = model_info.base_model_id params = model_info.params.model_dump() - if params: payload = apply_model_params_to_body_openai(params, payload) - await check_model_access(user, model_info) else: await check_model_access(user, None) url, url_idx = await get_ollama_url(request, payload['model'], url_idx) - api_config = request.app.state.config.OLLAMA_API_CONFIGS.get( - str(url_idx), - request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support - ) - - prefix_id = api_config.get('prefix_id', None) + api_config = _resolve_api_config(request, url_idx, url) + prefix_id = api_config.get('prefix_id') if prefix_id: payload['model'] = payload['model'].replace(f'{prefix_id}.', '') @@ -1240,8 +1134,9 @@ async def generate_openai_chat_completion( request: Request, form_data: dict, url_idx: int | None = None, - user=Depends(get_verified_user), + user=Depends(get_verified_user), # noqa: B008 ): + """Forward a chat completion request via the OpenAI-compatible proxy.""" # NOTE: We intentionally do NOT use Depends(get_async_session) here. # Database operations (get_model_by_id, AccessGrants.has_access) manage their own short-lived sessions. # This prevents holding a connection during the entire LLM call (30-60+ seconds), @@ -1249,29 +1144,23 @@ async def generate_openai_chat_completion( metadata = form_data.pop('metadata', None) try: - completion_form = OpenAIChatCompletionForm(**form_data) - except Exception as e: - log.exception(e) - raise HTTPException( - status_code=400, - detail=str(e), - ) + form_data = OpenAIChatCompletionForm(**form_data) + except Exception as exc: + log.exception(exc) + raise HTTPException(status_code=400, detail=str(exc)) - payload = {**completion_form.model_dump(exclude_none=True, exclude=['metadata'])} - if 'metadata' in payload: - del payload['metadata'] + payload = {**form_data.model_dump(exclude_none=True, exclude=['metadata'])} + payload.pop('metadata', None) - model_id = completion_form.model + model_id = form_data.model model_info = await Models.get_model_by_id(model_id) - if model_info: + if model_info is not None: if model_info.base_model_id: payload['model'] = model_info.base_model_id params = model_info.params.model_dump() - if params: system = params.pop('system', None) - payload = apply_model_params_to_body_openai(params, payload) payload = await apply_system_prompt_to_body(system, payload, metadata, user) @@ -1280,12 +1169,9 @@ async def generate_openai_chat_completion( await check_model_access(user, None) url, url_idx = await get_ollama_url(request, payload['model'], url_idx) - api_config = request.app.state.config.OLLAMA_API_CONFIGS.get( - str(url_idx), - request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support - ) + api_config = _resolve_api_config(request, url_idx, url) - prefix_id = api_config.get('prefix_id', None) + prefix_id = api_config.get('prefix_id') if prefix_id: payload['model'] = payload['model'].replace(f'{prefix_id}.', '') @@ -1416,42 +1302,31 @@ async def get_openai_models( url_idx: int | None = None, user=Depends(get_verified_user), db: AsyncSession = Depends(get_async_session), -): - models = [] +) -> dict: + """List models in the OpenAI-compatible format.""" if url_idx is None: model_list = await get_all_models(request, user=user) - models = [ - { - 'id': model['model'], - 'object': 'model', - 'created': int(time.time()), - 'owned_by': 'openai', - } - for model in model_list['models'] - ] - + raw_models = model_list['models'] else: url = request.app.state.config.OLLAMA_BASE_URLS[url_idx] model_list = await send_request(f'{url}/api/tags', 'GET') + raw_models = model_list.get('models', []) - models = [ - { - 'id': model['model'], - 'object': 'model', - 'created': int(time.time()), - 'owned_by': 'openai', - } - for model in model_list.get('models', []) - ] + now_ts = int(time.time()) + models = [ + {'id': m['model'], 'object': 'model', 'created': now_ts, 'owned_by': 'openai'} + for m in raw_models + ] if user.role == 'user' and not BYPASS_MODEL_ACCESS_CONTROL: - # Filter models based on user access control - model_ids = [model['id'] for model in models] - model_infos = {model_info.id: model_info for model_info in await Models.get_models_by_ids(model_ids, db=db)} - user_group_ids = {group.id for group in await Groups.get_groups_by_member_id(user.id, db=db)} - - # Batch-fetch accessible resource IDs in a single query instead of N has_access calls - accessible_model_ids = await AccessGrants.get_accessible_resource_ids( + model_ids = [m['id'] for m in models] + model_infos = { + mi.id: mi for mi in await Models.get_models_by_ids(model_ids, db=db) + } + user_group_ids = { + g.id for g in await Groups.get_groups_by_member_id(user.id, db=db) + } + accessible_ids = await AccessGrants.get_accessible_resource_ids( user_id=user.id, resource_type='model', resource_ids=list(model_infos.keys()), @@ -1459,98 +1334,84 @@ async def get_openai_models( user_group_ids=user_group_ids, db=db, ) + models = [ + m for m in models + if (mi := model_infos.get(m['id'])) and (user.id == mi.user_id or mi.id in accessible_ids) + ] - filtered_models = [] - for model in models: - model_info = model_infos.get(model['id']) - if model_info: - if user.id == model_info.user_id or model_info.id in accessible_model_ids: - filtered_models.append(model) - models = filtered_models - - return { - 'data': models, - 'object': 'list', - } + return {'data': models, 'object': 'list'} class UrlForm(BaseModel): + """Form carrying a single URL string.""" + url: str class UploadBlobForm(BaseModel): + """Form carrying a filename for blob uploads.""" + filename: str -def parse_huggingface_url(hf_url): +def parse_huggingface_url(hf_url: str) -> str | None: + """Extract the filename from a HuggingFace download URL.""" try: - # Parse the URL - parsed_url = urlparse(hf_url) - - # Get the path and split it into components - path_components = parsed_url.path.split('/') - - # Extract the desired output - model_file = path_components[-1] - - return model_file - except ValueError: + return urlparse(hf_url).path.split('/')[-1] + except (ValueError, IndexError): return None -async def download_file_stream(ollama_url, file_url, file_path, file_name, chunk_size=1024 * 1024): - done = False - - if os.path.exists(file_path): - current_size = os.path.getsize(file_path) - else: - current_size = 0 - +async def download_file_stream( + ollama_url: str, + file_url: str, + file_path: str, + file_name: str, + chunk_size: int = 1024 * 1024, +): + """Stream a model file download from *file_url*, then push the blob to Ollama.""" + current_size = os.path.getsize(file_path) if os.path.exists(file_path) else 0 headers = {'Range': f'bytes={current_size}-'} if current_size > 0 else {} - timeout = aiohttp.ClientTimeout(total=600) # Set the timeout + session = await get_session() + async with session.get( + file_url, + headers=headers, + ssl=AIOHTTP_CLIENT_SESSION_SSL, + timeout=aiohttp.ClientTimeout(total=600), + ) as response: + total_size = int(response.headers.get('content-length', 0)) + current_size - async with aiohttp.ClientSession(timeout=timeout, trust_env=True) as session: - async with session.get(file_url, headers=headers, ssl=AIOHTTP_CLIENT_SESSION_SSL) as response: - total_size = int(response.headers.get('content-length', 0)) + current_size + with open(file_path, 'ab+') as f: + async for data in response.content.iter_chunked(chunk_size): + current_size += len(data) + f.write(data) - with open(file_path, 'ab+') as file: - async for data in response.content.iter_chunked(chunk_size): - current_size += len(data) - file.write(data) + done = current_size == total_size + progress = round((current_size / total_size) * 100, 2) + yield f'data: {{"progress": {progress}, "completed": {current_size}, "total": {total_size}}}\n\n' - done = current_size == total_size - progress = round((current_size / total_size) * 100, 2) + if done: + f.close() + hashed = calculate_sha256(file_path, chunk_size) - yield f'data: {{"progress": {progress}, "completed": {current_size}, "total": {total_size}}}\n\n' + with open(file_path, 'rb') as blob_f: + blob_data = blob_f.read() - if done: - file.close() - hashed = calculate_sha256(file_path, chunk_size) - - with open(file_path, 'rb') as f: - blob_data = f.read() - - url = f'{ollama_url}/api/blobs/sha256:{hashed}' - blob_timeout = aiohttp.ClientTimeout(total=30) - async with aiohttp.ClientSession(timeout=blob_timeout, trust_env=True) as blob_session: - async with blob_session.post( - url, data=blob_data, ssl=AIOHTTP_CLIENT_SESSION_SSL - ) as blob_response: - if blob_response.ok: - res = { - 'done': done, - 'blob': f'sha256:{hashed}', - 'name': file_name, - } - os.remove(file_path) - - yield f'data: {json.dumps(res)}\n\n' - else: - raise RuntimeError('Ollama: Could not create blob, Please try again.') + blob_url = f'{ollama_url}/api/blobs/sha256:{hashed}' + async with session.post( + blob_url, + data=blob_data, + ssl=AIOHTTP_CLIENT_SESSION_SSL, + timeout=aiohttp.ClientTimeout(total=30), + ) as blob_resp: + if blob_resp.ok: + os.remove(file_path) + yield f'data: {json.dumps({"done": done, "blob": f"sha256:{hashed}", "name": file_name})}\n\n' + else: + raise RuntimeError('Ollama: Could not create blob, Please try again.') -# url = "https://huggingface.co/TheBloke/stablelm-zephyr-3b-GGUF/resolve/main/stablelm-zephyr-3b.Q2_K.gguf" @router.post('/models/download') @router.post('/models/download/{url_idx}') async def download_model( @@ -1559,31 +1420,26 @@ async def download_model( url_idx: int | None = None, user=Depends(get_admin_user), ): + """Download a GGUF model from HuggingFace or GitHub and register it with Ollama.""" allowed_hosts = ['https://huggingface.co/', 'https://github.com/'] - if not any(form_data.url.startswith(host) for host in allowed_hosts): raise HTTPException( status_code=400, detail='Invalid file_url. Only URLs from allowed hosts are permitted.', ) - if url_idx is None: - url_idx = 0 - url = request.app.state.config.OLLAMA_BASE_URLS[url_idx] - + url = request.app.state.config.OLLAMA_BASE_URLS[url_idx if url_idx is not None else 0] file_name = parse_huggingface_url(form_data.url) - if file_name: - file_path = os.path.join(UPLOAD_DIR, file_name) - - return StreamingResponse( - download_file_stream(url, form_data.url, file_path, file_name), - ) - else: + if not file_name: return None + file_path = os.path.join(UPLOAD_DIR, file_name) + return StreamingResponse( + download_file_stream(url, form_data.url, file_path, file_name), + ) + -# TODO: Progress bar does not reflect size & duration of upload. @router.post('/models/upload') @router.post('/models/upload/{url_idx}') async def upload_model( @@ -1592,20 +1448,18 @@ async def upload_model( url_idx: int | None = None, user=Depends(get_admin_user), ): - if url_idx is None: - url_idx = 0 - ollama_url = request.app.state.config.OLLAMA_BASE_URLS[url_idx] + """Upload a local model file, push it as a blob, and create the model in Ollama.""" + ollama_url = request.app.state.config.OLLAMA_BASE_URLS[url_idx if url_idx is not None else 0] filename = os.path.basename(file.filename) file_path = os.path.join(UPLOAD_DIR, filename) os.makedirs(UPLOAD_DIR, exist_ok=True) - # --- P1: save file locally --- - chunk_size = 1024 * 1024 * 2 # 2 MB chunks + # Stage 1: persist the uploaded file to disk + chunk_size = 1024 * 1024 * 2 # 2 MiB with open(file_path, 'wb') as out_f: while True: chunk = file.file.read(chunk_size) - # log.info(f"Chunk: {str(chunk)}") # DEBUG if not chunk: break out_f.write(chunk) @@ -1613,74 +1467,63 @@ async def upload_model( async def file_process_stream(): nonlocal ollama_url total_size = os.path.getsize(file_path) - log.info(f'Total Model Size: {str(total_size)}') # DEBUG + log.info(f'Total Model Size: {total_size}') - # --- P2: SSE progress + calculate sha256 hash --- + # Stage 2: hash the file and emit SSE progress file_hash = calculate_sha256(file_path, chunk_size) - log.info(f'Model Hash: {str(file_hash)}') # DEBUG + log.info(f'Model Hash: {file_hash}') + try: with open(file_path, 'rb') as f: bytes_read = 0 while chunk := f.read(chunk_size): bytes_read += len(chunk) progress = round(bytes_read / total_size * 100, 2) - data_msg = { - 'progress': progress, - 'total': total_size, - 'completed': bytes_read, - } - yield f'data: {json.dumps(data_msg)}\n\n' + yield f'data: {json.dumps({"progress": progress, "total": total_size, "completed": bytes_read})}\n\n' - # --- P3: Upload to ollama /api/blobs --- + # Stage 3: push blob to Ollama with open(file_path, 'rb') as f: blob_data = f.read() - url = f'{ollama_url}/api/blobs/sha256:{file_hash}' - upload_timeout = aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT) - async with aiohttp.ClientSession(timeout=upload_timeout, trust_env=True) as upload_session: - async with upload_session.post(url, data=blob_data, ssl=AIOHTTP_CLIENT_SESSION_SSL) as response: - if not response.ok: - raise Exception('Ollama: Could not create blob, Please try again.') + session = await get_session() + blob_url = f'{ollama_url}/api/blobs/sha256:{file_hash}' + async with session.post( + blob_url, + data=blob_data, + ssl=AIOHTTP_CLIENT_SESSION_SSL, + timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT), + ) as resp: + if not resp.ok: + raise Exception('Ollama: Could not create blob, Please try again.') - log.info(f'Uploaded to /api/blobs') # DEBUG - # Remove local file + log.info('Uploaded to /api/blobs') os.remove(file_path) - # Create model in ollama - model_name, ext = os.path.splitext(filename) - log.info(f'Created Model: {model_name}') # DEBUG + # Stage 4: create the model + model, _ext = os.path.splitext(filename) + log.info(f'Created Model: {model}') create_payload = { - 'model': model_name, - # Reference the file by its original name => the uploaded blob's digest + 'model': model, 'files': {filename: f'sha256:{file_hash}'}, } - log.info(f'Model Payload: {create_payload}') # DEBUG + log.info(f'Model Payload: {create_payload}') - # Call ollama /api/create - # https://github.com/ollama/ollama/blob/main/docs/api.md#create-a-model - async with aiohttp.ClientSession(timeout=upload_timeout, trust_env=True) as create_session: - async with create_session.post( - f'{ollama_url}/api/create', - headers={'Content-Type': 'application/json'}, - data=json.dumps(create_payload), - ssl=AIOHTTP_CLIENT_SESSION_SSL, - ) as create_resp: - if create_resp.ok: - log.info(f'API SUCCESS!') # DEBUG - done_msg = { - 'done': True, - 'blob': f'sha256:{file_hash}', - 'name': filename, - 'model_created': model_name, - } - yield f'data: {json.dumps(done_msg)}\n\n' - else: - resp_text = await create_resp.text() - raise Exception(f'Failed to create model in Ollama. {resp_text}') + async with session.post( + f'{ollama_url}/api/create', + headers={'Content-Type': 'application/json'}, + data=json.dumps(create_payload), + ssl=AIOHTTP_CLIENT_SESSION_SSL, + timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT), + ) as create_resp: + if create_resp.ok: + log.info('API SUCCESS!') + yield f'data: {json.dumps({"done": True, "blob": f"sha256:{file_hash}", "name": filename, "model_created": model})}\n\n' + else: + resp_text = await create_resp.text() + raise Exception(f'Failed to create model in Ollama. {resp_text}') - except Exception as e: - res = {'error': str(e)} - yield f'data: {json.dumps(res)}\n\n' + except Exception as exc: + yield f'data: {json.dumps({"error": str(exc)})}\n\n' return StreamingResponse(file_process_stream(), media_type='text/event-stream')