From 760652c4cc73f1b8839272dddf5846d9eb7b8923 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Tue, 12 Nov 2024 20:26:33 -0800 Subject: [PATCH] refac --- backend/open_webui/apps/ollama/main.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/backend/open_webui/apps/ollama/main.py b/backend/open_webui/apps/ollama/main.py index 3afe75866..e566e032b 100644 --- a/backend/open_webui/apps/ollama/main.py +++ b/backend/open_webui/apps/ollama/main.py @@ -300,7 +300,7 @@ async def get_all_models(): prefix_id = api_config.get("prefix_id", None) model_ids = api_config.get("model_ids", []) - if len(model_ids) != 0: + if len(model_ids) != 0 and "models" in response: response["models"] = list( filter( lambda model: model["model"] in model_ids, @@ -309,7 +309,7 @@ async def get_all_models(): ) if prefix_id: - for model in response["models"]: + for model in response.get("models", []): model["model"] = f"{prefix_id}.{model['model']}" print(responses) @@ -317,7 +317,8 @@ async def get_all_models(): models = { "models": merge_models_lists( map( - lambda response: response["models"] if response else None, responses + lambda response: response.get("models", []) if response else None, + responses, ) ) }