Spaces:
Sleeping
Sleeping
| import sqlite3 | |
| import threading | |
| import time | |
| import os | |
| from huggingface_hub import hf_hub_download, HfApi | |
| import config | |
| class DatabaseManager: | |
| def __init__(self): | |
| self.api = HfApi(token=config.HF_TOKEN) | |
| self.lock = threading.Lock() | |
| self._last_sync = 0 | |
| self._needs_sync = False | |
| self._sync_thread = None | |
| self._stop_event = threading.Event() | |
| def initialize(self): | |
| """啟動時下載資料庫或建立新的""" | |
| with self.lock: | |
| try: | |
| # 確保目錄存在 | |
| config.LOCAL_DB_PATH.parent.mkdir(parents=True, exist_ok=True) | |
| print(f"正在從 {config.DB_REPO_ID} 下載資料庫...") | |
| hf_hub_download( | |
| repo_id=config.DB_REPO_ID, | |
| filename=config.DB_FILENAME, | |
| repo_type="dataset", | |
| token=config.HF_TOKEN, | |
| local_dir=str(config.RUNTIME_ROOT), | |
| local_dir_use_symlinks=False | |
| ) | |
| print("資料庫下載成功。") | |
| except Exception as e: | |
| print(f"無法下載資料庫 (可能是首次建立): {e}") | |
| self._create_empty_db() | |
| self._ensure_tables() | |
| # 啟動背景同步執行緒 | |
| if not self._sync_thread or not self._sync_thread.is_alive(): | |
| self._stop_event.clear() | |
| self._sync_thread = threading.Thread(target=self._sync_worker, daemon=True) | |
| self._sync_thread.start() | |
| def _sync_worker(self): | |
| """背景同步工作執行緒，負責彙整變更並重試""" | |
| print("資料庫同步執行緒已啟動。") | |
| while not self._stop_event.is_set(): | |
| if self._needs_sync: | |
| # 等待一小段時間（例如 5 秒），彙整多個連續變更 | |
| time.sleep(5) | |
| self.sync_to_hf() | |
| else: | |
| time.sleep(1) | |
| def _create_empty_db(self): | |
| conn = sqlite3.connect(config.LOCAL_DB_PATH) | |
| conn.close() | |
| def _ensure_tables(self): | |
| conn = sqlite3.connect(config.LOCAL_DB_PATH) | |
| cursor = conn.cursor() | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS knowledge ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| key TEXT UNIQUE, | |
| content TEXT, | |
| updated_at DATETIME DEFAULT CURRENT_TIMESTAMP | |
| ) | |
| """) | |
| # 新增 personas 表，支援儲存多種人設 | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS personas ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| name TEXT UNIQUE, | |
| content TEXT, | |
| is_active INTEGER DEFAULT 0, | |
| updated_at DATETIME DEFAULT CURRENT_TIMESTAMP | |
| ) | |
| """) | |
| cursor.execute(""" | |
| CREATE TABLE IF NOT EXISTS chat_history ( | |
| id INTEGER PRIMARY KEY AUTOINCREMENT, | |
| role TEXT, | |
| content TEXT, | |
| timestamp DATETIME DEFAULT CURRENT_TIMESTAMP | |
| ) | |
| """) | |
| conn.commit() | |
| conn.close() | |
| def sync_to_hf(self, max_retries=3): | |
| if not config.HF_TOKEN: | |
| self._needs_sync = False | |
| return | |
| with self.lock: | |
| for attempt in range(max_retries): | |
| try: | |
| print(f"正在同步資料庫到 {config.DB_REPO_ID} (嘗試 {attempt + 1}/{max_retries})...") | |
| self.api.upload_file( | |
| path_or_fileobj=str(config.LOCAL_DB_PATH), | |
| path_in_repo=config.DB_FILENAME, | |
| repo_id=config.DB_REPO_ID, | |
| repo_type="dataset", | |
| commit_message=f"Auto-sync database at {time.strftime('%H:%M:%S')}" | |
| ) | |
| self._last_sync = time.time() | |
| self._needs_sync = False | |
| print("同步成功。") | |
| return | |
| except Exception as e: | |
| print(f"同步嘗試 {attempt + 1} 失敗: {e}") | |
| if "503" in str(e) or "Service Temporarily Unavailable" in str(e): | |
| # 針對 HF 503 錯誤，等待較長時間後重試 | |
| wait_time = (attempt + 1) * 10 | |
| print(f"HF 服務忙碌中，等待 {wait_time} 秒後重試...") | |
| time.sleep(wait_time) | |
| else: | |
| # 其他錯誤，稍微等待後重試 | |
| time.sleep(2) | |
| print("所有同步嘗試皆失敗，將於下一次週期再次嘗試。") | |
| def _request_sync(self): | |
| """標記需要同步，由背景執行緒處理""" | |
| self._needs_sync = True | |
| def query_knowledge(self, user_input): | |
| if not user_input or len(user_input.strip()) < 2: | |
| return [] | |
| conn = sqlite3.connect(config.LOCAL_DB_PATH) | |
| cursor = conn.cursor() | |
| # 1. 先嘗試原有的關鍵字比對 (User Input 包含 Key) | |
| cursor.execute(""" | |
| SELECT content FROM knowledge | |
| WHERE ? LIKE '%' || key || '%' | |
| AND length(key) > 1 | |
| ORDER BY length(key) DESC | |
| LIMIT 3 | |
| """, (user_input,)) | |
| results = cursor.fetchall() | |
| # 2. 如果結果不夠，嘗試「反向比對」(Key 包含 User Input 中的關鍵詞，或 Content 包含 User Input 中的關鍵詞) | |
| if len(results) < 5: | |
| # 簡單提取關鍵詞 (移除常見廢話) | |
| clean_input = user_input | |
| for stop in ["什麼是", "請問", "你知道", "嗎", "關於", "我想問", "的"]: | |
| clean_input = clean_input.replace(stop, "") | |
| clean_input = clean_input.strip("?？!！ ") | |
| if len(clean_input) >= 2: | |
| # 搜尋內容或 Key 中包含這個關鍵詞的內容 | |
| cursor.execute(""" | |
| SELECT content FROM knowledge | |
| WHERE (content LIKE '%' || ? || '%' OR key LIKE '%' || ? || '%') | |
| AND id NOT IN (SELECT id FROM ( | |
| SELECT id FROM knowledge WHERE ? LIKE '%' || key || '%' | |
| )) | |
| LIMIT ? | |
| """, (clean_input, clean_input, user_input, 5 - len(results))) | |
| results.extend(cursor.fetchall()) | |
| conn.close() | |
| # 確保回傳值是字串列表，且去除重複 | |
| final_results = [] | |
| seen = set() | |
| for r in results: | |
| content = r[0] | |
| if content not in seen: | |
| final_results.append(content) | |
| seen.add(content) | |
| return final_results[:5] | |
| def add_knowledge(self, key, content): | |
| with self.lock: | |
| conn = sqlite3.connect(config.LOCAL_DB_PATH) | |
| cursor = conn.cursor() | |
| cursor.execute("INSERT OR REPLACE INTO knowledge (key, content) VALUES (?, ?)", (key, content)) | |
| conn.commit() | |
| conn.close() | |
| self._request_sync() | |
| def add_persona(self, name, content, set_active=False): | |
| with self.lock: | |
| conn = sqlite3.connect(config.LOCAL_DB_PATH) | |
| cursor = conn.cursor() | |
| if set_active: | |
| cursor.execute("UPDATE personas SET is_active = 0") | |
| cursor.execute(""" | |
| INSERT OR REPLACE INTO personas (name, content, is_active) | |
| VALUES (?, ?, ?) | |
| """, (name, content, 1 if set_active else 0)) | |
| conn.commit() | |
| conn.close() | |
| self._request_sync() | |
| def get_all_personas(self): | |
| conn = sqlite3.connect(config.LOCAL_DB_PATH) | |
| cursor = conn.cursor() | |
| cursor.execute("SELECT name, is_active FROM personas") | |
| rows = cursor.fetchall() | |
| conn.close() | |
| return [{"name": r[0], "active": bool(r[1])} for r in rows] | |
| def switch_persona(self, name): | |
| with self.lock: | |
| conn = sqlite3.connect(config.LOCAL_DB_PATH) | |
| cursor = conn.cursor() | |
| cursor.execute("UPDATE personas SET is_active = 0") | |
| cursor.execute("UPDATE personas SET is_active = 1 WHERE name = ?", (name,)) | |
| conn.commit() | |
| conn.close() | |
| self._request_sync() | |
| def get_active_persona(self): | |
| try: | |
| conn = sqlite3.connect(config.LOCAL_DB_PATH) | |
| cursor = conn.cursor() | |
| cursor.execute("SELECT content FROM personas WHERE is_active = 1") | |
| row = cursor.fetchone() | |
| if not row: | |
| # 備用：檢查舊有的 knowledge 表 | |
| cursor.execute("SELECT content FROM knowledge WHERE key = 'persona'") | |
| row = cursor.fetchone() | |
| conn.close() | |
| return row[0] if row else None | |
| except Exception: | |
| return None | |
| db_mgr = DatabaseManager() | |