update gemini

rag_core/llm.py

@@ -1,40 +1,53 @@
 import requests
 import logging
 import time
+import os
+
+LLM_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent"
 
-LLM_ENDPOINT = "https://vietcat-gemma34b.hf.space/purechat"
 
 def generate_answer(prompt: str) -> str:
+    api_key = os.getenv("GEMINI_API_KEY")
+    if not api_key:
+        return "Thiếu biến môi trường GEMINI_API_KEY."
     max_retries = 3
-    timeout = 300  # ⏱️ timeout ban đầu: 5 phút
+    timeout = 60  # Gemini Flash rất nhanh, timeout 1 phút là đủ
+
+    payload = {
+        "contents": [
+            {
+                "parts": [
+                    {"text": prompt}
+                ]
+            }
+        ]
+    }
+
+    headers = {"Content-Type": "application/json"}
+    url = f"{LLM_ENDPOINT}?key={api_key}"
 
     for attempt in range(1, max_retries + 1):
         try:
-            logging.info(f"📡 Gửi request đến LLM tại {LLM_ENDPOINT}")
-            logging.info(f"📡 Gửi request đến LLM (lần {attempt}, timeout={timeout}s)...")
-            prompt_length = len(prompt.split())
-            response = requests.post(
-                LLM_ENDPOINT,
-                json={"prompt": prompt, "max_tokens": prompt_length},
-                timeout=timeout
-            )
+            logging.info(f"📡 Gửi request đến Gemini 2.5 Flash tại {LLM_ENDPOINT}")
+            response = requests.post(url, json=payload, headers=headers, timeout=timeout)
             response.raise_for_status()
-            return response.json()["response"]
-
+            data = response.json()
+            # Lấy text trả về
+            return data["candidates"][0]["content"]["parts"][0]["text"]
         except requests.exceptions.Timeout as e:
             logging.warning(f"⚠️ Timeout ở lần {attempt}: {e}")
             if attempt < max_retries:
-                timeout *= 2  # ⏫ tăng timeout gấp đôi
+                timeout *= 2
                 continue
             else:
                 logging.error("❌ Lỗi timeout sau 3 lần retry.")
-                return "Lỗi timeout khi gửi tới LLM."
-
+                return "Lỗi timeout khi gửi tới Gemini."
         except Exception as e:
-            logging.warning(f"⚠️ Lỗi khi gửi request tới LLM (lần {attempt}): {e}")
+            logging.warning(f"⚠️ Lỗi khi gửi request tới Gemini (lần {attempt}): {e}")
             if attempt < max_retries:
-                time.sleep(1)  # chờ 1s rồi thử lại
+                time.sleep(1)
                 continue
             else:
-                logging.error("❌ Lỗi khi gửi tới LLM sau 3 lần thử.")
-                return "Lỗi khi gửi tới LLM."
+                logging.error("❌ Lỗi khi gửi tới Gemini sau 3 lần thử.")
+                return "Lỗi khi gửi tới Gemini."
+    return "Lỗi không xác định khi gửi tới Gemini."