feat: switch to config proxy, remove config_get/set

This commit is contained in:
Jun Siang Cheah
2024-05-10 15:03:24 +08:00
parent f712c90019
commit 298e6848b3
11 changed files with 340 additions and 379 deletions

View File

@@ -46,8 +46,7 @@ from config import (
ENABLE_MODEL_FILTER,
MODEL_FILTER_LIST,
UPLOAD_DIR,
config_set,
config_get,
AppConfig,
)
from utils.misc import calculate_sha256
@@ -63,11 +62,12 @@ app.add_middleware(
allow_headers=["*"],
)
app.state.config = AppConfig()
app.state.ENABLE_MODEL_FILTER = ENABLE_MODEL_FILTER
app.state.MODEL_FILTER_LIST = MODEL_FILTER_LIST
app.state.OLLAMA_BASE_URLS = OLLAMA_BASE_URLS
app.state.config.OLLAMA_BASE_URLS = OLLAMA_BASE_URLS
app.state.MODELS = {}
@@ -98,7 +98,7 @@ async def get_status():
@app.get("/urls")
async def get_ollama_api_urls(user=Depends(get_admin_user)):
return {"OLLAMA_BASE_URLS": config_get(app.state.OLLAMA_BASE_URLS)}
return {"OLLAMA_BASE_URLS": app.state.config.OLLAMA_BASE_URLS}
class UrlUpdateForm(BaseModel):
@@ -107,10 +107,10 @@ class UrlUpdateForm(BaseModel):
@app.post("/urls/update")
async def update_ollama_api_url(form_data: UrlUpdateForm, user=Depends(get_admin_user)):
config_set(app.state.OLLAMA_BASE_URLS, form_data.urls)
app.state.config.OLLAMA_BASE_URLS = form_data.urls
log.info(f"app.state.OLLAMA_BASE_URLS: {app.state.OLLAMA_BASE_URLS}")
return {"OLLAMA_BASE_URLS": config_get(app.state.OLLAMA_BASE_URLS)}
log.info(f"app.state.config.OLLAMA_BASE_URLS: {app.state.config.OLLAMA_BASE_URLS}")
return {"OLLAMA_BASE_URLS": app.state.config.OLLAMA_BASE_URLS}
@app.get("/cancel/{request_id}")
@@ -155,9 +155,7 @@ def merge_models_lists(model_lists):
async def get_all_models():
log.info("get_all_models()")
tasks = [
fetch_url(f"{url}/api/tags") for url in config_get(app.state.OLLAMA_BASE_URLS)
]
tasks = [fetch_url(f"{url}/api/tags") for url in app.state.config.OLLAMA_BASE_URLS]
responses = await asyncio.gather(*tasks)
models = {
@@ -183,15 +181,14 @@ async def get_ollama_tags(
if user.role == "user":
models["models"] = list(
filter(
lambda model: model["name"]
in config_get(app.state.MODEL_FILTER_LIST),
lambda model: model["name"] in app.state.MODEL_FILTER_LIST,
models["models"],
)
)
return models
return models
else:
url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
try:
r = requests.request(method="GET", url=f"{url}/api/tags")
r.raise_for_status()
@@ -222,8 +219,7 @@ async def get_ollama_versions(url_idx: Optional[int] = None):
# returns lowest version
tasks = [
fetch_url(f"{url}/api/version")
for url in config_get(app.state.OLLAMA_BASE_URLS)
fetch_url(f"{url}/api/version") for url in app.state.config.OLLAMA_BASE_URLS
]
responses = await asyncio.gather(*tasks)
responses = list(filter(lambda x: x is not None, responses))
@@ -243,7 +239,7 @@ async def get_ollama_versions(url_idx: Optional[int] = None):
detail=ERROR_MESSAGES.OLLAMA_NOT_FOUND,
)
else:
url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
try:
r = requests.request(method="GET", url=f"{url}/api/version")
r.raise_for_status()
@@ -275,7 +271,7 @@ class ModelNameForm(BaseModel):
async def pull_model(
form_data: ModelNameForm, url_idx: int = 0, user=Depends(get_admin_user)
):
url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
log.info(f"url: {url}")
r = None
@@ -363,7 +359,7 @@ async def push_model(
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.name),
)
url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
log.debug(f"url: {url}")
r = None
@@ -425,7 +421,7 @@ async def create_model(
form_data: CreateModelForm, url_idx: int = 0, user=Depends(get_admin_user)
):
log.debug(f"form_data: {form_data}")
url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
log.info(f"url: {url}")
r = None
@@ -498,7 +494,7 @@ async def copy_model(
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.source),
)
url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
log.info(f"url: {url}")
try:
@@ -545,7 +541,7 @@ async def delete_model(
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.name),
)
url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
log.info(f"url: {url}")
try:
@@ -585,7 +581,7 @@ async def show_model_info(form_data: ModelNameForm, user=Depends(get_verified_us
)
url_idx = random.choice(app.state.MODELS[form_data.name]["urls"])
url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
log.info(f"url: {url}")
try:
@@ -642,7 +638,7 @@ async def generate_embeddings(
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model),
)
url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
log.info(f"url: {url}")
try:
@@ -692,7 +688,7 @@ def generate_ollama_embeddings(
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model),
)
url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
log.info(f"url: {url}")
try:
@@ -761,7 +757,7 @@ async def generate_completion(
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model),
)
url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
log.info(f"url: {url}")
r = None
@@ -864,7 +860,7 @@ async def generate_chat_completion(
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model),
)
url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
log.info(f"url: {url}")
r = None
@@ -973,7 +969,7 @@ async def generate_openai_chat_completion(
detail=ERROR_MESSAGES.MODEL_NOT_FOUND(form_data.model),
)
url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
log.info(f"url: {url}")
r = None
@@ -1072,7 +1068,7 @@ async def get_openai_models(
}
else:
url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
try:
r = requests.request(method="GET", url=f"{url}/api/tags")
r.raise_for_status()
@@ -1206,7 +1202,7 @@ async def download_model(
if url_idx == None:
url_idx = 0
url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
url = app.state.config.OLLAMA_BASE_URLS[url_idx]
file_name = parse_huggingface_url(form_data.url)
@@ -1225,7 +1221,7 @@ async def download_model(
def upload_model(file: UploadFile = File(...), url_idx: Optional[int] = None):
if url_idx == None:
url_idx = 0
ollama_url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
ollama_url = app.state.config.OLLAMA_BASE_URLS[url_idx]
file_path = f"{UPLOAD_DIR}/{file.filename}"
@@ -1290,7 +1286,7 @@ def upload_model(file: UploadFile = File(...), url_idx: Optional[int] = None):
# async def upload_model(file: UploadFile = File(), url_idx: Optional[int] = None):
# if url_idx == None:
# url_idx = 0
# url = config_get(app.state.OLLAMA_BASE_URLS)[url_idx]
# url = app.state.config.OLLAMA_BASE_URLS[url_idx]
# file_location = os.path.join(UPLOAD_DIR, file.filename)
# total_size = file.size
@@ -1327,7 +1323,7 @@ def upload_model(file: UploadFile = File(...), url_idx: Optional[int] = None):
async def deprecated_proxy(
path: str, request: Request, user=Depends(get_verified_user)
):
url = config_get(app.state.OLLAMA_BASE_URLS)[0]
url = app.state.config.OLLAMA_BASE_URLS[0]
target_url = f"{url}/{path}"
body = await request.body()