feat: add utils dir

Author: yuhexiong

app.py

@@ -1,22 +1,13 @@
-import os
+from config import DB_URL
 
 import streamlit as st
-from dotenv import load_dotenv
 from langchain.chains import create_sql_query_chain
 from langchain_community.utilities import SQLDatabase
 
-from llm_util import get_llm
-from rag_util import get_vector_store, run_rag
-from sql_util import clean_sql_response, convert_result_to_df
+from utils.llm_util import get_llm
+from utils.rag_util import get_vector_store, run_rag
+from utils.sql_util import clean_sql_response, convert_result_to_df
 
-# 讀取 .env 變數
-load_dotenv()
-
-
-# 取得資料庫連線字串
-DB_URL = os.getenv("DB_URL")
-if not DB_URL:
-    raise Exception("未在 .env 檔案中找到 DB_URL")
 
 MAX_RETRIES = 3  # 最多重試次數
 
@@ -43,7 +34,6 @@
 user_input = st.chat_input("請輸入您的問題...")
 
 if user_input:
-
     # 顯示使用者輸入
     with st.chat_message("user"):
         st.markdown(user_input)
@@ -52,6 +42,7 @@
     query_result = None
     memory = []
 
+    # 嘗試 MAX_RETRIES 次
     for retry in range(1, MAX_RETRIES + 1):
         try:
             # 生成 SQL 查詢
@@ -67,12 +58,16 @@
             break
 
         except Exception as e:
+            # 將錯誤轉為字串
             error_message = str(e)
+
+            # 記錄錯誤到 memory 中已供下次使用
             memory.append({
                 "sql": sql_query,
                 "error": error_message,
             })
 
+            # 顯示在 Streamlit 上
             if retry < MAX_RETRIES:
                 with st.chat_message("assistant"):
                     st.markdown(
@@ -90,7 +85,7 @@
 
     # 順利產生 sql 後嘗試執行
     if sql_query:
-
+        # 顯示在 Streamlit 上
         with st.chat_message("assistant"):
             st.markdown(f"**生成的 SQL 查詢：**\n```sql\n{sql_query}\n```")
 
@@ -105,9 +100,11 @@
                 with st.chat_message("table"):
                     st.dataframe(result_df)
             else:
+                # 沒有資料
                 with st.chat_message("table"):
                     st.markdown(f"⚠️ 沒有查詢結果。")
 
+        # 轉換錯誤
         except Exception as e:
             with st.chat_message("table"):
                 st.markdown(f"❌ 結果處理錯誤：{e}")

config.py

@@ -0,0 +1,41 @@
+import os
+from dotenv import load_dotenv
+
+
+
+# 讀取 .env 變數
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+load_dotenv(dotenv_path=os.path.join(BASE_DIR, ".env"))
+
+# 資料庫連線
+DB_URL = os.getenv("DB_URL")
+if not DB_URL:
+    raise Exception("未在 .env 檔案中找到 DB_URL。")
+
+# LLM 模型
+LLM_TYPE = os.getenv("LLM_TYPE", "OPENAI")  # 默認為 OPENAI
+
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+OPENAI_MODEL = os.getenv("OPENAI_MODEL")
+
+OLLAMA_URL = os.getenv("OLLAMA_URL")
+OLLAMA_MODEL = os.getenv("OLLAMA_MODEL")
+
+if LLM_TYPE == "OPENAI":
+    if not OPENAI_API_KEY:
+        raise Exception("未在 .env 檔案中找到 OPENAI_API_KEY。")
+    if not OPENAI_MODEL:
+        raise Exception("未在 .env 檔案中找到 OPENAI_MODEL")
+elif LLM_TYPE == "OLLAMA":
+    if not OLLAMA_URL:
+        raise Exception("未在 .env 檔案中找到 OLLAMA_URL")
+    if not OLLAMA_MODEL:
+        raise Exception("未在 .env 檔案中找到 OLLAMA_MODEL")
+
+# Embedding
+OLLAMA_EMBEDDING_URL = os.getenv("OLLAMA_EMBEDDING_URL")
+if not OLLAMA_EMBEDDING_URL:
+    raise Exception("未在 .env 檔案中找到 OLLAMA_EMBEDDING_URL。")
+OLLAMA_EMBEDDING_MODEL = os.getenv("OLLAMA_EMBEDDING_MODEL")
+if not OLLAMA_EMBEDDING_MODEL:
+    raise Exception("未在 .env 檔案中找到 OLLAMA_EMBEDDING_MODEL。")

llm_util.py

@@ -4,9 +4,9 @@
 from langchain.chains import create_sql_query_chain
 from langchain_community.utilities import SQLDatabase
 
-from llm_util import get_llm
-from prompt_util import get_prompt
-from sql_util import clean_sql_response, convert_result_to_df
+from utils.llm_util import get_llm
+from utils.prompt_util import get_prompt
+from utils.sql_util import clean_sql_response, convert_result_to_df
 
 # 讀取 .env 變數
 load_dotenv()

script.py

@@ -0,0 +1,17 @@
+from config import LLM_TYPE, OLLAMA_MODEL, OLLAMA_URL, OPENAI_API_KEY, OPENAI_MODEL
+
+from langchain_ollama.chat_models import ChatOllama
+from langchain_openai import ChatOpenAI
+
+
+
+def get_llm():
+    # 初始化 LLM 模型
+    if LLM_TYPE == "OPENAI":
+        return ChatOpenAI(model=OPENAI_MODEL, temperature=0, api_key=OPENAI_API_KEY)
+    elif LLM_TYPE == "OLLAMA":
+        return ChatOllama(model=OLLAMA_MODEL, base_url=OLLAMA_URL)
+
+    else:
+        raise Exception(f"未支援的 LLM_TYPE: {LLM_TYPE}")
+

utils/llm_util.py

@@ -1,27 +1,15 @@
-import json
-import os
+from config import OLLAMA_EMBEDDING_MODEL, OLLAMA_EMBEDDING_URL
 
+import json
 from langchain_ollama import OllamaEmbeddings
-from dotenv import load_dotenv
 from langchain_community.document_loaders import DirectoryLoader
 from langchain_core.documents import Document
 from langchain_core.vectorstores import InMemoryVectorStore
 
-from prompt_util import get_prompt
+from utils.prompt_util import get_prompt
 
 
-# 讀取 .env 變數
-load_dotenv()
-
 def get_vector_store():
-    # embedding
-    OLLAMA_EMBEDDING_URL = os.getenv("OLLAMA_EMBEDDING_URL")
-    if not OLLAMA_EMBEDDING_URL:
-        raise Exception("未在 .env 檔案中找到 OLLAMA_EMBEDDING_URL。")
-    OLLAMA_EMBEDDING_MODEL = os.getenv("OLLAMA_EMBEDDING_MODEL")
-    if not OLLAMA_EMBEDDING_MODEL:
-        raise Exception("未在 .env 檔案中找到 OLLAMA_EMBEDDING_MODEL。")
-
     embeddings = OllamaEmbeddings(model=OLLAMA_EMBEDDING_MODEL, base_url=OLLAMA_EMBEDDING_URL)
 
     vector_store = InMemoryVectorStore(embeddings)