Add support for auto-refreshing available models based on an API call (#3576)

backend/onyx/background/celery/tasks/llm_model_update/tasks.py

@@ -0,0 +1,105 @@
+from typing import Any
+
+import requests
+from celery import shared_task
+from celery import Task
+
+from onyx.background.celery.apps.app_base import task_logger
+from onyx.configs.app_configs import JOB_TIMEOUT
+from onyx.configs.app_configs import LLM_MODEL_UPDATE_API_URL
+from onyx.configs.constants import OnyxCeleryTask
+from onyx.db.engine import get_session_with_tenant
+from onyx.db.models import LLMProvider
+
+
+def _process_model_list_response(model_list_json: Any) -> list[str]:
+    # Handle case where response is wrapped in a "data" field
+    if isinstance(model_list_json, dict) and "data" in model_list_json:
+        model_list_json = model_list_json["data"]
+
+    if not isinstance(model_list_json, list):
+        raise ValueError(
+            f"Invalid response from API - expected list, got {type(model_list_json)}"
+        )
+
+    # Handle both string list and object list cases
+    model_names: list[str] = []
+    for item in model_list_json:
+        if isinstance(item, str):
+            model_names.append(item)
+        elif isinstance(item, dict) and "model_name" in item:
+            model_names.append(item["model_name"])
+        else:
+            raise ValueError(
+                f"Invalid item in model list - expected string or dict with model_name, got {type(item)}"
+            )
+
+    return model_names
+
+
+@shared_task(
+    name=OnyxCeleryTask.CHECK_FOR_LLM_MODEL_UPDATE,
+    soft_time_limit=JOB_TIMEOUT,
+    trail=False,
+    bind=True,
+)
+def check_for_llm_model_update(self: Task, *, tenant_id: str | None) -> bool | None:
+    if not LLM_MODEL_UPDATE_API_URL:
+        raise ValueError("LLM model update API URL not configured")
+
+    # First fetch the models from the API
+    try:
+        response = requests.get(LLM_MODEL_UPDATE_API_URL)
+        response.raise_for_status()
+        available_models = _process_model_list_response(response.json())
+        task_logger.info(f"Found available models: {available_models}")
+
+    except Exception:
+        task_logger.exception("Failed to fetch models from API.")
+        return None
+
+    # Then update the database with the fetched models
+    with get_session_with_tenant(tenant_id) as db_session:
+        # Get the default LLM provider
+        default_provider = (
+            db_session.query(LLMProvider)
+            .filter(LLMProvider.is_default_provider.is_(True))
+            .first()
+        )
+
+        if not default_provider:
+            task_logger.warning("No default LLM provider found")
+            return None
+
+        # log change if any
+        old_models = set(default_provider.model_names or [])
+        new_models = set(available_models)
+        added_models = new_models - old_models
+        removed_models = old_models - new_models
+
+        if added_models:
+            task_logger.info(f"Adding models: {sorted(added_models)}")
+        if removed_models:
+            task_logger.info(f"Removing models: {sorted(removed_models)}")
+
+        # Update the provider's model list
+        default_provider.model_names = available_models
+        # if the default model is no longer available, set it to the first model in the list
+        if default_provider.default_model_name not in available_models:
+            task_logger.info(
+                f"Default model {default_provider.default_model_name} not "
+                f"available, setting to first model in list: {available_models[0]}"
+            )
+            default_provider.default_model_name = available_models[0]
+        if default_provider.fast_default_model_name not in available_models:
+            task_logger.info(
+                f"Fast default model {default_provider.fast_default_model_name} "
+                f"not available, setting to first model in list: {available_models[0]}"
+            )
+            default_provider.fast_default_model_name = available_models[0]
+        db_session.commit()
+
+        if added_models or removed_models:
+            task_logger.info("Updated model list for default provider.")
+
+    return True