Update app.py

app.py

@@ -1,6 +1,6 @@
 # ==============================================================================
-#  Text-to-SQL 智能助手 - Hugging Face CPU 最终版 v6
-#  (融合模板引擎 + 强化 Prompt + 修复所有 Bug)
+#  Text-to-SQL 智能助手 - Hugging Face CPU 最终版 v7
+#  (修复所有已知 Bug)
 # ==============================================================================
 import gradio as gr
 import os
@@ -22,57 +22,32 @@ from transformers import AutoModel, AutoTokenizer
 import torch.nn.functional as F
 
 # ==================== 配置參數 ====================
-# --- Hugging Face CPU 部署配置 ---
 GGUF_FILENAME = "qwen2.5-coder-1.5b-sql-finetuned.q8_0.gguf"
-N_GPU_LAYERS = 0 # 在 Hugging Face CPU 环境下设置为 0
-
+N_GPU_LAYERS = 0 # CPU 环境
 DATASET_REPO_ID = "Paul720810/Text-to-SQL-Softline"
 GGUF_REPO_ID = "Paul720810/gguf-models"
 FEW_SHOT_EXAMPLES_COUNT = 1
-DEVICE = "cuda" if torch.cuda.is_available() and N_GPU_LAYERS != 0 else "cpu"
+DEVICE = "cpu"
 EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
-
 TEMP_DIR = tempfile.gettempdir()
 os.makedirs(os.path.join(TEMP_DIR, 'text_to_sql_cache'), exist_ok=True)
 
 print("=" * 60)
-print("🤖 Text-to-SQL 智能助手 v6.0 (Hugging Face CPU 版)...")
+print("🤖 Text-to-SQL 智能助手 v7.0 (Hugging Face CPU 版)...")
 print(f"🚀 模型: {GGUF_FILENAME}")
 print(f"💻 設備: {DEVICE} (GPU Layers: {N_GPU_LAYERS})")
 print("=" * 60)
 
 # ==================== 工具函數 ====================
-def get_current_time():
-    return datetime.now().strftime('%Y-%m-%d %H:%M:%S')
-
+def get_current_time(): return datetime.now().strftime('%Y-%m-%d %H:%M:%S')
 def format_log(message: str, level: str = "INFO") -> str:
-    log_entry = f"[{get_current_time()}] [{level.upper()}] {message}"
-    print(log_entry)
-    return log_entry
-
+    log_entry = f"[{get_current_time()}] [{level.upper()}] {message}"; print(log_entry); return log_entry
 def parse_sql_from_response(response_text: str) -> Optional[str]:
     if not response_text: return None
     response_text = response_text.strip()
     match = re.search(r"```sql\s*\n(.*?)\n```", response_text, re.DOTALL | re.IGNORECASE)
     if match: return match.group(1).strip()
-    match = re.search(r"```\s*\n?(.*?)\n?```", response_text, re.DOTALL)
-    if match:
-        sql_candidate = match.group(1).strip()
-        if sql_candidate.upper().startswith('SELECT'): return sql_candidate
-    match = re.search(r"(SELECT\s+.*?;)", response_text, re.DOTALL | re.IGNORECASE)
-    if match: return match.group(1).strip()
-    match = re.search(r"(SELECT\s+.*?)(?=\n\n|\n```|$|\n[^,\s])", response_text, re.DOTALL | re.IGNORECASE)
-    if match:
-        sql = match.group(1).strip()
-        if not sql.endswith(';'): sql += ';'
-        return sql
-    if 'SELECT' in response_text.upper():
-        for line in response_text.split('\n'):
-            line = line.strip()
-            if line.upper().startswith('SELECT'):
-                if not line.endswith(';'): line += ';'
-                return line
-    return None
+    return response_text if response_text.upper().startswith("SELECT") else None
 
 # ==================== Text-to-SQL 核心類 ====================
 class TextToSQLSystem:
@@ -84,21 +59,15 @@ class TextToSQLSystem:
             self._log(f"載入嵌入模型: {embed_model_name}")
             self.embed_tokenizer = AutoTokenizer.from_pretrained(embed_model_name)
             self.embed_model = AutoModel.from_pretrained(embed_model_name)
-            if DEVICE == "cuda":
-                self.embed_model.to(DEVICE)
-
             self.schema = self._load_schema()
             self.dataset, self.faiss_index = self._load_and_index_dataset()
             self._load_gguf_model()
             self._log("✅ 系統初始化完成")
         except Exception as e:
-            self._log(f"❌ 系統初始化過程中發生嚴重錯誤: {e}", "CRITICAL")
-            self._log(traceback.format_exc(), "DEBUG")
-            self.llm = None
-
-    def _log(self, message: str, level: str = "INFO"):
-        self.log_history.append(format_log(message, level))
+            self._log(f"❌ 系統初始化過程中發生嚴重錯誤: {e}", "CRITICAL"); self._log(traceback.format_exc(), "DEBUG"); self.llm = None
 
+    def _log(self, message: str, level: str = "INFO"): self.log_history.append(format_log(message, level))
+    
     def _load_gguf_model(self):
         try:
             model_path = hf_hub_download(repo_id=GGUF_REPO_ID, filename=GGUF_FILENAME, repo_type="dataset", cache_dir=TEMP_DIR)
@@ -107,54 +76,47 @@ class TextToSQLSystem:
             self.llm = Llama(model_path=model_path, n_ctx=2048, n_threads=4, n_batch=512, verbose=False, n_gpu_layers=N_GPU_LAYERS)
             self._log("✅ GGUF 模型成功載入")
         except Exception as e:
-            self._log(f"❌ GGUF 載入失敗: {e}", "CRITICAL")
-            self.llm = None
+            self._log(f"❌ GGUF 載入失敗: {e}", "CRITICAL"); self.llm = None
 
     def huggingface_api_call(self, prompt: str) -> str:
         if self.llm is None: return ""
         try:
             output = self.llm(prompt, max_tokens=150, temperature=0.1, top_p=0.9, echo=False, stop=["```", ";", "\n\n", "</s>", "###", "Q:"], repeat_penalty=1.1)
-            generated_text = output["choices"][0]["text"] if output and "choices" in output and len(output["choices"]) > 0 else ""
+            generated_text = output["choices"]["text"] if output and "choices" in output and len(output["choices"]) > 0 else ""
             self._log(f"🧠 模型原始輸出: {generated_text.strip()}", "DEBUG")
             return generated_text.strip()
         except Exception as e:
-            self._log(f"❌ 模型生成錯誤: {e}", "CRITICAL")
-            return ""
+            self._log(f"❌ 模型生成錯誤: {e}", "CRITICAL"); return ""
     
     def _load_schema(self) -> Dict:
         try:
             schema_path = hf_hub_download(repo_id=DATASET_REPO_ID, filename="sqlite_schema_FULL.json", repo_type="dataset")
-            with open(schema_path, "r", encoding="utf-8") as f:
-                schema_data = json.load(f)
+            with open(schema_path, "r", encoding="utf-8") as f: schema_data = json.load(f)
             self._log(f"📊 Schema 載入成功，包含 {len(schema_data)} 個表格。")
             return schema_data
         except Exception as e:
-            self._log(f"❌ 載入 schema 失敗: {e}", "ERROR")
-            return {}
+            self._log(f"❌ 載入 schema 失敗: {e}", "ERROR"); return {}
 
     def _encode_texts(self, texts):
         if isinstance(texts, str): texts = [texts]
-        inputs = self.embed_tokenizer(texts, padding=True, truncation=True, return_tensors="pt", max_length=512).to(DEVICE)
+        inputs = self.embed_tokenizer(texts, padding=True, truncation=True, return_tensors="pt", max_length=512)
         with torch.no_grad():
             outputs = self.embed_model(**inputs)
-        embeddings = outputs.last_hidden_state.mean(dim=1)
-        return embeddings.cpu()
+        return outputs.last_hidden_state.mean(dim=1).cpu()
 
     def _load_and_index_dataset(self):
         try:
             dataset = load_dataset(DATASET_REPO_ID, data_files="training_data.jsonl", split="train")
             dataset = dataset.filter(lambda ex: isinstance(ex.get("messages"), list) and len(ex["messages"]) >= 2)
-            corpus = [item['messages'][0]['content'] for item in dataset]
+            corpus = [item['messages']['content'] for item in dataset]
             self._log(f"正在編碼 {len(corpus)} 個問題...")
             all_embeddings = torch.cat([self._encode_texts(corpus[i:i+32]) for i in range(0, len(corpus), 32)], dim=0).numpy()
-            index = faiss.IndexFlatIP(all_embeddings.shape[1])
+            index = faiss.IndexFlatIP(all_embeddings.shape)
             index.add(all_embeddings.astype('float32'))
             self._log("✅ 向量索引建立完成")
             return dataset, index
         except Exception as e:
-            self._log(f"❌ 載入數據失敗: {e}", "ERROR")
-            self._log(traceback.format_exc(), "DEBUG")
-            return None, None
+            self._log(f"❌ 載入數據失敗: {e}", "ERROR"); self._log(traceback.format_exc(), "DEBUG"); return None, None
             
     def _identify_relevant_tables(self, question: str) -> List[str]:
         question_lower = question.lower()
@@ -185,31 +147,27 @@ class TextToSQLSystem:
             q_embedding = self._encode_texts([question]).numpy().astype('float32')
             distances, indices = self.faiss_index.search(q_embedding, min(top_k + 2, len(self.dataset)))
             results, seen_questions = [], set()
-            for i, idx in enumerate(indices[0]):
+            for i, idx in enumerate(indices):
                 if len(results) >= top_k: break
                 idx = int(idx)
                 if idx >= len(self.dataset): continue
                 item = self.dataset[idx]
                 if not (isinstance(item.get('messages'), list) and len(item['messages']) >= 2): continue
-                q_content = (item['messages'][0].get('content') or '').strip()
-                a_content = (item['messages'][1].get('content') or '').strip()
+                q_content = (item['messages'].get('content') or '').strip()
+                a_content = (item['messages'].get('content') or '').strip()
                 if not q_content or not a_content: continue
                 clean_q = re.sub(r"以下是一個SQL查詢任務：\s*指令:\s*", "", q_content).strip()
                 if clean_q in seen_questions: continue
                 seen_questions.add(clean_q)
                 sql = parse_sql_from_response(a_content) or "無法解析範例SQL"
-                results.append({"similarity": float(distances[0][i]), "question": clean_q, "sql": sql})
+                results.append({"similarity": float(distances[i]), "question": clean_q, "sql": sql})
             return results
         except Exception as e:
-            self._log(f"❌ 檢索失敗: {e}", "ERROR")
-            return []
+            self._log(f"❌ 檢索失敗: {e}", "ERROR"); return []
 
     def _build_prompt(self, user_q: str, examples: List[Dict]) -> str:
         schema_str = self._format_relevant_schema(self._identify_relevant_tables(user_q))
-        example_str = ""
-        if examples:
-            example_prompts = [f"Q: {ex['question']}\nA: ```sql\n{ex['sql']}\n```" for ex in examples]
-            example_str = "\n---\n".join(example_prompts)
+        example_str = "\n---\n".join([f"Q: {ex['question']}\nA: ```sql\n{ex['sql']}\n```" for ex in examples]) if examples else "No examples provided."
         prompt = f"""You are an expert SQLite programmer. Your task is to generate a SQL query based on the database schema and a user's question.
 
 ## Database Schema
@@ -238,12 +196,8 @@ A: ```sql
         entity_patterns = [
             {'pattern': r"(买家|buyer)\s*(?:id|代號|代码)\s*'\"?\b([A-Z]\d{4}[A-Z])\b'\"?", 'column': 'sd.BuyerID', 'type': '买家ID'},
             {'pattern': r"(申请方|申请厂商|applicant)\s*(?:id|代號|代码)\s*'\"?\b([A-Z]\d{4}[A-Z])\b'\"?", 'column': 'sd.ApplicantID', 'type': '申请方ID'},
-            {'pattern': r"(付款方|付款厂商|invoiceto)\s*(?:id|代號|代码)\s*'\"?\b([A-Z]\d{4}[A-Z])\b'\"?", 'column': 'sd.InvoiceToID', 'type': '付款方ID'},
-            {'pattern': r"(代理商|agent)\s*(?:id|代號|代码)\s*'\"?\b([A-Z]\d{4}[A-Z])\b'\"?", 'column': 'sd.AgentID', 'type': '代理商ID'},
             {'pattern': r"(买家|buyer|客戶)\s*'\"?([a-zA-Z0-9&.\s-]+?)(?:\s*的|\s+|$|有)", 'column': 'sd.BuyerName', 'type': '买家'},
             {'pattern': r"(申请方|申请厂商|applicant)\s*'\"?([a-zA-Z0-9&.\s-]+?)(?:\s*的|\s+|$|有)", 'column': 'sd.ApplicantName', 'type': '申请方'},
-            {'pattern': r"(付款方|付款厂商|invoiceto)\s*'\"?([a-zA-Z0-9&.\s-]+?)(?:\s*的|\s+|$|有)", 'column': 'sd.InvoiceToName', 'type': '付款方'},
-            {'pattern': r"(代理商|agent)\s*'\"?([a-zA-Z0-9&.\s-]+?)(?:\s*的|\s+|$|有)", 'column': 'sd.AgentName', 'type': '代理商'},
             {'pattern': r"\b([A-Z]\d{4}[A-Z])\b", 'column': 'sd.ApplicantID', 'type': 'ID'}
         ]
         for p in entity_patterns:
@@ -253,64 +207,38 @@ A: ```sql
                 entity_match_data = {"type": p['type'], "name": entity_value.strip().upper(), "column": p['column']}
                 break
 
-        if any(kw in q_lower for kw in ['報告號碼', '報告清單', '列出報告', 'report number', 'list of reports']):
-            year_match = re.search(r'(\d{4})\s*年?', question)
-            month_match = re.search(r'(\d{1,2})\s*月', question)
-            from_clause = "FROM JobTimeline AS jt"
-            select_clause = "SELECT jt.JobNo, jt.ReportAuthorization"
-            where_conditions = ["jt.ReportAuthorization IS NOT NULL"]
-            log_parts = []
-            if year_match: where_conditions.append(f"strftime('%Y', jt.ReportAuthorization) = '{year_match.group(1)}'"); log_parts.append(f"{year_match.group(1)}年")
-            if month_match: where_conditions.append(f"strftime('%m', jt.ReportAuthorization) = '{month_match.group(1).zfill(2)}'"); log_parts.append(f"{month_match.group(1)}月")
-            if 'fail' in q_lower or '失敗' in q_lower:
-                if "JOIN TSR53SampleDescription" not in from_clause: from_clause += " JOIN TSR53SampleDescription AS sd ON jt.JobNo = sd.JobNo"
-                where_conditions.append("sd.OverallRating = 'Fail'"); log_parts.append("Fail")
-            elif 'pass' in q_lower or '通過' in q_lower:
-                if "JOIN TSR53SampleDescription" not in from_clause: from_clause += " JOIN TSR53SampleDescription AS sd ON jt.JobNo = sd.JobNo"
-                where_conditions.append("sd.OverallRating = 'Pass'"); log_parts.append("Pass")
+        if any(kw in q_lower for kw in ['報告號碼', '報告清單', '列出報告']):
+            year_match, month_match = re.search(r'(\d{4})\s*年?', question), re.search(r'(\d{1,2})\s*月', question)
+            from_clause = "FROM JobTimeline AS jt JOIN TSR53SampleDescription AS sd ON jt.JobNo = sd.JobNo"
+            select_clause = "SELECT jt.JobNo, sd.BuyerName, jt.ReportAuthorization"
+            where, log_parts = ["jt.ReportAuthorization IS NOT NULL"], []
+            if year_match: where.append(f"strftime('%Y', jt.ReportAuthorization) = '{year_match.group(1)}'"); log_parts.append(f"{year_match.group(1)}年")
+            if month_match: where.append(f"strftime('%m', jt.ReportAuthorization) = '{month_match.group(1).zfill(2)}'"); log_parts.append(f"{month_match.group(1)}月")
+            if 'fail' in q_lower or '失敗' in q_lower: where.append("sd.OverallRating = 'Fail'"); log_parts.append("Fail")
+            elif 'pass' in q_lower or '通過' in q_lower: where.append("sd.OverallRating = 'Pass'"); log_parts.append("Pass")
             if entity_match_data:
-                entity_name, column_name = entity_match_data["name"], entity_match_data["column"]
-                if "JOIN TSR53SampleDescription" not in from_clause: from_clause += " JOIN TSR53SampleDescription AS sd ON jt.JobNo = sd.JobNo"
-                match_operator = "=" if column_name.endswith("ID") else "LIKE"
-                entity_value = f"'{entity_name}'" if match_operator == "=" else f"'%{entity_name}%'"
-                where_conditions.append(f"{column_name} {match_operator} {entity_value}")
-                log_parts.append(entity_name)
-                select_clause = "SELECT jt.JobNo, sd.BuyerName, jt.ReportAuthorization"
-            final_where_clause = "WHERE " + " AND ".join(where_conditions) if where_conditions else ""
-            time_log = " ".join(log_parts) if log_parts else "全部"
-            self._log(f"🔄 檢測到查詢【{time_log} 報告列表】意圖，啟用智能模板。", "INFO")
-            template_sql = f"{select_clause} {from_clause} {final_where_clause} ORDER BY jt.ReportAuthorization DESC;"
-            return self._finalize_sql(template_sql, f"模板覆寫: {time_log} 報告列表查詢")
+                op = "=" if entity_match_data["column"].endswith("ID") else "LIKE"
+                val = f"'{entity_match_data['name']}'" if op == "=" else f"'%{entity_match_data['name']}%'"
+                where.append(f"{entity_match_data['column']} {op} {val}"); log_parts.append(entity_match_data["name"])
+            final_where = "WHERE " + " AND ".join(where) if where else ""
+            log_msg = " ".join(log_parts) if log_parts else "全部"
+            self._log(f"🔄 檢測到【{log_msg} 報告列表】意圖，啟用模板。", "INFO")
+            template_sql = f"{select_clause} {from_clause} {final_where} ORDER BY jt.ReportAuthorization DESC;"
+            return self._finalize_sql(template_sql, f"模板覆寫: {log_msg}")
 
         if '報告' in q_lower and any(kw in q_lower for kw in ['幾份', '多少', '數量', '總數']) and not entity_match_data:
             year_match = re.search(r'(\d{4})\s*年?', question)
-            time_condition, time_log = "", "總"
-            if year_match:
-                time_condition = f"WHERE ReportAuthorization IS NOT NULL AND strftime('%Y', ReportAuthorization) = '{year_match.group(1)}'"
-                time_log = f"{year_match.group(1)}年"
-            else:
-                time_condition = "WHERE ReportAuthorization IS NOT NULL"
-            self._log(f"🔄 檢測到查詢【{time_log}全局報告總數】意圖，啟用模板。", "INFO")
-            template_sql = f"SELECT COUNT(DISTINCT JobNo) AS report_count FROM JobTimeline {time_condition};"
-            return self._finalize_sql(template_sql, f"模板覆寫: {time_log}全局報告總數查詢")
-
-        self._log("未觸發任何模板，嘗試解析並修正 AI 輸出...", "INFO")
+            time_cond = f"WHERE strftime('%Y', ReportAuthorization) = '{year_match.group(1)}'" if year_match else ""
+            time_log = f"{year_match.group(1)}年" if year_match else "总"
+            self._log(f"🔄 檢測到【{time_log}报告总数】意图，启用模板。", "INFO")
+            template_sql = f"SELECT COUNT(DISTINCT JobNo) AS report_count FROM JobTimeline {time_cond};"
+            return self._finalize_sql(template_sql, f"模板覆寫: {time_log}")
+        
+        # Fallback to AI
+        self._log("未觸發模板，由 AI 生成...", "INFO")
         parsed_sql = parse_sql_from_response(raw_response)
-        if not parsed_sql:
-            return None, f"無法解析SQL。原始回應:\n{raw_response}"
-        fixed_sql = " " + parsed_sql.strip() + " "
-        fixes_applied_fallback = []
-        dialect_corrections = {r'YEAR\s*\(([^)]+)\)': r"strftime('%Y', \1)"}
-        for p, r in dialect_corrections.items():
-            if re.search(p, fixed_sql, re.IGNORECASE):
-                fixed_sql = re.sub(p, r, fixed_sql, flags=re.IGNORECASE); fixes_applied_fallback.append(f"修正方言: {p}")
-        schema_corrections = {'TSR53Report':'TSR53SampleDescription', 'TSR53InvoiceReportNo':'JobNo', 'Status':'OverallRating'}
-        for w, c in schema_corrections.items():
-            pattern = r'\b' + re.escape(w) + r'\b'
-            if re.search(pattern, fixed_sql, re.IGNORECASE):
-                fixed_sql = re.sub(pattern, c, fixed_sql, flags=re.IGNORECASE); fixes_applied_fallback.append(f"映射 Schema: '{w}' -> '{c}'")
-        log_msg = "AI 生成並成功修正" if fixes_applied_fallback else "AI 生成且無需修正"
-        return self._finalize_sql(fixed_sql, log_msg)
+        if not parsed_sql: return None, f"無法解析SQL: {raw_response}"
+        return self._finalize_sql(parsed_sql, "AI 生成")
 
     def process_question(self, question: str) -> Tuple[str, str]:
         if question in self.query_cache: self._log("⚡ 使用緩存結果"); return self.query_cache[question]
@@ -323,42 +251,28 @@ A: ```sql
         self._log("🧠 開始生成 AI 回應...")
         response = self.huggingface_api_call(prompt)
         final_sql, status_message = self._validate_and_fix_sql(question, response)
-        if not final_sql: result = (status_message, "生成失敗")
-        else: result = (final_sql, status_message)
+        result = (final_sql, status_message) if final_sql else (status_message, "生成失敗")
         self.query_cache[question] = result
         return result
 
 # ==================== Gradio 介面 ====================
 text_to_sql_system = TextToSQLSystem()
-
 def process_query(q: str):
     if not q.strip(): return "", "等待輸入", "請輸入問題"
-    if text_to_sql_system.llm is None:
-        return "模型未能成功載入，請檢查終端日誌。", "模型載入失敗", "\n".join(text_to_sql_system.log_history)
+    if text_to_sql_system.llm is None: return "模型未能成功載入...", "模型載入失敗", "\n".join(text_to_sql_system.log_history)
     sql, status = text_to_sql_system.process_question(q)
-    logs = "\n".join(text_to_sql_system.log_history[-15:])
-    return sql, status, logs
+    return sql, status, "\n".join(text_to_sql_system.log_history[-15:])
 
-examples = [
-    "2024年7月買家 Gap 的 Fail 報告號碼",
-    "列出2023年所有失败的报告",
-    "找出总金额最高的10个工作单",
-    "哪些客户的工作单数量最多？",
-    "A組2024年完成了多少個測試項目？",
-    "2024年每月完成多少份報告？"
-]
+examples = ["2024年7月買家 Gap 的 Fail 報告號碼", "列出2023年所有失败的报告", "找出总金额最高的10个工作单", "哪些客户的工作单数量最多？", "A組2024年完成了多少個測試項目？", "2024年每月完成多少份報告？"]
 with gr.Blocks(theme=gr.themes.Soft(), title="Text-to-SQL 智能助手") as demo:
     gr.Markdown("# ⚡ Text-to-SQL 智能助手 (终极版)")
-    gr.Markdown("融合了模板引擎和 GGUF 模型的强大版本")
     with gr.Row():
         with gr.Column(scale=2):
             inp = gr.Textbox(lines=3, label="💬 您的問題", placeholder="例如：2024年每月完成多少份報告？")
             btn = gr.Button("🚀 生成 SQL", variant="primary")
             status = gr.Textbox(label="狀態", interactive=False)
-        with gr.Column(scale=3):
-            sql_out = gr.Code(label="🤖 生成的 SQL", language="sql", lines=8)
-    with gr.Accordion("📋 處理日誌", open=False):
-        logs = gr.Textbox(lines=10, label="日誌", interactive=False)
+        with gr.Column(scale=3): sql_out = gr.Code(label="🤖 生成的 SQL", language="sql", lines=8)
+    with gr.Accordion("📋 處理日誌", open=False): logs = gr.Textbox(lines=10, label="日誌", interactive=False)
     gr.Examples(examples=examples, inputs=inp, label="💡 點擊試用範例問題")
     btn.click(process_query, inputs=[inp], outputs=[sql_out, status, logs])
     inp.submit(process_query, inputs=[inp], outputs=[sql_out, status, logs])