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()