z90486091/litellm-sqlite

1#!/usr/bin/env python32import hashlib3import logging4import os5import time6from huggingface_hub import hf_hub_download, upload_file7 8logging.basicConfig(9    format="%(asctime)s [db_syncer] %(levelname)s %(message)s",10    level=logging.INFO,11)12log = logging.getLogger("db_syncer")13 14HF_TOKEN = os.environ.get("HF_TOKEN", "")15HF_DATASET_REPO = os.environ.get("HF_DATASET_REPO", "z90486091/LITELLM_DB")16DB_PATH = os.environ.get("DB_PATH", "/app/data/litellm.db")17REPO_DB_FILENAME = "litellm.db"18SYNC_INTERVAL = int(os.environ.get("SYNC_INTERVAL", "300"))19 20 21def file_hash(path: str) -> str:22    h = hashlib.sha256()23    with open(path, "rb") as f:24        for chunk in iter(lambda: f.read(65536), b""):25            h.update(chunk)26    return h.hexdigest()27 28 29def download_db() -> bool:30    try:31        log.info(f"Downloading {REPO_DB_FILENAME} from {HF_DATASET_REPO}...")32        hf_hub_download(33            repo_id=HF_DATASET_REPO,34            filename=REPO_DB_FILENAME,35            repo_type="dataset",36            token=HF_TOKEN,37            local_dir=os.path.dirname(DB_PATH),38        )39        log.info(f"Downloaded DB to {DB_PATH}")40        return True41    except Exception as e:42        log.warning(f"Download failed (will use existing or fresh DB): {e}")43        return False44 45 46def upload_db() -> bool:47    try:48        upload_file(49            path_or_fileobj=DB_PATH,50            path_in_repo=REPO_DB_FILENAME,51            repo_id=HF_DATASET_REPO,52            repo_type="dataset",53            token=HF_TOKEN,54            commit_message="db_syncer: auto-sync litellm.db",55        )56        log.info("Synced litellm.db to HF Dataset repo")57        return True58    except Exception as e:59        log.error(f"Upload failed: {e}")60        return False61 62 63def main():64    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)65    download_db()66    log.info(f"Starting db_syncer — watching {DB_PATH} every {SYNC_INTERVAL}s")67    last_hash = None68 69    while True:70        try:71            if not os.path.exists(DB_PATH):72                log.warning(f"{DB_PATH} not found — waiting...")73                time.sleep(30)74                continue75            current_hash = file_hash(DB_PATH)76            if current_hash != last_hash:77                log.info(f"Change detected (hash: {current_hash[:8]}...) — syncing")78                if upload_db():79                    last_hash = current_hash80            else:81                log.debug("No change — skipping")82        except Exception as e:83            log.error(f"Unexpected error: {e}")84        time.sleep(SYNC_INTERVAL)85 86 87if __name__ == "__main__":88    main()