Update PolyAgent/gradio_interface.py

PolyAgent/gradio_interface.py

@@ -1201,9 +1201,7 @@ def gpt_only_answer(state: Dict[str, Any], prompt: str) -> str:
 # ----------------------------- Other LLMs (Hugging Face Inference) ----------------------------- #
 def llm_only_answer(state: Dict[str, Any], model_name: str, prompt: str) -> str:
     """
-    LLM-only responses via huggingface_hub.InferenceClient.
-    - Forces provider to avoid unwanted auto-routing (e.g., fireworks-ai).
-    - Tries chat_completion first; if model/provider doesn't support chat, falls back to text_generation.
+    LLM-only responses using Hugging Face Inference API for non-GPT models.
     """
     ensure_orch(state)
 
@@ -1212,7 +1210,12 @@ def llm_only_answer(state: Dict[str, Any], model_name: str, prompt: str) -> str:
 
     HF_TOKEN = (os.getenv("HF_TOKEN") or "").strip()
     if not HF_TOKEN:
-        return pretty_json({"ok": False, "error": "HF_TOKEN is not set. Add HF_TOKEN=hf_... to Space Secrets."})
+        return pretty_json(
+            {
+                "ok": False,
+                "error": "HF_TOKEN is not set. Add HF_TOKEN=hf_... to your .env or env vars.",
+            }
+        )
 
     HF_MODEL_MAP = {
         "mixtral-8x22b-instruct": "mistralai/Mixtral-8x22B-Instruct-v0.1",
@@ -1221,6 +1224,7 @@ def llm_only_answer(state: Dict[str, Any], model_name: str, prompt: str) -> str:
 
     m = (model_name or "").strip()
     p = (prompt or "").strip()
+
     if not p:
         return "Please provide a prompt."
     if not m:
@@ -1228,60 +1232,34 @@ def llm_only_answer(state: Dict[str, Any], model_name: str, prompt: str) -> str:
 
     model_id = HF_MODEL_MAP.get(m)
     if not model_id:
-        return pretty_json({"ok": False, "error": f"Unsupported model selection: {m}", "supported": list(HF_MODEL_MAP.keys())})
-
-    # IMPORTANT: force provider so HF doesn't auto-route you to a provider that lacks the needed task
-    provider = (os.getenv("HF_PROVIDER") or "hf-inference").strip()
-
-    client = InferenceClient(
-        provider=provider,
-        model=model_id,
-        api_key=HF_TOKEN,  # api_key works for both HF token + provider keys
-    )
+        return pretty_json(
+            {
+                "ok": False,
+                "error": f"Unsupported model selection: {m}",
+                "supported": list(HF_MODEL_MAP.keys()),
+            }
+        )
 
-    system = (
-        "You are a polymer R&D assistant. Answer directly and clearly. "
-        "Do not call tools or run web searches. If you are uncertain, state uncertainty."
-    )
+    client = InferenceClient(model=model_id, token=HF_TOKEN)
 
-    # 1) Try chat (conversational)
     try:
         resp = client.chat_completion(
             messages=[
-                {"role": "system", "content": system},
+                {
+                    "role": "system",
+                    "content": (
+                        "You are a polymer R&D assistant. Answer directly and clearly. "
+                        "Do not call tools or run web searches. If you are uncertain, state uncertainty."
+                    ),
+                },
                 {"role": "user", "content": p},
             ],
             max_tokens=900,
             temperature=0.7,
         )
         return resp.choices[0].message.content or ""
-    except Exception as e_chat:
-        # 2) Fallback to plain text-generation (works on hf-inference; many providers don't support it)
-        try:
-            if provider != "hf-inference":
-                # text_generation is not universally supported across providers
-                raise RuntimeError(
-                    f"Chat failed and provider='{provider}' may not support text_generation. "
-                    f"Set HF_PROVIDER=hf-inference (recommended) or choose a compatible model/provider."
-                )
-
-            # A simple prompt wrapper for non-chat models / non-chat endpoints
-            wrapped = f"{system}\n\nUser: {p}\nAssistant:"
-            out = client.text_generation(
-                wrapped,
-                max_new_tokens=900,
-                temperature=0.7,
-                do_sample=True,
-                return_full_text=False,
-            )
-            return out if isinstance(out, str) else str(out)
-        except Exception as e_gen:
-            return pretty_json({
-                "ok": False,
-                "error": f"chat_completion failed: {e_chat}; text_generation failed: {e_gen}",
-                "model_id": model_id,
-                "provider": provider,
-            })
+    except Exception as e:
+        return pretty_json({"ok": False, "error": str(e), "model_id": model_id})
 
 
 def build_ui() -> gr.Blocks: