Update app.py

app.py

@@ -15,53 +15,40 @@ from typing import List, Dict, Tuple, Optional
 # ==================== 配置區 ====================
 HF_TOKEN = os.environ.get("HF_TOKEN", "您的_HuggingFace_Token")
 DATASET_REPO_ID = "Paul720810/Text-to-SQL-Softline"
-SIMILARITY_THRESHOLD = 0.75  # 相似度閾值
+SIMILARITY_THRESHOLD = 0.75
 
-# 多個備用LLM模型（保證可用性）
+# 多個備用LLM模型
 LLM_MODELS = [
     "https://api-inference.huggingface.co/models/gpt2",
     "https://api-inference.huggingface.co/models/distilgpt2", 
     "https://api-inference.huggingface.co/models/microsoft/DialoGPT-small"
 ]
 
-# 數據庫連接配置（可選）
-DB_CONFIG = {
-    "enabled": False,  # 設置為True啟用真實數據庫連接
-    "path": "您的數據庫路徑.db",
-    "test_queries": True  # 是否啟用SQL測試功能
-}
-
 print("=" * 50)
 print("🚀 智能 Text-to-SQL 系統啟動中...")
 print("=" * 50)
 
 # ==================== 工具函數 ====================
 def get_current_time():
-    """獲取當前時間字符串"""
     return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
 
-def safe_json_load(data, default=None):
-    """安全的JSON解析"""
-    try:
-        return json.loads(data) if isinstance(data, str) else data
-    except (json.JSONDecodeError, TypeError):
-        return default
-
 def validate_sql(sql_query: str) -> Dict:
     """驗證SQL語句的安全性"""
     security_issues = []
     
     # 檢查危險操作
     dangerous_keywords = ['DROP', 'DELETE', 'INSERT', 'UPDATE', 'ALTER', 'TRUNCATE', 'EXEC', 'EXECUTE']
+    sql_upper = sql_query.upper()
+    
     for keyword in dangerous_keywords:
-        if f" {keyword} " in sql_query.upper():
+        if f" {keyword} " in sql_upper:
             security_issues.append(f"發現危險操作: {keyword}")
     
     # 檢查基本語法
-    if "SELECT" not in sql_query.upper():
+    if "SELECT" not in sql_upper:
         security_issues.append("缺少SELECT語句")
     
-    if "FROM" not in sql_query.upper():
+    if "FROM" not in sql_upper:
         security_issues.append("缺少FROM子句")
     
     return {
@@ -70,30 +57,28 @@ def validate_sql(sql_query: str) -> Dict:
         "is_safe": len([i for i in security_issues if '危險' in i]) == 0
     }
 
-def execute_test_query(sql_query: str) -> Tuple[bool, str]:
-    """執行測試查詢（可選功能）"""
-    if not DB_CONFIG["enabled"]:
-        return False, "數據庫連接未啟用"
+def repair_sql(sql_query: str) -> str:
+    """修復有問題的SQL語句"""
+    if not sql_query or sql_query.strip() == "":
+        return "SELECT 'SQL語句為空' AS error;"
     
-    try:
-        validation = validate_sql(sql_query)
-        if not validation["is_safe"]:
-            return False, f"SQL安全檢查失敗: {', '.join(validation['issues'])}"
-        
-        # 連接數據庫並執行
-        conn = sqlite3.connect(DB_CONFIG["path"])
-        df = pd.read_sql_query(sql_query, conn)
-        conn.close()
-        
-        if len(df) == 0:
-            return True, "✅ SQL執行成功，但返回0條數據\n💡 可能原因: 條件太嚴格或數據不存在"
-        else:
-            sample_info = f"✅ SQL執行成功，返回 {len(df)} 條數據\n"
-            sample_info += f"📊 前3條數據:\n{df.head(3).to_string()}"
-            return True, sample_info
-            
-    except Exception as e:
-        return False, f"❌ SQL執行錯誤: {str(e)}"
+    # 清理SQL
+    sql_clean = re.sub(r'^```sql|```$', '', sql_query).strip()
+    
+    # 檢查是否已經是完整SQL
+    if "SELECT" in sql_clean.upper() and "FROM" in sql_clean.upper():
+        return sql_clean
+    
+    # 如果只有SELECT部分
+    if "SELECT" in sql_clean.upper() and "FROM" not in sql_clean.upper():
+        return sql_clean + " FROM appropriate_table WHERE 1=1;"
+    
+    # 如果只有FROM部分
+    if "FROM" in sql_clean.upper() and "SELECT" not in sql_clean.upper():
+        return "SELECT * " + sql_clean
+    
+    # 如果什麼都沒有，提供默認查詢
+    return "SELECT '請檢查SQL語法' AS status;"
 
 # ==================== 數據加載模塊 ====================
 class DataLoader:
@@ -126,8 +111,14 @@ class DataLoader:
                             sql_query = sql_match.group(1).strip()
                             sql_query = re.sub(r'^sql\s*', '', sql_query)
                             sql_query = re.sub(r'```sql|```', '', sql_query).strip()
+                            
+                            # 驗證並修復SQL
+                            validation = validate_sql(sql_query)
+                            if not validation["valid"]:
+                                print(f"發現有問題的SQL，將進行修復: {sql_query}")
+                                sql_query = repair_sql(sql_query)
                         else:
-                            sql_query = assistant_content
+                            sql_query = repair_sql(assistant_content)
                         
                         self.questions.append(question)
                         self.sql_answers.append(sql_query)
@@ -154,7 +145,7 @@ class DataLoader:
                 token=self.hf_token
             )
             with open(schema_file_path, 'r', encoding='utf-8') as f:
-                self.schema_data = safe_json_load(f.read(), {})
+                self.schema_data = json.load(f)
             print("Schema加載成功")
             return True
         except Exception as e:
@@ -184,7 +175,7 @@ class LLMClient:
         self.hf_token = hf_token
     
     def call_llm_api(self, prompt: str, model_urls: List[str] = LLM_MODELS) -> Optional[str]:
-        """調用LLM API（多模型備用）"""
+        """調用LLM API"""
         headers = {"Authorization": f"Bearer {self.hf_token}"}
         payload = {
             "inputs": prompt,
@@ -203,12 +194,10 @@ class LLMClient:
                     result = response.json()
                     if isinstance(result, list) and len(result) > 0:
                         generated_text = result[0]['generated_text'].strip()
-                        # 清理輸出
                         generated_text = re.sub(r'^```sql|```$', '', generated_text).strip()
                         return generated_text
                 
             except Exception as e:
-                print(f"模型 {model_url} 調用失敗: {e}")
                 continue
         
         return None
@@ -249,25 +238,23 @@ class TextToSQLSystem:
         self.llm_client = LLMClient(hf_token)
         self.retrieval_system = RetrievalSystem()
         
-        # 初始化組件
         self.initialize_system()
     
     def initialize_system(self):
         """初始化系統組件"""
         print("正在初始化系統組件...")
         
-        # 加載數據
         self.data_loader.load_dataset()
         self.data_loader.load_schema()
-        
-        # 初始化檢索系統
         self.retrieval_system.compute_embeddings(self.data_loader.questions)
-        
         self.schema_context = self.data_loader.build_schema_context()
+        
         print("系統初始化完成")
+        print(f"可用問題數量: {len(self.data_loader.questions)}")
+        print(f"Schema表數量: {len(self.data_loader.schema_data)}")
     
     def generate_sql(self, user_question: str) -> Tuple[str, str]:
-        """生成SQL查詢（主函數）"""
+        """生成SQL查詢"""
         log_messages = [f"🕒 開始處理: {get_current_time()}"]
         
         if not user_question or user_question.strip() == "":
@@ -281,20 +268,23 @@ class TextToSQLSystem:
                 best_hit = hits[0]
                 similarity_score = best_hit['score']
                 similar_question = self.data_loader.questions[best_hit['corpus_id']]
+                original_sql = self.data_loader.sql_answers[best_hit['corpus_id']]
                 
                 log_messages.append(f"🔍 檢索到相似問題: '{similar_question}'")
                 log_messages.append(f"📊 相似度: {similarity_score:.3f}")
                 
                 if similarity_score > SIMILARITY_THRESHOLD:
-                    sql_result = self.data_loader.sql_answers[best_hit['corpus_id']]
-                    log_messages.append(f"✅ 相似度 > {SIMILARITY_THRESHOLD}，直接返回預先SQL")
-                    
-                    # 驗證SQL安全性
-                    validation = validate_sql(sql_result)
-                    if not validation["is_safe"]:
-                        log_messages.append(f"⚠️ 安全警告: {', '.join(validation['issues'])}")
-                    
-                    return sql_result, "\n".join(log_messages)
+                    # 驗證並可能修復SQL
+                    validation = validate_sql(original_sql)
+                    if not validation["valid"]:
+                        log_messages.append(f"⚠️ 原始SQL有問題: {', '.join(validation['issues'])}")
+                        log_messages.append("🛠️ 正在修復SQL...")
+                        repaired_sql = repair_sql(original_sql)
+                        log_messages.append(f"✅ 修復完成")
+                        return repaired_sql, "\n".join(log_messages)
+                    else:
+                        log_messages.append(f"✅ 相似度 > {SIMILARITY_THRESHOLD}，直接返回")
+                        return original_sql, "\n".join(log_messages)
                 else:
                     log_messages.append(f"ℹ️ 相似度低於閾值 {SIMILARITY_THRESHOLD}")
         
@@ -305,17 +295,14 @@ class TextToSQLSystem:
         generated_sql = self.llm_client.call_llm_api(prompt)
         
         if generated_sql:
-            # 清理和驗證生成的SQL
-            generated_sql = re.sub(r'^```sql|```$', '', generated_sql).strip()
+            # 驗證生成的SQL
             validation = validate_sql(generated_sql)
+            if not validation["valid"]:
+                log_messages.append(f"⚠️ LLM生成的SQL有問題: {', '.join(validation['issues'])}")
+                log_messages.append("🛠️ 正在修復SQL...")
+                generated_sql = repair_sql(generated_sql)
             
-            if validation["valid"]:
-                log_messages.append("✅ LLM生成成功")
-                if validation["issues"]:
-                    log_messages.append(f"ℹ️ 驗證提示: {', '.join(validation['issues'])}")
-            else:
-                log_messages.append("⚠️ LLM生成可能存在问题")
-                
+            log_messages.append("✅ SQL生成完成")
             return generated_sql, "\n".join(log_messages)
         else:
             # 3. 備用方案
@@ -334,7 +321,7 @@ class TextToSQLSystem:
 
 要求：
 1. 只輸出SQL語句
-2. 不要任何解釋
+2. 必須包含SELECT和FROM
 3. 使用正確的語法
 
 SQL查詢："""
@@ -346,11 +333,9 @@ SQL查詢："""
         if any(kw in user_question_lower for kw in ['銷售', '業績', '金額', '收入']):
             return "SELECT product_name, SUM(sales_amount) as total_sales FROM sales GROUP BY product_name ORDER BY total_sales DESC LIMIT 10;"
         elif any(kw in user_question_lower for kw in ['客戶', '買家', '用戶']):
-            return "SELECT customer_name, COUNT(*) as order_count, SUM(order_amount) as total_spent FROM orders GROUP BY customer_name ORDER BY total_spent DESC;"
+            return "SELECT customer_name, COUNT(*) as order_count FROM orders GROUP BY customer_name ORDER BY order_count DESC;"
         elif any(kw in user_question_lower for kw in ['時間', '日期', '最近', '月份']):
-            return "SELECT strftime('%Y-%m', order_date) as month, COUNT(*) as orders, SUM(order_amount) as revenue FROM orders WHERE order_date >= date('now', '-6 months') GROUP BY month ORDER BY month DESC;"
-        elif any(kw in user_question_lower for kw in ['產品', '商品', '項目']):
-            return "SELECT product_name, category, stock_quantity, price FROM products WHERE stock_quantity > 0 ORDER BY price DESC;"
+            return "SELECT strftime('%Y-%m', order_date) as month, COUNT(*) as orders FROM orders GROUP BY month ORDER BY month DESC;"
         else:
             return "SELECT '請重試或提供更詳細的問題' AS status;"
 
@@ -359,70 +344,52 @@ print("正在初始化Text-to-SQL系統...")
 text_to_sql_system = TextToSQLSystem(HF_TOKEN)
 
 # ==================== Gradio界面 ====================
-def process_query(user_question: str, test_query: bool = False) -> Tuple[str, str, str]:
+def process_query(user_question: str) -> Tuple[str, str]:
     """處理用戶查詢"""
     sql_result, log_message = text_to_sql_system.generate_sql(user_question)
     
-    # SQL調試信息
+    # 最終驗證
+    final_validation = validate_sql(sql_result)
     debug_info = ""
-    validation = validate_sql(sql_result)
     
-    if not validation["valid"]:
-        debug_info = "❌ SQL驗證失敗:\n" + "\n".join(validation["issues"])
+    if not final_validation["valid"]:
+        debug_info = "❌ 最終SQL驗證失敗:\n" + "\n".join(final_validation["issues"])
+        debug_info += "\n🛠️ 已嘗試自動修復，但仍存在问题"
     else:
-        debug_info = "✅ SQL語法驗證通過"
-        
-        if validation["issues"]:
-            debug_info += "\nℹ️ 提示: " + ", ".join(validation["issues"])
-        
-        # 如果啟用測試功能
-        if test_query and DB_CONFIG["test_queries"]:
-            success, test_result = execute_test_query(sql_result)
-            debug_info += f"\n\n🔧 測試結果:\n{test_result}"
+        debug_info = "✅ 最終SQL驗證通過"
+        if final_validation["issues"]:
+            debug_info += "\nℹ️ 提示: " + ", ".join(final_validation["issues"])
     
     return sql_result, debug_info, log_message
 
 # 創建界面
-with gr.Blocks(
-    title="智能Text-to-SQL系統",
-    theme=gr.themes.Soft(),
-    css="""
-    .gradio-container { max-width: 1000px; margin: 0 auto; }
-    .success { color: green; }
-    .warning { color: orange; }
-    .error { color: red; }
-    """
-) as demo:
+with gr.Blocks(title="智能Text-to-SQL系統", theme=gr.themes.Soft()) as demo:
     
     gr.Markdown("# 🚀 智能 Text-to-SQL 系統")
     gr.Markdown("輸入自然語言問題，自動生成並驗證SQL查詢")
     
     with gr.Row():
-        with gr.Column(scale=3):
-            question_input = gr.Textbox(
-                label="📝 您的問題",
-                placeholder="例如：查詢2024年銷售額最高的產品",
-                lines=2,
-                max_lines=4
-            )
-            
-            with gr.Row():
-                submit_btn = gr.Button("🚀 生成SQL", variant="primary")
-                test_btn = gr.Button("🔧 測試SQL", variant="secondary")
-                clear_btn = gr.Button("🗑️ 清除", variant="secondary")
+        question_input = gr.Textbox(
+            label="📝 您的問題",
+            placeholder="例如：查詢2024年銷售額最高的產品",
+            lines=2
+        )
+    
+    with gr.Row():
+        submit_btn = gr.Button("🚀 生成SQL", variant="primary")
+        clear_btn = gr.Button("🗑️ 清除", variant="secondary")
     
     with gr.Row():
         sql_output = gr.Code(
             label="📊 生成的SQL",
             language="sql",
-            lines=6,
-            interactive=True
+            lines=6
         )
     
     with gr.Row():
         debug_output = gr.Textbox(
-            label="🔍 SQL調試信息",
-            lines=4,
+            label="🔍 SQL驗證信息",
+            lines=3,
             interactive=False
         )
     
@@ -433,28 +400,9 @@ with gr.Blocks(
             interactive=False
         )
     
-    # 示例問題
-    gr.Examples(
-        examples=[
-            "2024年銷售額最高的5個產品",
-            "最近30天每個客戶的訂單數量",
-            "庫存不足的商品列表",
-            "比較2023年和2024年的月度銷售額",
-            "付款不及時的客戶統計"
-        ],
-        inputs=question_input,
-        label="💡 示例問題"
-    )
-    
     # 事件處理
     submit_btn.click(
-        fn=lambda q: process_query(q, False),
-        inputs=question_input,
-        outputs=[sql_output, debug_output, log_output]
-    )
-    
-    test_btn.click(
-        fn=lambda q: process_query(q, True),
+        fn=process_query,
         inputs=question_input,
         outputs=[sql_output, debug_output, log_output]
     )
@@ -470,12 +418,6 @@ if __name__ == "__main__":
     print("=" * 50)
     print("🌐 啟動Gradio Web界面...")
     print("📍 本地訪問: http://localhost:7860")
-    print("🔄 如果需要公網訪問，設置 share=True")
     print("=" * 50)
     
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        show_error=True
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)