Update PolyAgent/gradio_interface.py

PolyAgent/gradio_interface.py

@@ -1240,34 +1240,23 @@ def llm_only_answer(state: Dict[str, Any], model_name: str, prompt: str) -> str:
             }
         )
 
-    client = InferenceClient(model=model_id, token=HF_TOKEN)
+    client = InferenceClient(model=model_id, token=HF_TOKEN, provider="fireworks-ai")
 
     try:
-        if model_id.startswith("mistralai/"):
-            # Mixtral: use text-generation, not chat
-            prompt_text = (
-                "You are a polymer R&D assistant. Answer directly and clearly.\n\n"
-                f"User: {p}\nAssistant:"
-            )
-            resp = client.text_generation(
-                prompt_text,
-                max_new_tokens=900,
-                temperature=0.7,
-                top_p=0.95,
-                return_full_text=False,
-            )
-            return resp
-        else:
-            # Llama: chat endpoint works
-            resp = client.chat_completion(
-                messages=[
-                    {"role": "system", "content": "You are a polymer R&D assistant..."},
-                    {"role": "user", "content": p},
-                ],
-                max_tokens=900,
-                temperature=0.7,
-            )
-            return resp.choices[0].message.content or ""
+        resp = client.chat_completion(
+            messages=[
+                {
+                    "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:
         return pretty_json({"ok": False, "error": str(e), "model_id": model_id})