fix: model status in ui

etalab-ia · Oct 9, 2024 · 5c8dda5 · 5c8dda5
1 parent d1e6368
commit 5c8dda5
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/app/main.py b/app/main.py
@@ -28,6 +28,7 @@ def health(user: User = Security(check_api_key)):
 
 app.add_middleware(ContentSizeLimitMiddleware)
 
+# Core
 app.include_router(models.router, tags=["Core"], prefix="/v1")
 app.include_router(chat.router, tags=["Core"], prefix="/v1")
 app.include_router(completions.router, tags=["Core"], prefix="/v1")

diff --git a/ui/utils.py b/ui/utils.py
@@ -73,8 +73,8 @@ def get_models(api_key: str):
     assert response.status_code == 200
     models = response.json()["data"]
 
-    embeddings_models = [model["id"] for model in models if model["type"] == EMBEDDINGS_MODEL_TYPE]
-    language_models = [model["id"] for model in models if model["type"] == LANGUAGE_MODEL_TYPE]
+    embeddings_models = [model["id"] for model in models if model["type"] == EMBEDDINGS_MODEL_TYPE and model["status"] == "available"]
+    language_models = [model["id"] for model in models if model["type"] == LANGUAGE_MODEL_TYPE and model["status"] == "available"]
 
     return language_models, embeddings_models