Update app/llm.py

app/llm.py

@@ -1,88 +1,89 @@
-import requests
 import os
 import random
+import time
+import requests
+from typing import List
+
 API_URL = "https://router.huggingface.co/v1/chat/completions"
 
-def load_tokens():
-    tokens = []
-    for k, v in os.environ.items():
-        if k.startswith("HF_TOKEN_") and v:
-            tokens.append(v)
+MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"
+
+TIMEOUT_SECONDS = 30
+MAX_RETRIES_PER_TOKEN = 3
+
+
+def load_tokens() -> List[str]:
+    """
+    Load all Hugging Face tokens that start with HF_TOKEN_
+    from environment variables.
+    """
+    tokens = [
+        v for k, v in os.environ.items()
+        if k.startswith("HF_TOKEN_") and v
+    ]
+
     if not tokens:
-        raise RuntimeError("No HF tokens found in environment variables")
+        raise RuntimeError(
+            "No HF_TOKEN_* variables found. "
+            "Add at least one token in Hugging Face Space settings."
+        )
+
     return tokens
 
+
 HF_TOKENS = load_tokens()
 
+
 def generate_answer(context: str, query: str) -> str:
-    random.shuffle(HF_TOKENS)
+ 
+    tokens = HF_TOKENS[:]
+    random.shuffle(tokens)
+
+    context = context[:3000]
 
-    for token in HF_TOKENS:
+    for token in tokens:
         headers = {
             "Authorization": f"Bearer {token}",
             "Content-Type": "application/json",
         }
 
-        response = requests.post(API_URL, headers=headers, json={
-            "model": "deepseek-ai/DeepSeek-V3.2",
+        payload = {
+            "model": MODEL_NAME,
             "messages": [
-                {"role": "system", "content": "You are a Harry Potter knowledge assistant."},
-                {"role": "user", "content": f"Context:\n{context}\nQuestion:\n{query}\nAnswer:"},
+                {
+                    "role": "system",
+                    "content": "You are a Harry Potter knowledge assistant created by Ishmeet Kaur, who is awesome. Answer in short concise sentences, stay grounded and faithful."
+                },
+                {
+                    "role": "user",
+                    "content": f"Context:\n{context}\n\nQuestion:\n{query}\n\nAnswer:"
+                },
             ],
             "temperature": 0.3,
             "max_tokens": 500,
-        })
-
-        if response.status_code == 200:
-            return response.json()["choices"][0]["message"]["content"]
-
-        if response.status_code != 429:
-            break  # real error, don't retry blindly
-
-    return "All magical channels are busy right now. Please try again shortly."
-
-# def generate_answer(context: str, query: str) -> str:
-#     headers = {
-#         "Authorization": f"Bearer {HF_TOKEN}",
-#         "Content-Type": "application/json",
-#     }
-
-#     payload = {
-#         "model": "deepseek-ai/DeepSeek-V3.2",
-#         "messages": [
-#             {
-#                 "role": "system",
-#                 "content": (
-#                     """You are a Harry Potter knowledge assistant.
-
-#                     Answer the question concisely in 2–3 sentences.
-#                     Do not use bullet points.
-#                     Do not add extra background.
-#                     Be clear and direct."""
-#                 ),
-#             },
-#             {
-#                 "role": "user",
-#                 "content": f"""Context:
-# {context}
-
-# Question:
-# {query}
-
-# Answer:
-# """,
-#             },
-#         ],
-#         "temperature": 0.3,
-#         "max_tokens": 500,
-#     }
-
-#     response = requests.post(API_URL, headers=headers, json=payload)
-
-#     if response.status_code != 200:
-#         return (
-#         "The magic is unstable right now.\n\n"
-#         f"Model response: {response.status_code}\n\n"
-#         "Try again in a moment."
-#     )
-#     return response.json()["choices"][0]["message"]["content"]
+        }
+
+        for attempt in range(MAX_RETRIES_PER_TOKEN):
+            try:
+                response = requests.post(
+                    API_URL,
+                    headers=headers,
+                    json=payload,
+                    timeout=TIMEOUT_SECONDS,
+                )
+            except requests.RequestException:
+                break
+
+            if response.status_code == 200:
+                return response.json()["choices"][0]["message"]["content"]
+
+            if response.status_code == 429:
+                time.sleep(2 ** attempt)
+                continue
+
+            break
+
+    return (
+        "The library is a bit crowded right now. "
+        "Please try again in a moment."
+    )