123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812 |
- # 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
- import os
- import random
- import re
- import time
- from datetime import datetime
- from typing import Optional, Union
- from urllib.parse import urlparse
- import aiohttp
- from aiocache import cached
- import requests
- from open_webui.models.chats import Chats
- from open_webui.models.users import UserModel
- from open_webui.env import (
- ENABLE_FORWARD_USER_INFO_HEADERS,
- )
- from fastapi import (
- Depends,
- FastAPI,
- File,
- HTTPException,
- Request,
- UploadFile,
- APIRouter,
- )
- from fastapi.middleware.cors import CORSMiddleware
- from fastapi.responses import StreamingResponse
- from pydantic import BaseModel, ConfigDict, validator
- from starlette.background import BackgroundTask
- from open_webui.models.models import Models
- 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_model_system_prompt_to_body,
- )
- from open_webui.utils.auth import get_admin_user, get_verified_user
- from open_webui.utils.access_control import has_access
- from open_webui.config import (
- UPLOAD_DIR,
- )
- from open_webui.env import (
- ENV,
- SRC_LOG_LEVELS,
- AIOHTTP_CLIENT_SESSION_SSL,
- AIOHTTP_CLIENT_TIMEOUT,
- AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST,
- BYPASS_MODEL_ACCESS_CONTROL,
- )
- from open_webui.constants import ERROR_MESSAGES
- log = logging.getLogger(__name__)
- log.setLevel(SRC_LOG_LEVELS["OLLAMA"])
- ##########################################
- #
- # Utility functions
- #
- ##########################################
- async def send_get_request(url, key=None, user: UserModel = None):
- timeout = aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST)
- try:
- async with aiohttp.ClientSession(timeout=timeout, trust_env=True) as session:
- async with session.get(
- url,
- headers={
- "Content-Type": "application/json",
- **({"Authorization": f"Bearer {key}"} if key else {}),
- **(
- {
- "X-OpenWebUI-User-Name": user.name,
- "X-OpenWebUI-User-Id": user.id,
- "X-OpenWebUI-User-Email": user.email,
- "X-OpenWebUI-User-Role": user.role,
- }
- if ENABLE_FORWARD_USER_INFO_HEADERS and user
- else {}
- ),
- },
- 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}")
- return None
- async def cleanup_response(
- response: Optional[aiohttp.ClientResponse],
- session: Optional[aiohttp.ClientSession],
- ):
- if response:
- response.close()
- if session:
- await session.close()
- async def send_post_request(
- url: str,
- payload: Union[str, bytes],
- stream: bool = True,
- key: Optional[str] = None,
- content_type: Optional[str] = None,
- user: UserModel = None,
- ):
- r = None
- try:
- session = aiohttp.ClientSession(
- trust_env=True, timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT)
- )
- r = await session.post(
- url,
- data=payload,
- headers={
- "Content-Type": "application/json",
- **({"Authorization": f"Bearer {key}"} if key else {}),
- **(
- {
- "X-OpenWebUI-User-Name": user.name,
- "X-OpenWebUI-User-Id": user.id,
- "X-OpenWebUI-User-Email": user.email,
- "X-OpenWebUI-User-Role": user.role,
- }
- if ENABLE_FORWARD_USER_INFO_HEADERS and user
- else {}
- ),
- },
- ssl=AIOHTTP_CLIENT_SESSION_SSL,
- )
- if r.ok is False:
- try:
- res = await r.json()
- await cleanup_response(r, session)
- if "error" in res:
- raise HTTPException(status_code=r.status, detail=res["error"])
- except HTTPException as e:
- raise e # Re-raise HTTPException to be handled by FastAPI
- except Exception as e:
- log.error(f"Failed to parse error response: {e}")
- raise HTTPException(
- status_code=r.status,
- detail=f"Open WebUI: Server Connection Error",
- )
- r.raise_for_status() # Raises an error for bad responses (4xx, 5xx)
- if stream:
- response_headers = dict(r.headers)
- if content_type:
- response_headers["Content-Type"] = content_type
- return StreamingResponse(
- r.content,
- status_code=r.status,
- headers=response_headers,
- background=BackgroundTask(
- cleanup_response, response=r, session=session
- ),
- )
- else:
- res = await r.json()
- await cleanup_response(r, session)
- return res
- except HTTPException as e:
- raise e # Re-raise HTTPException to be handled by FastAPI
- except Exception as e:
- detail = f"Ollama: {e}"
- raise HTTPException(
- status_code=r.status if r else 500,
- detail=detail if e else "Open WebUI: Server Connection Error",
- )
- 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()
- @router.head("/")
- @router.get("/")
- async def get_status():
- return {"status": True}
- class ConnectionVerificationForm(BaseModel):
- url: str
- key: Optional[str] = None
- @router.post("/verify")
- async def verify_connection(
- form_data: ConnectionVerificationForm, user=Depends(get_admin_user)
- ):
- url = form_data.url
- key = form_data.key
- async with aiohttp.ClientSession(
- trust_env=True,
- timeout=aiohttp.ClientTimeout(total=AIOHTTP_CLIENT_TIMEOUT_MODEL_LIST),
- ) as session:
- try:
- async with session.get(
- f"{url}/api/version",
- headers={
- **({"Authorization": f"Bearer {key}"} if key else {}),
- **(
- {
- "X-OpenWebUI-User-Name": user.name,
- "X-OpenWebUI-User-Id": user.id,
- "X-OpenWebUI-User-Email": user.email,
- "X-OpenWebUI-User-Role": user.role,
- }
- if ENABLE_FORWARD_USER_INFO_HEADERS and user
- else {}
- ),
- },
- 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="Open WebUI: 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)
- @router.get("/config")
- async def get_config(request: Request, user=Depends(get_admin_user)):
- return {
- "ENABLE_OLLAMA_API": request.app.state.config.ENABLE_OLLAMA_API,
- "OLLAMA_BASE_URLS": request.app.state.config.OLLAMA_BASE_URLS,
- "OLLAMA_API_CONFIGS": request.app.state.config.OLLAMA_API_CONFIGS,
- }
- class OllamaConfigForm(BaseModel):
- ENABLE_OLLAMA_API: Optional[bool] = 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)
- ):
- 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))))
- 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
- }
- return {
- "ENABLE_OLLAMA_API": request.app.state.config.ENABLE_OLLAMA_API,
- "OLLAMA_BASE_URLS": request.app.state.config.OLLAMA_BASE_URLS,
- "OLLAMA_API_CONFIGS": request.app.state.config.OLLAMA_API_CONFIGS,
- }
- def merge_ollama_models_lists(model_lists):
- merged_models = {}
- for idx, model_list in enumerate(model_lists):
- if model_list is not None:
- for model in model_list:
- id = model["model"]
- 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())
- @cached(ttl=1)
- async def get_all_models(request: Request, user: UserModel = None):
- 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 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)))
- responses = await asyncio.gather(*request_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
- )
- connection_type = api_config.get("connection_type", "local")
- prefix_id = api_config.get("prefix_id", None)
- tags = api_config.get("tags", [])
- model_ids = api_config.get("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 model in response.get("models", []):
- if prefix_id:
- model["model"] = f"{prefix_id}.{model['model']}"
- if tags:
- model["tags"] = tags
- 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["name"]: m["expires_at"]
- for m in loaded_models["models"]
- if "expires_at" in m
- }
- for m in models["models"]:
- if m["name"] in expires_map:
- # Parse ISO8601 datetime with offset, get unix timestamp as int
- dt = datetime.fromisoformat(expires_map[m["name"]])
- 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
- async def get_filtered_models(models, user):
- # Filter models based on user access control
- filtered_models = []
- for model in models.get("models", []):
- model_info = Models.get_model_by_id(model["model"])
- if model_info:
- if user.id == model_info.user_id or has_access(
- user.id, type="read", access_control=model_info.access_control
- ):
- filtered_models.append(model)
- return filtered_models
- @router.get("/api/tags")
- @router.get("/api/tags/{url_idx}")
- async def get_ollama_tags(
- request: Request, url_idx: Optional[int] = None, user=Depends(get_verified_user)
- ):
- models = []
- if url_idx is None:
- models = 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)
- r = None
- try:
- r = requests.request(
- method="GET",
- url=f"{url}/api/tags",
- headers={
- **({"Authorization": f"Bearer {key}"} if key else {}),
- **(
- {
- "X-OpenWebUI-User-Name": user.name,
- "X-OpenWebUI-User-Id": user.id,
- "X-OpenWebUI-User-Email": user.email,
- "X-OpenWebUI-User-Role": user.role,
- }
- if ENABLE_FORWARD_USER_INFO_HEADERS and user
- else {}
- ),
- },
- )
- r.raise_for_status()
- models = r.json()
- except Exception as e:
- log.exception(e)
- detail = None
- if r is not None:
- try:
- res = r.json()
- if "error" in res:
- detail = f"Ollama: {res['error']}"
- except Exception:
- detail = f"Ollama: {e}"
- raise HTTPException(
- status_code=r.status_code if r else 500,
- detail=detail if detail else "Open WebUI: Server Connection Error",
- )
- if user.role == "user" and not BYPASS_MODEL_ACCESS_CONTROL:
- models["models"] = await get_filtered_models(models, user)
- return models
- @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
- )
- enable = api_config.get("enable", True)
- key = api_config.get("key", 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(*request_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
- )
- 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
- @router.get("/api/version")
- @router.get("/api/version/{url_idx}")
- async def get_ollama_versions(request: Request, url_idx: Optional[int] = 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]
- r = None
- try:
- r = requests.request(method="GET", url=f"{url}/api/version")
- r.raise_for_status()
- return r.json()
- except Exception as e:
- log.exception(e)
- detail = None
- if r is not None:
- try:
- res = r.json()
- if "error" in res:
- detail = f"Ollama: {res['error']}"
- except Exception:
- detail = f"Ollama: {e}"
- raise HTTPException(
- status_code=r.status_code if r else 500,
- detail=detail if detail else "Open WebUI: Server Connection Error",
- )
- else:
- return {"version": False}
- class ModelNameForm(BaseModel):
- name: str
- @router.post("/api/unload")
- async def unload_model(
- request: Request,
- form_data: ModelNameForm,
- user=Depends(get_admin_user),
- ):
- model_name = form_data.name
- if not model_name:
- raise HTTPException(
- status_code=400, detail="Missing 'name' of 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
- # Canonicalize model name (if not supplied with version)
- if ":" not in model_name:
- model_name = f"{model_name}:latest"
- 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"]
- # Send unload to ALL url_indices
- results = []
- errors = []
- for idx in url_indices:
- 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, {})
- )
- 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}.") :]
- payload = {"model": model_name, "keep_alive": 0, "prompt": ""}
- try:
- res = await send_post_request(
- url=f"{url}/api/generate",
- payload=json.dumps(payload),
- stream=False,
- 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),
- ):
- 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.model_dump(exclude_none=True), "insecure": True}
- return await send_post_request(
- url=f"{url}/api/pull",
- payload=json.dumps(payload),
- key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
- user=user,
- )
- class PushModelForm(BaseModel):
- name: str
- insecure: Optional[bool] = None
- stream: Optional[bool] = None
- @router.delete("/api/push")
- @router.delete("/api/push/{url_idx}")
- async def push_model(
- request: Request,
- form_data: PushModelForm,
- url_idx: Optional[int] = None,
- user=Depends(get_admin_user),
- ):
- if url_idx is None:
- await get_all_models(request, user=user)
- models = request.app.state.OLLAMA_MODELS
- if form_data.name in models:
- url_idx = models[form_data.name]["urls"][0]
- else:
- raise HTTPException(
- status_code=400,
- detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.name),
- )
- url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
- log.debug(f"url: {url}")
- return await send_post_request(
- url=f"{url}/api/push",
- payload=form_data.model_dump_json(exclude_none=True).encode(),
- key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
- user=user,
- )
- class CreateModelForm(BaseModel):
- model: Optional[str] = None
- stream: Optional[bool] = None
- path: Optional[str] = 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),
- ):
- log.debug(f"form_data: {form_data}")
- url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
- return await send_post_request(
- url=f"{url}/api/create",
- payload=form_data.model_dump_json(exclude_none=True).encode(),
- key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
- user=user,
- )
- class CopyModelForm(BaseModel):
- 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: Optional[int] = None,
- user=Depends(get_admin_user),
- ):
- 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),
- )
- url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
- key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS)
- try:
- r = requests.request(
- method="POST",
- url=f"{url}/api/copy",
- headers={
- "Content-Type": "application/json",
- **({"Authorization": f"Bearer {key}"} if key else {}),
- **(
- {
- "X-OpenWebUI-User-Name": user.name,
- "X-OpenWebUI-User-Id": user.id,
- "X-OpenWebUI-User-Email": user.email,
- "X-OpenWebUI-User-Role": user.role,
- }
- if ENABLE_FORWARD_USER_INFO_HEADERS and user
- else {}
- ),
- },
- data=form_data.model_dump_json(exclude_none=True).encode(),
- )
- r.raise_for_status()
- log.debug(f"r.text: {r.text}")
- return True
- except Exception as e:
- log.exception(e)
- detail = None
- if r is not None:
- try:
- res = r.json()
- if "error" in res:
- detail = f"Ollama: {res['error']}"
- except Exception:
- detail = f"Ollama: {e}"
- raise HTTPException(
- status_code=r.status_code if r else 500,
- detail=detail if detail else "Open WebUI: Server Connection Error",
- )
- @router.delete("/api/delete")
- @router.delete("/api/delete/{url_idx}")
- async def delete_model(
- request: Request,
- form_data: ModelNameForm,
- url_idx: Optional[int] = None,
- user=Depends(get_admin_user),
- ):
- if url_idx is None:
- await get_all_models(request, user=user)
- models = request.app.state.OLLAMA_MODELS
- if form_data.name in models:
- url_idx = models[form_data.name]["urls"][0]
- else:
- raise HTTPException(
- status_code=400,
- detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.name),
- )
- url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
- key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS)
- try:
- r = requests.request(
- method="DELETE",
- url=f"{url}/api/delete",
- data=form_data.model_dump_json(exclude_none=True).encode(),
- headers={
- "Content-Type": "application/json",
- **({"Authorization": f"Bearer {key}"} if key else {}),
- **(
- {
- "X-OpenWebUI-User-Name": user.name,
- "X-OpenWebUI-User-Id": user.id,
- "X-OpenWebUI-User-Email": user.email,
- "X-OpenWebUI-User-Role": user.role,
- }
- if ENABLE_FORWARD_USER_INFO_HEADERS and user
- else {}
- ),
- },
- )
- r.raise_for_status()
- log.debug(f"r.text: {r.text}")
- return True
- except Exception as e:
- log.exception(e)
- detail = None
- if r is not None:
- try:
- res = r.json()
- if "error" in res:
- detail = f"Ollama: {res['error']}"
- except Exception:
- detail = f"Ollama: {e}"
- raise HTTPException(
- status_code=r.status_code if r else 500,
- detail=detail if detail else "Open WebUI: Server Connection Error",
- )
- @router.post("/api/show")
- async def show_model_info(
- request: Request, form_data: ModelNameForm, user=Depends(get_verified_user)
- ):
- await get_all_models(request, user=user)
- models = request.app.state.OLLAMA_MODELS
- if form_data.name not in models:
- raise HTTPException(
- status_code=400,
- detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.name),
- )
- url_idx = random.choice(models[form_data.name]["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)
- try:
- r = requests.request(
- method="POST",
- url=f"{url}/api/show",
- headers={
- "Content-Type": "application/json",
- **({"Authorization": f"Bearer {key}"} if key else {}),
- **(
- {
- "X-OpenWebUI-User-Name": user.name,
- "X-OpenWebUI-User-Id": user.id,
- "X-OpenWebUI-User-Email": user.email,
- "X-OpenWebUI-User-Role": user.role,
- }
- if ENABLE_FORWARD_USER_INFO_HEADERS and user
- else {}
- ),
- },
- data=form_data.model_dump_json(exclude_none=True).encode(),
- )
- r.raise_for_status()
- return r.json()
- except Exception as e:
- log.exception(e)
- detail = None
- if r is not None:
- try:
- res = r.json()
- if "error" in res:
- detail = f"Ollama: {res['error']}"
- except Exception:
- detail = f"Ollama: {e}"
- raise HTTPException(
- status_code=r.status_code if r else 500,
- detail=detail if detail else "Open WebUI: Server Connection Error",
- )
- class GenerateEmbedForm(BaseModel):
- model: str
- input: list[str] | str
- truncate: Optional[bool] = None
- options: Optional[dict] = None
- keep_alive: Optional[Union[int, str]] = None
- @router.post("/api/embed")
- @router.post("/api/embed/{url_idx}")
- async def embed(
- request: Request,
- form_data: GenerateEmbedForm,
- url_idx: Optional[int] = None,
- user=Depends(get_verified_user),
- ):
- log.info(f"generate_ollama_batch_embeddings {form_data}")
- if url_idx is None:
- await get_all_models(request, user=user)
- models = request.app.state.OLLAMA_MODELS
- model = form_data.model
- if ":" not in model:
- model = f"{model}:latest"
- 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),
- )
- 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
- )
- key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS)
- prefix_id = api_config.get("prefix_id", None)
- if prefix_id:
- form_data.model = form_data.model.replace(f"{prefix_id}.", "")
- try:
- r = requests.request(
- method="POST",
- url=f"{url}/api/embed",
- headers={
- "Content-Type": "application/json",
- **({"Authorization": f"Bearer {key}"} if key else {}),
- **(
- {
- "X-OpenWebUI-User-Name": user.name,
- "X-OpenWebUI-User-Id": user.id,
- "X-OpenWebUI-User-Email": user.email,
- "X-OpenWebUI-User-Role": user.role,
- }
- if ENABLE_FORWARD_USER_INFO_HEADERS and user
- else {}
- ),
- },
- data=form_data.model_dump_json(exclude_none=True).encode(),
- )
- r.raise_for_status()
- data = r.json()
- return data
- except Exception as e:
- log.exception(e)
- detail = None
- if r is not None:
- try:
- res = r.json()
- if "error" in res:
- detail = f"Ollama: {res['error']}"
- except Exception:
- detail = f"Ollama: {e}"
- raise HTTPException(
- status_code=r.status_code if r else 500,
- detail=detail if detail else "Open WebUI: Server Connection Error",
- )
- class GenerateEmbeddingsForm(BaseModel):
- model: str
- prompt: str
- options: Optional[dict] = None
- keep_alive: Optional[Union[int, str]] = None
- @router.post("/api/embeddings")
- @router.post("/api/embeddings/{url_idx}")
- async def embeddings(
- request: Request,
- form_data: GenerateEmbeddingsForm,
- url_idx: Optional[int] = None,
- user=Depends(get_verified_user),
- ):
- log.info(f"generate_ollama_embeddings {form_data}")
- if url_idx is None:
- await get_all_models(request, user=user)
- models = request.app.state.OLLAMA_MODELS
- model = form_data.model
- if ":" not in model:
- model = f"{model}:latest"
- 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),
- )
- 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
- )
- key = get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS)
- prefix_id = api_config.get("prefix_id", None)
- if prefix_id:
- form_data.model = form_data.model.replace(f"{prefix_id}.", "")
- try:
- r = requests.request(
- method="POST",
- url=f"{url}/api/embeddings",
- headers={
- "Content-Type": "application/json",
- **({"Authorization": f"Bearer {key}"} if key else {}),
- **(
- {
- "X-OpenWebUI-User-Name": user.name,
- "X-OpenWebUI-User-Id": user.id,
- "X-OpenWebUI-User-Email": user.email,
- "X-OpenWebUI-User-Role": user.role,
- }
- if ENABLE_FORWARD_USER_INFO_HEADERS and user
- else {}
- ),
- },
- data=form_data.model_dump_json(exclude_none=True).encode(),
- )
- r.raise_for_status()
- data = r.json()
- return data
- except Exception as e:
- log.exception(e)
- detail = None
- if r is not None:
- try:
- res = r.json()
- if "error" in res:
- detail = f"Ollama: {res['error']}"
- except Exception:
- detail = f"Ollama: {e}"
- raise HTTPException(
- status_code=r.status_code if r else 500,
- detail=detail if detail else "Open WebUI: Server Connection Error",
- )
- class GenerateCompletionForm(BaseModel):
- model: str
- prompt: str
- suffix: Optional[str] = None
- images: Optional[list[str]] = None
- format: Optional[Union[dict, str]] = None
- options: Optional[dict] = None
- system: Optional[str] = None
- template: Optional[str] = None
- context: Optional[list[int]] = None
- stream: Optional[bool] = True
- raw: Optional[bool] = None
- keep_alive: Optional[Union[int, str]] = None
- @router.post("/api/generate")
- @router.post("/api/generate/{url_idx}")
- async def generate_completion(
- request: Request,
- form_data: GenerateCompletionForm,
- url_idx: Optional[int] = None,
- user=Depends(get_verified_user),
- ):
- if url_idx is None:
- await get_all_models(request, user=user)
- models = request.app.state.OLLAMA_MODELS
- model = form_data.model
- if ":" not in model:
- model = f"{model}:latest"
- 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),
- )
- 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
- )
- prefix_id = api_config.get("prefix_id", None)
- if prefix_id:
- form_data.model = form_data.model.replace(f"{prefix_id}.", "")
- return await send_post_request(
- url=f"{url}/api/generate",
- payload=form_data.model_dump_json(exclude_none=True).encode(),
- key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
- user=user,
- )
- class ChatMessage(BaseModel):
- role: str
- content: Optional[str] = None
- tool_calls: Optional[list[dict]] = None
- images: Optional[list[str]] = None
- @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):
- model: str
- messages: list[ChatMessage]
- format: Optional[Union[dict, str]] = None
- options: Optional[dict] = None
- template: Optional[str] = None
- stream: Optional[bool] = True
- keep_alive: Optional[Union[int, str]] = None
- tools: Optional[list[dict]] = None
- model_config = ConfigDict(
- extra="allow",
- )
- async def get_ollama_url(request: Request, model: str, url_idx: Optional[int] = None):
- if url_idx is None:
- 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 = request.app.state.config.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: Optional[int] = None,
- user=Depends(get_verified_user),
- bypass_filter: Optional[bool] = False,
- ):
- if BYPASS_MODEL_ACCESS_CONTROL:
- bypass_filter = True
- 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),
- )
- if isinstance(form_data, BaseModel):
- payload = {**form_data.model_dump(exclude_none=True)}
- if "metadata" in payload:
- del payload["metadata"]
- model_id = payload["model"]
- model_info = Models.get_model_by_id(model_id)
- if model_info:
- 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_ollama(params, payload)
- payload = apply_model_system_prompt_to_body(system, payload, metadata, user)
- # Check if user has access to the model
- if not bypass_filter and user.role == "user":
- if not (
- user.id == model_info.user_id
- or has_access(
- user.id, type="read", access_control=model_info.access_control
- )
- ):
- raise HTTPException(
- status_code=403,
- detail="Model not found",
- )
- elif not bypass_filter:
- if user.role != "admin":
- raise HTTPException(
- status_code=403,
- detail="Model not found",
- )
- if ":" not in payload["model"]:
- payload["model"] = f"{payload['model']}:latest"
- 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)
- if prefix_id:
- payload["model"] = payload["model"].replace(f"{prefix_id}.", "")
- return await send_post_request(
- url=f"{url}/api/chat",
- payload=json.dumps(payload),
- stream=form_data.stream,
- key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
- content_type="application/x-ndjson",
- user=user,
- )
- # TODO: we should update this part once Ollama supports other types
- class OpenAIChatMessageContent(BaseModel):
- type: str
- model_config = ConfigDict(extra="allow")
- class OpenAIChatMessage(BaseModel):
- role: str
- content: Union[Optional[str], list[OpenAIChatMessageContent]]
- model_config = ConfigDict(extra="allow")
- class OpenAIChatCompletionForm(BaseModel):
- model: str
- messages: list[OpenAIChatMessage]
- model_config = ConfigDict(extra="allow")
- class OpenAICompletionForm(BaseModel):
- model: str
- prompt: str
- 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: Optional[int] = None,
- user=Depends(get_verified_user),
- ):
- try:
- form_data = OpenAICompletionForm(**form_data)
- except Exception as e:
- log.exception(e)
- raise HTTPException(
- status_code=400,
- detail=str(e),
- )
- payload = {**form_data.model_dump(exclude_none=True, exclude=["metadata"])}
- if "metadata" in payload:
- del payload["metadata"]
- model_id = form_data.model
- if ":" not in model_id:
- model_id = f"{model_id}:latest"
- model_info = Models.get_model_by_id(model_id)
- if model_info:
- 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)
- # Check if user has access to the model
- if user.role == "user":
- if not (
- user.id == model_info.user_id
- or has_access(
- user.id, type="read", access_control=model_info.access_control
- )
- ):
- raise HTTPException(
- status_code=403,
- detail="Model not found",
- )
- else:
- if user.role != "admin":
- raise HTTPException(
- status_code=403,
- detail="Model not found",
- )
- if ":" not in payload["model"]:
- payload["model"] = f"{payload['model']}:latest"
- 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)
- if prefix_id:
- payload["model"] = payload["model"].replace(f"{prefix_id}.", "")
- return await send_post_request(
- url=f"{url}/v1/completions",
- payload=json.dumps(payload),
- stream=payload.get("stream", False),
- key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
- user=user,
- )
- @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: Optional[int] = None,
- user=Depends(get_verified_user),
- ):
- 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),
- )
- payload = {**completion_form.model_dump(exclude_none=True, exclude=["metadata"])}
- if "metadata" in payload:
- del payload["metadata"]
- model_id = completion_form.model
- if ":" not in model_id:
- model_id = f"{model_id}:latest"
- model_info = Models.get_model_by_id(model_id)
- if model_info:
- 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 = apply_model_system_prompt_to_body(system, payload, metadata, user)
- # Check if user has access to the model
- if user.role == "user":
- if not (
- user.id == model_info.user_id
- or has_access(
- user.id, type="read", access_control=model_info.access_control
- )
- ):
- raise HTTPException(
- status_code=403,
- detail="Model not found",
- )
- else:
- if user.role != "admin":
- raise HTTPException(
- status_code=403,
- detail="Model not found",
- )
- if ":" not in payload["model"]:
- payload["model"] = f"{payload['model']}:latest"
- 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)
- if prefix_id:
- payload["model"] = payload["model"].replace(f"{prefix_id}.", "")
- return await send_post_request(
- url=f"{url}/v1/chat/completions",
- payload=json.dumps(payload),
- stream=payload.get("stream", False),
- key=get_api_key(url_idx, url, request.app.state.config.OLLAMA_API_CONFIGS),
- user=user,
- )
- @router.get("/v1/models")
- @router.get("/v1/models/{url_idx}")
- async def get_openai_models(
- request: Request,
- url_idx: Optional[int] = None,
- user=Depends(get_verified_user),
- ):
- models = []
- 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"]
- ]
- else:
- url = request.app.state.config.OLLAMA_BASE_URLS[url_idx]
- try:
- r = requests.request(method="GET", url=f"{url}/api/tags")
- r.raise_for_status()
- model_list = r.json()
- models = [
- {
- "id": model["model"],
- "object": "model",
- "created": int(time.time()),
- "owned_by": "openai",
- }
- for model in models["models"]
- ]
- except Exception as e:
- log.exception(e)
- error_detail = "Open WebUI: Server Connection Error"
- if r is not None:
- try:
- res = r.json()
- if "error" in res:
- error_detail = f"Ollama: {res['error']}"
- except Exception:
- error_detail = f"Ollama: {e}"
- raise HTTPException(
- status_code=r.status_code if r else 500,
- detail=error_detail,
- )
- if user.role == "user" and not BYPASS_MODEL_ACCESS_CONTROL:
- # Filter models based on user access control
- filtered_models = []
- for model in models:
- model_info = Models.get_model_by_id(model["id"])
- if model_info:
- if user.id == model_info.user_id or has_access(
- user.id, type="read", access_control=model_info.access_control
- ):
- filtered_models.append(model)
- models = filtered_models
- return {
- "data": models,
- "object": "list",
- }
- class UrlForm(BaseModel):
- url: str
- class UploadBlobForm(BaseModel):
- filename: str
- def parse_huggingface_url(hf_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 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
- headers = {"Range": f"bytes={current_size}-"} if current_size > 0 else {}
- timeout = aiohttp.ClientTimeout(total=600) # Set the timeout
- 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 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'
- if done:
- file.seek(0)
- chunk_size = 1024 * 1024 * 2
- hashed = calculate_sha256(file, chunk_size)
- file.seek(0)
- url = f"{ollama_url}/api/blobs/sha256:{hashed}"
- response = requests.post(url, data=file)
- if 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 "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(
- request: Request,
- form_data: UrlForm,
- url_idx: Optional[int] = None,
- user=Depends(get_admin_user),
- ):
- 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]
- file_name = parse_huggingface_url(form_data.url)
- if file_name:
- file_path = f"{UPLOAD_DIR}/{file_name}"
- return StreamingResponse(
- download_file_stream(url, form_data.url, file_path, file_name),
- )
- else:
- return None
- # TODO: Progress bar does not reflect size & duration of upload.
- @router.post("/models/upload")
- @router.post("/models/upload/{url_idx}")
- async def upload_model(
- request: Request,
- file: UploadFile = File(...),
- url_idx: Optional[int] = 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]
- 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
- 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)
- 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
- # --- P2: SSE progress + calculate sha256 hash ---
- file_hash = calculate_sha256(file_path, chunk_size)
- log.info(f"Model Hash: {str(file_hash)}") # DEBUG
- 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"
- # --- P3: Upload to ollama /api/blobs ---
- with open(file_path, "rb") as f:
- url = f"{ollama_url}/api/blobs/sha256:{file_hash}"
- response = requests.post(url, data=f)
- if response.ok:
- log.info(f"Uploaded to /api/blobs") # DEBUG
- # Remove local file
- os.remove(file_path)
- # Create model in ollama
- model_name, ext = os.path.splitext(filename)
- log.info(f"Created Model: {model_name}") # DEBUG
- create_payload = {
- "model": model_name,
- # Reference the file by its original name => the uploaded blob's digest
- "files": {filename: f"sha256:{file_hash}"},
- }
- log.info(f"Model Payload: {create_payload}") # DEBUG
- # Call ollama /api/create
- # https://github.com/ollama/ollama/blob/main/docs/api.md#create-a-model
- create_resp = requests.post(
- url=f"{ollama_url}/api/create",
- headers={"Content-Type": "application/json"},
- data=json.dumps(create_payload),
- )
- 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:
- raise Exception(
- f"Failed to create model in Ollama. {create_resp.text}"
- )
- else:
- raise Exception("Ollama: Could not create blob, Please try again.")
- except Exception as e:
- res = {"error": str(e)}
- yield f"data: {json.dumps(res)}\n\n"
- return StreamingResponse(file_process_stream(), media_type="text/event-stream")
|