Update app.py

app.py

@@ -6,11 +6,13 @@ from datasets import load_dataset, Dataset
 from sentence_transformers import SentenceTransformer, util
 import torch
 from huggingface_hub import hf_hub_download
+import re
 
 # --- 配置區 ---
 HF_TOKEN = os.environ.get("HF_TOKEN") 
 DATASET_REPO_ID = "Paul720810/Text-to-SQL-Softline"
-LLM_API_URL = "https://api-inference.huggingface.co/models/codellama/CodeLlama-34b-Instruct-hf"
+# 使用正確的模型名稱（7B版本更適合免費使用）
+LLM_API_URL = "https://api-inference.huggingface.co/models/codellama/CodeLlama-7b-hf"
 SIMILARITY_THRESHOLD = 0.90 
 
 print("--- [1/5] 開始初始化應用 ---")
@@ -18,57 +20,90 @@ print("--- [1/5] 開始初始化應用 ---")
 # --- 1. 載入知識庫 ---
 qa_dataset = None
 schema_data = {}
+questions = []
+sql_answers = []
+
 try:
     print(f"--- [2/5] 正在從 '{DATASET_REPO_ID}' 載入知識庫... ---")
     raw_dataset = load_dataset(DATASET_REPO_ID, token=HF_TOKEN)['train']
     
-    # *** 關鍵修正：智能解析 Dataset ***
-    # 檢查第一條數據的結構來判斷格式
-    if raw_dataset and len(raw_dataset) > 0:
-        first_item = raw_dataset[0]
-        if 'text' in first_item and 'question' not in first_item:
-            # 這是舊的 {'text': '...'} 格式，需要解析
-            print("--- > 檢測到 'text' 格式，正在解析JSON...")
-            parsed_qa_data = []
-            for item in raw_dataset:
-                try:
-                    line_dict = json.loads(item['text'])
-                    parsed_qa_data.append(line_dict)
-                except (json.JSONDecodeError, KeyError):
-                    continue # 跳過錯誤行
-            qa_dataset = Dataset.from_list(parsed_qa_data)
-        elif 'question' in first_item and 'sql' in first_item:
-            # 這已經是正確的 {'question': ..., 'sql': ...} 格式
-            print("--- > 檢測到已解析的 'question'/'sql' 格式，直接使用。")
-            qa_dataset = raw_dataset
-        else:
-            raise ValueError(f"未知的Dataset格式: {first_item}")
+    # 解析新的 messages 格式
+    print("--- > 檢測到 'messages' 格式，正在解析...")
+    
+    for item in raw_dataset:
+        try:
+            if 'messages' in item and len(item['messages']) >= 2:
+                user_content = item['messages'][0]['content']
+                assistant_content = item['messages'][1]['content']
+                
+                # 從用戶消息中提取問題
+                question_match = re.search(r'指令:\s*(.*?)(?:\n|$)', user_content)
+                if question_match:
+                    question = question_match.group(1).strip()
+                else:
+                    # 如果沒有找到指令，使用整個內容
+                    question = user_content
+                
+                # 從助手消息中提取SQL
+                sql_match = re.search(r'SQL查詢:\s*(.*?)(?:\n|$)', assistant_content, re.DOTALL)
+                if sql_match:
+                    sql_query = sql_match.group(1).strip()
+                    # 清理SQL語句
+                    sql_query = re.sub(r'^sql\s*', '', sql_query)  # 移除開頭的sql
+                    sql_query = re.sub(r'```sql|```', '', sql_query).strip()  # 移除代碼塊標記
+                else:
+                    sql_query = assistant_content
+                
+                questions.append(question)
+                sql_answers.append(sql_query)
+                
+        except (KeyError, IndexError, TypeError) as e:
+            print(f"解析錯誤，跳過該條目: {e}")
+            continue
+    
+    # 創建問答數據集
+    if questions:
+        qa_dataset = Dataset.from_dict({
+            'question': questions,
+            'sql': sql_answers
+        })
     else:
-        # 數據集為空
-        raise ValueError("載入的Dataset為空。")
-
+        raise ValueError("沒有成功解析出任何問答對")
+    
     # 載入並解析 Schema JSON
     schema_file_path = "sqlite_schema_FULL.json"
-    hf_hub_download(repo_id=DATASET_REPO_ID, filename=schema_file_path, repo_type='dataset', local_dir='.', token=HF_TOKEN)
-    
-    with open(schema_file_path, 'r', encoding='utf-8') as f:
-        schema_data = json.load(f)
+    try:
+        hf_hub_download(repo_id=DATASET_REPO_ID, filename=schema_file_path, 
+                       repo_type='dataset', local_dir='.', token=HF_TOKEN)
+        
+        with open(schema_file_path, 'r', encoding='utf-8') as f:
+            schema_data = json.load(f)
+    except Exception as e:
+        print(f"警告: 無法載入Schema文件: {e}")
+        schema_data = {}
         
-    print(f"--- > 成功載入 {len(qa_dataset)} 條問答範例和 Schema。 ---")
+    print(f"--- > 成功解析 {len(questions)} 條問答範例。 ---")
 
 except Exception as e:
-    print(f"!!! 致命錯誤: 無法載入或解析 Dataset '{DATASET_REPO_ID}'.")
-    print(f"詳細錯誤: {e}")
-    qa_dataset = Dataset.from_dict({"question": ["示例問題"], "sql": ["SELECT 'Dataset failed to load'"]})
+    print(f"!!! 錯誤: 處理Dataset時發生問題: {e}")
+    # 創建備用數據集
+    questions = ["示例問題"]
+    sql_answers = ["SELECT '請檢查數據集格式' AS error;"]
+    qa_dataset = Dataset.from_dict({"question": questions, "sql": sql_answers})
 
 # --- 2. 構建 DDL 和初始化檢索模型 ---
 def load_schema_as_ddl(schema_dict: dict) -> str:
-    # (此函式無需修改)
-    ddl_string = ""
+    ddl_string = "/* 數據庫結構 */\n"
     for table_name, columns in schema_dict.items():
-        if not isinstance(columns, list): continue
+        if not isinstance(columns, list): 
+            continue
         ddl_string += f"CREATE TABLE `{table_name}` (\n"
-        ddl_cols = [f"  `{col.get('name', '')}` {col.get('type', '')} -- {col.get('description', '')}" for col in columns]
+        ddl_cols = []
+        for col in columns:
+            col_name = col.get('name', 'unknown')
+            col_type = col.get('type', 'TEXT')
+            col_desc = col.get('description', '')
+            ddl_cols.append(f"  `{col_name}` {col_type} -- {col_desc}")
         ddl_string += ",\n".join(ddl_cols) + "\n);\n\n"
     return ddl_string
 
@@ -77,113 +112,156 @@ SCHEMA_DDL = load_schema_as_ddl(schema_data)
 print("--- [3/5] 正在載入句向量模型 (all-MiniLM-L6-v2)... ---")
 embedder = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
 
-questions = [item['question'] for item in qa_dataset]
-sql_answers = [item['sql'] for item in qa_dataset]
-
-# 只有在 questions 列表不為空時才進行計算
+# 計算問題向量
 if questions:
     print(f"--- [4/5] 正在為 {len(questions)} 個問題計算向量... ---")
     question_embeddings = embedder.encode(questions, convert_to_tensor=True, show_progress_bar=True)
     print("--- > 向量計算完成！ ---")
 else:
-    print("--- [4/5] 警告：沒有可用的問題來計算向量。檢索功能將不可用。---")
+    print("--- [4/5] 警告：沒有可用的問題來計算向量。 ---")
     question_embeddings = torch.Tensor([])
 
-
 # --- 3. 混合系統核心邏輯 ---
 def get_sql_query(user_question: str):
-    # (此函式剩餘部分幾乎無需修改)
     if not user_question:
         return "請輸入您的問題。", "日誌：用戶未輸入問題。"
     
-    # 增加一個檢查，確保知識庫不是空的
-    if len(questions) == 0:
-        log_message = "錯誤：知識庫為空，無法進行檢索。"
-        return "系統錯誤：知識庫未成功載入。", log_message
-
-    question_embedding = embedder.encode(user_question, convert_to_tensor=True)
-    hits = util.semantic_search(question_embedding, question_embeddings, top_k=5)
+    log_messages = []
     
-    if not hits or not hits[0]:
-        log_message = "檢索失敗：找不到任何相似的問題。"
-        # 即使檢索失敗，也應該嘗試調用 LLM
-    else:
-        hits = hits[0]
-        most_similar_hit = hits[0]
-        similarity_score = most_similar_hit['score']
-        log_message = f"檢索到最相似問題: '{questions[most_similar_hit['corpus_id']]}' (相似度: {similarity_score:.4f})"
+    # 檢索相似問題
+    if len(questions) > 0:
+        question_embedding = embedder.encode(user_question, convert_to_tensor=True)
+        hits = util.semantic_search(question_embedding, question_embeddings, top_k=3)
         
-        if similarity_score > SIMILARITY_THRESHOLD:
-            sql_result = sql_answers[most_similar_hit['corpus_id']]
-            log_message += f"\n相似度 > {SIMILARITY_THRESHOLD}，[模式: 直接返回]。"
-            return sql_result, log_message
-
-    log_message += f"\n相似度低於閾值或檢索失敗，[模式: LLM生成]。正在構建 Prompt..."
+        if hits and hits[0]:
+            most_similar_hit = hits[0][0]
+            similarity_score = most_similar_hit['score']
+            similar_question = questions[most_similar_hit['corpus_id']]
+            
+            log_messages.append(f"檢索到相似問題: '{similar_question}' (相似度: {similarity_score:.4f})")
+            
+            if similarity_score > SIMILARITY_THRESHOLD:
+                sql_result = sql_answers[most_similar_hit['corpus_id']]
+                log_messages.append(f"相似度 > {SIMILARITY_THRESHOLD}，直接返回預先SQL")
+                return sql_result, "\n".join(log_messages)
+        else:
+            log_messages.append("檢索失敗：找不到相似問題")
+    else:
+        log_messages.append("知識庫為空，跳過檢索")
+    
+    # LLM生成模式
+    log_messages.append("進入LLM生成模式...")
+    
+    # 構建示例上下文
     examples_context = ""
-    if hits: # 只有在檢索到結果時才添加範例
-        for hit in hits[:3]:
-             examples_context += f"### A user asks: {questions[hit['corpus_id']]}\n{sql_answers[hit['corpus_id']]}\n\n"
-
-    prompt = f"""### Task
-Generate a SQLite SQL query that answers the following user question.
-Your response must contain ONLY the SQL query. Do not add any explanation.
+    if 'hits' in locals() and hits and hits[0]:
+        for i, hit in enumerate(hits[0][:2]):
+            examples_context += f"問題: {questions[hit['corpus_id']]}\nSQL: {sql_answers[hit['corpus_id']]}\n\n"
+    
+    # 構建提示詞
+    prompt = f"""你是一個SQL專家。請根據數據庫結構生成SQL查詢。
 
-### Database Schema
+數據庫結構：
 {SCHEMA_DDL}
-### Examples
+
+參考示例：
 {examples_context}
-### Question
+
+請為以下問題生成SQL查詢：
 {user_question}
 
-### SQL Query
+只輸出SQL語句，不要其他內容：
+
 """
-    log_message += "\n正在請求雲端 LLM..."
+    
+    log_messages.append("正在請求雲端LLM...")
+    
     headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 512, "temperature": 0.1, "return_full_text": False}}
-    response_text = ""
+    payload = {
+        "inputs": prompt,
+        "parameters": {
+            "max_new_tokens": 300,
+            "temperature": 0.1,
+            "do_sample": False
+        }
+    }
+    
     try:
-        response = requests.post(LLM_API_URL, headers=headers, json=payload)
-        response_text = response.text
-        response.raise_for_status() 
-        generated_text = response.json()[0]['generated_text'].strip()
-        if "```sql" in generated_text:
-            generated_text = generated_text.split("```sql")[1].split("```").strip()
-        if "```" in generated_text:
-            generated_text = generated_text.replace("```", "").strip()
-        log_message += f"\nLLM 生成成功！"
-        return generated_text, log_message
+        response = requests.post(LLM_API_URL, headers=headers, json=payload, timeout=30)
+        
+        if response.status_code == 200:
+            result = response.json()
+            if isinstance(result, list) and len(result) > 0:
+                generated_text = result[0]['generated_text'].strip()
+                
+                # 清理輸出，只保留SQL
+                if "```sql" in generated_text:
+                    generated_text = generated_text.split("```sql")[1].split("```")[0].strip()
+                elif "```" in generated_text:
+                    generated_text = generated_text.split("```")[1].strip() if len(generated_text.split("```")) > 2 else generated_text
+                
+                log_messages.append("LLM生成成功！")
+                return generated_text, "\n".join(log_messages)
+            else:
+                raise Exception(f"API返回格式異常: {result}")
+        else:
+            raise Exception(f"API錯誤: {response.status_code} - {response.text}")
+            
     except Exception as e:
-        error_msg = f"LLM API 調用失敗: {e}\nAPI 原始回應: {response_text}"
-        log_message += f"\n{error_msg}"
-        return "抱歉，調用雲端 AI 時發生錯誤。", log_message
+        error_msg = f"LLM API調用失敗: {str(e)}"
+        log_messages.append(error_msg)
+        
+        # 提供備用答案
+        backup_sql = "SELECT 'AI服務暫時不可用，請稍後重試' AS status;"
+        return backup_sql, "\n".join(log_messages)
 
 # --- 4. 創建 Gradio Web 界面 ---
 print("--- [5/5] 正在創建 Gradio Web 界面... ---")
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    # (此部分無需修改)
-    gr.Markdown("# 智能 Text-to-SQL 系統 (混合模式)")
-    # ... (Gradio界面代碼與之前相同)
-    gr.Markdown("輸入您的自然語言問題，系統將首先嘗試從知識庫中快速檢索答案。如果問題較新穎，則會調用雲端大語言模型生成SQL。")
+    gr.Markdown("# 🚀 智能 Text-to-SQL 系統 (混合模式)")
+    gr.Markdown("輸入自然語言問題，系統會智能生成SQL查詢")
+    
     with gr.Row():
-        question_input = gr.Textbox(label="輸入您的問題", placeholder="例如：去年Nike的總業績是多少？", scale=4)
+        question_input = gr.Textbox(
+            label="輸入您的問題", 
+            placeholder="例如：查詢去年的銷售總額", 
+            lines=2,
+            scale=4
+        )
         submit_button = gr.Button("生成SQL", variant="primary", scale=1)
-    sql_output = gr.Code(label="生成的 SQL 查詢", language="sql")
-    log_output = gr.Textbox(label="系統日誌 (執行過程)", lines=4, interactive=False)
+    
+    with gr.Row():
+        sql_output = gr.Code(
+            label="生成的 SQL 查詢", 
+            language="sql",
+            lines=6
+        )
+    
+    with gr.Row():
+        log_output = gr.Textbox(
+            label="系統日誌", 
+            lines=4, 
+            interactive=False
+        )
+    
     submit_button.click(
         fn=get_sql_query,
         inputs=question_input,
         outputs=[sql_output, log_output]
     )
+    
     gr.Examples(
         examples=[
-            "2024 最好的5個客人以及業績",
+            "2024年最好的5個客戶以及業績",
             "比較2023年跟2024年的業績",
-            "上禮拜C組 完成幾份報告",
-            "有沒��快到期的單子？",
+            "上週C組完成了幾份報告",
+            "有沒有快到期的訂單？",
             "哪個客戶的付款最不及時？"
         ],
-        inputs=question_input
+        inputs=question_input,
+        label="示例問題"
     )
 
 print("--- 應用準備啟動 ---")
-demo.launch()
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)