feat: backend model load balancing support (#4927)

api/core/helper/model_provider_cache.py

@@ -9,6 +9,7 @@ from extensions.ext_redis import redis_client
 class ProviderCredentialsCacheType(Enum):
     PROVIDER = "provider"
     MODEL = "provider_model"
+    LOAD_BALANCING_MODEL = "load_balancing_provider_model"
 
 
 class ProviderCredentialsCache:

api/core/helper/module_import_helper.py

@@ -0,0 +1,62 @@
+import importlib.util
+import logging
+import sys
+from types import ModuleType
+from typing import AnyStr
+
+
+def import_module_from_source(
+        module_name: str,
+        py_file_path: AnyStr,
+        use_lazy_loader: bool = False
+) -> ModuleType:
+    """
+    Importing a module from the source file directly
+    """
+    try:
+        existed_spec = importlib.util.find_spec(module_name)
+        if existed_spec:
+            spec = existed_spec
+        else:
+            # Refer to: https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly
+            spec = importlib.util.spec_from_file_location(module_name, py_file_path)
+            if use_lazy_loader:
+                # Refer to: https://docs.python.org/3/library/importlib.html#implementing-lazy-imports
+                spec.loader = importlib.util.LazyLoader(spec.loader)
+        module = importlib.util.module_from_spec(spec)
+        if not existed_spec:
+            sys.modules[module_name] = module
+        spec.loader.exec_module(module)
+        return module
+    except Exception as e:
+        logging.exception(f'Failed to load module {module_name} from {py_file_path}: {str(e)}')
+        raise e
+
+
+def get_subclasses_from_module(mod: ModuleType, parent_type: type) -> list[type]:
+    """
+    Get all the subclasses of the parent type from the module
+    """
+    classes = [x for _, x in vars(mod).items()
+               if isinstance(x, type) and x != parent_type and issubclass(x, parent_type)]
+    return classes
+
+
+def load_single_subclass_from_source(
+        module_name: str,
+        script_path: AnyStr,
+        parent_type: type,
+        use_lazy_loader: bool = False,
+) -> type:
+    """
+    Load a single subclass from the source
+    """
+    module = import_module_from_source(module_name, script_path, use_lazy_loader)
+    subclasses = get_subclasses_from_module(module, parent_type)
+    match len(subclasses):
+        case 1:
+            return subclasses[0]
+        case 0:
+            raise Exception(f'Missing subclass of {parent_type.__name__} in {script_path}')
+        case _:
+            raise Exception(f'Multiple subclasses of {parent_type.__name__} in {script_path}')

api/core/helper/position_helper.py

@@ -0,0 +1,63 @@
+import os
+from collections import OrderedDict
+from collections.abc import Callable
+from typing import Any, AnyStr
+
+from core.tools.utils.yaml_utils import load_yaml_file
+
+
+def get_position_map(
+        folder_path: AnyStr,
+        file_name: str = '_position.yaml',
+) -> dict[str, int]:
+    """
+    Get the mapping from name to index from a YAML file
+    :param folder_path:
+    :param file_name: the YAML file name, default to '_position.yaml'
+    :return: a dict with name as key and index as value
+    """
+    position_file_name = os.path.join(folder_path, file_name)
+    positions = load_yaml_file(position_file_name, ignore_error=True)
+    position_map = {}
+    index = 0
+    for _, name in enumerate(positions):
+        if name and isinstance(name, str):
+            position_map[name.strip()] = index
+            index += 1
+    return position_map
+
+
+def sort_by_position_map(
+        position_map: dict[str, int],
+        data: list[Any],
+        name_func: Callable[[Any], str],
+) -> list[Any]:
+    """
+    Sort the objects by the position map.
+    If the name of the object is not in the position map, it will be put at the end.
+    :param position_map: the map holding positions in the form of {name: index}
+    :param name_func: the function to get the name of the object
+    :param data: the data to be sorted
+    :return: the sorted objects
+    """
+    if not position_map or not data:
+        return data
+
+    return sorted(data, key=lambda x: position_map.get(name_func(x), float('inf')))
+
+
+def sort_to_dict_by_position_map(
+        position_map: dict[str, int],
+        data: list[Any],
+        name_func: Callable[[Any], str],
+) -> OrderedDict[str, Any]:
+    """
+    Sort the objects into a ordered dict by the position map.
+    If the name of the object is not in the position map, it will be put at the end.
+    :param position_map: the map holding positions in the form of {name: index}
+    :param name_func: the function to get the name of the object
+    :param data: the data to be sorted
+    :return: an OrderedDict with the sorted pairs of name and object
+    """
+    sorted_items = sort_by_position_map(position_map, data, name_func)
+    return OrderedDict([(name_func(item), item) for item in sorted_items])