Update app.py

app.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 """
 Multi-Model AI API — HuggingFace Spaces Edition
-Unified API gateway for multiple AI models via Hugging Face Spaces.
+With load balancing (multiple provider instances per model) and 10 req/s rate limiting.
 """
 
 import re, os, json, uuid, time, random, string, logging, threading
@@ -23,7 +23,7 @@ except ImportError:
 #  CONFIG & CONSTANTS
 # ═══════════════════════════════════════════════════════════════
 
-VERSION = "2.2.0-hf"
+VERSION = "2.3.0-hf-lb"
 APP_NAME = "Multi-Model-AI-API"
 DEFAULT_SYSTEM_PROMPT = "You are a helpful, friendly AI assistant."
 DEFAULT_MODEL = "gpt-oss-120b"
@@ -34,6 +34,8 @@ log = logging.getLogger(APP_NAME)
 USER_AGENTS = [
     "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 Chrome/144.0.0.0 Safari/537.36",
     "Mozilla/5.0 (Macintosh; Intel Mac OS X 14_5) AppleWebKit/605.1.15 Safari/605.1.15",
+    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 Chrome/143.0.0.0 Safari/537.36",
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:128.0) Gecko/20100101 Firefox/128.0",
 ]
 
 # ═══════════════════════════════════════════════════════════════
@@ -61,12 +63,17 @@ class ModelDef:
     api_name: Optional[str] = None
     extra_params: Dict[str, Any] = field(default_factory=dict)
     clean_analysis: bool = False
+    # Load balancing config per model
+    lb_pool_size: int = 2  # number of provider instances for load balancing
+    lb_enabled: bool = True  # whether load balancing is enabled
 
 MODEL_REGISTRY: Dict[str, ModelDef] = {}
 
+
 def register_model(m: ModelDef):
     MODEL_REGISTRY[m.model_id] = m
 
+
 def _init_registry():
     register_model(ModelDef(
         model_id="gpt-oss-120b", display_name="AMD GPT-OSS-120B",
@@ -74,6 +81,7 @@ def _init_registry():
         owned_by="amd", description="AMD open-source 120B model",
         fn_index=8, clean_analysis=True, default_temperature=0.0,
         supports_vision=False, supports_thinking=False,
+        lb_pool_size=3, lb_enabled=True,
     ))
     register_model(ModelDef(
         model_id="command-a-vision", display_name="Cohere Command-A Vision",
@@ -83,6 +91,7 @@ def _init_registry():
         supports_temperature=False, supports_streaming=False, supports_history=False,
         supports_thinking=False, max_tokens_default=700,
         extra_params={"max_new_tokens": 700},
+        lb_pool_size=2, lb_enabled=True,
     ))
     register_model(ModelDef(
         model_id="command-a-translate", display_name="Cohere Command-A Translate",
@@ -92,6 +101,7 @@ def _init_registry():
         supports_temperature=False, supports_streaming=False, supports_history=False,
         supports_thinking=False, max_tokens_default=700,
         extra_params={"max_new_tokens": 700},
+        lb_pool_size=1, lb_enabled=False,  # NO load balancing for translate
     ))
     register_model(ModelDef(
         model_id="minimax-vl-01", display_name="MiniMax VL-01",
@@ -101,6 +111,7 @@ def _init_registry():
         supports_temperature=True, supports_streaming=False, supports_history=False,
         supports_thinking=False, max_tokens_default=12800, default_temperature=0.1,
         extra_params={"max_tokens": 12800, "top_p": 0.9},
+        lb_pool_size=2, lb_enabled=True,
     ))
     register_model(ModelDef(
         model_id="glm-4.5", display_name="GLM-4.5 (ZhipuAI)",
@@ -110,6 +121,7 @@ def _init_registry():
         supports_temperature=True, supports_streaming=False, supports_history=False,
         supports_thinking=True, thinking_default=True, default_temperature=1.0,
         extra_params={"thinking_enabled": True},
+        lb_pool_size=2, lb_enabled=True,
     ))
     register_model(ModelDef(
         model_id="chatgpt", display_name="ChatGPT (Community)",
@@ -119,6 +131,7 @@ def _init_registry():
         supports_temperature=True, supports_streaming=False, supports_history=True,
         supports_thinking=False, default_temperature=1.0,
         extra_params={"top_p": 1.0},
+        lb_pool_size=2, lb_enabled=True,
     ))
     register_model(ModelDef(
         model_id="qwen3-vl", display_name="Qwen3-VL (Alibaba)",
@@ -127,8 +140,10 @@ def _init_registry():
         api_name="/add_message", supports_vision=True, supports_system_prompt=False,
         supports_temperature=False, supports_streaming=False, supports_history=False,
         supports_thinking=False, max_tokens_default=4096,
+        lb_pool_size=2, lb_enabled=True,
     ))
 
+
 _init_registry()
 
 # ═══════════════════════════════════════════════════════════════
@@ -143,8 +158,8 @@ class Config:
     max_retries: int = 3
     retry_backoff_base: float = 1.5
     retry_jitter: float = 0.5
-    rate_limit_rpm: int = 10
-    rate_limit_burst: int = 3
+    rate_limit_rps: int = 10  # requests per SECOND (changed from RPM)
+    rate_limit_burst: int = 15  # burst capacity
     pool_size: int = 2
     max_history_messages: int = 50
     max_message_length: int = 10000
@@ -158,12 +173,14 @@ class Config:
         env_map = {
             "MMAI_TIMEOUT": ("timeout_stream", int),
             "MMAI_MAX_RETRIES": ("max_retries", int),
-            "MMAI_RATE_LIMIT": ("rate_limit_rpm", int),
+            "MMAI_RATE_LIMIT_RPS": ("rate_limit_rps", int),
+            "MMAI_RATE_LIMIT_BURST": ("rate_limit_burst", int),
             "MMAI_POOL_SIZE": ("pool_size", int),
             "MMAI_SYSTEM_PROMPT": ("default_system_prompt", str),
             "MMAI_TEMPERATURE": ("default_temperature", float),
             "MMAI_DEFAULT_MODEL": ("default_model", str),
-            "MMAI_INCLUDE_THINKING": ("include_thinking", lambda x: x.lower() in ("1", "true")),
+            "MMAI_INCLUDE_THINKING": ("include_thinking",
+                                      lambda x: x.lower() in ("1", "true")),
         }
         for env_key, (attr, conv) in env_map.items():
             val = os.environ.get(env_key)
@@ -183,12 +200,17 @@ class APIError(Exception):
         super().__init__(message)
         self.code = code
         self.status = status
+
     def to_dict(self):
         return {"error": str(self), "code": self.code}
 
+
 class ModelNotFoundError(APIError):
     def __init__(self, model_id: str):
-        super().__init__(f"Model '{model_id}' not found. Available: {list(MODEL_REGISTRY.keys())}", "MODEL_NOT_FOUND", 404)
+        super().__init__(
+            f"Model '{model_id}' not found. Available: {list(MODEL_REGISTRY.keys())}",
+            "MODEL_NOT_FOUND", 404,
+        )
 
 # ═══════════════════════════════════════════════════════════════
 #  RESPONSE CLEANER
@@ -232,7 +254,8 @@ class ResponseCleaner:
         }
         for entity, char in entities.items():
             text = text.replace(entity, char)
-        text = re.sub(r'&#x([0-9a-fA-F]+);', lambda m: chr(int(m.group(1), 16)), text)
+        text = re.sub(r'&#x([0-9a-fA-F]+);',
+                      lambda m: chr(int(m.group(1), 16)), text)
         text = re.sub(r'&#(\d+);', lambda m: chr(int(m.group(1))), text)
         return text
 
@@ -249,12 +272,25 @@ class ResponseCleaner:
         if '<details' not in text and '<div' not in text:
             return text.strip()
         thinking_text = ""
-        thinking_match = re.search(r'<details[^>]*>.*?<div[^>]*>(.*?)</div>\s*</details>', text, re.DOTALL | re.IGNORECASE)
+        thinking_match = re.search(
+            r'<details[^>]*>.*?<div[^>]*>(.*?)</div>\s*</details>',
+            text, re.DOTALL | re.IGNORECASE,
+        )
         if thinking_match:
             thinking_text = cls._strip_html(thinking_match.group(1)).strip()
-        text_without_details = re.sub(r'<details[^>]*>.*?</details>', '', text, flags=re.DOTALL | re.IGNORECASE).strip()
-        div_match = re.search(r"<div[^>]*>\s*(.*?)\s*</div>", text_without_details, re.DOTALL | re.IGNORECASE)
-        response_text = cls._strip_html(div_match.group(1)).strip() if div_match else cls._strip_html(text_without_details).strip()
+        text_without_details = re.sub(
+            r'<details[^>]*>.*?</details>', '', text,
+            flags=re.DOTALL | re.IGNORECASE,
+        ).strip()
+        div_match = re.search(
+            r"<div[^>]*>\s*(.*?)\s*</div>",
+            text_without_details, re.DOTALL | re.IGNORECASE,
+        )
+        response_text = (
+            cls._strip_html(div_match.group(1)).strip()
+            if div_match
+            else cls._strip_html(text_without_details).strip()
+        )
         if thinking_text and include_thinking:
             return f"<thinking>\n{thinking_text}\n</thinking>\n{response_text}"
         return response_text
@@ -307,7 +343,8 @@ class ResponseCleaner:
         return str(result)
 
     @classmethod
-    def clean(cls, text: str, model_id: str = "", include_thinking: bool = True) -> str:
+    def clean(cls, text: str, model_id: str = "",
+              include_thinking: bool = True) -> str:
         if not text:
             return text
         text = text.strip()
@@ -326,7 +363,10 @@ class ResponseCleaner:
 class ThinkingParser:
     @staticmethod
     def split(text: str) -> Tuple[Optional[str], str]:
-        match = re.match(r'\s*<thinking>\s*\n?(.*?)\n?\s*</thinking>\s*\n?(.*)', text, re.DOTALL | re.IGNORECASE)
+        match = re.match(
+            r'\s*<thinking>\s*\n?(.*?)\n?\s*</thinking>\s*\n?(.*)',
+            text, re.DOTALL | re.IGNORECASE,
+        )
         if match:
             thinking = match.group(1).strip()
             response = match.group(2).strip()
@@ -351,6 +391,7 @@ class Message:
     timestamp: float = field(default_factory=time.time)
     message_id: str = field(default_factory=lambda: str(uuid.uuid4()))
 
+
 @dataclass
 class Conversation:
     conversation_id: str = field(default_factory=lambda: str(uuid.uuid4()))
@@ -361,7 +402,9 @@ class Conversation:
     system_prompt: str = DEFAULT_SYSTEM_PROMPT
     model_id: str = DEFAULT_MODEL
 
-    def add_message(self, role: str, content: str, max_messages: int = 50, thinking: Optional[str] = None) -> Message:
+    def add_message(self, role: str, content: str,
+                    max_messages: int = 50,
+                    thinking: Optional[str] = None) -> Message:
         msg = Message(role=role, content=content, thinking=thinking)
         self.messages.append(msg)
         self.updated_at = time.time()
@@ -378,7 +421,9 @@ class Conversation:
         non_system = [m for m in self.messages if m.role != "system"]
         i = 0
         while i < len(non_system) - 1:
-            if non_system[i].role == "user" and i + 1 < len(non_system) and non_system[i + 1].role == "assistant":
+            if (non_system[i].role == "user"
+                    and i + 1 < len(non_system)
+                    and non_system[i + 1].role == "assistant"):
                 history.append([non_system[i].content, non_system[i + 1].content])
                 i += 2
             else:
@@ -390,13 +435,16 @@ class Conversation:
 
     def to_dict(self) -> Dict:
         return {
-            "conversation_id": self.conversation_id, "title": self.title,
-            "model": self.model_id, "message_count": len(self.messages),
-            "created_at": self.created_at, "updated_at": self.updated_at,
+            "conversation_id": self.conversation_id,
+            "title": self.title,
+            "model": self.model_id,
+            "message_count": len(self.messages),
+            "created_at": self.created_at,
+            "updated_at": self.updated_at,
         }
 
 # ═══════════════════════════════════════════════════════════════
-#  METRICS & RATE LIMITER
+#  METRICS
 # ═══════════════════════════════════════════════════════════════
 
 @dataclass
@@ -411,8 +459,12 @@ class Metrics:
     requests_per_model: Dict[str, int] = field(default_factory=dict)
     _latencies: deque = field(default_factory=lambda: deque(maxlen=1000), repr=False)
     started_at: float = field(default_factory=time.time)
+    # Load balancer metrics
+    lb_total_dispatches: int = 0
+    lb_failovers: int = 0
 
-    def record_request(self, success: bool, duration_ms: float, chars: int = 0, model: str = ""):
+    def record_request(self, success: bool, duration_ms: float,
+                       chars: int = 0, model: str = ""):
         with self._lock:
             self.total_requests += 1
             if success:
@@ -422,48 +474,85 @@ class Metrics:
                 self.failed_requests += 1
             self._latencies.append(duration_ms)
             if model:
-                self.requests_per_model[model] = self.requests_per_model.get(model, 0) + 1
+                self.requests_per_model[model] = (
+                    self.requests_per_model.get(model, 0) + 1
+                )
 
     def record_retry(self):
         with self._lock:
             self.total_retries += 1
 
+    def record_lb_dispatch(self, failover: bool = False):
+        with self._lock:
+            self.lb_total_dispatches += 1
+            if failover:
+                self.lb_failovers += 1
+
     def to_dict(self) -> Dict:
         with self._lock:
-            avg = sum(self._latencies) / len(self._latencies) if self._latencies else 0
-            rate = self.successful_requests / self.total_requests if self.total_requests else 1
+            avg = (sum(self._latencies) / len(self._latencies)
+                   if self._latencies else 0)
+            rate = (self.successful_requests / self.total_requests
+                    if self.total_requests else 1)
             return {
-                "total_requests": self.total_requests, "successful": self.successful_requests,
-                "failed": self.failed_requests, "success_rate": round(rate, 4),
-                "retries": self.total_retries, "chars_received": self.total_chars_received,
-                "avg_latency_ms": round(avg, 1), "active_streams": self.active_streams,
+                "total_requests": self.total_requests,
+                "successful": self.successful_requests,
+                "failed": self.failed_requests,
+                "success_rate": round(rate, 4),
+                "retries": self.total_retries,
+                "chars_received": self.total_chars_received,
+                "avg_latency_ms": round(avg, 1),
+                "active_streams": self.active_streams,
                 "uptime_s": round(time.time() - self.started_at, 1),
                 "per_model": dict(self.requests_per_model),
+                "load_balancer": {
+                    "total_dispatches": self.lb_total_dispatches,
+                    "failovers": self.lb_failovers,
+                },
             }
 
+
 metrics = Metrics()
 
+# ═══════════════════════════════════════════════════════════════
+#  RATE LIMITER — 10 requests per SECOND (token bucket)
+# ═══════════════════════════════════════════════════════════════
+
 class RateLimiter:
-    def __init__(self, rpm: int = 10, burst: int = 3):
-        self.rate = rpm / 60.0
+    """Token-bucket rate limiter. Default: 10 requests/second with burst."""
+
+    def __init__(self, rps: int = 10, burst: int = 15):
+        self.rate = float(rps)           # tokens per second
         self.max_tokens = float(burst)
         self.tokens = float(burst)
         self.last_refill = time.monotonic()
         self._lock = threading.Lock()
 
-    def acquire(self, timeout: float = 30.0) -> bool:
+    def acquire(self, timeout: float = 10.0) -> bool:
         deadline = time.monotonic() + timeout
         while True:
             with self._lock:
                 now = time.monotonic()
-                self.tokens = min(self.max_tokens, self.tokens + (now - self.last_refill) * self.rate)
+                elapsed = now - self.last_refill
+                self.tokens = min(
+                    self.max_tokens,
+                    self.tokens + elapsed * self.rate,
+                )
                 self.last_refill = now
                 if self.tokens >= 1.0:
                     self.tokens -= 1.0
                     return True
             if time.monotonic() >= deadline:
                 return False
-            time.sleep(0.1)
+            time.sleep(0.05)  # short sleep for per-second limiting
+
+    def get_info(self) -> Dict:
+        with self._lock:
+            return {
+                "rate_rps": self.rate,
+                "burst": self.max_tokens,
+                "available_tokens": round(self.tokens, 2),
+            }
 
 # ═══════════════════════════════════════════════════════════════
 #  CIRCUIT BREAKER
@@ -514,7 +603,8 @@ class CircuitBreaker:
 
 class GradioSSEParser:
     @staticmethod
-    def parse_sse(response: requests.Response, log_raw: bool = False) -> Generator[Dict, None, None]:
+    def parse_sse(response: requests.Response,
+                  log_raw: bool = False) -> Generator[Dict, None, None]:
         buffer = ""
         for chunk in response.iter_content(chunk_size=None, decode_unicode=True):
             if chunk is None:
@@ -554,11 +644,19 @@ class GradioSSEParser:
 # ═══════════════════════════════════════════════════════════════
 
 class ModelProvider(ABC):
-    def __init__(self, model_def: ModelDef, config: Config):
+    def __init__(self, model_def: ModelDef, config: Config, instance_id: int = 0):
         self.model_def = model_def
         self.config = config
+        self.instance_id = instance_id
         self.ready = False
         self._lock = threading.Lock()
+        # Per-instance health tracking
+        self._consecutive_failures = 0
+        self._last_success_time = 0.0
+        self._last_failure_time = 0.0
+        self._total_requests = 0
+        self._total_failures = 0
+        self._latencies: deque = deque(maxlen=50)
 
     @abstractmethod
     def initialize(self) -> bool: ...
@@ -570,10 +668,58 @@ class ModelProvider(ABC):
     def generate_stream(self, message: str, **kwargs) -> Generator[str, None, None]:
         yield self.generate(message, **kwargs)
 
+    def record_success(self, latency_ms: float):
+        self._consecutive_failures = 0
+        self._last_success_time = time.time()
+        self._total_requests += 1
+        self._latencies.append(latency_ms)
+
+    def record_failure(self):
+        self._consecutive_failures += 1
+        self._last_failure_time = time.time()
+        self._total_requests += 1
+        self._total_failures += 1
+
+    @property
+    def avg_latency(self) -> float:
+        return sum(self._latencies) / len(self._latencies) if self._latencies else 0.0
+
+    @property
+    def health_score(self) -> float:
+        """0.0 (worst) to 1.0 (best). Used by load balancer to pick instance."""
+        if not self.ready:
+            return 0.0
+        score = 1.0
+        # Penalise consecutive failures
+        score -= min(self._consecutive_failures * 0.2, 0.8)
+        # Penalise high avg latency (>10s = bad)
+        if self._latencies:
+            avg = self.avg_latency
+            if avg > 10000:
+                score -= 0.3
+            elif avg > 5000:
+                score -= 0.15
+        # Penalise high failure rate
+        if self._total_requests > 5:
+            fail_rate = self._total_failures / self._total_requests
+            score -= fail_rate * 0.4
+        return max(0.0, min(1.0, score))
+
+    def get_instance_info(self) -> Dict:
+        return {
+            "instance_id": self.instance_id,
+            "ready": self.ready,
+            "health_score": round(self.health_score, 3),
+            "consecutive_failures": self._consecutive_failures,
+            "total_requests": self._total_requests,
+            "total_failures": self._total_failures,
+            "avg_latency_ms": round(self.avg_latency, 1),
+        }
+
 
 class GptOssProvider(ModelProvider):
-    def __init__(self, model_def, config):
-        super().__init__(model_def, config)
+    def __init__(self, model_def, config, instance_id=0):
+        super().__init__(model_def, config, instance_id)
         self._session = requests.Session()
         self._rotate()
 
@@ -594,31 +740,48 @@ class GptOssProvider(ModelProvider):
                 return True
             self._rotate()
             try:
-                r = self._session.get(f"{self.model_def.space_id}/gradio_api/info", timeout=15)
+                r = self._session.get(
+                    f"{self.model_def.space_id}/gradio_api/info", timeout=15,
+                )
                 self.ready = r.status_code == 200
                 return self.ready
-            except:
+            except Exception:
                 return False
 
-    def generate(self, message, history=None, system_prompt=None, temperature=None, max_tokens=None, **kw):
+    def generate(self, message, history=None, system_prompt=None,
+                 temperature=None, max_tokens=None, **kw):
         if not self.ready:
             self.initialize()
         sys_p = system_prompt or self.config.default_system_prompt
-        temp = temperature if temperature is not None else self.model_def.default_temperature
+        temp = (temperature if temperature is not None
+                else self.model_def.default_temperature)
         h = self._hash()
-        payload = {"data": [message, history or [], sys_p, temp], "event_data": None,
-                   "fn_index": self.model_def.fn_index, "trigger_id": None, "session_hash": h}
-        r = self._session.post(f"{self.model_def.space_id}/gradio_api/queue/join?",
-                               json=payload, headers={"Content-Type": "application/json"}, timeout=30)
+        payload = {
+            "data": [message, history or [], sys_p, temp],
+            "event_data": None,
+            "fn_index": self.model_def.fn_index,
+            "trigger_id": None,
+            "session_hash": h,
+        }
+        r = self._session.post(
+            f"{self.model_def.space_id}/gradio_api/queue/join?",
+            json=payload,
+            headers={"Content-Type": "application/json"},
+            timeout=30,
+        )
         if r.status_code != 200:
             raise APIError(f"Queue join failed: {r.status_code}")
         data = r.json()
         if not data.get("event_id"):
-            raise APIError(f"No event_id")
-
-        resp = self._session.get(f"{self.model_def.space_id}/gradio_api/queue/data",
-                                 params={"session_hash": h}, headers={"Accept": "text/event-stream"},
-                                 timeout=self.config.timeout_stream, stream=True)
+            raise APIError("No event_id")
+
+        resp = self._session.get(
+            f"{self.model_def.space_id}/gradio_api/queue/data",
+            params={"session_hash": h},
+            headers={"Accept": "text/event-stream"},
+            timeout=self.config.timeout_stream,
+            stream=True,
+        )
         full = ""
         for d in GradioSSEParser.parse_sse(resp):
             msg = d.get("msg", "")
@@ -635,21 +798,37 @@ class GptOssProvider(ModelProvider):
                 break
         if not full.strip():
             raise APIError("Empty response", "EMPTY")
-        return ResponseCleaner.clean_analysis(full) if self.model_def.clean_analysis else full
+        return (ResponseCleaner.clean_analysis(full)
+                if self.model_def.clean_analysis else full)
 
-    def generate_stream(self, message, history=None, system_prompt=None, temperature=None, max_tokens=None, **kw):
+    def generate_stream(self, message, history=None, system_prompt=None,
+                        temperature=None, max_tokens=None, **kw):
         if not self.ready:
             self.initialize()
         sys_p = system_prompt or self.config.default_system_prompt
-        temp = temperature if temperature is not None else self.model_def.default_temperature
+        temp = (temperature if temperature is not None
+                else self.model_def.default_temperature)
         h = self._hash()
-        payload = {"data": [message, history or [], sys_p, temp], "event_data": None,
-                   "fn_index": self.model_def.fn_index, "trigger_id": None, "session_hash": h}
-        self._session.post(f"{self.model_def.space_id}/gradio_api/queue/join?",
-                          json=payload, headers={"Content-Type": "application/json"}, timeout=30)
-        resp = self._session.get(f"{self.model_def.space_id}/gradio_api/queue/data",
-                                 params={"session_hash": h}, headers={"Accept": "text/event-stream"},
-                                 timeout=self.config.timeout_stream, stream=True)
+        payload = {
+            "data": [message, history or [], sys_p, temp],
+            "event_data": None,
+            "fn_index": self.model_def.fn_index,
+            "trigger_id": None,
+            "session_hash": h,
+        }
+        self._session.post(
+            f"{self.model_def.space_id}/gradio_api/queue/join?",
+            json=payload,
+            headers={"Content-Type": "application/json"},
+            timeout=30,
+        )
+        resp = self._session.get(
+            f"{self.model_def.space_id}/gradio_api/queue/data",
+            params={"session_hash": h},
+            headers={"Accept": "text/event-stream"},
+            timeout=self.config.timeout_stream,
+            stream=True,
+        )
         metrics.active_streams += 1
         last = ""
         try:
@@ -658,7 +837,7 @@ class GptOssProvider(ModelProvider):
                 if msg in ("process_generating", "process_completed"):
                     output = d.get("output", {})
                     if not output.get("success", True):
-                        raise APIError(f"Gradio error")
+                        raise APIError("Gradio error")
                     raw = GradioSSEParser.extract_text(output)
                     if raw:
                         if self.model_def.clean_analysis:
@@ -680,27 +859,35 @@ class GptOssProvider(ModelProvider):
 
 class GradioClientProvider(ModelProvider):
     """Generic provider for all gradio_client based models."""
-    def __init__(self, model_def, config):
-        super().__init__(model_def, config)
+
+    def __init__(self, model_def, config, instance_id=0):
+        super().__init__(model_def, config, instance_id)
         self._client = None
         self._chat_counter = 0
 
     def initialize(self) -> bool:
         if not HAS_GRADIO_CLIENT:
-            raise APIError(f"gradio_client not installed", "MISSING_DEP")
+            raise APIError("gradio_client not installed", "MISSING_DEP")
         with self._lock:
             if self.ready:
                 return True
             try:
-                log.info(f"Connecting to {self.model_def.space_id}...")
+                log.info(
+                    f"[Instance {self.instance_id}] Connecting to "
+                    f"{self.model_def.space_id}..."
+                )
                 self._client = GradioClient(self.model_def.space_id)
                 self.ready = True
                 return True
             except Exception as e:
-                log.error(f"Init failed for {self.model_def.model_id}: {e}")
+                log.error(
+                    f"[Instance {self.instance_id}] Init failed for "
+                    f"{self.model_def.model_id}: {e}"
+                )
                 return False
 
-    def generate(self, message, history=None, system_prompt=None, temperature=None, max_tokens=None, **kw):
+    def generate(self, message, history=None, system_prompt=None,
+                 temperature=None, max_tokens=None, **kw):
         if not self.ready:
             self.initialize()
         if not self._client:
@@ -709,50 +896,73 @@ class GradioClientProvider(ModelProvider):
         mid = self.model_def.model_id
         try:
             if mid == "command-a-vision":
-                max_new = max_tokens or self.model_def.extra_params.get("max_new_tokens", 700)
-                result = self._client.predict(message={"text": message, "files": []},
-                                              max_new_tokens=max_new, api_name=self.model_def.api_name)
+                max_new = (max_tokens
+                           or self.model_def.extra_params.get("max_new_tokens", 700))
+                result = self._client.predict(
+                    message={"text": message, "files": []},
+                    max_new_tokens=max_new,
+                    api_name=self.model_def.api_name,
+                )
             elif mid == "command-a-translate":
-                max_new = max_tokens or self.model_def.extra_params.get("max_new_tokens", 700)
-                result = self._client.predict(message=message, max_new_tokens=max_new,
-                                              api_name=self.model_def.api_name)
+                max_new = (max_tokens
+                           or self.model_def.extra_params.get("max_new_tokens", 700))
+                result = self._client.predict(
+                    message=message,
+                    max_new_tokens=max_new,
+                    api_name=self.model_def.api_name,
+                )
             elif mid == "minimax-vl-01":
-                temp = temperature if temperature is not None else self.model_def.default_temperature
-                max_tok = max_tokens or self.model_def.extra_params.get("max_tokens", 12800)
-                top_p = kw.get("top_p", self.model_def.extra_params.get("top_p", 0.9))
-                result = self._client.predict(message={"text": message, "files": []},
-                                              max_tokens=max_tok, temperature=temp, top_p=top_p,
-                                              api_name=self.model_def.api_name)
+                temp = (temperature if temperature is not None
+                        else self.model_def.default_temperature)
+                max_tok = (max_tokens
+                           or self.model_def.extra_params.get("max_tokens", 12800))
+                top_p = kw.get("top_p",
+                               self.model_def.extra_params.get("top_p", 0.9))
+                result = self._client.predict(
+                    message={"text": message, "files": []},
+                    max_tokens=max_tok, temperature=temp, top_p=top_p,
+                    api_name=self.model_def.api_name,
+                )
             elif mid == "glm-4.5":
                 sys_p = system_prompt or self.config.default_system_prompt
-                temp = temperature if temperature is not None else self.model_def.default_temperature
-                thinking = kw.get("thinking_enabled", self.model_def.thinking_default)
-                include = kw.get("include_thinking", self.config.include_thinking)
-                result = self._client.predict(msg=message, sys_prompt=sys_p,
-                                              thinking_enabled=thinking, temperature=temp,
-                                              api_name=self.model_def.api_name)
+                temp = (temperature if temperature is not None
+                        else self.model_def.default_temperature)
+                thinking = kw.get("thinking_enabled",
+                                  self.model_def.thinking_default)
+                include = kw.get("include_thinking",
+                                 self.config.include_thinking)
+                result = self._client.predict(
+                    msg=message, sys_prompt=sys_p,
+                    thinking_enabled=thinking, temperature=temp,
+                    api_name=self.model_def.api_name,
+                )
                 return self._extract_glm(result, include)
             elif mid == "chatgpt":
-                temp = temperature if temperature is not None else self.model_def.default_temperature
-                top_p = kw.get("top_p", self.model_def.extra_params.get("top_p", 1.0))
+                temp = (temperature if temperature is not None
+                        else self.model_def.default_temperature)
+                top_p = kw.get("top_p",
+                               self.model_def.extra_params.get("top_p", 1.0))
                 chat_hist = []
                 if history:
                     for pair in history:
                         if isinstance(pair, (list, tuple)) and len(pair) == 2:
                             chat_hist.append([str(pair[0]), str(pair[1])])
-                result = self._client.predict(inputs=message, top_p=top_p, temperature=temp,
-                                              chat_counter=self._chat_counter, chatbot=chat_hist,
-                                              api_name=self.model_def.api_name)
+                result = self._client.predict(
+                    inputs=message, top_p=top_p, temperature=temp,
+                    chat_counter=self._chat_counter, chatbot=chat_hist,
+                    api_name=self.model_def.api_name,
+                )
                 self._chat_counter += 1
                 return ResponseCleaner.extract_chatgpt_text(result)
             elif mid == "qwen3-vl":
-                result = self._client.predict(input_value={"files": None, "text": message},
-                                              api_name="/add_message")
+                result = self._client.predict(
+                    input_value={"files": None, "text": message},
+                    api_name="/add_message",
+                )
                 return ResponseCleaner.extract_qwen_text(result)
             else:
                 raise APIError(f"Unknown model handler: {mid}")
 
-            # Default extraction for simple results
             if isinstance(result, str):
                 return result.strip()
             if isinstance(result, dict):
@@ -786,28 +996,268 @@ class GradioClientProvider(ModelProvider):
         return ResponseCleaner.clean_glm(str(result), include_thinking)
 
 
-# Factory
-def create_provider(model_id: str, config: Config) -> ModelProvider:
+# Factory — creates a single provider instance
+def create_provider(model_id: str, config: Config,
+                    instance_id: int = 0) -> ModelProvider:
     if model_id not in MODEL_REGISTRY:
         raise ModelNotFoundError(model_id)
     mdef = MODEL_REGISTRY[model_id]
     if model_id == "gpt-oss-120b":
-        return GptOssProvider(mdef, config)
-    return GradioClientProvider(mdef, config)
+        return GptOssProvider(mdef, config, instance_id)
+    return GradioClientProvider(mdef, config, instance_id)
+
+# ═══════════════════════════════════════════════════════════════
+#  LOAD BALANCER — Per-model provider pool with health-aware
+#                  round-robin + failover
+# ═══════════════════════════════════════════════════════════════
+
+class LoadBalancedProviderPool:
+    """
+    Manages multiple provider instances for a single model.
+    Selects the best instance based on health score with
+    weighted-random selection (healthier instances chosen more).
+    Falls back through all instances on failure.
+    """
+
+    def __init__(self, model_id: str, config: Config):
+        self.model_id = model_id
+        self.config = config
+        self.mdef = MODEL_REGISTRY[model_id]
+        pool_size = self.mdef.lb_pool_size if self.mdef.lb_enabled else 1
+        self._instances: List[ModelProvider] = []
+        self._rr_index = 0
+        self._lock = threading.Lock()
+
+        for i in range(pool_size):
+            self._instances.append(create_provider(model_id, config, instance_id=i))
+
+        log.info(
+            f"[LB] Created pool for '{model_id}' with {len(self._instances)} "
+            f"instance(s), lb_enabled={self.mdef.lb_enabled}"
+        )
+
+    @property
+    def pool_size(self) -> int:
+        return len(self._instances)
+
+    def initialize_all(self) -> int:
+        """Initialize all instances, return count of successful ones."""
+        ok = 0
+        for inst in self._instances:
+            try:
+                if inst.initialize():
+                    ok += 1
+            except Exception as e:
+                log.warning(
+                    f"[LB] Failed to init {self.model_id} "
+                    f"instance {inst.instance_id}: {e}"
+                )
+        return ok
+
+    def initialize_one(self) -> bool:
+        """Initialize at least one instance."""
+        for inst in self._instances:
+            try:
+                if inst.initialize():
+                    return True
+            except Exception:
+                continue
+        return False
+
+    def _select_instance(self) -> ModelProvider:
+        """
+        Select best available instance.
+        Strategy: weighted random by health score.
+        If all have equal scores, falls back to round-robin.
+        """
+        if len(self._instances) == 1:
+            return self._instances[0]
+
+        with self._lock:
+            # Collect health scores
+            scored = []
+            for inst in self._instances:
+                score = inst.health_score
+                # Give a minimum weight so unhealthy instances can still recover
+                scored.append((inst, max(score, 0.05)))
+
+            total_weight = sum(s for _, s in scored)
+            if total_weight <= 0:
+                # All dead, just round-robin
+                inst = self._instances[self._rr_index % len(self._instances)]
+                self._rr_index += 1
+                return inst
+
+            # Weighted random selection
+            r = random.uniform(0, total_weight)
+            cumulative = 0.0
+            for inst, weight in scored:
+                cumulative += weight
+                if r <= cumulative:
+                    return inst
+
+            # Fallback
+            return scored[-1][0]
+
+    def _get_ordered_instances(self) -> List[ModelProvider]:
+        """Return instances ordered by health score (best first)."""
+        return sorted(self._instances, key=lambda p: p.health_score, reverse=True)
+
+    def execute(self, fn_name: str, **kwargs) -> Any:
+        """
+        Execute a provider method with automatic failover.
+        Tries the best instance first, fails over to others.
+        """
+        primary = self._select_instance()
+        metrics.record_lb_dispatch()
+
+        # Ensure primary is ready
+        if not primary.ready:
+            try:
+                primary.initialize()
+            except Exception:
+                pass
+
+        # Try primary
+        start = time.monotonic()
+        try:
+            result = self._call_provider(primary, fn_name, **kwargs)
+            latency = (time.monotonic() - start) * 1000
+            primary.record_success(latency)
+            return result
+        except Exception as primary_err:
+            primary.record_failure()
+            log.warning(
+                f"[LB] Primary instance {primary.instance_id} for "
+                f"'{self.model_id}' failed: {primary_err}"
+            )
+
+        # Failover through remaining instances
+        for inst in self._get_ordered_instances():
+            if inst is primary:
+                continue
+            if not inst.ready:
+                try:
+                    inst.initialize()
+                except Exception:
+                    continue
+
+            metrics.record_lb_dispatch(failover=True)
+            start = time.monotonic()
+            try:
+                result = self._call_provider(inst, fn_name, **kwargs)
+                latency = (time.monotonic() - start) * 1000
+                inst.record_success(latency)
+                log.info(
+                    f"[LB] Failover to instance {inst.instance_id} "
+                    f"for '{self.model_id}' succeeded"
+                )
+                return result
+            except Exception as e:
+                inst.record_failure()
+                log.warning(
+                    f"[LB] Failover instance {inst.instance_id} "
+                    f"for '{self.model_id}' also failed: {e}"
+                )
+
+        raise APIError(
+            f"All {len(self._instances)} instances for '{self.model_id}' failed",
+            "ALL_INSTANCES_FAILED",
+        )
+
+    def execute_stream(self, **kwargs) -> Generator[str, None, None]:
+        """
+        Execute streaming with failover.
+        Since generators can't easily be retried mid-stream,
+        we do failover only on initial connection failure.
+        """
+        primary = self._select_instance()
+        metrics.record_lb_dispatch()
+
+        if not primary.ready:
+            try:
+                primary.initialize()
+            except Exception:
+                pass
+
+        # Try primary
+        try:
+            yield from self._call_provider_stream(primary, **kwargs)
+            return
+        except Exception as primary_err:
+            primary.record_failure()
+            log.warning(
+                f"[LB] Stream primary instance {primary.instance_id} "
+                f"for '{self.model_id}' failed: {primary_err}"
+            )
+
+        # Failover
+        for inst in self._get_ordered_instances():
+            if inst is primary:
+                continue
+            if not inst.ready:
+                try:
+                    inst.initialize()
+                except Exception:
+                    continue
+
+            metrics.record_lb_dispatch(failover=True)
+            try:
+                yield from self._call_provider_stream(inst, **kwargs)
+                return
+            except Exception as e:
+                inst.record_failure()
+                log.warning(
+                    f"[LB] Stream failover instance {inst.instance_id} "
+                    f"for '{self.model_id}' failed: {e}"
+                )
+
+        raise APIError(
+            f"All streaming instances for '{self.model_id}' failed",
+            "ALL_INSTANCES_FAILED",
+        )
+
+    def _call_provider(self, provider: ModelProvider, fn_name: str,
+                       **kwargs) -> Any:
+        if not provider.ready:
+            provider.initialize()
+        fn = getattr(provider, fn_name)
+        return fn(**kwargs)
+
+    def _call_provider_stream(self, provider: ModelProvider,
+                               **kwargs) -> Generator[str, None, None]:
+        if not provider.ready:
+            provider.initialize()
+        start = time.monotonic()
+        try:
+            yield from provider.generate_stream(**kwargs)
+            latency = (time.monotonic() - start) * 1000
+            provider.record_success(latency)
+        except Exception:
+            provider.record_failure()
+            raise
+
+    def get_pool_info(self) -> Dict:
+        return {
+            "model_id": self.model_id,
+            "lb_enabled": self.mdef.lb_enabled,
+            "pool_size": len(self._instances),
+            "instances": [inst.get_instance_info() for inst in self._instances],
+        }
 
 # ═══════════════════════════════════════════════════════════════
-#  MULTI-MODEL CLIENT
+#  MULTI-MODEL CLIENT (with load balancing)
 # ═══════════════════════════════════════════════════════════════
 
 class MultiModelClient:
     def __init__(self, config: Config):
         self.config = config
-        self._providers: Dict[str, ModelProvider] = {}
+        self._lb_pools: Dict[str, LoadBalancedProviderPool] = {}
         self._lock = threading.Lock()
         self._conversations: Dict[str, Conversation] = {}
         self._active_conv_id: Optional[str] = None
         self._current_model = config.default_model
-        self.rate_limiter = RateLimiter(config.rate_limit_rpm, config.rate_limit_burst)
+        self.rate_limiter = RateLimiter(config.rate_limit_rps, config.rate_limit_burst)
         self.circuit_breaker = CircuitBreaker()
 
     @property
@@ -820,45 +1270,74 @@ class MultiModelClient:
             raise ModelNotFoundError(m)
         self._current_model = m
 
-    def _get_provider(self, model_id: str) -> ModelProvider:
-        if model_id not in self._providers:
+    def _get_lb_pool(self, model_id: str) -> LoadBalancedProviderPool:
+        if model_id not in self._lb_pools:
             with self._lock:
-                if model_id not in self._providers:
-                    self._providers[model_id] = create_provider(model_id, self.config)
-        return self._providers[model_id]
-
-    def _ensure_ready(self, model_id: str) -> ModelProvider:
-        p = self._get_provider(model_id)
-        if not p.ready:
-            if not p.initialize():
-                raise APIError(f"Cannot init {model_id}", "INIT_FAILED")
-        return p
+                if model_id not in self._lb_pools:
+                    self._lb_pools[model_id] = LoadBalancedProviderPool(
+                        model_id, self.config
+                    )
+        return self._lb_pools[model_id]
+
+    def _ensure_ready(self, model_id: str) -> LoadBalancedProviderPool:
+        pool = self._get_lb_pool(model_id)
+        # Make sure at least one instance is ready
+        has_ready = any(inst.ready for inst in pool._instances)
+        if not has_ready:
+            if not pool.initialize_one():
+                raise APIError(f"Cannot init any instance for {model_id}",
+                               "INIT_FAILED")
+        return pool
 
     @property
     def active_conversation(self) -> Conversation:
         if self._active_conv_id not in self._conversations:
-            conv = Conversation(system_prompt=self.config.default_system_prompt, model_id=self._current_model)
+            conv = Conversation(
+                system_prompt=self.config.default_system_prompt,
+                model_id=self._current_model,
+            )
             self._conversations[conv.conversation_id] = conv
             self._active_conv_id = conv.conversation_id
         return self._conversations[self._active_conv_id]
 
-    def new_conversation(self, system_prompt=None, model_id=None) -> Conversation:
-        conv = Conversation(system_prompt=system_prompt or self.config.default_system_prompt,
-                            model_id=model_id or self._current_model)
+    def new_conversation(self, system_prompt=None,
+                         model_id=None) -> Conversation:
+        conv = Conversation(
+            system_prompt=system_prompt or self.config.default_system_prompt,
+            model_id=model_id or self._current_model,
+        )
         self._conversations[conv.conversation_id] = conv
         self._active_conv_id = conv.conversation_id
         return conv
 
     def init_model(self, model_id: str) -> bool:
         try:
-            return self._get_provider(model_id).initialize()
-        except:
+            pool = self._get_lb_pool(model_id)
+            return pool.initialize_one()
+        except Exception:
             return False
 
-    def send_message(self, message: str, *, stream: bool = False, model: Optional[str] = None,
-                     conversation_id: Optional[str] = None, system_prompt: Optional[str] = None,
-                     temperature: Optional[float] = None, max_tokens: Optional[int] = None,
-                     include_thinking: Optional[bool] = None, **kwargs) -> Union[str, Generator]:
+    def init_model_all(self, model_id: str) -> int:
+        """Init all instances in the pool, return count of ready ones."""
+        try:
+            pool = self._get_lb_pool(model_id)
+            return pool.initialize_all()
+        except Exception:
+            return 0
+
+    def send_message(
+        self,
+        message: str,
+        *,
+        stream: bool = False,
+        model: Optional[str] = None,
+        conversation_id: Optional[str] = None,
+        system_prompt: Optional[str] = None,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        include_thinking: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[str, Generator]:
         model_id = model or self._current_model
         if model_id not in MODEL_REGISTRY:
             raise ModelNotFoundError(model_id)
@@ -871,9 +1350,10 @@ class MultiModelClient:
         if not self.circuit_breaker.can_execute():
             raise APIError("Circuit breaker open", "CIRCUIT_OPEN", 503)
         if not self.rate_limiter.acquire(timeout=10.0):
-            raise APIError("Rate limited", "RATE_LIMITED", 429)
+            raise APIError("Rate limited (10 req/s max)", "RATE_LIMITED", 429)
 
-        conv = self._conversations.get(conversation_id, self.active_conversation) if conversation_id else self.active_conversation
+        conv = (self._conversations.get(conversation_id, self.active_conversation)
+                if conversation_id else self.active_conversation)
         conv.model_id = model_id
         if system_prompt:
             conv.system_prompt = system_prompt
@@ -881,9 +1361,11 @@ class MultiModelClient:
         history = conv.build_gradio_history() if mdef.supports_history else None
         conv.add_message("user", message, self.config.max_history_messages)
 
-        eff_temp = temperature if temperature is not None else mdef.default_temperature
+        eff_temp = (temperature if temperature is not None
+                    else mdef.default_temperature)
         eff_sys = conv.system_prompt if mdef.supports_system_prompt else None
-        eff_thinking = include_thinking if include_thinking is not None else self.config.include_thinking
+        eff_thinking = (include_thinking if include_thinking is not None
+                        else self.config.include_thinking)
 
         extra = dict(kwargs)
         if mdef.supports_thinking:
@@ -894,21 +1376,39 @@ class MultiModelClient:
         for attempt in range(self.config.max_retries + 1):
             try:
                 if attempt > 0:
-                    time.sleep(self.config.retry_backoff_base ** attempt + random.uniform(0, self.config.retry_jitter))
+                    time.sleep(
+                        self.config.retry_backoff_base ** attempt
+                        + random.uniform(0, self.config.retry_jitter)
+                    )
                     metrics.record_retry()
 
-                provider = self._ensure_ready(model_id)
+                lb_pool = self._ensure_ready(model_id)
 
                 if stream and mdef.supports_streaming:
-                    gen = provider.generate_stream(message, history=history, system_prompt=eff_sys,
-                                                   temperature=eff_temp, max_tokens=max_tokens, **extra)
+                    gen = lb_pool.execute_stream(
+                        message=message,
+                        history=history,
+                        system_prompt=eff_sys,
+                        temperature=eff_temp,
+                        max_tokens=max_tokens,
+                        **extra,
+                    )
                     return self._wrap_stream(gen, conv, start, model_id)
 
-                result = provider.generate(message, history=history, system_prompt=eff_sys,
-                                           temperature=eff_temp, max_tokens=max_tokens, **extra)
+                result = lb_pool.execute(
+                    "generate",
+                    message=message,
+                    history=history,
+                    system_prompt=eff_sys,
+                    temperature=eff_temp,
+                    max_tokens=max_tokens,
+                    **extra,
+                )
                 dur = (time.monotonic() - start) * 1000
                 thinking, response = ThinkingParser.split(result)
-                conv.add_message("assistant", response, self.config.max_history_messages, thinking=thinking)
+                conv.add_message("assistant", response,
+                                 self.config.max_history_messages,
+                                 thinking=thinking)
                 metrics.record_request(True, dur, len(result), model_id)
                 self.circuit_breaker.record_success()
                 return result
@@ -933,31 +1433,46 @@ class MultiModelClient:
                 full += chunk
                 yield chunk
             thinking, response = ThinkingParser.split(full)
-            conv.add_message("assistant", response, self.config.max_history_messages, thinking=thinking)
-            metrics.record_request(True, (time.monotonic() - start) * 1000, len(full), model_id)
+            conv.add_message("assistant", response,
+                             self.config.max_history_messages,
+                             thinking=thinking)
+            metrics.record_request(
+                True, (time.monotonic() - start) * 1000,
+                len(full), model_id,
+            )
             self.circuit_breaker.record_success()
         except Exception:
-            metrics.record_request(False, (time.monotonic() - start) * 1000, model=model_id)
+            metrics.record_request(
+                False, (time.monotonic() - start) * 1000, model=model_id,
+            )
             self.circuit_breaker.record_failure()
             raise
 
     def get_status(self) -> Dict:
+        lb_info = {}
+        for model_id, pool in self._lb_pools.items():
+            lb_info[model_id] = pool.get_pool_info()
+
         return {
-            "version": VERSION, "current_model": self._current_model,
+            "version": VERSION,
+            "current_model": self._current_model,
             "models": list(MODEL_REGISTRY.keys()),
-            "providers": {m: "READY" if p.ready else "NOT_READY" for m, p in self._providers.items()},
+            "load_balancer": lb_info,
             "conversations": len(self._conversations),
             "circuit_breaker": self.circuit_breaker.state,
+            "rate_limiter": self.rate_limiter.get_info(),
         }
 
 # ═══════════════════════════════════════════════════════════════
-#  SESSION POOL
+#  SESSION POOL (top-level pool of MultiModelClients)
 # ═══════════════════════════════════════════════════════════════
 
 class SessionPool:
     def __init__(self, config: Config):
         self.config = config
-        self._clients = [MultiModelClient(config) for _ in range(config.pool_size)]
+        self._clients = [
+            MultiModelClient(config) for _ in range(config.pool_size)
+        ]
         self._idx = 0
         self._lock = threading.Lock()
 
@@ -966,7 +1481,10 @@ class SessionPool:
             c.init_model(self.config.default_model)
 
     def init_model(self, model_id: str) -> int:
-        return sum(1 for c in self._clients if c.init_model(model_id))
+        total = 0
+        for c in self._clients:
+            total += c.init_model_all(model_id)
+        return total
 
     def acquire(self) -> MultiModelClient:
         with self._lock:
@@ -980,14 +1498,17 @@ class SessionPool:
 
 ALIASES = {
     "gpt-oss": "gpt-oss-120b", "gptoss": "gpt-oss-120b", "amd": "gpt-oss-120b",
-    "command-a": "command-a-vision", "command-vision": "command-a-vision", "cohere-vision": "command-a-vision",
-    "command-translate": "command-a-translate", "cohere-translate": "command-a-translate", "translate": "command-a-translate",
+    "command-a": "command-a-vision", "command-vision": "command-a-vision",
+    "cohere-vision": "command-a-vision",
+    "command-translate": "command-a-translate",
+    "cohere-translate": "command-a-translate", "translate": "command-a-translate",
     "minimax": "minimax-vl-01", "minimax-vl": "minimax-vl-01",
     "glm": "glm-4.5", "glm4": "glm-4.5", "glm-4": "glm-4.5", "zhipu": "glm-4.5",
     "gpt": "chatgpt", "gpt-3.5": "chatgpt", "gpt3": "chatgpt", "openai": "chatgpt",
     "qwen": "qwen3-vl", "qwen3": "qwen3-vl", "qwen-vl": "qwen3-vl",
 }
 
+
 def resolve_alias(model_id: str) -> str:
     return ALIASES.get(model_id.lower(), model_id)
 
@@ -1001,6 +1522,7 @@ pool.init_default()
 
 app = Flask(APP_NAME)
 
+
 @app.after_request
 def cors(response):
     response.headers["Access-Control-Allow-Origin"] = "*"
@@ -1008,15 +1530,19 @@ def cors(response):
     response.headers["Access-Control-Allow-Methods"] = "GET, POST, OPTIONS"
     return response
 
+
 @app.errorhandler(APIError)
 def handle_api_error(e: APIError):
     return jsonify({"ok": False, **e.to_dict()}), e.status
 
+
 @app.route("/")
 def index():
     return jsonify({
-        "name": APP_NAME, "version": VERSION,
+        "name": APP_NAME,
+        "version": VERSION,
         "default_model": config.default_model,
+        "features": ["load_balancing", "10_req_per_second_limit", "failover"],
         "models": list(MODEL_REGISTRY.keys()),
         "endpoints": {
             "POST /chat": "Chat with any model",
@@ -1024,11 +1550,13 @@ def index():
             "POST /v1/chat/completions": "OpenAI-compatible",
             "GET /v1/models": "List models",
             "POST /models/init": "Init a model",
-            "GET /health": "Health check",
+            "GET /health": "Health check (incl. LB status)",
             "GET /metrics": "Metrics",
+            "GET /lb/status": "Load balancer detailed status",
         },
     })
 
+
 @app.route("/chat", methods=["POST"])
 def chat():
     data = freq.get_json(force=True, silent=True) or {}
@@ -1040,17 +1568,26 @@ def chat():
     client = pool.acquire()
     if data.get("new_conversation"):
         client.new_conversation(data.get("system_prompt"), model_id)
-    result = client.send_message(message, model=model_id, system_prompt=data.get("system_prompt"),
-                                  temperature=data.get("temperature"), max_tokens=data.get("max_tokens"),
-                                  include_thinking=include_thinking)
+    result = client.send_message(
+        message, model=model_id,
+        system_prompt=data.get("system_prompt"),
+        temperature=data.get("temperature"),
+        max_tokens=data.get("max_tokens"),
+        include_thinking=include_thinking,
+    )
     thinking, clean = ThinkingParser.split(result)
-    resp = {"ok": True, "response": clean, "model": model_id,
-            "conversation_id": client.active_conversation.conversation_id,
-            "history_size": len(client.active_conversation.messages)}
+    resp = {
+        "ok": True,
+        "response": clean,
+        "model": model_id,
+        "conversation_id": client.active_conversation.conversation_id,
+        "history_size": len(client.active_conversation.messages),
+    }
     if thinking:
         resp["thinking"] = thinking
     return jsonify(resp)
 
+
 @app.route("/chat/stream", methods=["POST"])
 def chat_stream():
     data = freq.get_json(force=True, silent=True) or {}
@@ -1068,40 +1605,57 @@ def chat_stream():
     def generate():
         try:
             if use_stream:
-                for chunk in client.send_message(message, stream=True, model=model_id,
-                                                  system_prompt=data.get("system_prompt"),
-                                                  temperature=data.get("temperature"),
-                                                  max_tokens=data.get("max_tokens"),
-                                                  include_thinking=include_thinking):
+                for chunk in client.send_message(
+                    message, stream=True, model=model_id,
+                    system_prompt=data.get("system_prompt"),
+                    temperature=data.get("temperature"),
+                    max_tokens=data.get("max_tokens"),
+                    include_thinking=include_thinking,
+                ):
                     yield f"data: {json.dumps({'chunk': chunk})}\n\n"
             else:
-                result = client.send_message(message, model=model_id,
-                                              system_prompt=data.get("system_prompt"),
-                                              temperature=data.get("temperature"),
-                                              max_tokens=data.get("max_tokens"),
-                                              include_thinking=include_thinking)
+                result = client.send_message(
+                    message, model=model_id,
+                    system_prompt=data.get("system_prompt"),
+                    temperature=data.get("temperature"),
+                    max_tokens=data.get("max_tokens"),
+                    include_thinking=include_thinking,
+                )
                 yield f"data: {json.dumps({'chunk': result})}\n\n"
             yield "data: [DONE]\n\n"
         except APIError as e:
             yield f"data: {json.dumps(e.to_dict())}\n\n"
 
-    return Response(stream_with_context(generate()), content_type="text/event-stream")
+    return Response(stream_with_context(generate()),
+                    content_type="text/event-stream")
+
 
 @app.route("/v1/models", methods=["GET"])
 def list_models():
     models = []
     for mid, mdef in MODEL_REGISTRY.items():
         models.append({
-            "id": mid, "object": "model", "owned_by": mdef.owned_by, "created": 0,
+            "id": mid,
+            "object": "model",
+            "owned_by": mdef.owned_by,
+            "created": 0,
             "description": mdef.description,
             "capabilities": {
-                "vision": mdef.supports_vision, "streaming": mdef.supports_streaming,
-                "system_prompt": mdef.supports_system_prompt, "temperature": mdef.supports_temperature,
-                "history": mdef.supports_history, "thinking": mdef.supports_thinking,
+                "vision": mdef.supports_vision,
+                "streaming": mdef.supports_streaming,
+                "system_prompt": mdef.supports_system_prompt,
+                "temperature": mdef.supports_temperature,
+                "history": mdef.supports_history,
+                "thinking": mdef.supports_thinking,
+            },
+            "load_balancing": {
+                "enabled": mdef.lb_enabled,
+                "pool_size": mdef.lb_pool_size,
             },
         })
     return jsonify({"object": "list", "data": models})
 
+
 @app.route("/v1/chat/completions", methods=["POST", "OPTIONS"])
 def openai_compat():
     if freq.method == "OPTIONS":
@@ -1115,7 +1669,12 @@ def openai_compat():
     include_thinking = data.get("include_thinking", config.include_thinking)
 
     if model_id not in MODEL_REGISTRY:
-        return jsonify({"error": {"message": f"Model '{model_id}' not found", "type": "invalid_request_error"}}), 404
+        return jsonify({
+            "error": {
+                "message": f"Model '{model_id}' not found",
+                "type": "invalid_request_error",
+            }
+        }), 404
     if not messages:
         return jsonify({"error": {"message": "messages required"}}), 400
 
@@ -1145,68 +1704,160 @@ def openai_compat():
     if do_stream:
         def generate():
             try:
-                yield f"data: {json.dumps({'id': rid, 'object': 'chat.completion.chunk', 'created': created, 'model': model_id, 'choices': [{'index': 0, 'delta': {'role': 'assistant'}, 'finish_reason': None}]})}\n\n"
+                yield (
+                    f"data: {json.dumps({'id': rid, 'object': 'chat.completion.chunk', "
+                    f"'created': created, 'model': model_id, 'choices': ["
+                    f"{{'index': 0, 'delta': {{'role': 'assistant'}}, "
+                    f"'finish_reason': None}}]})}\n\n"
+                )
                 if mdef.supports_streaming:
-                    for chunk in client.send_message(user_msg, stream=True, model=model_id,
-                                                      temperature=temperature, max_tokens=max_tokens,
-                                                      include_thinking=include_thinking):
-                        yield f"data: {json.dumps({'id': rid, 'object': 'chat.completion.chunk', 'created': created, 'model': model_id, 'choices': [{'index': 0, 'delta': {'content': chunk}, 'finish_reason': None}]})}\n\n"
+                    for chunk in client.send_message(
+                        user_msg, stream=True, model=model_id,
+                        temperature=temperature, max_tokens=max_tokens,
+                        include_thinking=include_thinking,
+                    ):
+                        yield (
+                            f"data: {json.dumps({'id': rid, "
+                            f"'object': 'chat.completion.chunk', "
+                            f"'created': created, 'model': model_id, "
+                            f"'choices': [{{'index': 0, "
+                            f"'delta': {{'content': chunk}}, "
+                            f"'finish_reason': None}}]})}\n\n"
+                        )
                 else:
-                    result = client.send_message(user_msg, model=model_id, temperature=temperature,
-                                                  max_tokens=max_tokens, include_thinking=include_thinking)
-                    yield f"data: {json.dumps({'id': rid, 'object': 'chat.completion.chunk', 'created': created, 'model': model_id, 'choices': [{'index': 0, 'delta': {'content': result}, 'finish_reason': None}]})}\n\n"
-                yield f"data: {json.dumps({'id': rid, 'object': 'chat.completion.chunk', 'created': created, 'model': model_id, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]})}\n\n"
+                    result = client.send_message(
+                        user_msg, model=model_id, temperature=temperature,
+                        max_tokens=max_tokens,
+                        include_thinking=include_thinking,
+                    )
+                    yield (
+                        f"data: {json.dumps({'id': rid, "
+                        f"'object': 'chat.completion.chunk', "
+                        f"'created': created, 'model': model_id, "
+                        f"'choices': [{{'index': 0, "
+                        f"'delta': {{'content': result}}, "
+                        f"'finish_reason': None}}]})}\n\n"
+                    )
+                yield (
+                    f"data: {json.dumps({'id': rid, "
+                    f"'object': 'chat.completion.chunk', "
+                    f"'created': created, 'model': model_id, "
+                    f"'choices': [{{'index': 0, 'delta': {{}}, "
+                    f"'finish_reason': 'stop'}}]})}\n\n"
+                )
                 yield "data: [DONE]\n\n"
             except Exception as e:
                 yield f"data: {json.dumps({'error': {'message': str(e)}})}\n\n"
-        return Response(stream_with_context(generate()), content_type="text/event-stream")
 
-    result = client.send_message(user_msg, model=model_id, temperature=temperature,
-                                  max_tokens=max_tokens, include_thinking=include_thinking)
+        return Response(stream_with_context(generate()),
+                        content_type="text/event-stream")
+
+    result = client.send_message(
+        user_msg, model=model_id, temperature=temperature,
+        max_tokens=max_tokens, include_thinking=include_thinking,
+    )
     return jsonify({
-        "id": rid, "object": "chat.completion", "created": created, "model": model_id,
-        "choices": [{"index": 0, "message": {"role": "assistant", "content": result}, "finish_reason": "stop"}],
-        "usage": {"prompt_tokens": len(user_msg) // 4, "completion_tokens": len(result) // 4,
-                  "total_tokens": (len(user_msg) + len(result)) // 4},
+        "id": rid,
+        "object": "chat.completion",
+        "created": created,
+        "model": model_id,
+        "choices": [{
+            "index": 0,
+            "message": {"role": "assistant", "content": result},
+            "finish_reason": "stop",
+        }],
+        "usage": {
+            "prompt_tokens": len(user_msg) // 4,
+            "completion_tokens": len(result) // 4,
+            "total_tokens": (len(user_msg) + len(result)) // 4,
+        },
     })
 
+
 @app.route("/new", methods=["POST"])
 def new_conv():
     data = freq.get_json(force=True, silent=True) or {}
     model_id = resolve_alias(data.get("model", config.default_model))
     client = pool.acquire()
     conv = client.new_conversation(data.get("system_prompt"), model_id)
-    return jsonify({"ok": True, "conversation_id": conv.conversation_id, "model": model_id})
+    return jsonify({
+        "ok": True,
+        "conversation_id": conv.conversation_id,
+        "model": model_id,
+    })
+
 
 @app.route("/health", methods=["GET"])
 def health():
     client = pool.acquire()
     return jsonify(client.get_status())
 
+
 @app.route("/metrics", methods=["GET"])
 def metrics_endpoint():
     return jsonify(metrics.to_dict())
 
+
+@app.route("/lb/status", methods=["GET"])
+def lb_status():
+    """Detailed load balancer status for all models across all clients."""
+    all_pools = {}
+    for client in pool._clients:
+        for model_id, lb_pool in client._lb_pools.items():
+            key = f"{model_id}"
+            if key not in all_pools:
+                all_pools[key] = []
+            all_pools[key].append(lb_pool.get_pool_info())
+    return jsonify({
+        "ok": True,
+        "version": VERSION,
+        "rate_limit": f"{config.rate_limit_rps} req/s",
+        "models": all_pools,
+    })
+
+
 @app.route("/conversations", methods=["GET"])
 def conversations():
     client = pool.acquire()
-    return jsonify({"conversations": [c.to_dict() for c in client._conversations.values()]})
+    return jsonify({
+        "conversations": [c.to_dict() for c in client._conversations.values()]
+    })
+
 
 @app.route("/models/init", methods=["POST"])
 def init_model_ep():
     data = freq.get_json(force=True, silent=True) or {}
     model_id = resolve_alias(data.get("model", ""))
     if not model_id or model_id not in MODEL_REGISTRY:
-        return jsonify({"ok": False, "error": f"Unknown model. Available: {list(MODEL_REGISTRY.keys())}"}), 400
+        return jsonify({
+            "ok": False,
+            "error": f"Unknown model. Available: {list(MODEL_REGISTRY.keys())}",
+        }), 400
     count = pool.init_model(model_id)
-    return jsonify({"ok": True, "model": model_id, "initialized_clients": count})
+    mdef = MODEL_REGISTRY[model_id]
+    return jsonify({
+        "ok": True,
+        "model": model_id,
+        "initialized_instances": count,
+        "lb_enabled": mdef.lb_enabled,
+        "pool_size_per_client": mdef.lb_pool_size,
+    })
+
 
 # ═══════════════════════════════════════════════════════════════
-#  ENTRY POINT (for HuggingFace Spaces)
+#  ENTRY POINT
 # ═══════════════════════════════════════════════════════════════
 
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 7860))
     log.info(f"Starting Multi-Model AI API v{VERSION} on port {port}")
     log.info(f"Models: {list(MODEL_REGISTRY.keys())}")
+    log.info(f"Rate limit: {config.rate_limit_rps} req/s (burst: {config.rate_limit_burst})")
+    for mid, mdef in MODEL_REGISTRY.items():
+        lb_status_str = (
+            f"LB ON (pool={mdef.lb_pool_size})"
+            if mdef.lb_enabled
+            else "LB OFF (single instance)"
+        )
+        log.info(f"  {mid}: {lb_status_str}")
     app.run(host="0.0.0.0", port=port, threaded=True)