Update app.py

app.py

@@ -231,7 +231,7 @@ class TextToSQLSystem:
             output = self.llm(
                 prompt,
                 max_tokens=150,
-                temperature=0.1,
+                temperature=0.05,
                 top_p=0.9,
                 echo=False,
                 # --- 將 stop 參數加回來 ---
@@ -240,13 +240,25 @@ class TextToSQLSystem:
 
             self._log(f"🧠 模型原始輸出 (Raw Output): {output}", "DEBUG")
 
+
             if output and "choices" in output and len(output["choices"]) > 0:
-                generated_text = output["choices"][0]["text"]
+                generated_text = output["choices"]["text"]
                 self._log(f"📝 提取出的生成文本: {generated_text.strip()}", "DEBUG")
-                return generated_text.strip()
+
+                # --- 新增的清理邏輯 ---
+                lines = generated_text.strip().split('\n')
+                # 過濾掉所有以 '--' 開頭的註解行
+                non_comment_lines = [line for line in lines if not line.strip().startswith('--')]
+                cleaned_text = "\n".join(non_comment_lines).strip()
+
+                if cleaned_text != generated_text.strip():
+                    self._log(f"🧹 清理掉註解後的文本: {cleaned_text}", "DEBUG")
+
+                return cleaned_text  # <-- 返回清理後的文本
+                # --- 清理邏輯結束 ---
             else:
                 self._log("❌ 模型的原始輸出格式不正確或為空。", "ERROR")
-                return ""
+                return ""                
 
         except Exception as e:
             self._log(f"❌ 模型生成過程中發生嚴重錯誤: {e}", "CRITICAL")
@@ -693,11 +705,10 @@ class TextToSQLSystem:
 
         # 使用強分隔符和清晰的標題來構建 prompt
         prompt = f"""### INSTRUCTIONS ###
-You are a SQLite expert. Your only job is to generate a single, valid SQLite query based on the provided schema and question.
-- ONLY use the tables and columns from the schema below.
-- ALWAYS use SQLite syntax (e.g., `strftime('%Y', date_column)` for years).
-- The report completion date is the `ReportAuthorization` column in the `JobTimeline` table.
-- Your output MUST be ONLY the SQL query inside a ```sql code block.
+You are a SQLite code generation machine. Your SOLE PURPOSE is to generate a single, valid SQLite query.
+- DO NOT provide any explanations, comments, or introductory text.
+- Your response MUST start directly with ```sql and end with ```.
+- The output must be ONLY the SQL code. Any other text is a failure.
 
 ### SCHEMA ###
 {schema_str}
@@ -771,30 +782,46 @@ SQL:
         self.log_history = []
         self._log(f"⏰ 處理問題: {question}")
 
-        # 1. 檢索相似範例
-        self._log("🔍 尋找相似範例...")
-        examples = self.find_most_similar(question, FEW_SHOT_EXAMPLES_COUNT)
-        if examples: self._log(f"✅ 找到 {len(examples)} 個相似範例")
 
-        # 2. 建立提示詞
-        self._log("📝 建立 Prompt...")
-        prompt = self._build_prompt(question, examples)
+        for attempt in range(2): # --- 新增：最多嘗試 2 次 ---
+            self._log(f"🚀 開始第 {attempt + 1} 次嘗試...")
 
-        # 3. 生成 AI 回應
-        self._log("🧠 開始生成 AI 回應...")
-        response = self.huggingface_api_call(prompt)
+            # 1. 檢索相似範例 (第二次嘗試時不再重複)
+            if attempt == 0:
+                self._log("🔍 尋找相似範例...")
+                examples = self.find_most_similar(question, FEW_SHOT_EXAMPLES_COUNT)
+                if examples: self._log(f"✅ 找到 {len(examples)} 個相似範例")
 
-        # 4. **新的核心步驟**: 呼叫決策引擎來生成最終 SQL
-        final_sql, status_message = self._validate_and_fix_sql(question, response)
+            # 2. 建立提示詞
+            self._log("📝 建立 Prompt...")
+            prompt = self._build_prompt(question, examples)
+            
+            # --- 新增：如果是第二次嘗試，加入修正指令 ---
+            if attempt > 0:
+                correction_prompt = "\nYour previous attempt failed because you did not provide a valid SQL query. REMEMBER: ONLY output the SQL code inside a ```sql block. DO NOT write comments or explanations.\nSQL:\n```sql\n"
+                # 將原本 prompt 的結尾替換成我們的修正指令
+                prompt = prompt.rsplit("SQL:\n```sql", 1)[0] + correction_prompt
 
-        if final_sql:
-            result = (final_sql, status_message)
-        else:
-            result = (status_message, "生成失敗")
 
-        # 緩存結果
-        self.query_cache[question] = result
-        return result
+            # 3. 生成 AI 回應
+            self._log("🧠 開始生成 AI 回應...")
+            response = self.huggingface_api_call(prompt)
+
+            # 4. 驗證與生成
+            final_sql, status_message = self._validate_and_fix_sql(question, response)
+
+            if final_sql:
+                self._log(f"✅ 在第 {attempt + 1} 次嘗試成功！", "INFO")
+                result = (final_sql, status_message)
+                self.query_cache[question] = result # 緩存成功結果
+                return result
+
+            self._log(f"❌ 第 {attempt + 1} 次嘗試失敗。原因: {status_message}", "WARNING")
+
+        # --- 如果兩次都失敗 ---
+        self._log("❌ 所有嘗試均失敗，返回錯誤訊息。", "ERROR")
+        final_fallback_message = "模型多次嘗試後仍無法生成有效的SQL。"
+        return (final_fallback_message, "生成失敗")
 
 # ==================== Gradio 介面 ====================
 text_to_sql_system = TextToSQLSystem()