scripts/bootstrap_faiss.py

import os
import sys
import shutil
import urllib.request
import tarfile
import zipfile
from pathlib import Path

# Configuração via variáveis de ambiente (podem ser definidas no docker-compose.prod.yaml)
BOOTSTRAP_URL = os.getenv("BOOTSTRAP_FAISS_URL")
FAISS_DB_DIR = Path(os.getenv("BASE_FAISS_DIR", "/app/faiss_db"))

def _sincronizar_faiss_do_hub() -> None:
    """
    Descarrega o faiss_db/ do Hugging Face Dataset repo para o disco local.
    Se HF_TOKEN ou HF_DATASET_REPO não estiverem definidos, salta silenciosamente.
    Nunca bloqueia o arranque nem lança excepções não tratadas.
    """
    hf_token = os.getenv("HF_TOKEN")
    hf_repo = os.getenv("HF_DATASET_REPO")

    if not hf_token or not hf_repo:
        print("[BOOTSTRAP] HF_TOKEN ou HF_DATASET_REPO não configurados — a usar faiss_db local.")
        return

    faiss_base_dir = os.getenv("FAISS_BASE_DIR", "faiss_db")

    try:
        from huggingface_hub import snapshot_download

        print(f"[BOOTSTRAP] A sincronizar faiss_db/ de {hf_repo}...")
        snapshot_download(
            repo_id=hf_repo,
            repo_type="dataset",
            token=hf_token,
            local_dir=faiss_base_dir,
            local_dir_use_symlinks=False,
            ignore_patterns=[".gitattributes", "README.md", ".git*"],
        )
        print(f"[BOOTSTRAP] faiss_db/ sincronizado com sucesso de {hf_repo}.")
    except Exception as exc:
        print(f"[BOOTSTRAP] ⚠️  Falha ao sincronizar do HF Hub ({type(exc).__name__}: {exc}). A continuar com faiss_db local...")

def bootstrap():
    # Sincronização com HF Hub (best-effort)
    _sincronizar_faiss_do_hub()

    if not BOOTSTRAP_URL:
        print("[BOOTSTRAP] BOOTSTRAP_FAISS_URL não definida. A saltar download.")
        return

    if any(FAISS_DB_DIR.iterdir()):
        print(f"[BOOTSTRAP] Directório {FAISS_DB_DIR} não está vazio. A saltar bootstrap.")
        return

    print(f"[BOOTSTRAP] A iniciar download de {BOOTSTRAP_URL}...")
    temp_file = "/tmp/faiss_bootstrap.archive"
    
    try:
        urllib.request.urlretrieve(BOOTSTRAP_URL, temp_file)
        print("[BOOTSTRAP] Download concluído. A extrair...")

        if tarfile.is_tarfile(temp_file):
            with tarfile.open(temp_file) as tar:
                tar.extractall(path=FAISS_DB_DIR)
        elif zipfile.is_zipfile(temp_file):
            with zipfile.open(temp_file) as zip_ref:
                zip_ref.extractall(path=FAISS_DB_DIR)
        else:
            print("[BOOTSTRAP] Erro: Formato de arquivo desconhecido.")
            sys.exit(1)

        print(f"[BOOTSTRAP] Extração concluída para {FAISS_DB_DIR}")
    except Exception as e:
        print(f"[BOOTSTRAP] Erro durante o bootstrap: {e}")
        # Em produção, podemos decidir se falhamos ou se continuamos vazio
        # sys.exit(1) 

if __name__ == "__main__":
    bootstrap()