Update app.py

app.py

@@ -144,12 +144,18 @@ class OptionalHFInferenceClient:
     detailed errors instead of hiding provider/model failures.
     """
 
-    def __init__(self, model_id: str, token: str | None = None) -> None:
+    def __init__(
+        self,
+        model_id: str,
+        token: str | None = None,
+        provider: str | None = None,
+    ) -> None:
         """Initialize the optional Hugging Face inference adapter."""
 
         from huggingface_hub import InferenceClient
 
         self.model_id = model_id.strip().strip('"').strip("'").strip()
+        self.provider = None if provider is None or not str(provider).strip() else str(provider).strip()
 
         api_key = (
             token
@@ -165,10 +171,17 @@ class OptionalHFInferenceClient:
                 "and add it as a Space secret named HF_TOKEN."
             )
 
-        self.client = InferenceClient(
-            model=self.model_id,
-            api_key=api_key,
-        )
+        if self.provider:
+            self.client = InferenceClient(
+                model=self.model_id,
+                provider=self.provider,
+                api_key=api_key,
+            )
+        else:
+            self.client = InferenceClient(
+                model=self.model_id,
+                api_key=api_key,
+            )
 
     def generate(
         self,
@@ -376,6 +389,12 @@ def build_optional_model_client() -> tuple[Any | None, AppModelStatus]:
             or os.getenv("HUGGINGFACEHUB_API_TOKEN")
             )
 
+    provider = (
+            os.getenv("WORLDSMITHAI_PROVIDER")
+            or os.getenv("HF_INFERENCE_PROVIDER")
+            or None
+        )
+
     if not raw_model_id:
         return None, AppModelStatus(
             enabled=False,
@@ -399,7 +418,11 @@ def build_optional_model_client() -> tuple[Any | None, AppModelStatus]:
         )
 
     try:
-        client = OptionalHFInferenceClient(model_id=model_id, token=token)
+        client = OptionalHFInferenceClient(
+            model_id=model_id,
+            token=token,
+            provider=provider,
+        )
         return client, AppModelStatus(
             enabled=True,
             model_id=model_id,
@@ -430,6 +453,13 @@ def model_health_check_ui() -> str:
             "Check that WORLDSMITHAI_MODEL_ID is set in Space settings."
         )
 
+    provider_label = (
+        getattr(MODEL_CLIENT, "provider", None)
+        or os.getenv("WORLDSMITHAI_PROVIDER")
+        or os.getenv("HF_INFERENCE_PROVIDER")
+        or "auto"
+    )
+
     try:
         response = MODEL_CLIENT.generate(
             prompt='Return exactly this JSON object and nothing else: {"ok": true}',
@@ -450,6 +480,7 @@ def model_health_check_ui() -> str:
         return (
             "Model client is configured and returned text.\n\n"
             f"Model id: {MODEL_STATUS.model_id}\n\n"
+            f"Provider: {provider_label}\n\n"
             f"Response:\n{response}"
         )
     except Exception as exc:
@@ -457,6 +488,7 @@ def model_health_check_ui() -> str:
         return (
             "Model client is configured but inference failed.\n\n"
             f"Model id: {MODEL_STATUS.model_id}\n\n"
+            f"Provider: {provider_label}\n\n"
             f"Error type: {exc.__class__.__name__}\n"
             f"Error:\n{exc}"
         )