File size: 5,069 Bytes
abb807c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
import streamlit as st
from langchain_community.utilities import SQLDatabase
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
import os
import re
# === Konfigurasi API Hugging Face Inference API ===
HF_TOKEN = os.environ.get("HF_TOKEN", "") # simpan secret via HF Spaces
os.environ["OPENAI_API_KEY"] = HF_TOKEN
HF_BASE_URL = "https://api-inference.huggingface.com/v1"
# Fungsi koneksi database
def connectDatabase(username, port, host, password, database):
try:
db_uri = f"postgresql://{username}:{password}@{host}:{port}/{database}"
st.session_state.db = SQLDatabase.from_uri(db_uri)
st.success("✅ Database connected successfully.")
except Exception as e:
st.session_state.db = None
st.error(f"❌ Failed to connect to database:\n{e}")
# Fungsi menjalankan query SQL
def runQuery(query):
if "db" in st.session_state and st.session_state.db:
try:
return st.session_state.db.run(query)
except Exception as e:
return f"❌ Query execution failed: {e}"
return "❌ Please connect to database first."
# Ambil skema database
def getDatabaseSchema():
if "db" in st.session_state and st.session_state.db:
try:
return st.session_state.db.get_table_info()
except Exception as e:
return f"❌ Failed to fetch schema: {e}"
return "❌ Please connect to database first."
# Prompt untuk menghasilkan SQL
def getQueryFromLLM(question):
template = """
Below is the schema of a PostgreSQL database. Read the schema carefully and answer the user's question using a valid SQL query. Be careful with table and column names (case-sensitive). Only provide the SQL query, and nothing else.
Schema:
{schema}
Sekarang giliran Anda:
Question: {question}
SQL query:
"""
try:
schema = getDatabaseSchema()[:3000]
prompt = ChatPromptTemplate.from_template(template)
chain = prompt | llm
response = chain.invoke({"question": question, "schema": schema})
query_only = re.sub(r"<think>.*?</think>", "", response.content, flags=re.DOTALL).strip()
return query_only
except Exception as e:
st.error(f"❌ Model gagal merespons: {e}")
return "Gagal membuat query SQL."
# Prompt untuk menjelaskan hasil
def getResponseForQueryResult(question, query, result):
template = """
Berdasarkan pertanyaan, SQL query, dan hasilnya, berikan jawaban dalam bahasa alami. Gunakan bahasa Indonesia jika pertanyaannya dalam bahasa Indonesia.
Schema:
{schema}
Pertanyaan: {question}
SQL query: {query}
Hasil: {result}
Jawaban:
"""
prompt2 = ChatPromptTemplate.from_template(template)
chain2 = prompt2 | llm
response = chain2.invoke({
"question": question,
"schema": getDatabaseSchema(),
"query": query,
"result": result
})
return response.content.strip()
# Konfigurasi halaman
st.set_page_config(page_title="Chat with PostgreSQL DB", page_icon="🧠", layout="centered")
st.title("🧠 Chat with your Database")
# Inisialisasi state
if "chat" not in st.session_state:
st.session_state.chat = []
# Sidebar
with st.sidebar:
st.subheader("🔗 Koneksi Database")
host = st.text_input("Host", value="aws-0-ap-southeast-1.pooler.supabase.com")
port = st.text_input("Port", value="6543")
username = st.text_input("Username", value="postgres")
password = st.text_input("Password", type="password", value="password")
database = st.text_input("Database", value="postgres")
if st.button("Connect"):
connectDatabase(username, port, host, password, database)
# Inisialisasi LLM
llm = ChatOpenAI(
model="qwen/qwen3-1.7b",
base_url=HF_BASE_URL,
api_key=HF_TOKEN,
temperature=0
)
# Input pertanyaan
question = st.chat_input("Tanyakan sesuatu tentang database Anda...")
# Proses pertanyaan
if question:
if "db" not in st.session_state or st.session_state.db is None:
st.error("❌ Please connect to database first.")
else:
st.session_state.chat.append({"role": "user", "content": question})
with st.spinner("💡 Sedang berpikir..."):
query = getQueryFromLLM(question)
result = runQuery(query)
response = getResponseForQueryResult(question, query, result)
st.session_state.chat.append({"role": "assistant", "content": response})
with st.expander("🧾 Generated SQL Query"):
st.code(query, language="sql")
with st.expander("📊 Raw Result"):
st.write(result)
# Tampilkan chat
for chat in st.session_state.chat:
st.chat_message(chat["role"]).markdown(chat["content"])
|