Update app.py

app.py

@@ -83,19 +83,21 @@ def setup_model(model_key, intial=False):
     model_config = get_model_config(model_key)
     log_debug(f"update_model() --> Model config: {model_config}")
 
-    url_list = (model_config.get('VLLM_API_URL_LIST') or "").split(",")
-    if USE_RANDOM_ENDPOINT:
-        base_url = random.choice(url_list) if len(url_list) > 0 else model_config.get('VLLM_API_URL')
-    else:
-        base_url = url_list[endpoint_rotation_count % len(url_list)]
-        endpoint_rotation_count += 1
+    # ENVIRONMENT VARIABLE CONFIGURATION
+    base_url = os.environ.get("API_BASE_URL")
+    api_key = os.environ.get("API_KEY")
+
+    if not base_url:
+        raise ValueError("API_BASE_URL environment variable not set")
+    if not api_key:
+        raise ValueError("API_KEY environment variable not set")
 
     openai_client = OpenAI(
-        api_key=model_config.get('AUTH_TOKEN'),
+        api_key=api_key,
         base_url=base_url
     )
     model_config['base_url'] = base_url
-    log_debug(f"Switched to model {model_key} using endpoint {base_url}")
+    log_debug(f"Switched to model {model_key} using endpoint {base_url} (ENV VARS)")
 
     _model_hf_name = model_config.get("MODEL_HF_URL").split('https://huggingface.co/')[1]
     _link = f"<a href='{model_config.get('MODEL_HF_URL')}{COMMUNITY_POSTFIX_URL}' target='_blank'>{_model_hf_name}</a>"
@@ -134,7 +136,11 @@ def run_chat_inference(history, message, state):
     state["is_streaming"] = True
     state["stop_flag"] = False
     error = None
-    model_name = model_config.get('MODEL_NAME')
+    # ENVIRONMENT VARIABLE MODEL
+    model_name = os.environ.get("API_MODEL")
+    if not model_name:
+         raise ValueError("API_MODEL environment variable not set")
+    # model_name = model_config.get('MODEL_NAME')
     temperature = model_config.get('TEMPERATURE', DEFAULT_MODEL_TEMPERATURE)
     output_tag_start = model_config.get('OUTPUT_TAG_START', "[BEGIN FINAL RESPONSE]")
     output_tag_end = model_config.get('OUTPUT_TAG_END', "[END FINAL RESPONSE]")
@@ -659,4 +665,4 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
     )
 
 demo.queue(default_concurrency_limit=2).launch(ssr_mode=False, show_api=False, max_file_size="10mb")
-log_info("Gradio app launched")
+log_info("Gradio app launched")