Merge branch 'features/lb'

Author: Paweł Kędzia

llm_router_api/base/lb/balanced.py

@@ -1,4 +1,4 @@
-from typing import List, Dict
+from typing import List, Dict, Optional, Any
 from collections import defaultdict
 
 from llm_router_api.base.lb.strategy import ChooseProviderStrategyI
@@ -13,7 +13,12 @@ def __init__(self, models_config_path: str) -> None:
             lambda: defaultdict(int)
         )
 
-    def get_provider(self, model_name: str, providers: List[Dict]) -> Dict:
+    def get_provider(
+        self,
+        model_name: str,
+        providers: List[Dict],
+        options: Optional[Dict[str, Any]] = None,
+    ) -> Dict:
         if not providers:
             raise ValueError(f"No providers configured for model '{model_name}'")
 

llm_router_api/base/lb/chooser.py

@@ -128,7 +128,9 @@ def __strategy_from_name(
 
         return _cls(models_config_path=models_config_path)
 
-    def get_provider(self, model_name: str, providers: List[Dict]) -> Dict:
+    def get_provider(
+        self, model_name: str, providers: List[Dict], options: Optional[Dict] = None
+    ) -> Dict:
         """
         Choose a provider for *model_name* from *providers* using the configured strategy.
 
@@ -141,6 +143,8 @@ def get_provider(self, model_name: str, providers: List[Dict]) -> Dict:
             The name of the model for which a provider is required.
         providers : List[Dict]
             A list of provider configuration dictionaries.
+        options: Optional[Dict], Default is ``None``.
+            Additional options to pass to ``self.strategy.choose``.
 
         Returns
         -------
@@ -154,7 +158,13 @@ def get_provider(self, model_name: str, providers: List[Dict]) -> Dict:
         """
         if not providers:
             raise RuntimeError(f"{model_name} does not have any providers!")
-        return self.strategy.get_provider(model_name, providers)
+        return self.strategy.get_provider(
+            model_name=model_name, providers=providers, options=options
+        )
 
-    def put_provider(self, model_name: str, provider: Dict) -> None:
-        self.strategy.put_provider(model_name, provider)
+    def put_provider(
+        self, model_name: str, provider: Dict, options: Optional[Dict] = None
+    ) -> None:
+        self.strategy.put_provider(
+            model_name=model_name, provider=provider, options=options
+        )

llm_router_api/base/lb/first_available.py

@@ -22,6 +22,7 @@
 
 """
 
+import random
 import time
 
 try:
@@ -31,7 +32,7 @@
 except ImportError:
     REDIS_IS_AVAILABLE = False
 
-from typing import List, Dict
+from typing import List, Dict, Optional, Any
 
 from llm_router_api.base.constants import REDIS_PORT, REDIS_HOST
 from llm_router_api.base.lb.strategy import ChooseProviderStrategyI
@@ -120,37 +121,58 @@ def __init__(
 
         self._clear_buffers()
 
-    def get_provider(self, model_name: str, providers: List[Dict]) -> Dict:
+    def get_provider(
+        self,
+        model_name: str,
+        providers: List[Dict],
+        options: Optional[Dict[str, Any]] = None,
+    ) -> Dict | None:
         """
-        Acquire the first available provider for *model_name*.
+        Acquire a provider for *model_name* from the supplied ``providers`` list.
 
-        The method repeatedly attempts to acquire a lock on each provider in the
-        order supplied by *providers*.  If a provider is successfully marked as
-        chosen in Redis, the provider dictionary is returned with an additional
-        ``"__chosen_field"`` entry that records the Redis hash field used for the
-        lock.  The call blocks until a provider is obtained or *self.timeout*
-        seconds have elapsed, in which case a :class:`TimeoutError` is raised.
+        The method attempts to lock a provider using a Redis‑backed
+        atomic Lua script.  If ``options`` contains ``{\"random_choice\": True}``,
+        the selection is performed on a shuffled copy of ``providers``; otherwise
+        providers are examined in the order they appear in the list.
 
         Parameters
         ----------
         model_name : str
-            The name of the model for which a provider is required.
+            Identifier of the model for which a provider is required.
         providers : List[Dict]
-            A list of provider configuration dictionaries.
+            A list of provider configuration dictionaries.  Each dictionary must
+            contain the information required by :meth:`_provider_field` to build a
+            unique Redis hash field name.
+        options : dict, optional
+            Additional flags that influence the acquisition strategy.  Currently
+            supported keys:
+            ``random_choice`` (bool) – when ``True`` the provider is chosen at
+            random; defaults to ``False``.
 
         Returns
         -------
-        Dict
-            The selected provider configuration, augmented with a
-            ``"__chosen_field"`` key.
+        dict | None
+            The chosen provider dictionary with an extra ``"__chosen_field"``
+            entry indicating the Redis hash field that was locked.  Returns
+            ``None`` if ``providers`` is empty.
 
         Raises
         ------
         TimeoutError
-            If no provider becomes available within the configured timeout.
-        RuntimeError
-            If Redis is not available.
+            Raised when no provider can be locked within the ``timeout`` period
+            configured for the strategy instance.
+
+        Notes
+        -----
+        * The method creates the Redis hash (``model:<model_name>``) and initial
+          ``false`` fields if they do not already exist.
+        * The lock is represented by the value ``'true'`` in the hash field.
+        * Call :meth:`put_provider` to release the lock once the provider is no
+          longer needed.
         """
+        if not providers:
+            return None
+
         redis_key = self._get_redis_key(model_name)
         start_time = time.time()
 
@@ -159,29 +181,47 @@ def get_provider(self, model_name: str, providers: List[Dict]) -> Dict:
             for p in providers:
                 self.redis_client.hset(redis_key, self._provider_field(p), "false")
 
+        # self._print_provider_status(redis_key, providers)
+
+        is_random = options and options.get("random_choice", False)
+
         while True:
             if time.time() - start_time > self.timeout:
                 raise TimeoutError(
                     f"No available provider found for model '{model_name}' "
                     f"within {self.timeout} seconds"
                 )
 
-            for provider in providers:
-                provider_field = self._provider_field(provider)
-                try:
-                    ok = int(
-                        self._acquire_script(keys=[redis_key], args=[provider_field])
-                    )
-                    if ok == 1:
-                        provider["__chosen_field"] = provider_field
-                        return provider
-                except Exception:
-                    time.sleep(self.check_interval)
-                    continue
-
+            if is_random:
+                provider = self._try_acquire_random_provider(
+                    redis_key=redis_key, providers=providers
+                )
+                if provider:
+                    return provider
+            else:
+                for provider in providers:
+                    provider_field = self._provider_field(provider)
+                    try:
+                        ok = int(
+                            self._acquire_script(
+                                keys=[redis_key], args=[provider_field]
+                            )
+                        )
+                        if ok == 1:
+                            provider["__chosen_field"] = provider_field
+                            return provider
+                    except Exception:
+                        time.sleep(self.check_interval)
+                        continue
+            # -------------------------------------------------------------
             time.sleep(self.check_interval)
 
-    def put_provider(self, model_name: str, provider: Dict) -> None:
+    def put_provider(
+        self,
+        model_name: str,
+        provider: Dict,
+        options: Optional[Dict[str, Any]] = None,
+    ) -> None:
         """
         Release a previously acquired provider back to the pool.
 
@@ -196,6 +236,8 @@ def put_provider(self, model_name: str, provider: Dict) -> None:
             The model name associated with the provider.
         provider : Dict
             The provider dictionary that was returned by :meth:`get_provider`.
+        options: Dict[str, Any], default: None
+            Additional options passed to the chosen provider.
         """
         redis_key = self._get_redis_key(model_name)
         provider_field = self._provider_field(provider)
@@ -206,6 +248,78 @@ def put_provider(self, model_name: str, provider: Dict) -> None:
 
         provider.pop("__chosen_field", None)
 
+    def _try_acquire_random_provider(
+        self, redis_key: str, providers: List[Dict]
+    ) -> Optional[Dict]:
+        """
+        Attempt to lock a provider chosen at random.
+
+        The method works in three stages:
+
+        1. **Shuffle** – a shallow copy of ``providers`` is shuffled so that each
+           provider has an equal probability of being tried first.  The original
+           list is left untouched.
+        2. **Atomic acquisition** – each shuffled provider is passed to the
+           ``_acquire_script`` Lua script which atomically sets the corresponding
+           Redis hash field to ``'true'`` *only if* it is currently ``'false'`` or
+           missing.  The first provider for which the script returns ``1`` is
+           considered successfully acquired.
+        3. **Fallback** – if none of the providers can be locked (e.g., all are
+           currently in use), the method falls back to the *first* provider in the
+           original ``providers`` list, marks its ``"__chosen_field"`` for
+           consistency, and returns it.  This fallback mirrors the behaviour of
+           the non‑random acquisition path and ensures the caller always receives
+           a provider dictionary (or ``None`` when ``providers`` is empty).
+
+        Parameters
+        ----------
+        redis_key : str
+            The Redis hash key associated with the model (e.g., ``model:<name>``).
+        providers : List[Dict]
+            A list of provider configuration dictionaries.  Each dictionary must
+            contain sufficient information for :meth:`_provider_field` to generate
+            a unique field name within the Redis hash.
+
+        Returns
+        -------
+        Optional[Dict]
+            The selected provider dictionary with an additional ``"__chosen_field"``
+            entry indicating the Redis hash field that was locked.  Returns ``None``
+            only when the input ``providers`` list is empty.
+
+        Raises
+        ------
+        Exception
+            Propagates any unexpected exceptions raised by the Lua script execution;
+            callers may catch these to implement retry or logging logic.
+
+        Notes
+        -----
+        * The random selection is *non‑deterministic* on each call; however, the
+          fallback to the first provider ensures deterministic behaviour when
+          all providers are currently busy.
+        * The method does **not** block; it returns immediately after trying all
+          shuffled providers.
+        """
+        shuffled = providers[:]
+        random.shuffle(shuffled)
+        for provider in shuffled:
+            provider_field = self._provider_field(provider)
+            try:
+                ok = int(
+                    self._acquire_script(keys=[redis_key], args=[provider_field])
+                )
+                if ok == 1:
+                    provider["__chosen_field"] = provider_field
+                    return provider
+            except Exception:
+                continue
+
+        provider = providers[0]
+        provider_field = self._provider_field(provider)
+        provider["__chosen_field"] = provider_field
+        return provider
+
     def _get_redis_key(self, model_name: str) -> str:
         """
         Return Redis key prefix for a given model.
@@ -319,3 +433,31 @@ def _clear_buffers(self) -> None:
                 self._initialize_providers(
                     model_name=model_name, providers=providers
                 )
+
+    def _print_provider_status(self, redis_key: str, providers: List[Dict]) -> None:
+        """
+        Print the lock status of each provider stored in the Redis hash
+        ``redis_key``.  Uses emojis for a quick visual cue:
+
+        * 🟢 – provider is free (`'false'` or missing)
+        * 🔴 – provider is currently taken (`'true'`)
+
+        The output is formatted in a table‑like layout for readability.
+        """
+        try:
+            # Retrieve the entire hash; missing fields default to None
+            hash_data = self.redis_client.hgetall(redis_key)
+        except Exception as exc:
+            print(f"[⚠️] Could not read Redis key '{redis_key}': {exc}")
+            return
+
+        print("\nProvider lock status:")
+        print("-" * 40)
+        for provider in providers:
+            field = self._provider_field(provider)
+            status = hash_data.get(field, "false")
+            icon = "🔴" if status == "true" else "🟢"
+            # Show a short identifier for the provider (fallback to field)
+            provider_id = provider.get("id") or provider.get("name") or field
+            print(f"{icon}  {provider_id:<30} [{field}]")
+        print("-" * 40)

llm_router_api/base/lb/strategy.py

@@ -1,4 +1,4 @@
-from typing import List, Dict
+from typing import List, Dict, Any, Optional
 from abc import ABC, abstractmethod
 
 from llm_router_api.base.model_config import ApiModelConfig
@@ -34,7 +34,12 @@ def _provider_key(self, provider_cfg: Dict) -> str:
         return _pk
 
     @abstractmethod
-    def get_provider(self, model_name: str, providers: List[Dict]) -> Dict:
+    def get_provider(
+        self,
+        model_name: str,
+        providers: List[Dict],
+        options: Optional[Dict[str, Any]] = None,
+    ) -> Dict:
         """
         Choose a provider for the given model.
 
@@ -44,6 +49,8 @@ def get_provider(self, model_name: str, providers: List[Dict]) -> Dict:
             Name of the model for which a provider is required.
         providers: List[Dict]
             List of provider configuration dictionaries.
+        options: Dict[str, Any], default: None
+            Options passed to the chosen provider.
 
         Returns
         -------
@@ -52,7 +59,12 @@ def get_provider(self, model_name: str, providers: List[Dict]) -> Dict:
         """
         raise NotImplementedError
 
-    def put_provider(self, model_name: str, provider: Dict) -> None:
+    def put_provider(
+        self,
+        model_name: str,
+        provider: Dict,
+        options: Optional[Dict[str, Any]] = None,
+    ) -> None:
         """
         Notify the strategy that a provider has been used.
 
@@ -70,5 +82,7 @@ def put_provider(self, model_name: str, provider: Dict) -> None:
             Name of the model for which the provider was used.
         provider : Dict
             The provider configuration dictionary that was used.
+        options: Dict[str, Any], default: None
+            Options passed to the chosen provider.
         """
         pass