Update app.py

app.py

@@ -3,8 +3,6 @@ import requests
 import json
 import os
 import re
-import sqlite3
-import pandas as pd
 from datetime import datetime
 from datasets import load_dataset
 from sentence_transformers import SentenceTransformer, util
@@ -34,12 +32,14 @@ def get_current_time():
 
 def validate_sql(sql_query: str) -> Dict:
     """驗證SQL語句的安全性"""
+    if not sql_query or sql_query.strip() == "":
+        return {"valid": False, "issues": ["SQL語句為空"], "is_safe": False}
+    
     security_issues = []
+    sql_upper = sql_query.upper()
     
     # 檢查危險操作
     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_upper:
             security_issues.append(f"發現危險操作: {keyword}")
@@ -57,28 +57,54 @@ def validate_sql(sql_query: str) -> Dict:
         "is_safe": len([i for i in security_issues if '危險' in i]) == 0
     }
 
-def repair_sql(sql_query: str) -> str:
-    """修復有問題的SQL語句"""
-    if not sql_query or sql_query.strip() == "":
-        return "SELECT 'SQL語句為空' AS error;"
+def intelligent_sql_repair(original_sql: str, user_question: str, similar_question: str) -> str:
+    """智能修復SQL語句"""
+    if not original_sql or original_sql.strip() == "":
+        # 根據問題內容生成有意義的SQL
+        user_question_lower = user_question.lower()
+        similar_question_lower = similar_question.lower()
+        
+        # 分析問題類型
+        if any(kw in user_question_lower for kw in ['報告', '完成', '份']):
+            return "SELECT strftime('%Y-%m', completion_date) as month, COUNT(*) as report_count FROM reports WHERE strftime('%Y', completion_date) = '2023' GROUP BY month ORDER BY month;"
+        
+        elif 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;"
+        
+        elif any(kw in user_question_lower for kw in ['客戶', '買家']):
+            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 FROM orders GROUP BY month ORDER BY month DESC;"
+        
+        else:
+            return "SELECT '請提供更詳細的查詢條件' AS status;"
     
-    # 清理SQL
-    sql_clean = re.sub(r'^```sql|```$', '', sql_query).strip()
+    # 如果SQL不為空但缺少關鍵字
+    sql_upper = original_sql.upper()
     
-    # 檢查是否已經是完整SQL
-    if "SELECT" in sql_clean.upper() and "FROM" in sql_clean.upper():
-        return sql_clean
+    if "SELECT" not in sql_upper and "FROM" not in sql_upper:
+        # 嘗試從問題推斷
+        if "count" in user_question_lower or "多少" in user_question_lower:
+            return f"SELECT COUNT(*) as count FROM appropriate_table WHERE condition; -- 原始SQL: {original_sql}"
+        else:
+            return f"SELECT * FROM appropriate_table WHERE condition; -- 原始SQL: {original_sql}"
     
-    # 如果只有SELECT部分
-    if "SELECT" in sql_clean.upper() and "FROM" not in sql_clean.upper():
-        return sql_clean + " FROM appropriate_table WHERE 1=1;"
+    elif "SELECT" not in sql_upper and "FROM" in sql_upper:
+        return "SELECT * " + original_sql
     
-    # 如果只有FROM部分
-    if "FROM" in sql_clean.upper() and "SELECT" not in sql_clean.upper():
-        return "SELECT * " + sql_clean
+    elif "SELECT" in sql_upper and "FROM" not in sql_upper:
+        # 嘗試找到合適的FROM子句
+        if "customer" in user_question_lower or "客戶" in user_question_lower:
+            return original_sql + " FROM customers WHERE 1=1;"
+        elif "product" in user_question_lower or "產品" in user_question_lower:
+            return original_sql + " FROM products WHERE 1=1;"
+        elif "sale" in user_question_lower or "銷售" in user_question_lower:
+            return original_sql + " FROM sales WHERE 1=1;"
+        else:
+            return original_sql + " FROM appropriate_table WHERE 1=1;"
     
-    # 如果什麼都沒有，提供默認查詢
-    return "SELECT '請檢查SQL語法' AS status;"
+    return original_sql  # 如果不需要修復
 
 # ==================== 數據加載模塊 ====================
 class DataLoader:
@@ -95,6 +121,9 @@ class DataLoader:
             raw_dataset = load_dataset(DATASET_REPO_ID, token=self.hf_token)['train']
             
             print("正在解析 messages 格式...")
+            valid_count = 0
+            invalid_count = 0
+            
             for item in raw_dataset:
                 try:
                     if 'messages' in item and len(item['messages']) >= 2:
@@ -111,22 +140,24 @@ 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 = repair_sql(assistant_content)
+                            sql_query = assistant_content
+                        
+                        # 驗證SQL
+                        validation = validate_sql(sql_query)
+                        if not validation["valid"]:
+                            invalid_count += 1
+                            print(f"發現無效SQL [{invalid_count}]: {sql_query}")
+                            # 暫時不修復，等待使用時再智能修復
                         
                         self.questions.append(question)
                         self.sql_answers.append(sql_query)
+                        valid_count += 1
                         
                 except Exception as e:
                     continue
             
-            print(f"成功解析 {len(self.questions)} 條問答範例")
+            print(f"成功解析 {valid_count} 條問答範例，其中 {invalid_count} 條需要修復")
             return True
             
         except Exception as e:
@@ -169,73 +200,12 @@ class DataLoader:
                     context += f"--   {col_name} ({col_type}) - {col_desc}\n"
         return context
 
-# ==================== LLM模塊 ====================
-class LLMClient:
-    def __init__(self, hf_token: str):
-        self.hf_token = hf_token
-    
-    def call_llm_api(self, prompt: str, model_urls: List[str] = LLM_MODELS) -> Optional[str]:
-        """調用LLM API"""
-        headers = {"Authorization": f"Bearer {self.hf_token}"}
-        payload = {
-            "inputs": prompt,
-            "parameters": {
-                "max_new_tokens": 200,
-                "temperature": 0.1,
-                "do_sample": False
-            }
-        }
-        
-        for model_url in model_urls:
-            try:
-                response = requests.post(model_url, headers=headers, json=payload, timeout=20)
-                
-                if response.status_code == 200:
-                    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:
-                continue
-        
-        return None
-
-# ==================== 檢索模塊 ====================
-class RetrievalSystem:
-    def __init__(self):
-        self.embedder = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
-        self.question_embeddings = None
-    
-    def compute_embeddings(self, questions: List[str]) -> None:
-        """計算問題向量"""
-        if questions:
-            print(f"正在為 {len(questions)} 個問題計算向量...")
-            self.question_embeddings = self.embedder.encode(questions, convert_to_tensor=True, show_progress_bar=False)
-            print("向量計算完成")
-        else:
-            self.question_embeddings = torch.Tensor([])
-    
-    def retrieve_similar(self, user_question: str, top_k: int = 3) -> List[Dict]:
-        """檢索相似問題"""
-        if self.question_embeddings is None or len(self.question_embeddings) == 0:
-            return []
-        
-        try:
-            question_embedding = self.embedder.encode(user_question, convert_to_tensor=True)
-            hits = util.semantic_search(question_embedding, self.question_embeddings, top_k=top_k)
-            return hits[0] if hits and hits[0] else []
-        except Exception as e:
-            print(f"檢索失敗: {e}")
-            return []
-
 # ==================== 主系統 ====================
 class TextToSQLSystem:
     def __init__(self, hf_token: str):
         self.hf_token = hf_token
         self.data_loader = DataLoader(hf_token)
-        self.llm_client = LLMClient(hf_token)
+        self.llm_client = None  # 延遲加載
         self.retrieval_system = RetrievalSystem()
         
         self.initialize_system()
@@ -251,7 +221,12 @@ class TextToSQLSystem:
         
         print("系統初始化完成")
         print(f"可用問題數量: {len(self.data_loader.questions)}")
-        print(f"Schema表數量: {len(self.data_loader.schema_data)}")
+    
+    def get_llm_client(self):
+        """延遲加載LLM客戶端"""
+        if self.llm_client is None:
+            self.llm_client = LLMClient(self.hf_token)
+        return self.llm_client
     
     def generate_sql(self, user_question: str) -> Tuple[str, str]:
         """生成SQL查詢"""
@@ -274,13 +249,23 @@ class TextToSQLSystem:
                 log_messages.append(f"📊 相似度: {similarity_score:.3f}")
                 
                 if similarity_score > SIMILARITY_THRESHOLD:
-                    # 驗證並可能修復SQL
+                    # 驗證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("🛠️ 正在智能修復SQL...")
+                        
+                        # 智能修復
+                        repaired_sql = intelligent_sql_repair(original_sql, user_question, similar_question)
                         log_messages.append(f"✅ 修復完成")
+                        
+                        # 驗證修復後的SQL
+                        final_validation = validate_sql(repaired_sql)
+                        if not final_validation["valid"]:
+                            log_messages.append(f"❌ 修復後仍有問題: {', '.join(final_validation['issues'])}")
+                        else:
+                            log_messages.append("✅ 修復後SQL驗證通過")
+                        
                         return repaired_sql, "\n".join(log_messages)
                     else:
                         log_messages.append(f"✅ 相似度 > {SIMILARITY_THRESHOLD}，直接返回")
@@ -291,16 +276,25 @@ class TextToSQLSystem:
         # 2. LLM生成模式
         log_messages.append("🤖 進入LLM生成模式...")
         
-        prompt = self.build_llm_prompt(user_question)
-        generated_sql = self.llm_client.call_llm_api(prompt)
+        prompt = f"""你是一個SQL專家。請為以下問題生成SQL查詢：
+
+問題：{user_question}
+
+要求：
+1. 只輸出SQL語句
+2. 必須包含SELECT和FROM
+3. 使用正確的語法
+
+SQL查詢："""
+        
+        generated_sql = self.get_llm_client().call_llm_api(prompt)
         
         if generated_sql:
             # 驗證生成的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)
+                generated_sql = intelligent_sql_repair(generated_sql, user_question, user_question)
             
             log_messages.append("✅ SQL生成完成")
             return generated_sql, "\n".join(log_messages)
@@ -310,34 +304,74 @@ class TextToSQLSystem:
             backup_sql = self.generate_backup_sql(user_question)
             return backup_sql, "\n".join(log_messages)
     
-    def build_llm_prompt(self, user_question: str) -> str:
-        """構建LLM提示詞"""
-        return f"""你是一個SQL專家。請根據以下數據庫結構生成SQL查詢。
-
-{self.schema_context}
-
-請為以下問題生成準確的SQL查詢：
-{user_question}
-
-要求：
-1. 只輸出SQL語句
-2. 必須包含SELECT和FROM
-3. 使用正確的語法
-
-SQL查詢："""
-    
     def generate_backup_sql(self, user_question: str) -> str:
         """生成備用SQL"""
         user_question_lower = user_question.lower()
         
-        if any(kw in user_question_lower for kw in ['銷售', '業績', '金額', '收入']):
+        if any(kw in user_question_lower for kw in ['報告', '完成', '份', 'report']):
+            return "SELECT strftime('%Y-%m', completion_date) as month, COUNT(*) as report_count FROM reports GROUP BY month ORDER BY month;"
+        
+        elif any(kw in user_question_lower for kw in ['銷售', '業績', '金額', 'sale']):
             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 ['客戶', '買家', '用戶']):
+        
+        elif any(kw in user_question_lower for kw in ['客戶', '買家', 'customer']):
             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 ['時間', '日期', '最近', '月份']):
+        
+        elif any(kw in user_question_lower for kw in ['時間', '日期', '月份', 'month']):
             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;"
+            return "SELECT '請提供更詳細的查詢條件' AS status;"
+
+# ==================== 其他類定義（保持不變） ====================
+class LLMClient:
+    def __init__(self, hf_token: str):
+        self.hf_token = hf_token
+    
+    def call_llm_api(self, prompt: str, model_urls: List[str] = LLM_MODELS) -> Optional[str]:
+        headers = {"Authorization": f"Bearer {self.hf_token}"}
+        payload = {
+            "inputs": prompt,
+            "parameters": {
+                "max_new_tokens": 200,
+                "temperature": 0.1,
+                "do_sample": False
+            }
+        }
+        
+        for model_url in model_urls:
+            try:
+                response = requests.post(model_url, headers=headers, json=payload, timeout=20)
+                if response.status_code == 200:
+                    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:
+                continue
+        return None
+
+class RetrievalSystem:
+    def __init__(self):
+        self.embedder = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
+        self.question_embeddings = None
+    
+    def compute_embeddings(self, questions: List[str]) -> None:
+        if questions:
+            self.question_embeddings = self.embedder.encode(questions, convert_to_tensor=True, show_progress_bar=False)
+        else:
+            self.question_embeddings = torch.Tensor([])
+    
+    def retrieve_similar(self, user_question: str, top_k: int = 3) -> List[Dict]:
+        if self.question_embeddings is None or len(self.question_embeddings) == 0:
+            return []
+        try:
+            question_embedding = self.embedder.encode(user_question, convert_to_tensor=True)
+            hits = util.semantic_search(question_embedding, self.question_embeddings, top_k=top_k)
+            return hits[0] if hits and hits[0] else []
+        except:
+            return []
 
 # ==================== 初始化系統 ====================
 print("正在初始化Text-to-SQL系統...")
@@ -345,79 +379,31 @@ text_to_sql_system = TextToSQLSystem(HF_TOKEN)
 
 # ==================== Gradio界面 ====================
 def process_query(user_question: str) -> Tuple[str, str]:
-    """處理用戶查詢"""
     sql_result, log_message = text_to_sql_system.generate_sql(user_question)
-    
-    # 最終驗證
     final_validation = validate_sql(sql_result)
-    debug_info = ""
     
     if not final_validation["valid"]:
         debug_info = "❌ 最終SQL驗證失敗:\n" + "\n".join(final_validation["issues"])
-        debug_info += "\n🛠️ 已嘗試自動修復，但仍存在问题"
     else:
         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()) as demo:
-    
     gr.Markdown("# 🚀 智能 Text-to-SQL 系統")
-    gr.Markdown("輸入自然語言問題，自動生成並驗證SQL查詢")
-    
-    with gr.Row():
-        question_input = gr.Textbox(
-            label="📝 您的問題",
-            placeholder="例如：查詢2024年銷售額最高的產品",
-            lines=2
-        )
     
     with gr.Row():
+        question_input = gr.Textbox(label="📝 您的問題", placeholder="例如：查詢2023年每月報告數量", lines=2)
         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
-        )
+        sql_output = gr.Code(label="📊 生成的SQL", language="sql", lines=6)
     
     with gr.Row():
-        debug_output = gr.Textbox(
-            label="🔍 SQL驗證信息",
-            lines=3,
-            interactive=False
-        )
+        debug_output = gr.Textbox(label="🔍 驗證信息", lines=2, interactive=False)
+        log_output = gr.Textbox(label="📋 執行日誌", lines=4, interactive=False)
     
-    with gr.Row():
-        log_output = gr.Textbox(
-            label="📋 執行日誌",
-            lines=4,
-            interactive=False
-        )
-    
-    # 事件處理
-    submit_btn.click(
-        fn=process_query,
-        inputs=question_input,
-        outputs=[sql_output, debug_output, log_output]
-    )
-    
-    clear_btn.click(
-        fn=lambda: ["", "", ""],
-        inputs=[],
-        outputs=[sql_output, debug_output, log_output]
-    )
+    submit_btn.click(process_query, inputs=question_input, outputs=[sql_output, debug_output, log_output])
 
-# ==================== 啟動應用 ====================
 if __name__ == "__main__":
-    print("=" * 50)
-    print("🌐 啟動Gradio Web界面...")
-    print("📍 本地訪問: http://localhost:7860")
-    print("=" * 50)
-    
     demo.launch(server_name="0.0.0.0", server_port=7860, share=False)