from __future__ import annotations import asyncio import logging import os import random import re import time from datetime import datetime from typing import Optional, Union from urllib.parse import urlparse import aiofiles import aiohttp from aiocache import cached from fastapi import APIRouter, Depends, File, HTTPException, Request, UploadFile from fastapi.responses import StreamingResponse from open_webui.config import UPLOAD_DIR from open_webui.constants import ERROR_MESSAGES from open_webui.env import ( AIOHTTP_CLIENT_SESSION_SSL, AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST, AIOHTTP_FILE_STREAM_CHUNK_SIZE, BYPASS_MODEL_ACCESS_CONTROL, ENABLE_FORWARD_USER_INFO_HEADERS, FORWARD_SESSION_INFO_HEADER_CHAT_ID, MODELS_CACHE_TTL, REDIS_KEY_PREFIX, ) from open_webui.events import EVENTS, publish_event, publish_model_provider_request_failed from open_webui.internal.db import get_async_session from open_webui.models.access_grants import AccessGrants from open_webui.models.config import Config 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 get_custom_headers, include_user_info_headers from open_webui.utils.json_codec import JSONCodec from open_webui.utils.misc import calculate_sha256 from open_webui.utils.model_ids import strip_provider_model_prefix 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_client_timeout, get_session, stream_wrapper from pydantic import BaseModel, ConfigDict, validator from sqlalchemy.ext.asyncio import AsyncSession log = logging.getLogger(__name__) # 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 # in ZlibError. See https://github.com/aio-libs/aiohttp/issues/4462. _STRIP_PROXY_HEADERS = frozenset({'Content-Encoding', 'Content-Length', 'Transfer-Encoding'}) _MODEL_LIST_TIMEOUT = aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST) BASE_MODELS_CACHE_KEY = f'{REDIS_KEY_PREFIX}:models:base' def _clean_proxy_headers(raw_headers) -> dict: """Return a copy of *raw_headers* with stale encoding headers removed.""" return {k: v for k, v in raw_headers.items() if k not in _STRIP_PROXY_HEADERS} 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: 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) async with session.get( url, headers=headers, ssl=AIOHTTP_CLIENT_SESSION_SSL, timeout=_MODEL_LIST_TIMEOUT, ) as r: return await r.json(loads=JSONCodec.loads) except Exception as exc: log.error(f'Connection error: {exc}') return None async def send_request( url: str, method: str = 'POST', *, payload: Union[str, bytes | None] = None, key: str | None = None, user: UserModel = None, stream: bool = False, # passthrough must stay False for /api/chat: middleware parses it per line passthrough: bool = False, content_type: str | None = None, metadata: dict | None = None, api_config: dict | None = None, request: Request | None = None, ): r = None streaming = False try: session = await get_session() headers = { 'Content-Type': 'application/json', **({'Authorization': f'Bearer {key}'} if key else {}), } if ENABLE_FORWARD_USER_INFO_HEADERS and user: headers = include_user_info_headers(headers, user) if metadata and metadata.get('chat_id'): headers[FORWARD_SESSION_INFO_HEADER_CHAT_ID] = metadata.get('chat_id') # Custom per-connection headers last so admin-set headers take precedence. if api_config and api_config.get('headers'): headers.update(await get_custom_headers(api_config['headers'], user, metadata, request=request)) r = await session.request( method, url, data=payload, headers=headers, ssl=AIOHTTP_CLIENT_SESSION_SSL, timeout=get_client_timeout(stream=stream), ) if not r.ok: try: res = await r.json(loads=JSONCodec.loads) await publish_model_provider_request_failed( request, actor=user, provider='ollama', base_url=url, status=r.status, upstream_error=res, ) if 'error' in res: raise HTTPException(status_code=r.status, detail=res['error']) except HTTPException: raise except Exception as e: log.error(f'Failed to parse error response: {e}') await publish_model_provider_request_failed( request, actor=user, provider='ollama', base_url=url, status=r.status, ) raise HTTPException( status_code=r.status, detail=ERROR_MESSAGES.SERVER_CONNECTION_ERROR, ) r.raise_for_status() if stream: response_headers = _clean_proxy_headers(r.headers) if content_type: response_headers['Content-Type'] = content_type streaming = True return StreamingResponse( stream_wrapper(r, passthrough=passthrough), status_code=r.status, headers=response_headers, ) else: try: return await r.json(loads=JSONCodec.loads) except Exception: return None except HTTPException: raise except Exception as e: raise HTTPException( status_code=r.status if r else 500, detail=f'Ollama: {e}' if str(e) else ERROR_MESSAGES.SERVER_CONNECTION_ERROR, ) finally: if not streaming: await cleanup_response(r) def get_api_key(idx, url, configs): parsed_url = urlparse(url) base_url = f'{parsed_url.scheme}://{parsed_url.netloc}' return configs.get(str(idx), configs.get(base_url, {})).get('key', None) # Legacy support ########################################## # # API routes # ########################################## router = APIRouter() OLLAMA_CONFIG_KEYS = { 'ENABLE_OLLAMA_API': 'ollama.enable', 'OLLAMA_BASE_URLS': 'ollama.base_urls', 'OLLAMA_API_CONFIGS': 'ollama.api_configs', } async def get_ollama_config_values() -> dict: values = await Config.get_many(*OLLAMA_CONFIG_KEYS.values()) return {field: values[storage_key] for field, storage_key in OLLAMA_CONFIG_KEYS.items() if storage_key in values} async def get_ollama_runtime_config() -> tuple[bool, list[str], dict]: values = await Config.get_many('ollama.enable', 'ollama.base_urls', 'ollama.api_configs') return ( values.get('ollama.enable'), values.get('ollama.base_urls') or [], values.get('ollama.api_configs') or {}, ) async def get_ollama_connection(idx: int) -> tuple[str, dict, str | None]: _, base_urls, api_configs = await get_ollama_runtime_config() url = base_urls[idx] return url, resolve_api_config(api_configs, idx, url), get_api_key(idx, url, api_configs) @router.head('/') @router.get('/') async def get_status() -> dict: """Health-check endpoint.""" return {'status': True} class ConnectionVerificationForm(BaseModel): url: str key: str | None = None @router.post('/verify') 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 session.get( f'{form_data.url}/api/version', headers=headers, ssl=AIOHTTP_CLIENT_SESSION_SSL, timeout=_MODEL_LIST_TIMEOUT, ) as r: if r.status != 200: detail = f'HTTP Error: {r.status}' res = await r.json(loads=JSONCodec.loads) if 'error' in res: detail = f'External Error: {res["error"]}' raise Exception(detail) return await r.json(loads=JSONCodec.loads) 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), ) -> dict: """Return the current Ollama connection configuration.""" return await get_ollama_config_values() 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), ) -> dict: """Persist updated Ollama connection settings.""" valid_keys = {str(i) for i in range(len(form_data.OLLAMA_BASE_URLS))} api_configs = {k: v for k, v in form_data.OLLAMA_API_CONFIGS.items() if k in valid_keys} await Config.upsert( { 'ollama.enable': form_data.ENABLE_OLLAMA_API, 'ollama.base_urls': form_data.OLLAMA_BASE_URLS, 'ollama.api_configs': api_configs, } ) await get_all_models.cache.clear() redis = getattr(request.app.state, 'redis', None) if redis is not None: await redis.delete(BASE_MODELS_CACHE_KEY) request.app.state.BASE_MODELS = [] request.app.state.OLLAMA_MODELS = {} models = getattr(request.app.state, 'MODELS', None) if hasattr(models, 'clear'): models.clear() else: request.app.state.MODELS = {} await publish_event( request, EVENTS.MODEL_PROVIDER_CONFIG_UPDATED, actor=user, subject_id='ollama', subject_type='model.provider_config', data={ 'provider': 'ollama', 'enabled': form_data.ENABLE_OLLAMA_API, 'base_url_count': len(form_data.OLLAMA_BASE_URLS), }, ) return { 'ENABLE_OLLAMA_API': form_data.ENABLE_OLLAMA_API, 'OLLAMA_BASE_URLS': form_data.OLLAMA_BASE_URLS, 'OLLAMA_API_CONFIGS': api_configs, } 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()) def resolve_api_config(api_configs: dict, idx: int, url: str) -> dict: """Look up the API config for a backend by numeric index, falling back to URL key (legacy).""" return api_configs.get(str(idx), api_configs.get(url, {})) async def get_ollama_connection_config() -> tuple[list, dict]: """Base URLs and per-connection API configs in one batched SELECT.""" config = await Config.get_many('ollama.base_urls', 'ollama.api_configs') return config.get('ollama.base_urls', []), config.get('ollama.api_configs', {}) @cached( ttl=MODELS_CACHE_TTL, # key_builder (not key) is the per-call hook in aiocache 0.12; `key=` is a # static key, so a `key=lambda` collapsed every caller to one shared entry. key_builder=lambda _func, request, user=None: f'ollama_all_models_{user.id}' if user else 'ollama_all_models', ) 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 not await Config.get('ollama.enable'): models_dict: dict = {'models': []} request.app.state.OLLAMA_MODELS = {} return models_dict # Fan-out tag requests to every backend tasks = [] base_urls = await Config.get('ollama.base_urls', []) api_configs = await Config.get('ollama.api_configs', {}) for idx, url in enumerate(base_urls): api_config = resolve_api_config(api_configs, 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(*tasks) # Track which backends failed so we can skip them for /api/ps failed_idxs: set[int] = set() # Post-process each response: apply prefix_id, tags, model filtering for idx, response in enumerate(responses): if not response: failed_idxs.add(idx) continue url = base_urls[idx] api_config = resolve_api_config(api_configs, idx, url) 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', []) if allowed_model_ids and 'models' in response: response['models'] = [m for m in response['models'] if m['model'] in allowed_model_ids] for m in response.get('models', []): if prefix_id: m['model'] = f'{prefix_id}.{m["model"]}' if m.get('name'): m['name'] = f'{prefix_id}.{m["name"]}' if allowed_tags: m['tags'] = allowed_tags if connection_type: m['connection_type'] = connection_type models_dict = {'models': merge_models_lists(r.get('models', []) if r else None for r in responses)} # Annotate with expiry info from loaded-model state try: loaded = await get_ollama_loaded_models(request, user=user, skip_idxs=failed_idxs) 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('Failed to get loaded models: %s', exc) 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): """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)} accessible_ids = await AccessGrants.get_accessible_resource_ids( user_id=user.id, resource_type='model', resource_ids=list(model_infos.keys()), permission='read', user_group_ids=user_group_ids, db=db, ) 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), ): """List Ollama model tags, optionally from a specific backend.""" if url_idx is not None and user.role != 'admin': raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED) if not await Config.get('ollama.enable'): raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) if url_idx is None: result = await get_all_models(request, user=user) else: url = (await Config.get('ollama.base_urls', []))[url_idx] key = get_api_key(url_idx, url, (await Config.get('ollama.api_configs', {}))) result = await send_request(f'{url}/api/tags', 'GET', key=key, user=user) if user.role == 'user' and not BYPASS_MODEL_ACCESS_CONTROL: result['models'] = await get_filtered_models(result, user) return result @router.get('/api/ps') async def get_ollama_loaded_models( request: Request, user=Depends(get_admin_user), skip_idxs: set[int] | None = None, ) -> dict: """List models currently loaded in Ollama memory across all backends.""" if not await Config.get('ollama.enable'): return {'models': []} tasks = [] base_urls = await Config.get('ollama.base_urls', []) api_configs = await Config.get('ollama.api_configs', {}) for idx, url in enumerate(base_urls): if skip_idxs and idx in skip_idxs: tasks.append(asyncio.ensure_future(asyncio.sleep(0, None))) continue api_config = resolve_api_config(api_configs, 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))) responses = await asyncio.gather(*tasks) for idx, response in enumerate(responses): if not response: continue api_config = resolve_api_config(api_configs, idx, 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"]}' 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, user=Depends(get_verified_user), url_idx: int | None = None, ): """Return the lowest Ollama version across all configured backends.""" if url_idx is not None and user.role != 'admin': raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED) if not await Config.get('ollama.enable'): return {'version': False} if url_idx is not None: url = (await Config.get('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(await Config.get('ollama.base_urls', [])): api_config = (await Config.get('ollama.api_configs', {})).get( str(idx), (await Config.get('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') @router.post('/api/unload') async def unload_model( request: Request, form_data: ModelNameForm, user=Depends(get_admin_user), ): form_data = form_data.model_dump(exclude_none=True) model = form_data.get('model', form_data.get('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 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 = [] errors = [] for idx in url_indices: url = (await Config.get('ollama.base_urls', []))[idx] api_config = (await Config.get('ollama.api_configs', {})).get( str(idx), (await Config.get('ollama.api_configs', {})).get(url, {}) ) key = get_api_key(idx, url, (await Config.get('ollama.api_configs', {}))) prefix_id = api_config.get('prefix_id', None) if prefix_id and model.startswith(f'{prefix_id}.'): model = model[len(f'{prefix_id}.') :] payload = {'model': model, 'keep_alive': 0, 'prompt': ''} try: res = await send_request( f'{url}/api/generate', payload=JSONCodec.dumps(payload), key=key, user=user, ) results.append({'url_idx': idx, 'success': True, 'response': res}) except Exception as e: log.exception(f'Failed to unload model on node {idx}: {e}') errors.append({'url_idx': idx, 'success': False, 'error': str(e)}) if len(errors) > 0: raise HTTPException( status_code=500, detail=f'Failed to unload model on {len(errors)} nodes: {errors}', ) return {'status': True} @router.post('/api/pull') @router.post('/api/pull/{url_idx}') async def pull_model( request: Request, form_data: ModelNameForm, url_idx: int = 0, user=Depends(get_admin_user), ): if not await Config.get('ollama.enable'): 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')) url = (await Config.get('ollama.base_urls', []))[url_idx] log.info('url: %s', url) # Admins may pull from any registry return await send_request( f'{url}/api/pull', payload=JSONCodec.dumps({**form_data, 'insecure': True}), key=get_api_key(url_idx, url, (await Config.get('ollama.api_configs', {}))), user=user, stream=True, passthrough=True, ) class PushModelForm(BaseModel): """Payload for pushing a model to a registry.""" model: str insecure: bool | None = None stream: bool | None = None @router.delete('/api/push') @router.delete('/api/push/{url_idx}') async def push_model( request: Request, form_data: PushModelForm, url_idx: int | None = None, user=Depends(get_admin_user), ): """Push a local model to a remote registry.""" if not await Config.get('ollama.enable'): 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 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 = (await Config.get('ollama.base_urls', []))[url_idx] log.debug('url: %s', url) return await send_request( f'{url}/api/push', payload=form_data.model_dump_json(exclude_none=True).encode(), key=get_api_key(url_idx, url, (await Config.get('ollama.api_configs', {}))), user=user, stream=True, passthrough=True, ) 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') @router.post('/api/create') @router.post('/api/create/{url_idx}') async def create_model( request: Request, form_data: CreateModelForm, url_idx: int = 0, user=Depends(get_admin_user), ): if not await Config.get('ollama.enable'): raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) log.debug('form_data: %s', form_data) url = (await Config.get('ollama.base_urls', []))[url_idx] return await send_request( f'{url}/api/create', payload=form_data.model_dump_json(exclude_none=True).encode(), key=get_api_key(url_idx, url, (await Config.get('ollama.api_configs', {}))), user=user, stream=True, passthrough=True, ) class CopyModelForm(BaseModel): """Payload for duplicating an existing model under a new name.""" source: str destination: str @router.post('/api/copy') @router.post('/api/copy/{url_idx}') async def copy_model( request: Request, form_data: CopyModelForm, url_idx: int | None = None, user=Depends(get_admin_user), ): """Duplicate an existing model under a new name.""" if not await Config.get('ollama.enable'): 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 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 = (await Config.get('ollama.base_urls', []))[url_idx] key = get_api_key(url_idx, url, (await Config.get('ollama.api_configs', {}))) await send_request( f'{url}/api/copy', payload=form_data.model_dump_json(exclude_none=True).encode(), key=key, user=user, ) await publish_event( request, EVENTS.MODEL_PROVIDER_MODEL_CREATED, actor=user, subject_id=form_data.destination, data={'provider': 'ollama', 'source': form_data.source, 'url_idx': url_idx}, ) return True @router.delete('/api/delete') @router.delete('/api/delete/{url_idx}') async def delete_model( request: Request, form_data: ModelNameForm, url_idx: int | None = None, user=Depends(get_admin_user), ): """Remove a model from an Ollama backend.""" if not await Config.get('ollama.enable'): raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) 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 not in models: raise HTTPException(status_code=400, detail=ERROR_MESSAGES.MODEL_NOT_FOUND(model)) url_idx = models[model]['urls'][0] url = (await Config.get('ollama.base_urls', []))[url_idx] key = get_api_key(url_idx, url, (await Config.get('ollama.api_configs', {}))) await send_request( f'{url}/api/delete', 'DELETE', payload=JSONCodec.dumps(payload), key=key, user=user, ) await publish_event( request, EVENTS.MODEL_PROVIDER_MODEL_DELETED, actor=user, subject_id=model, data={'provider': 'ollama', 'url_idx': url_idx}, ) return True @router.post('/api/show') async def show_model_info( request: Request, form_data: ModelNameForm, user=Depends(get_verified_user), ): """Retrieve model metadata from the Ollama backend.""" if not await Config.get('ollama.enable'): raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) payload = form_data.model_dump(exclude_none=True) payload['model'] = payload.get('model', payload.get('name')) model = payload.get('model') 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)) url_idx = random.choice(models[model]['urls']) url = (await Config.get('ollama.base_urls', []))[url_idx] key = get_api_key(url_idx, url, (await Config.get('ollama.api_configs', {}))) return await send_request( f'{url}/api/show', payload=JSONCodec.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') @router.post('/api/embed') @router.post('/api/embed/{url_idx}') async def embed( request: Request, form_data: GenerateEmbedForm, url_idx: int | None = None, user=Depends(get_verified_user), ): """Generate embeddings via the Ollama /api/embed endpoint.""" if not await Config.get('ollama.enable'): raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) log.info('generate_ollama_batch_embeddings %s', form_data) await check_model_access(user, await Models.get_model_by_id(form_data.model), BYPASS_MODEL_ACCESS_CONTROL) await validate_ollama_backend_idx(request, form_data.model, url_idx, user) if url_idx is None: model = form_data.model 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 not in models: raise HTTPException(status_code=400, detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model)) url_idx = random.choice(models[model]['urls']) base_urls, api_configs = await get_ollama_connection_config() url = base_urls[url_idx] api_config = api_configs.get(str(url_idx), api_configs.get(url, {})) key = get_api_key(url_idx, url, api_configs) prefix_id = api_config.get('prefix_id') form_data.model = strip_provider_model_prefix(form_data.model, prefix_id) return await send_request( f'{url}/api/embed', payload=form_data.model_dump_json(exclude_none=True).encode(), key=key, user=user, ) class GenerateEmbeddingsForm(BaseModel): """Payload for the legacy /api/embeddings endpoint (single-prompt).""" model: str prompt: str options: dict | None = None keep_alive: Union[int, str | None] = None @router.post('/api/embeddings') @router.post('/api/embeddings/{url_idx}') async def embeddings( request: Request, form_data: GenerateEmbeddingsForm, url_idx: int | None = None, user=Depends(get_verified_user), ): """Generate embeddings via the legacy Ollama /api/embeddings endpoint.""" if not await Config.get('ollama.enable'): raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) log.info('generate_ollama_embeddings %s', form_data) await check_model_access(user, await Models.get_model_by_id(form_data.model), BYPASS_MODEL_ACCESS_CONTROL) await validate_ollama_backend_idx(request, form_data.model, url_idx, user) if url_idx is None: model = form_data.model 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 not in models: raise HTTPException(status_code=400, detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model)) url_idx = random.choice(models[model]['urls']) base_urls, api_configs = await get_ollama_connection_config() url = base_urls[url_idx] api_config = api_configs.get(str(url_idx), api_configs.get(url, {})) key = get_api_key(url_idx, url, api_configs) prefix_id = api_config.get('prefix_id') form_data.model = strip_provider_model_prefix(form_data.model, prefix_id) return await send_request( f'{url}/api/embeddings', payload=form_data.model_dump_json(exclude_none=True).encode(), key=key, user=user, ) class GenerateCompletionForm(BaseModel): """Payload for the Ollama /api/generate endpoint.""" model: str prompt: str | None = None suffix: str | None = None images: list[str] | None = None format: Union[dict, str | None] = None options: dict | None = None system: str | None = None template: str | None = None context: list[int] | None = None stream: bool | None = True raw: bool | None = None keep_alive: Union[int, str | None] = None @router.post('/api/generate') @router.post('/api/generate/{url_idx}') async def generate_completion( request: Request, form_data: GenerateCompletionForm, url_idx: int | None = None, user=Depends(get_verified_user), ): """Run text completion via Ollama /api/generate.""" if not await Config.get('ollama.enable'): raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) await check_model_access(user, await Models.get_model_by_id(form_data.model), BYPASS_MODEL_ACCESS_CONTROL) await validate_ollama_backend_idx(request, form_data.model, url_idx, user) if url_idx is None: await get_all_models(request, user=user) models = request.app.state.OLLAMA_MODELS model = 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']) base_urls, api_configs = await get_ollama_connection_config() url = base_urls[url_idx] api_config = api_configs.get(str(url_idx), api_configs.get(url, {})) prefix_id = api_config.get('prefix_id') form_data.model = strip_provider_model_prefix(form_data.model, prefix_id) return await send_request( f'{url}/api/generate', payload=form_data.model_dump_json(exclude_none=True).encode(), key=get_api_key(url_idx, url, api_configs), user=user, stream=True, passthrough=True, ) 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): 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 options: dict | None = None template: str | None = None stream: bool | None = True keep_alive: Union[int, str | None] = None tools: list[dict] | None = None model_config = ConfigDict(extra='allow') async def validate_ollama_backend_idx(request: Request, model: str, url_idx: int | None, user) -> None: # A caller-supplied url_idx must point to a backend the model is actually # served from; the None path is already constrained to that allow-list. if url_idx is None or user is None or getattr(user, 'role', None) == 'admin': return 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 url_idx not in (models.get(model) or {}).get('urls', []): raise HTTPException(status_code=403, detail=ERROR_MESSAGES.ACCESS_PROHIBITED) async def get_ollama_url(request: Request, model: str, url_idx: int | None = None, user=None): await validate_ollama_backend_idx(request, model, url_idx, user) if url_idx is None: models = request.app.state.OLLAMA_MODELS if not models or model not in models: await get_all_models.cache.clear() 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), ) url_idx = random.choice(models[model].get('urls', [])) url = (await Config.get('ollama.base_urls', []))[url_idx] return url, url_idx @router.post('/api/chat') @router.post('/api/chat/{url_idx}') async def generate_chat_completion( request: Request, form_data: dict, url_idx: int | None = None, user=Depends(get_verified_user), # noqa: B008 ): """Forward a chat completion request to an Ollama backend.""" if not await Config.get('ollama.enable'): raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) # 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), # which would exhaust the connection pool under concurrent load. # bypass_filter and bypass_system_prompt are read from request.state to prevent # external clients from setting them via query parameter. Only internal # server-side callers (e.g. utils/chat.py) should set # request.state.bypass_filter / request.state.bypass_system_prompt = True. bypass_filter = getattr(request.state, 'bypass_filter', False) if BYPASS_MODEL_ACCESS_CONTROL: bypass_filter = True bypass_system_prompt = getattr(request.state, 'bypass_system_prompt', False) metadata = form_data.pop('metadata', None) try: form_data = GenerateChatCompletionForm(**form_data) 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) payload.pop('metadata', None) model_id = payload['model'] model_info = await Models.get_model_by_id(model_id) 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 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) await check_model_access(user, model_info, bypass_filter) else: await check_model_access(user, None, bypass_filter) url, url_idx = await get_ollama_url(request, payload['model'], url_idx, user) api_configs = await Config.get('ollama.api_configs', {}) api_config = resolve_api_config(api_configs, url_idx, url) prefix_id = api_config.get('prefix_id') payload['model'] = strip_provider_model_prefix(payload['model'], prefix_id) return await send_request( f'{url}/api/chat', payload=JSONCodec.dumps(payload), key=get_api_key(url_idx, url, api_configs), user=user, stream=form_data.stream, content_type='application/x-ndjson', metadata=metadata, api_config=api_config, request=request, ) # 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') class OpenAIEmbeddingsForm(BaseModel): """Payload for the OpenAI-compatible /v1/embeddings proxy.""" model: str input: object model_config = ConfigDict(extra='allow') @router.post('/v1/completions') @router.post('/v1/completions/{url_idx}') async def generate_openai_completion( request: Request, form_data: dict, url_idx: int | None = None, 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), # which would exhaust the connection pool under concurrent load. metadata = form_data.pop('metadata', None) try: form_data = OpenAICompletionForm(**form_data) 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']) model_id = form_data.model model_info = await Models.get_model_by_id(model_id) 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, user) api_configs = await Config.get('ollama.api_configs', {}) api_config = resolve_api_config(api_configs, url_idx, url) prefix_id = api_config.get('prefix_id') payload['model'] = strip_provider_model_prefix(payload['model'], prefix_id) return await send_request( f'{url}/v1/completions', payload=JSONCodec.dumps(payload), key=get_api_key(url_idx, url, api_configs), user=user, stream=payload.get('stream', False), passthrough=True, metadata=metadata, api_config=api_config, request=request, ) @router.post('/v1/embeddings') @router.post('/v1/embeddings/{url_idx}') async def generate_openai_embeddings( request: Request, form_data: dict, url_idx: int | None = None, user=Depends(get_verified_user), # noqa: B008 ): """Forward an embeddings request via the OpenAI-compatible proxy.""" if not await Config.get('ollama.enable'): raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) metadata = form_data.pop('metadata', None) try: form_data = OpenAIEmbeddingsForm(**form_data) except Exception as exc: log.exception(exc) raise HTTPException(status_code=400, detail=str(exc)) payload = form_data.model_dump(exclude_none=True) payload.pop('metadata', None) model_id = form_data.model model_info = await Models.get_model_by_id(model_id) if model_info is not None: if model_info.base_model_id: payload['model'] = model_info.base_model_id 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, user) api_config = resolve_api_config((await Config.get('ollama.api_configs', {})), url_idx, url) prefix_id = api_config.get('prefix_id') payload['model'] = strip_provider_model_prefix(payload['model'], prefix_id) return await send_request( f'{url}/v1/embeddings', payload=JSONCodec.dumps(payload), key=get_api_key(url_idx, url, (await Config.get('ollama.api_configs', {}))), user=user, metadata=metadata, api_config=api_config, request=request, ) @router.post('/v1/chat/completions') @router.post('/v1/chat/completions/{url_idx}') async def generate_openai_chat_completion( request: Request, form_data: dict, url_idx: int | None = None, 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), # which would exhaust the connection pool under concurrent load. metadata = form_data.pop('metadata', None) try: form_data = OpenAIChatCompletionForm(**form_data) 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']) model_id = form_data.model model_info = await Models.get_model_by_id(model_id) 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) 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, user) api_configs = await Config.get('ollama.api_configs', {}) api_config = resolve_api_config(api_configs, url_idx, url) prefix_id = api_config.get('prefix_id') payload['model'] = strip_provider_model_prefix(payload['model'], prefix_id) return await send_request( f'{url}/v1/chat/completions', payload=JSONCodec.dumps(payload), key=get_api_key(url_idx, url, api_configs), user=user, stream=payload.get('stream', False), passthrough=True, metadata=metadata, api_config=api_config, request=request, ) @router.post('/v1/messages') @router.post('/v1/messages/{url_idx}') async def generate_anthropic_messages( request: Request, form_data: dict, url_idx: int | None = None, user=Depends(get_verified_user), ): """ Proxy for Ollama's Anthropic-compatible /v1/messages endpoint. Forwards the request as-is to the Ollama backend, applying the same model resolution, access control, and prefix_id handling used by the OpenAI-compatible /v1/chat/completions proxy. See https://docs.ollama.com/api/anthropic-compatibility """ if not await Config.get('ollama.enable'): raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) payload = {**form_data} model_id = payload.get('model', '') model_info = await Models.get_model_by_id(model_id) if model_info: if model_info.base_model_id: payload['model'] = model_info.base_model_id 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, user) api_configs = await Config.get('ollama.api_configs', {}) api_config = api_configs.get(str(url_idx), api_configs.get(url, {})) # Legacy support prefix_id = api_config.get('prefix_id', None) payload['model'] = strip_provider_model_prefix(payload['model'], prefix_id) return await send_request( f'{url}/v1/messages', payload=JSONCodec.dumps(payload), key=get_api_key(url_idx, url, api_configs), user=user, stream=payload.get('stream', False), passthrough=True, content_type='text/event-stream' if payload.get('stream', False) else None, api_config=api_config, request=request, ) class ResponsesForm(BaseModel): model: str model_config = ConfigDict(extra='allow') @router.post('/v1/responses') @router.post('/v1/responses/{url_idx}') async def generate_responses( request: Request, form_data: ResponsesForm, url_idx: int | None = None, user=Depends(get_verified_user), ): """ Proxy for Ollama's OpenAI-compatible /v1/responses endpoint. Forwards the request as-is to the Ollama backend, applying the same model resolution, access control, and prefix_id handling used by the OpenAI-compatible /v1/chat/completions proxy. See https://ollama.com/blog/responses-api """ if not await Config.get('ollama.enable'): raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED) payload = form_data.model_dump() model_id = form_data.model model_info = await Models.get_model_by_id(model_id) if model_info: if model_info.base_model_id: payload['model'] = model_info.base_model_id 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, user) api_configs = await Config.get('ollama.api_configs', {}) api_config = api_configs.get(str(url_idx), api_configs.get(url, {})) # Legacy support prefix_id = api_config.get('prefix_id', None) payload['model'] = strip_provider_model_prefix(payload['model'], prefix_id) return await send_request( f'{url}/v1/responses', payload=JSONCodec.dumps(payload), key=get_api_key(url_idx, url, api_configs), user=user, stream=payload.get('stream', False), passthrough=True, content_type='text/event-stream' if payload.get('stream', False) else None, api_config=api_config, request=request, ) @router.get('/v1/models') @router.get('/v1/models/{url_idx}') async def get_openai_models( request: Request, url_idx: int | None = None, user=Depends(get_verified_user), db: AsyncSession = Depends(get_async_session), ) -> dict: """List models in the OpenAI-compatible format.""" if url_idx is not None and user.role != 'admin': raise HTTPException(status_code=401, detail=ERROR_MESSAGES.ACCESS_PROHIBITED) if url_idx is None: model_list = await get_all_models(request, user=user) raw_models = model_list['models'] else: url = (await Config.get('ollama.base_urls', []))[url_idx] model_list = await send_request(f'{url}/api/tags', 'GET') raw_models = 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: 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()), permission='read', 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) ] 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: str) -> str | None: """Extract the filename from a HuggingFace download URL.""" try: return urlparse(hf_url).path.split('/')[-1] except (ValueError, IndexError): return None async def download_file_stream( ollama_url: str, file_url: str, file_path: str, file_name: str, chunk_size: int = AIOHTTP_FILE_STREAM_CHUNK_SIZE, ): """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 {} 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 aiofiles.open(file_path, 'ab') as f: async for data in response.content.iter_chunked(chunk_size): current_size += len(data) await f.write(data) progress_total = total_size or current_size progress = round((current_size / progress_total) * 100, 2) yield f'data: {{"progress": {progress}, "completed": {current_size}, "total": {total_size}}}\n\n' done = True hashed = await asyncio.to_thread(calculate_sha256, file_path, chunk_size) blob_url = f'{ollama_url}/api/blobs/sha256:{hashed}' blob_size = await asyncio.to_thread(os.path.getsize, file_path) async def blob_chunks(): async with aiofiles.open(file_path, 'rb') as blob_file: while chunk := await blob_file.read(chunk_size): yield chunk async with session.post( blob_url, data=blob_chunks(), headers={'Content-Length': str(blob_size)}, ssl=AIOHTTP_CLIENT_SESSION_SSL, timeout=aiohttp.ClientTimeout(total=30), ) as blob_resp: if blob_resp.ok: await asyncio.to_thread(os.remove, file_path) yield f'data: {JSONCodec.dumps({"done": done, "blob": f"sha256:{hashed}", "name": file_name})}\n\n' else: raise RuntimeError('Ollama: Could not create blob, Please try again.') @router.post('/models/download') @router.post('/models/download/{url_idx}') async def download_model( request: Request, form_data: UrlForm, 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.', ) url = (await Config.get('ollama.base_urls', []))[url_idx if url_idx is not None else 0] file_name = parse_huggingface_url(form_data.url) 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), ) @router.post('/models/upload') @router.post('/models/upload/{url_idx}') async def upload_model( request: Request, file: UploadFile = File(...), url_idx: int | None = None, 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] filename = os.path.basename(file.filename) file_path = os.path.join(UPLOAD_DIR, filename) os.makedirs(UPLOAD_DIR, exist_ok=True) # Stage 1: persist the uploaded file to disk chunk_size = AIOHTTP_FILE_STREAM_CHUNK_SIZE async with aiofiles.open(file_path, 'wb') as out_f: while chunk := await file.read(chunk_size): await out_f.write(chunk) async def file_process_stream(): nonlocal ollama_url total_size = os.path.getsize(file_path) log.info('Total Model Size: %s', total_size) # Stage 2: hash the file and emit SSE progress file_hash = await asyncio.to_thread(calculate_sha256, file_path, chunk_size) log.info('Model Hash: %s', file_hash) try: bytes_read = 0 async with aiofiles.open(file_path, 'rb') as f: while chunk := await f.read(chunk_size): bytes_read += len(chunk) progress = round(bytes_read / total_size * 100, 2) event = JSONCodec.dumps({'progress': progress, 'total': total_size, 'completed': bytes_read}) yield f'data: {event}\n\n' session = await get_session() blob_url = f'{ollama_url}/api/blobs/sha256:{file_hash}' async def blob_chunks(): async with aiofiles.open(file_path, 'rb') as blob_file: while chunk := await blob_file.read(chunk_size): yield chunk async with session.post( blob_url, data=blob_chunks(), headers={'Content-Length': str(total_size)}, ssl=AIOHTTP_CLIENT_SESSION_SSL, timeout=get_client_timeout(), ) as resp: if not resp.ok: raise Exception('Ollama: Could not create blob, Please try again.') log.info('Uploaded to /api/blobs') await asyncio.to_thread(os.remove, file_path) # Stage 4: create the model model, _ext = os.path.splitext(filename) log.info('Created Model: %s', model) create_payload = { 'model': model, 'files': {filename: f'sha256:{file_hash}'}, } log.info('Model Payload: %s', create_payload) async with session.post( f'{ollama_url}/api/create', headers={'Content-Type': 'application/json'}, data=JSONCodec.dumps(create_payload), ssl=AIOHTTP_CLIENT_SESSION_SSL, timeout=get_client_timeout(), ) as create_resp: if create_resp.ok: log.info('API SUCCESS!') event = JSONCodec.dumps( {'done': True, 'blob': f'sha256:{file_hash}', 'name': filename, 'model_created': model} ) yield f'data: {event}\n\n' else: resp_text = await create_resp.text() raise Exception(f'Failed to create model in Ollama. {resp_text}') except Exception as exc: yield f'data: {JSONCodec.dumps({"error": str(exc)})}\n\n' return StreamingResponse(file_process_stream(), media_type='text/event-stream')