Update PolyAgent/gradio_interface.py

PolyAgent/gradio_interface.py

@@ -1243,20 +1243,31 @@ def llm_only_answer(state: Dict[str, Any], model_name: str, prompt: str) -> str:
     client = InferenceClient(model=model_id, token=HF_TOKEN)
 
     try:
-        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,
-        )
+        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 ""
         return resp.choices[0].message.content or ""
     except Exception as e:
         return pretty_json({"ok": False, "error": str(e), "model_id": model_id})