fixup! ai: Better handling of load balancing.

98abcc71727c051b4892144f797138cfe8d020c9.

jarvis.py

@@ -19,6 +19,7 @@ import io
 import uuid
 import concurrent.futures
 import itertools
+import threading
 
 from openai import OpenAI
 
@@ -37,14 +38,19 @@ from pptx import Presentation
 os.system("apt-get update -q -y && apt-get install -q -y tesseract-ocr tesseract-ocr-eng tesseract-ocr-ind libleptonica-dev libtesseract-dev")
 
 LINUX_SERVER_HOSTS = [host for host in json.loads(os.getenv("LINUX_SERVER_HOST", "[]")) if host]
+LINUX_SERVER_HOSTS_MARKED = set()
+LINUX_SERVER_HOSTS_ATTEMPTS = {}
+
 LINUX_SERVER_PROVIDER_KEYS = [key for key in json.loads(os.getenv("LINUX_SERVER_PROVIDER_KEY", "[]")) if key]
+LINUX_SERVER_PROVIDER_KEYS_MARKED = set()
+LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS = {}
 
 AI_TYPES = {f"AI_TYPE_{i}": os.getenv(f"AI_TYPE_{i}") for i in range(1, 7)}
 RESPONSES = {f"RESPONSE_{i}": os.getenv(f"RESPONSE_{i}") for i in range(1, 10)}
 
 MODEL_MAPPING = json.loads(os.getenv("MODEL_MAPPING", "{}"))
 MODEL_CONFIG = json.loads(os.getenv("MODEL_CONFIG", "{}"))
-MODEL_CHOICES = list(MODEL_MAPPING.values())
+MODEL_CHOICES = list(MODEL_MAPPING.values()) if MODEL_MAPPING else []
 DEFAULT_CONFIG = json.loads(os.getenv("DEFAULT_CONFIG", "{}"))
 
 META_TAGS = os.getenv("META_TAGS")
@@ -53,6 +59,21 @@ ALLOWED_EXTENSIONS = json.loads(os.getenv("ALLOWED_EXTENSIONS"))
 
 ACTIVE_CANDIDATE = None
 
+def get_available_items(items, marked):
+    available = [item for item in items if item not in marked]
+    random.shuffle(available)
+    return available
+
+def marked_item(item, marked, attempts):
+    marked.add(item)
+    attempts[item] = attempts.get(item, 0) + 1
+    if attempts[item] >= 3:
+        def remove_fail():
+            marked.discard(item)
+            if item in attempts:
+                del attempts[item]
+        threading.Timer(300, remove_fail).start()
+
 class SessionWithID(requests.Session):
     def __init__(self):
         super().__init__()
@@ -121,15 +142,20 @@ def process_ai_response(ai_text):
         return ai_text
 
 def fetch_response(host, provider_key, selected_model, messages, model_config, session_id):
-    client = OpenAI(base_url=host, api_key=provider_key)
-    data = {"model": selected_model, "messages": messages, **model_config}
-    response = client.chat.completions.create(extra_body={"optillm_approach": "rto|re2|cot_reflection|self_consistency|plansearch|leap|z3|bon|moa|mcts|mcp|router|privacy|executecode|json", "session_id": session_id}, **data)
-    ai_text = response.choices[0].message.content if response.choices and response.choices[0].message and response.choices[0].message.content else RESPONSES["RESPONSE_2"]
-    return process_ai_response(ai_text)
+    try:
+        client = OpenAI(base_url=host, api_key=provider_key)
+        data = {"model": selected_model, "messages": messages, **model_config}
+        response = client.chat.completions.create(extra_body={"optillm_approach": "rto|re2|cot_reflection|self_consistency|plansearch|leap|z3|bon|moa|mcts|mcp|router|privacy|executecode|json", "session_id": session_id}, **data)
+        ai_text = response.choices[0].message.content if response.choices and response.choices[0].message and response.choices[0].message.content else RESPONSES["RESPONSE_2"]
+        return process_ai_response(ai_text)
+    except Exception:
+        marked_item(provider_key, LINUX_SERVER_PROVIDER_KEYS_MARKED, LINUX_SERVER_PROVIDER_KEYS_ATTEMPTS)
+        marked_item(host, LINUX_SERVER_HOSTS_MARKED, LINUX_SERVER_HOSTS_ATTEMPTS)
+        raise
 
 def chat_with_model(history, user_input, selected_model_display, sess):
     global ACTIVE_CANDIDATE
-    if not LINUX_SERVER_PROVIDER_KEYS or not LINUX_SERVER_HOSTS:
+    if not get_available_items(LINUX_SERVER_PROVIDER_KEYS, LINUX_SERVER_PROVIDER_KEYS_MARKED) or not get_available_items(LINUX_SERVER_HOSTS, LINUX_SERVER_HOSTS_MARKED):
         return RESPONSES["RESPONSE_3"]
     if not hasattr(sess, "session_id"):
         sess.session_id = str(uuid.uuid4())
@@ -143,7 +169,9 @@ def chat_with_model(history, user_input, selected_model_display, sess):
             return fetch_response(ACTIVE_CANDIDATE[0], ACTIVE_CANDIDATE[1], selected_model, messages, model_config, sess.session_id)
         except Exception:
             ACTIVE_CANDIDATE = None
-    candidates = [(host, key) for host in LINUX_SERVER_HOSTS for key in LINUX_SERVER_PROVIDER_KEYS]
+    available_keys = get_available_items(LINUX_SERVER_PROVIDER_KEYS, LINUX_SERVER_PROVIDER_KEYS_MARKED)
+    available_servers = get_available_items(LINUX_SERVER_HOSTS, LINUX_SERVER_HOSTS_MARKED)
+    candidates = [(host, key) for host in available_servers for key in available_keys]
     random.shuffle(candidates)
     with concurrent.futures.ThreadPoolExecutor(max_workers=len(candidates)) as executor:
         futures = {executor.submit(fetch_response, host, key, selected_model, messages, model_config, sess.session_id): (host, key) for (host, key) in candidates}