rag_core/embedder.py

import requests
import time
import logging
import os
from rag_core.utils import log_timed

@log_timed("gửi API tạo embedding")
def get_embedding(text: str, retries: int = 3, base_timeout: int = 30):
    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key:
        raise ValueError("Thiếu biến môi trường GEMINI_API_KEY.")

    # Giới hạn độ dài text gửi lên API (ví dụ 8000 ký tự)
    MAX_LEN = 8000
    if not text or not text.strip():
        raise ValueError("Chunk rỗng, không thể tạo embedding.")
    if len(text) > MAX_LEN:
        logging.warning(f"Chunk quá dài ({len(text)} ký tự), sẽ cắt còn {MAX_LEN} ký tự đầu.")
        text = text[:MAX_LEN]

    url = f"https://generativelanguage.googleapis.com/v1/models/text-embedding-004:embedContent?key={api_key}"

    payload = {
        "content": {
            "parts": [
                { "text": text }
            ]
        }
    }

    for i in range(retries):
        try:
            current_timeout = base_timeout * (i + 1)
            response = requests.post(url, json=payload, timeout=current_timeout)
            response.raise_for_status()

            # Kiểm tra phản hồi chứa trường embedding không
            data = response.json()
            if "embedding" not in data:
                raise ValueError(f"Phản hồi không có 'embedding': {data}")
            return data["embedding"]['values']

        except requests.exceptions.RequestException as e:
            logging.warning(f"Lỗi embedding (lần {i+1}/{retries}, timeout={current_timeout}s): {e}")
            if i < retries - 1:
                time.sleep(2)
            else:
                raise