core/database.py

import os
import json
import tempfile
from datetime import datetime
from typing import List, Dict
from huggingface_hub import HfApi, upload_file, hf_hub_download, login

from config import REPO_ID, MASTER_SECRET

HF_TOKEN = os.getenv("HF_TOKEN", "")
if HF_TOKEN:
    try:
        login(token=HF_TOKEN)
        print("✅ Hugging Face login success")
    except Exception as e:
        print(f"⚠️ Hugging Face login failed: {e}")

hf_api = HfApi(token=HF_TOKEN if HF_TOKEN else None)

def save_results_to_hf(results_data: List[Dict]) -> bool:
    if not results_data:
        return True
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"results/result_{timestamp}.jsonl"
    content = "\n".join(json.dumps(r, ensure_ascii=False) for r in results_data)
    try:
        with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=False) as f:
            f.write(content)
            tmp_path = f.name
        try:
            hf_api.repo_info(repo_id=REPO_ID, repo_type="dataset")
        except:
            hf_api.create_repo(repo_id=REPO_ID, repo_type="dataset", exist_ok=True)
        upload_file(
            path_or_fileobj=tmp_path,
            path_in_repo=filename,
            repo_id=REPO_ID,
            repo_type="dataset",
            token=HF_TOKEN if HF_TOKEN else None
        )
        os.unlink(tmp_path)
        print(f"✅ Saved results: {filename} ({len(results_data)} entries)")
        return True
    except Exception as e:
        print(f"❌ Save failed: {e}")
        return False

def load_results_from_hf(limit: int = 100) -> List[Dict]:
    try:
        files = hf_api.list_repo_files(REPO_ID, repo_type="dataset")
        result_files = [f for f in files if f.startswith("results/") and f.endswith(".jsonl")]
        if not result_files:
            return []
        result_files.sort(reverse=True)
        all_results = []
        for fname in result_files[:5]:
            try:
                with tempfile.NamedTemporaryFile(suffix='.jsonl', delete=False) as tmp:
                    hf_hub_download(
                        repo_id=REPO_ID,
                        filename=fname,
                        repo_type="dataset",
                        local_dir=os.path.dirname(tmp.name),
                        local_dir_use_symlinks=False,
                        token=HF_TOKEN if HF_TOKEN else None
                    )
                    actual_path = os.path.join(os.path.dirname(tmp.name), os.path.basename(fname))
                    with open(actual_path, 'r') as f2:
                        for line in f2:
                            line = line.strip()
                            if line:
                                try:
                                    all_results.append(json.loads(line))
                                except:
                                    pass
                    os.unlink(actual_path)
            except Exception as e:
                print(f"⚠️ Load {fname} failed: {e}")
                continue
        return all_results[-limit:]
    except Exception as e:
        print(f"❌ Load results failed: {e}")
        return []