Spaces:
Running
Running
| # ./storage.py | |
| """ | |
| Persistence Layer - Handles the "Save/Load" functionality using Hugging Face Dataset as a database | |
| """ | |
| import json | |
| import os | |
| from datetime import datetime | |
| from huggingface_hub import HfApi, hf_hub_download | |
| # --- CONFIGURATION --- | |
| REPO_ID = "prashantmatlani/chathistorycoderg" | |
| HISTORY_DIR = "./chathistory" | |
| # Initialize the API with your token | |
| api = HfApi(token=os.getenv("HF_TOKEN")) | |
| def save_chat(chat_id, history): | |
| """Saves chat to local subdirectory and syncs to Hugging Face Dataset.""" | |
| if not os.path.exists(HISTORY_DIR): | |
| os.makedirs(HISTORY_DIR) | |
| # Generate a unique ID if none exists (e.g., for a brand new chat) | |
| if not chat_id: | |
| chat_id = datetime.now().strftime("%Y%m%d_%H%M%S") | |
| filename = f"{chat_id}.json" | |
| local_path = os.path.join(HISTORY_DIR, filename) | |
| # 1. Save Locally | |
| with open(local_path, "w", encoding="utf-8") as f: | |
| json.dump(history, f, indent=4) | |
| # 2. Sync to Hugging Face Dataset (Master Stroke Persistence) | |
| try: | |
| api.upload_file( | |
| path_or_fileobj=local_path, | |
| path_in_repo=f"chats/{filename}", | |
| repo_id=REPO_ID, | |
| repo_type="dataset" | |
| ) | |
| except Exception as e: | |
| print(f"Cloud Sync Warning: {e}") | |
| return chat_id | |
| def load_history(): | |
| """Retrieves list of chat IDs from the Hub to populate the sidebar.""" | |
| try: | |
| # We pull the list from the Hub so the sidebar reflects all saved sessions | |
| files = api.list_repo_files(repo_id=REPO_ID, repo_type="dataset") | |
| chat_files = [f.split("/")[-1].replace(".json", "") for f in files if f.startswith("chats/")] | |
| # IMPORTANT: Sorted by newest first; return as list of lists for Gradio Dataset component | |
| return [[f] for f in sorted(chat_files, reverse=True)] | |
| except: | |
| return [] | |
| def get_chat_content(chat_id): | |
| """Loads a specific chat's content from the Hub or local cache.""" | |
| filename = f"chats/{chat_id}.json" | |
| local_path = os.path.join(HISTORY_DIR, f"{chat_id}.json") | |
| try: | |
| # Ensure local dir exists | |
| if not os.path.exists(HISTORY_DIR): | |
| os.makedirs(HISTORY_DIR) | |
| # Download from Hub to keep local state fresh | |
| from huggingface_hub import hf_hub_download | |
| downloaded_path = hf_hub_download( | |
| repo_id=REPO_ID, | |
| repo_type="dataset", | |
| filename=filename, | |
| token=os.getenv("HF_TOKEN") | |
| ) | |
| with open(downloaded_path, "r", encoding="utf-8") as f: | |
| return json.load(f) | |
| except Exception: | |
| # Fallback to local if Hub is unreachable | |
| if os.path.exists(local_path): | |
| with open(local_path, "r", encoding="utf-8") as f: | |
| return json.load(f) | |
| return [] |