feat: add memory to prompt

Author: yuhexiong

README-CH.md

@@ -43,7 +43,7 @@ DB_URL='postgresql://user:password@host:5432/database'
 ```bash
 pip install streamlit==1.42.2 pandas==2.2.3 python-dotenv==1.0.1 \
 langchain-community==0.3.19 langchain-openai==0.3.7 langchain==0.3.20 \
-langchain-ollama==0.2.3
+langchain-ollama==0.2.3 python-magic-bin==0.4.14
 ```
 
 

README.md

@@ -42,7 +42,7 @@ Modify the prompt in the code based on your database information. Some templates
 ```bash
 pip install streamlit==1.42.2 pandas==2.2.3 python-dotenv==1.0.1 \
 langchain-community==0.3.19 langchain-openai==0.3.7 langchain==0.3.20 \
-langchain-ollama==0.2.3
+langchain-ollama==0.2.3 python-magic-bin==0.4.14
 ```
 
 

app.py

@@ -6,7 +6,6 @@
 from langchain_community.utilities import SQLDatabase
 
 from llm_util import get_llm
-from prompt_util import get_prompt
 from rag_util import get_vector_store, run_rag
 from sql_util import clean_sql_response, convert_result_to_df
 
@@ -51,11 +50,12 @@
 
     sql_query = None
     query_result = None
+    memory = []
 
     for retry in range(1, MAX_RETRIES + 1):
         try:
             # 生成 SQL 查詢
-            sql_query = run_rag(llm, vector_store, user_input, table_info)
+            sql_query = run_rag(llm, vector_store, user_input, table_info, memory)
 
             # 清理 SQL 查詢字串
             sql_query = clean_sql_response(sql_query)
@@ -67,18 +67,28 @@
             break
 
         except Exception as e:
+            error_message = str(e)
+            memory.append({
+                "sql": sql_query,
+                "error": error_message,
+            })
+
             if retry < MAX_RETRIES:
                 with st.chat_message("assistant"):
                     st.markdown(
-                        f"⚠️ SQL 執行失敗：`{sql_query}`，正在重新嘗試 ({retry}/{MAX_RETRIES})...")
+                        f"❌ SQL 執行失敗：\n```sql\n{sql_query}\n```\n"
+                        f"❗ 錯誤訊息：`{error_message}`，正在重新嘗試 ({retry}/{MAX_RETRIES})...")
             else:
                 with st.chat_message("assistant"):
-                    st.markdown(f"❌ SQL 執行失敗：{e}")
+                    st.markdown(
+                        f"❌ SQL 執行失敗：\n```sql\n{sql_query}\n```\n"
+                        f"❗ 錯誤訊息：`{error_message}`")
 
                 # 移除錯誤的 SQL 語法
                 sql_query = None
                 break
 
+    # 順利產生 sql 後嘗試執行
     if sql_query:
 
         with st.chat_message("assistant"):

prompt_util.py

@@ -2,7 +2,7 @@
 
 
 
-def get_prompt(example: str | None):
+def get_prompt(example: str | None, memory: str | None):
     """
     自訂產生 SQL 的 Prompt
     """
@@ -19,31 +19,24 @@ def get_prompt(example: str | None):
         ### 1. SQL 語法規則  
 
         - **表格名稱必須加上雙引號**  
-        - ✅ 正確：`SELECT 名稱 FROM "table_name"`  
-        - ❌ 錯誤：`SELECT 名稱 FROM table_name`（**表格名稱沒有雙引號，錯誤！**）  
-        - ❌ 錯誤：`SELECT 名稱 FROM 'table_name'`（**單引號錯誤！**）  
-
-        - **欄位名稱不可使用雙引號**  
-        - ✅ 正確：`SELECT 名稱 FROM "table_name"`  
-        - ❌ 錯誤：`SELECT "名稱" FROM "table_name"`（**欄位名稱不應加雙引號！**）  
+        - **欄位名稱不可使用任何符號**  
+        - 範例：`SELECT 名稱 FROM "table_name"`  
 
         - **聚合函數必須加括號**  
-        - ✅ 正確：`SELECT MIN(日期) FROM "table_name"`  
-        - ❌ 錯誤：`SELECT MIN 日期 FROM "table_name"`  
+        - 範例：`SELECT MIN(日期) FROM "table_name"`  
 
         ### 2. `GROUP BY` 使用規則  
 
         - **當問題涉及「最高」「最低」「平均」「總和」時，一定要 `GROUP BY`**  
         - **當問題要求「每個對象」、「每個類別」、「每個項目」時，一定要 `GROUP BY`**  
-        - ❌ 錯誤：`SELECT AVG(數值) FROM "table_name"`  
-        - ✅ 正確：`SELECT 類別, AVG(數值) FROM "table_name" GROUP BY 類別`  
+        - 範例：`SELECT 類別, AVG(數值) FROM "table_name" GROUP BY 類別`  
+
         - **選擇正確的 `GROUP BY` 屬性**  
         - 若問題涉及某項目（如：「哪個項目的值最高？」）➡ `GROUP BY 項目名稱`  
         - 若問題涉及某類別（如：「哪個類別的平均值最高？」）➡ `GROUP BY 類別名稱`  
         - 若問題未明確說明，預設使用 `GROUP BY 項目名稱`  
         - **`SELECT` 中必須包含 `GROUP BY` 的欄位**  
-        - ❌ 錯誤：`SELECT MAX(數值) FROM "table_name" GROUP BY 類別`  
-        - ✅ 正確：`SELECT 類別, MAX(數值) FROM "table_name" GROUP BY 類別`  
+        - 範例：`SELECT 類別, MAX(數值) FROM "table_name" GROUP BY 類別`  
 
         """
 
@@ -54,4 +47,11 @@ def get_prompt(example: str | None):
             {example}
         """
 
+    if memory is not None and memory != "":
+        prompt_template += f"""
+
+        以下是目前為止的嘗試及錯誤訊息：
+            {memory}
+        """
+
     return PromptTemplate.from_template(prompt_template)

rag_util.py

@@ -1,16 +1,9 @@
-import ast
 import json
 import os
-import re
 
 from langchain_ollama import OllamaEmbeddings
-import pandas as pd
-import streamlit as st
 from dotenv import load_dotenv
-import matplotlib
-import matplotlib.pyplot as plt
 from langchain_community.document_loaders import DirectoryLoader
-from langchain.prompts import PromptTemplate
 from langchain_core.documents import Document
 from langchain_core.vectorstores import InMemoryVectorStore
 
@@ -52,7 +45,7 @@ def get_vector_store():
     return vector_store
 
 
-def run_rag(llm, vector_store, user_input, table_info):
+def run_rag(llm, vector_store, user_input, table_info, memory):
     """
     執行 RAG 流程，先檢索相似內容，再生成 SQL 查詢
     """
@@ -70,13 +63,22 @@ def run_rag(llm, vector_store, user_input, table_info):
         回答: `{best_match.metadata["response"]}`
     """ if best_match else None
 
-    prompt = get_prompt(example)
+    # 整理記憶
+    memory_str = ""
+    if memory:
+        for item in memory:
+            memory_str += f"""
+                SQL： {item['sql']}
+                錯誤訊息：`{item['error']}`
+            """
+
+
+    prompt = get_prompt(example, memory_str)
 
     # 產生 SQL 查詢
     messages = prompt.invoke({
         "input": user_input, 
         "table_info": table_info,
-        "example": example,
         "top_k": 20
     })