"""Streamlit-интерфейс утилиты Ru2SQL.
Архитектурно — клиент REST API на FastAPI. Соответствует разделу 3.5
пояснительной записки: все обращения к модели и базе данных идут через
HTTP к ``src.api.main:app``. Запуск двух процессов:
uvicorn src.api.main:app --reload # на 127.0.0.1:8000
streamlit run streamlit_app.py # на 127.0.0.1:8501
"""
from __future__ import annotations
import logging
import os
import sys
import warnings
from pathlib import Path
# ──────────────────────────────────────────────
# Глушим шумные warning'и
# ──────────────────────────────────────────────
# Streamlit-watcher ходит по всему пакету transformers (image-processors)
# и спамит ModuleNotFoundError про torchvision. На работу это не влияет —
# Qwen2.5-Coder text-only, torchvision не нужен.
warnings.filterwarnings("ignore", message=".*torchvision.*")
logging.getLogger("transformers").setLevel(logging.ERROR)
logging.getLogger("streamlit.watcher.local_sources_watcher").setLevel(logging.ERROR)
import httpx
import streamlit as st
ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))
# Бизнес-словарь парсим локально — он не требует обращения к серверу
from src.business.vocabulary import BusinessVocabulary
API_URL = os.environ.get("RU2SQL_API_URL", "http://127.0.0.1:8000")
QUERY_TIMEOUT = 1800.0 # 30 минут — фактически безлимит
SHORT_TIMEOUT = 10.0 # для /health, /schema
# ──────────────────────────────────────────────
# Конфигурация страницы
# ──────────────────────────────────────────────
st.set_page_config(
page_title="Ru2SQL",
layout="wide",
initial_sidebar_state="expanded",
)
# ──────────────────────────────────────────────
# CSS — оформление в стиле тёмной темы GitHub
# ──────────────────────────────────────────────
st.markdown("""
""", unsafe_allow_html=True)
# ──────────────────────────────────────────────
# Session state
# ──────────────────────────────────────────────
def _default_vocab_yaml() -> str:
example = ROOT / "configs" / "example_vocabulary.yaml"
if example.exists():
return example.read_text(encoding="utf-8")
return "company: Моя компания\n\nterms: {}\nfilters: {}\nnotes: []\n"
def _init_state():
defaults = {
"history": [],
"api_health": None, # dict | None
"api_error": None, # str | None
"connection_string": "",
"schema_tables": None, # list[TablePayload-like dict] | None
"schema_error": None,
"vocabulary": None, # BusinessVocabulary | None
"vocab_yaml": _default_vocab_yaml(),
"db_mode": None,
"warmup_done": False,
}
for k, v in defaults.items():
if k not in st.session_state:
st.session_state[k] = v
_init_state()
# ──────────────────────────────────────────────
# Обёртки над API
# ──────────────────────────────────────────────
def _api_get_health() -> dict | None:
"""GET /health. None если API недоступен."""
try:
r = httpx.get(f"{API_URL}/health", timeout=SHORT_TIMEOUT)
r.raise_for_status()
return r.json()
except Exception as e:
st.session_state.api_error = str(e)
return None
def _api_get_schema(cs: str) -> tuple[list[dict] | None, str | None]:
"""POST /schema. Возвращает (tables, error)."""
try:
r = httpx.post(
f"{API_URL}/schema",
json={"connection_string": cs, "include_samples": True},
timeout=SHORT_TIMEOUT,
)
if r.status_code != 200:
try:
return None, r.json().get("detail", r.text)
except Exception:
return None, r.text
return r.json().get("tables", []), None
except Exception as e:
return None, str(e)
def _api_query(question: str, cs: str, vocab: BusinessVocabulary | None) -> dict:
"""POST /query — генерация SQL + опциональное исполнение."""
payload = {
"question": question,
"connection_string": cs,
"execute": True,
}
if vocab is not None and bool(vocab):
payload["vocabulary"] = {
"company": vocab.company,
"terms": vocab.terms,
"filters": vocab.filters,
"notes": vocab.notes,
}
r = httpx.post(f"{API_URL}/query", json=payload, timeout=QUERY_TIMEOUT)
if r.status_code != 200:
try:
detail = r.json().get("detail", r.text)
except Exception:
detail = r.text
raise RuntimeError(f"API вернул {r.status_code}: {detail}")
return r.json()
def _api_warmup() -> tuple[bool, str | None]:
"""POST /warmup — короткий прогон для прогрева модели на CPU."""
try:
r = httpx.post(f"{API_URL}/warmup", timeout=QUERY_TIMEOUT)
if r.status_code == 200:
return True, None
return False, r.text
except Exception as e:
return False, str(e)
def _load_vocab_from_yaml(yaml_text: str) -> BusinessVocabulary:
import tempfile
tmp = Path(tempfile.mktemp(suffix=".yaml"))
tmp.write_text(yaml_text, encoding="utf-8")
try:
return BusinessVocabulary.from_yaml(tmp)
finally:
tmp.unlink(missing_ok=True)
# ──────────────────────────────────────────────
# Диалог редактирования бизнес-словаря
# ──────────────────────────────────────────────
@st.dialog("Бизнес-словарь", width="large")
def vocab_dialog():
st.caption(
"Опишите термины и метрики компании в формате YAML. "
"Модель будет учитывать их при генерации SQL."
)
yaml_text = st.text_area(
"YAML-конфигурация",
value=st.session_state.vocab_yaml,
height=480,
label_visibility="collapsed",
)
c1, c2 = st.columns(2)
with c1:
if st.button("Применить", type="primary", width='stretch'):
try:
st.session_state.vocabulary = _load_vocab_from_yaml(yaml_text)
st.session_state.vocab_yaml = yaml_text
st.rerun()
except Exception as e:
st.error(f"Ошибка синтаксиса YAML: {e}")
with c2:
if st.button("Отмена", width='stretch'):
st.rerun()
# ──────────────────────────────────────────────
# Sidebar
# ──────────────────────────────────────────────
with st.sidebar:
# ── API ──
st.markdown('API
', unsafe_allow_html=True)
health = _api_get_health()
st.session_state.api_health = health
if health is None:
st.markdown('API недоступен', unsafe_allow_html=True)
st.caption(f"Адрес: {API_URL}")
st.caption("Запусти в отдельной консоли: `uvicorn src.api.main:app --reload`")
if st.session_state.api_error:
st.caption(f"Причина: {st.session_state.api_error[:160]}")
else:
if health.get("model_loaded"):
st.markdown(
f'✅ {health.get("base_model", "модель")}',
unsafe_allow_html=True,
)
else:
st.markdown(
'⏳ Модель ещё загружается',
unsafe_allow_html=True,
)
st.caption("Подождите несколько минут — модель ещё инициализируется.")
st.markdown('
', unsafe_allow_html=True)
# ── База данных ──
st.markdown('База данных
', unsafe_allow_html=True)
modes = ["Демо-база", "Загрузить файл", "Строка подключения"]
prev = st.session_state.db_mode
db_mode = st.radio(
"Источник данных", modes,
index=modes.index(prev) if prev in modes else None,
label_visibility="collapsed",
)
if db_mode != prev:
st.session_state.schema_tables = None
st.session_state.connection_string = ""
st.session_state.db_mode = db_mode
cs = ""
if db_mode == "Демо-база":
st.caption("Встроенная база: интернет-магазин электроники, 120 заказов.")
cs = str(ROOT / "data" / "demo" / "sales.sqlite")
elif db_mode == "Загрузить файл":
uploaded = st.file_uploader(
"SQLite-файл базы данных", type=["sqlite", "db"],
label_visibility="collapsed",
)
if uploaded:
import tempfile
tmp_db = Path(tempfile.mktemp(suffix=".sqlite"))
tmp_db.write_bytes(uploaded.read())
cs = str(tmp_db)
else:
st.caption("Перетащите .sqlite или .db файл сюда")
else:
cs = st.text_input(
"Строка подключения",
placeholder="postgresql://user:pass@host:5432/db",
value=st.session_state.connection_string,
label_visibility="collapsed",
)
st.caption("PostgreSQL · MySQL · SQLite (sqlite:///path)")
if cs and st.button("Подключиться", width='stretch', type="primary"):
with st.spinner("Чтение схемы…"):
tables, err = _api_get_schema(cs)
if err:
st.error(f"Ошибка подключения: {err}")
st.session_state.schema_tables = None
else:
st.session_state.schema_tables = tables
st.session_state.connection_string = cs
st.session_state.schema_error = None
if not st.session_state.get("warmup_done", False):
with st.spinner("Прогрев модели (запускается один раз за сессию)…"):
ok, _err = _api_warmup()
if ok:
st.session_state.warmup_done = True
# Автозагрузка словаря для демо-базы
if "sales" in cs and st.session_state.vocabulary is None:
vp = ROOT / "configs" / "example_vocabulary.yaml"
if vp.exists():
try:
st.session_state.vocabulary = _load_vocab_from_yaml(
vp.read_text(encoding="utf-8")
)
except Exception:
pass
st.success(f"Подключено. Таблиц: {len(tables)}")
if st.session_state.schema_tables is not None:
n = len(st.session_state.schema_tables)
st.markdown(
'✅ База данных подключена',
unsafe_allow_html=True,
)
with st.expander(f"Таблицы ({n})"):
for t in st.session_state.schema_tables:
st.code(t.get("name", ""), language=None)
st.markdown('
', unsafe_allow_html=True)
# ── Бизнес-словарь ──
st.markdown('Бизнес-словарь
', unsafe_allow_html=True)
if st.session_state.vocabulary:
v = st.session_state.vocabulary
label = v.company if v.company else "Загружен"
st.markdown(f'✅ {label}', unsafe_allow_html=True)
if v.terms:
st.caption(f"Терминов: {len(v.terms)}")
else:
st.caption("Словарь не применён")
if st.button("Редактировать словарь", width='stretch'):
vocab_dialog()
# ──────────────────────────────────────────────
# Шапка
# ──────────────────────────────────────────────
st.markdown("""
""", unsafe_allow_html=True)
tab_query, tab_schema, tab_history = st.tabs(["Запрос", "Схема базы данных", "История"])
# ──────────── Tab: Запрос ────────────
with tab_query:
api_ready = (
st.session_state.api_health is not None
and st.session_state.api_health.get("model_loaded", False)
)
db_ready = st.session_state.schema_tables is not None
ready = api_ready and db_ready
if not ready:
missing = []
if not api_ready:
missing.append("API/модель не готовы")
if not db_ready:
missing.append("база данных не подключена")
st.warning("Система не готова: " + ", ".join(missing) + ". Используйте панель слева.")
question = st.text_area(
"Вопрос на естественном языке",
placeholder="Например: Какая выручка за январь этого года?",
height=100,
disabled=not ready,
)
col_btn, _ = st.columns([1, 5])
with col_btn:
run_btn = st.button(
"Выполнить",
type="primary",
disabled=not ready or not question.strip(),
width='stretch',
)
# Примеры для демо-базы
if (
st.session_state.connection_string
and "sales" in st.session_state.connection_string
):
st.markdown('Примеры запросов
', unsafe_allow_html=True)
ex_cols = st.columns(3)
examples = [
"Какая выручка за 2026 год?",
"Топ-5 клиентов по сумме заказов",
"Сколько заказов у каждого менеджера?",
]
for i, ex in enumerate(examples):
with ex_cols[i]:
if st.button(ex, key=f"ex_{i}", width='stretch'):
question = ex
run_btn = True
if run_btn and question.strip():
cs = st.session_state.connection_string
vocab = st.session_state.vocabulary
with st.spinner("Запрос к API. Это может занять несколько минут"):
try:
resp = _api_query(question, cs, vocab)
except Exception as e:
st.error(f"Ошибка: {e}")
st.stop()
st.markdown("**Сгенерированный SQL**")
st.markdown(f'{resp.get("sql", "")}
', unsafe_allow_html=True)
gen_time = resp.get("gen_time_seconds", 0.0)
execution = resp.get("execution")
err = resp.get("error")
c1, c2, c3 = st.columns(3)
c1.metric("Время генерации", f"{gen_time:.1f} с")
if execution:
c2.metric("Строк получено", execution.get("row_count", 0))
c3.metric("Статус", "Успешно")
elif err:
c2.metric("Строк получено", "—")
c3.metric("Статус", "Ошибка")
if execution and execution.get("rows"):
import pandas as pd
st.markdown("**Результат**")
df = pd.DataFrame(execution["rows"], columns=execution["columns"])
st.dataframe(df, width='stretch')
elif execution and not execution.get("rows"):
st.info("Запрос выполнен успешно. Результат пустой.")
elif err:
st.error(f"Ошибка выполнения SQL: {err}")
st.session_state.history.append({
"question": question,
"sql": resp.get("sql", ""),
"success": bool(execution),
"rows": execution.get("row_count", 0) if execution else 0,
"time": gen_time,
})
# ──────────── Tab: Схема БД ────────────
with tab_schema:
if st.session_state.schema_tables is None:
st.info("Подключитесь к базе данных через панель слева.")
else:
show_samples = st.toggle("Показывать примеры данных", value=True)
for t in st.session_state.schema_tables:
with st.expander(f"{t['name']} — {len(t['columns'])} колонок"):
st.code(t.get("ddl", ""), language="sql")
if show_samples and t.get("sample_rows"):
import pandas as pd
cols = [c["name"] for c in t["columns"]]
st.caption("Примеры данных:")
st.dataframe(
pd.DataFrame(t["sample_rows"], columns=cols),
width='stretch',
)
# ──────────── Tab: История ────────────
with tab_history:
history = st.session_state.history
if not history:
st.info("История пуста. Выполните первый запрос на вкладке «Запрос».")
else:
col_h, col_clr = st.columns([5, 1])
with col_h:
st.markdown(f"**Запросов в сессии: {len(history)}**")
with col_clr:
if st.button("Очистить", width='stretch'):
st.session_state.history = []
st.rerun()
for i, item in enumerate(reversed(history)):
icon = "✅" if item["success"] else "❌"
with st.expander(f"{icon} {item['question']}", expanded=(i == 0)):
st.markdown(
f'{item["sql"]}
',
unsafe_allow_html=True,
)
c1, c2, c3 = st.columns(3)
c1.metric("Время генерации", f"{item['time']:.1f} с")
c2.metric("Строк", item["rows"])
c3.metric("Статус", "Успешно" if item["success"] else "Ошибка")