rag_core/llm.py

import requests
import logging
import time
import os

LLM_ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent"


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 = 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 Gemini 2.5 Flash tại {LLM_ENDPOINT}")
            response = requests.post(url, json=payload, headers=headers, timeout=timeout)
            response.raise_for_status()
            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
                continue
            else:
                logging.error("❌ Lỗi timeout sau 3 lần retry.")
                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 Gemini (lần {attempt}): {e}")
            if attempt < max_retries:
                time.sleep(1)
                continue
            else:
                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."