Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2024-11-16 04:41:07 -08:00
parent 931e03bd9e
commit 932de8f1e2
9 changed files with 309 additions and 214 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/apps/ollama/main.py
+46 -37
View File
@@ -68,25 +68,12 @@ app.state.config.ENABLE_OLLAMA_API = ENABLE_OLLAMA_API
app.state.config.OLLAMA_BASE_URLS = OLLAMA_BASE_URLS
app.state.config.OLLAMA_API_CONFIGS = OLLAMA_API_CONFIGS
app.state.MODELS = {}
# 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.
@app.middleware("http")
async def check_url(request: Request, call_next):
if len(app.state.MODELS) == 0:
await get_all_models()
else:
pass
response = await call_next(request)
return response
@app.head("/")
@app.get("/")
async def get_status():
@@ -321,8 +308,6 @@ async def get_all_models():
else:
models = {"models": []}
app.state.MODELS = {model["model"]: model for model in models["models"]}
return models
@@ -470,8 +455,11 @@ async def push_model(
user=Depends(get_admin_user),
):
if url_idx is None:
if form_data.name in app.state.MODELS:
url_idx = app.state.MODELS[form_data.name]["urls"][0]
model_list = await get_all_models()
models = {model["model"]: model for model in model_list["models"]}
if form_data.name in models:
url_idx = models[form_data.name]["urls"][0]
else:
raise HTTPException(
status_code=400,
@@ -520,8 +508,11 @@ async def copy_model(
user=Depends(get_admin_user),
):
if url_idx is None:
if form_data.source in app.state.MODELS:
url_idx = app.state.MODELS[form_data.source]["urls"][0]
model_list = await get_all_models()
models = {model["model"]: model for model in model_list["models"]}
if form_data.source in models:
url_idx = models[form_data.source]["urls"][0]
else:
raise HTTPException(
status_code=400,
@@ -576,8 +567,11 @@ async def delete_model(
user=Depends(get_admin_user),
):
if url_idx is None:
if form_data.name in app.state.MODELS:
url_idx = app.state.MODELS[form_data.name]["urls"][0]
model_list = await get_all_models()
models = {model["model"]: model for model in model_list["models"]}
if form_data.name in models:
url_idx = models[form_data.name]["urls"][0]
else:
raise HTTPException(
status_code=400,
@@ -625,13 +619,16 @@ async def delete_model(
@app.post("/api/show")
async def show_model_info(form_data: ModelNameForm, user=Depends(get_verified_user)):
if form_data.name not in app.state.MODELS:
model_list = await get_all_models()
models = {model["model"]: model for model in model_list["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(app.state.MODELS[form_data.name]["urls"])
url_idx = random.choice(models[form_data.name]["urls"])
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
log.info(f"url: {url}")
@@ -701,23 +698,26 @@ async def generate_embeddings(
url_idx: Optional[int] = None,
user=Depends(get_verified_user),
):
return generate_ollama_embeddings(form_data=form_data, url_idx=url_idx)
return await generate_ollama_embeddings(form_data=form_data, url_idx=url_idx)
def generate_ollama_embeddings(
async def generate_ollama_embeddings(
form_data: GenerateEmbeddingsForm,
url_idx: Optional[int] = None,
):
log.info(f"generate_ollama_embeddings {form_data}")
if url_idx is None:
model_list = await get_all_models()
models = {model["model"]: model for model in model_list["models"]}
model = form_data.model
if ":" not in model:
model = f"{model}:latest"
if model in app.state.MODELS:
url_idx = random.choice(app.state.MODELS[model]["urls"])
if model in models:
url_idx = random.choice(models[model]["urls"])
else:
raise HTTPException(
status_code=400,
@@ -768,20 +768,23 @@ def generate_ollama_embeddings(
)
def generate_ollama_batch_embeddings(
async def generate_ollama_batch_embeddings(
form_data: GenerateEmbedForm,
url_idx: Optional[int] = None,
):
log.info(f"generate_ollama_batch_embeddings {form_data}")
if url_idx is None:
model_list = await get_all_models()
models = {model["model"]: model for model in model_list["models"]}
model = form_data.model
if ":" not in model:
model = f"{model}:latest"
if model in app.state.MODELS:
url_idx = random.choice(app.state.MODELS[model]["urls"])
if model in models:
url_idx = random.choice(models[model]["urls"])
else:
raise HTTPException(
status_code=400,
@@ -851,13 +854,16 @@ async def generate_completion(
user=Depends(get_verified_user),
):
if url_idx is None:
model_list = await get_all_models()
models = {model["model"]: model for model in model_list["models"]}
model = form_data.model
if ":" not in model:
model = f"{model}:latest"
if model in app.state.MODELS:
url_idx = random.choice(app.state.MODELS[model]["urls"])
if model in models:
url_idx = random.choice(models[model]["urls"])
else:
raise HTTPException(
status_code=400,
@@ -892,14 +898,17 @@ class GenerateChatCompletionForm(BaseModel):
keep_alive: Optional[Union[int, str]] = None
def get_ollama_url(url_idx: Optional[int], model: str):
async def get_ollama_url(url_idx: Optional[int], model: str):
if url_idx is None:
if model not in app.state.MODELS:
model_list = await get_all_models()
models = {model["model"]: model for model in model_list["models"]}
if model not in models:
raise HTTPException(
status_code=400,
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(model),
)
url_idx = random.choice(app.state.MODELS[model]["urls"])
url_idx = random.choice(models[model]["urls"])
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
return url
@@ -948,7 +957,7 @@ async def generate_chat_completion(
if ":" not in payload["model"]:
payload["model"] = f"{payload['model']}:latest"
url = get_ollama_url(url_idx, payload["model"])
url = await get_ollama_url(url_idx, payload["model"])
log.info(f"url: {url}")
log.debug(f"generate_chat_completion() - 2.payload = {payload}")
@@ -1030,7 +1039,7 @@ async def generate_openai_chat_completion(
if ":" not in payload["model"]:
payload["model"] = f"{payload['model']}:latest"
url = get_ollama_url(url_idx, payload["model"])
url = await get_ollama_url(url_idx, payload["model"])
log.info(f"url: {url}")
api_config = app.state.config.OLLAMA_API_CONFIGS.get(url, {})
backend/open_webui/apps/openai/main.py
+59 -60
View File
@@ -36,7 +36,7 @@ from open_webui.utils.payload import (
apply_model_system_prompt_to_body,
)
from open_webui.utils.utils import get_admin_user, get_verified_user
from open_webui.utils.utils import get_admin_user, get_verified_user, has_access
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["OPENAI"])
@@ -64,17 +64,6 @@ app.state.config.OPENAI_API_BASE_URLS = OPENAI_API_BASE_URLS
app.state.config.OPENAI_API_KEYS = OPENAI_API_KEYS
app.state.config.OPENAI_API_CONFIGS = OPENAI_API_CONFIGS
app.state.MODELS = {}
@app.middleware("http")
async def check_url(request: Request, call_next):
if len(app.state.MODELS) == 0:
await get_all_models()
response = await call_next(request)
return response
@app.get("/config")
async def get_config(user=Depends(get_admin_user)):
@@ -259,7 +248,7 @@ def merge_models_lists(model_lists):
return merged_list
async def get_all_models_raw() -> list:
async def get_all_models_responses() -> list:
if not app.state.config.ENABLE_OPENAI_API:
return []
@@ -330,22 +319,13 @@ async def get_all_models_raw() -> list:
return responses
@overload
async def get_all_models(raw: Literal[True]) -> list: ...
@overload
async def get_all_models(raw: Literal[False] = False) -> dict[str, list]: ...
async def get_all_models(raw=False) -> dict[str, list] | list:
async def get_all_models() -> dict[str, list]:
log.info("get_all_models()")
if not app.state.config.ENABLE_OPENAI_API:
return [] if raw else {"data": []}
responses = await get_all_models_raw()
if raw:
return responses
if not app.state.config.ENABLE_OPENAI_API:
return {"data": []}
responses = await get_all_models_responses()
def extract_data(response):
if response and "data" in response:
@@ -355,9 +335,7 @@ async def get_all_models(raw=False) -> dict[str, list] | list:
return None
models = {"data": merge_models_lists(map(extract_data, responses))}
log.debug(f"models: {models}")
app.state.MODELS = {model["id"]: model for model in models["data"]}
return models
@@ -365,21 +343,12 @@ async def get_all_models(raw=False) -> dict[str, list] | list:
@app.get("/models")
@app.get("/models/{url_idx}")
async def get_models(url_idx: Optional[int] = None, user=Depends(get_verified_user)):
models = {
"data": [],
}
if url_idx is None:
models = await get_all_models()
# TODO: Check User Group and Filter Models
# if app.state.config.ENABLE_MODEL_FILTER:
# if user.role == "user":
# models["data"] = list(
# filter(
# lambda model: model["id"] in app.state.config.MODEL_FILTER_LIST,
# models["data"],
# )
# )
# return models
return models
else:
url = app.state.config.OPENAI_API_BASE_URLS[url_idx]
key = app.state.config.OPENAI_API_KEYS[url_idx]
@@ -387,6 +356,7 @@ async def get_models(url_idx: Optional[int] = None, user=Depends(get_verified_us
headers = {}
headers["Authorization"] = f"Bearer {key}"
headers["Content-Type"] = "application/json"
if ENABLE_FORWARD_USER_INFO_HEADERS:
headers["X-OpenWebUI-User-Name"] = user.name
headers["X-OpenWebUI-User-Id"] = user.id
@@ -428,8 +398,7 @@ async def get_models(url_idx: Optional[int] = None, user=Depends(get_verified_us
)
]
return response_data
models = response_data
except aiohttp.ClientError as e:
# ClientError covers all aiohttp requests issues
log.exception(f"Client error: {str(e)}")
@@ -443,6 +412,22 @@ async def get_models(url_idx: Optional[int] = None, user=Depends(get_verified_us
error_detail = f"Unexpected error: {str(e)}"
raise HTTPException(status_code=500, detail=error_detail)
if user.role == "user":
# Filter models based on user access control
filtered_models = []
for model in models.get("data", []):
model_info = Models.get_model_by_id(model["id"])
if model_info:
if has_access(
user.id, type="read", access_control=model_info.access_control
):
filtered_models.append(model)
else:
filtered_models.append(model)
models["data"] = filtered_models
return models
class ConnectionVerificationForm(BaseModel):
url: str
@@ -502,18 +487,9 @@ async def generate_chat_completion(
del payload["metadata"]
model_id = form_data.get("model")
# TODO: Check User Group and Filter Models
# if not bypass_filter:
# if app.state.config.ENABLE_MODEL_FILTER:
# if user.role == "user" and model_id not in app.state.config.MODEL_FILTER_LIST:
# raise HTTPException(
# status_code=403,
# detail="Model not found",
# )
model_info = Models.get_model_by_id(model_id)
# Check model info and override the payload
if model_info:
if model_info.base_model_id:
payload["model"] = model_info.base_model_id
@@ -522,12 +498,33 @@ async def generate_chat_completion(
payload = apply_model_params_to_body_openai(params, payload)
payload = apply_model_system_prompt_to_body(params, payload, user)
try:
model = app.state.MODELS[payload.get("model")]
idx = model["urlIdx"]
except Exception as e:
raise HTTPException(status_code=404, detail="Model not found")
# Check if user has access to the model
if user.role == "user" and not has_access(
user.id, type="read", access_control=model_info.access_control
):
raise HTTPException(
status_code=403,
detail="Model not found",
)
# Attemp to get urlIdx from the model
models = await get_all_models()
# Find the model from the list
model = next(
(model for model in models["data"] if model["id"] == payload.get("model")),
None,
)
if model:
idx = model["urlIdx"]
else:
raise HTTPException(
status_code=404,
detail="Model not found",
)
# Get the API config for the model
api_config = app.state.config.OPENAI_API_CONFIGS.get(
app.state.config.OPENAI_API_BASE_URLS[idx], {}
)
@@ -536,6 +533,7 @@ async def generate_chat_completion(
if prefix_id:
payload["model"] = payload["model"].replace(f"{prefix_id}.", "")
# Add user info to the payload if the model is a pipeline
if "pipeline" in model and model.get("pipeline"):
payload["user"] = {
"name": user.name,
@@ -546,8 +544,9 @@ async def generate_chat_completion(
url = app.state.config.OPENAI_API_BASE_URLS[idx]
key = app.state.config.OPENAI_API_KEYS[idx]
is_o1 = payload["model"].lower().startswith("o1-")
# Fix: O1 does not support the "max_tokens" parameter, Modify "max_tokens" to "max_completion_tokens"
is_o1 = payload["model"].lower().startswith("o1-")
# Change max_completion_tokens to max_tokens (Backward compatible)
if "api.openai.com" not in url and not is_o1:
if "max_completion_tokens" in payload:
backend/open_webui/apps/retrieval/utils.py
+16 -7
View File
@@ -3,6 +3,7 @@ import os
import uuid
from typing import Optional, Union
import asyncio
import requests
from huggingface_hub import snapshot_download
@@ -291,7 +292,13 @@ def get_embedding_function(
if embedding_engine == "":
return lambda query: embedding_function.encode(query).tolist()
elif embedding_engine in ["ollama", "openai"]:
func = lambda query: generate_embeddings(
# Wrapper to run the async generate_embeddings synchronously.
def sync_generate_embeddings(*args, **kwargs):
return asyncio.run(generate_embeddings(*args, **kwargs))
# Semantic expectation from the original version (using sync wrapper).
func = lambda query: sync_generate_embeddings(
engine=embedding_engine,
model=embedding_model,
text=query,
@@ -469,7 +476,7 @@ def get_model_path(model: str, update_model: bool = False):
return model
def generate_openai_batch_embeddings(
async def generate_openai_batch_embeddings(
model: str, texts: list[str], key: str, url: str = "https://api.openai.com/v1"
) -> Optional[list[list[float]]]:
try:
@@ -492,14 +499,16 @@ def generate_openai_batch_embeddings(
return None
def generate_embeddings(engine: str, model: str, text: Union[str, list[str]], **kwargs):
async def generate_embeddings(
engine: str, model: str, text: Union[str, list[str]], **kwargs
):
if engine == "ollama":
if isinstance(text, list):
embeddings = generate_ollama_batch_embeddings(
embeddings = await generate_ollama_batch_embeddings(
GenerateEmbedForm(**{"model": model, "input": text})
)
else:
embeddings = generate_ollama_batch_embeddings(
embeddings = await generate_ollama_batch_embeddings(
GenerateEmbedForm(**{"model": model, "input": [text]})
)
return (
@@ -512,9 +521,9 @@ def generate_embeddings(engine: str, model: str, text: Union[str, list[str]], **
url = kwargs.get("url", "https://api.openai.com/v1")
if isinstance(text, list):
embeddings = generate_openai_batch_embeddings(model, text, key, url)
embeddings = await generate_openai_batch_embeddings(model, text, key, url)
else:
embeddings = generate_openai_batch_embeddings(model, [text], key, url)
embeddings = await generate_openai_batch_embeddings(model, [text], key, url)
return embeddings[0] if isinstance(text, str) else embeddings
backend/open_webui/apps/webui/main.py
+2 -3
View File
@@ -142,7 +142,6 @@ app.state.config.LDAP_USE_TLS = LDAP_USE_TLS
app.state.config.LDAP_CA_CERT_FILE = LDAP_CA_CERT_FILE
app.state.config.LDAP_CIPHERS = LDAP_CIPHERS
app.state.MODELS = {}
app.state.TOOLS = {}
app.state.FUNCTIONS = {}
@@ -369,7 +368,7 @@ def get_function_params(function_module, form_data, user, extra_params=None):
return params
async def generate_function_chat_completion(form_data, user):
async def generate_function_chat_completion(form_data, user, models: dict = {}):
model_id = form_data.get("model")
model_info = Models.get_model_by_id(model_id)
@@ -412,7 +411,7 @@ async def generate_function_chat_completion(form_data, user):
user,
{
**extra_params,
"__model__": app.state.MODELS[form_data["model"]],
"__model__": models.get(form_data["model"], None),
"__messages__": form_data["messages"],
"__files__": files,
},