"""Streamlit-интерфейс утилиты Ru2SQL. Запуск: streamlit run streamlit_app.py Что умеет: - Подключиться к любой SQLite/PostgreSQL/MySQL базе данных - Загрузить бизнес-словарь компании из YAML-файла или редактировать прямо в браузере - Принять вопрос на русском → сгенерировать SQL → выполнить → показать результат - Хранить историю запросов в текущей сессии """ from __future__ import annotations import sys import time from pathlib import Path import streamlit as st # Путь к src/ ROOT = Path(__file__).resolve().parent sys.path.insert(0, str(ROOT)) # ────────────────────────────────────────────── # Конфигурация страницы # ────────────────────────────────────────────── st.set_page_config( page_title="Ru2SQL — Natural Language → SQL", page_icon="🗄️", layout="wide", initial_sidebar_state="expanded", ) # ────────────────────────────────────────────── # CSS # ────────────────────────────────────────────── 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\n" "terms:\n" " выручка: SUM(orders.amount) WHERE status = 'paid'\n\n" "filters:\n" " только_оплаченные: orders.status = 'paid'\n\n" "notes: []\n" ) def _init_state(): defaults = { "history": [], "model_loaded": False, "engine": None, "db_connector": None, "db_executor": None, "vocabulary": None, "db_connection_string": "", "vocab_yaml": _default_vocab_yaml(), } for k, v in defaults.items(): if k not in st.session_state: st.session_state[k] = v _init_state() # ────────────────────────────────────────────── # Вспомогательные функции # ────────────────────────────────────────────── @st.cache_resource(show_spinner="Загружаю модель… (~30 с на первый раз)") def _load_engine(): from src.models.inference import InferenceEngine engine = InferenceEngine() engine.load() return engine def _connect_db(cs: str): from src.db.connector import DbConnector from src.db.executor import SqlExecutor connector = DbConnector(cs) executor = SqlExecutor(cs) return connector, executor def _load_vocab_from_yaml(yaml_text: str): import tempfile from src.business.vocabulary import BusinessVocabulary tmp = Path(tempfile.mktemp(suffix=".yaml")) tmp.write_text(yaml_text, encoding="utf-8") vocab = BusinessVocabulary.from_yaml(tmp) tmp.unlink(missing_ok=True) return vocab # ────────────────────────────────────────────── # Боковая панель # ────────────────────────────────────────────── with st.sidebar: st.title("⚙️ Настройки") # ── Модель — загружается автоматически при старте ── st.subheader("🤖 Модель") if not st.session_state.model_loaded: with st.spinner("Загружаю модель…"): try: st.session_state.engine = _load_engine() st.session_state.model_loaded = True except Exception as e: st.error(f"Ошибка загрузки модели: {e}") if st.session_state.model_loaded: st.markdown('✅ Модель готова', unsafe_allow_html=True) else: st.markdown('⚠️ Модель не загружена', unsafe_allow_html=True) st.divider() # ── База данных ── st.subheader("🗄️ База данных") db_type = st.radio("Тип подключения", ["SQLite файл", "Строка подключения"], horizontal=True) if db_type == "SQLite файл": uploaded = st.file_uploader("Загрузить .sqlite файл", type=["sqlite", "db"]) use_demo = st.checkbox("Использовать демо-базу", value=True) if use_demo: demo_path = ROOT / "data" / "demo" / "sales.sqlite" cs = str(demo_path) elif uploaded: import tempfile tmp_db = Path(tempfile.mktemp(suffix=".sqlite")) tmp_db.write_bytes(uploaded.read()) cs = str(tmp_db) else: cs = "" else: cs = st.text_input( "Строка подключения", placeholder="postgresql://user:pass@localhost/mydb", value=st.session_state.db_connection_string, ) if cs and st.button("Подключиться к БД", use_container_width=True): try: connector, executor = _connect_db(cs) tables = connector.list_tables() st.session_state.db_connector = connector st.session_state.db_executor = executor st.session_state.db_connection_string = cs st.success(f"Подключено! Таблиц: {len(tables)}") except Exception as e: st.error(f"Ошибка подключения: {e}") if st.session_state.db_connector: tables = st.session_state.db_connector.list_tables() st.markdown('✅ БД подключена', unsafe_allow_html=True) with st.expander("Таблицы"): for t in tables: st.code(t) st.divider() # ── Бизнес-словарь ── st.subheader("📖 Бизнес-словарь") vocab_yaml = st.text_area( "YAML-конфигурация", value=st.session_state.vocab_yaml, height=260, help="Определите термины вашей компании — модель будет их учитывать при генерации SQL", ) st.session_state.vocab_yaml = vocab_yaml if st.button("Применить словарь", use_container_width=True): try: st.session_state.vocabulary = _load_vocab_from_yaml(vocab_yaml) st.success("Словарь применён!") except Exception as e: st.error(f"Ошибка в YAML: {e}") if st.session_state.vocabulary: v = st.session_state.vocabulary st.markdown(f'✅ Словарь: {v.company or "загружен"}', unsafe_allow_html=True) terms_count = len(v.terms) if terms_count: st.caption(f"{terms_count} терминов определено") # ────────────────────────────────────────────── # Основная область # ────────────────────────────────────────────── st.title("🗄️ Ru2SQL — Бизнес-аналитика на русском языке") st.caption("Задайте вопрос на русском → получите SQL и данные из вашей базы") tab_query, tab_schema, tab_history = st.tabs(["💬 Запрос", "📐 Схема БД", "🕓 История"]) # ──────────── Вкладка: Запрос ──────────── with tab_query: ready = st.session_state.model_loaded and st.session_state.db_connector is not None if not ready: cols = st.columns(2) with cols[0]: if not st.session_state.model_loaded: st.warning("⚠️ Загрузите модель в левой панели") with cols[1]: if st.session_state.db_connector is None: st.warning("⚠️ Подключитесь к базе данных в левой панели") question = st.text_area( "Ваш вопрос", placeholder="Например: Какая выручка за январь этого года?", height=100, disabled=not ready, ) col_btn, col_hint = st.columns([1, 4]) with col_btn: run_btn = st.button("▶ Выполнить", type="primary", disabled=not ready or not question.strip(), use_container_width=True) with col_hint: if ready: st.caption("Модель сгенерирует SQL и выполнит его на вашей БД") # Быстрые примеры if st.session_state.db_connection_string and "sales" in st.session_state.db_connection_string: st.caption("💡 Попробуйте:") example_cols = st.columns(3) examples = [ "Какая выручка за 2026 год?", "Топ-5 клиентов по сумме заказов", "Сколько заказов по каждому менеджеру?", ] for i, ex in enumerate(examples): with example_cols[i]: if st.button(ex, key=f"ex_{i}", use_container_width=True): question = ex run_btn = True if run_btn and question.strip(): engine = st.session_state.engine connector = st.session_state.db_connector executor = st.session_state.db_executor vocab = st.session_state.vocabulary # Обогащаем вопрос бизнес-словарём enriched_question = vocab.enrich_prompt(question) if vocab else question # Получаем схему schema = connector.render_schema(include_samples=True) with st.spinner("Генерирую SQL…"): t0 = time.time() result = engine.generate(schema, enriched_question) gen_time = time.time() - t0 st.subheader("Сгенерированный SQL") st.markdown(f'
{result.sql}
', unsafe_allow_html=True) col1, col2 = st.columns(2) col1.metric("Время генерации", f"{gen_time:.1f} с") # Выполняем SQL if result.sql.strip(): with st.spinner("Выполняю запрос…"): qr = executor.run(result.sql) if qr.success: col2.metric("Строк в результате", qr.row_count) st.subheader("Результат") if qr.rows: import pandas as pd df = pd.DataFrame(qr.rows, columns=qr.columns) st.dataframe(df, use_container_width=True) else: st.info("Запрос выполнен успешно, результат пустой") else: col2.error("Ошибка выполнения") st.error(f"SQL ошибка: {qr.error}") # Добавляем в историю st.session_state.history.append({ "question": question, "sql": result.sql, "success": qr.success if result.sql.strip() else False, "rows": qr.row_count if result.sql.strip() and qr.success else 0, "time": gen_time, }) # ──────────── Вкладка: Схема БД ──────────── with tab_schema: if st.session_state.db_connector is None: st.info("Подключитесь к базе данных в левой панели") else: connector = st.session_state.db_connector st.subheader("Структура базы данных") show_samples = st.toggle("Показывать примеры строк", value=True) schema_text = connector.render_schema(include_samples=show_samples) for table in connector.get_schema(include_samples=show_samples): with st.expander(f"📋 {table.name} ({len(table.columns)} колонок)"): st.code(table.to_ddl(), language="sql") if show_samples and table.sample_rows: import pandas as pd cols = [c.name for c in table.columns] st.caption("Примеры строк:") st.dataframe( pd.DataFrame(table.sample_rows, columns=cols), use_container_width=True, ) # ──────────── Вкладка: История ──────────── with tab_history: history = st.session_state.history if not history: st.info("История запросов пуста. Задайте первый вопрос на вкладке «Запрос».") else: st.subheader(f"История запросов ({len(history)})") if st.button("Очистить историю"): st.session_state.history = [] st.rerun() for i, item in enumerate(reversed(history)): status = "✅" if item["success"] else "❌" with st.expander(f"{status} {item['question']}", expanded=(i == 0)): st.markdown(f'
{item["sql"]}
', unsafe_allow_html=True) cols = st.columns(3) cols[0].metric("Время генерации", f"{item['time']:.1f} с") cols[1].metric("Строк", item["rows"]) cols[2].metric("Статус", "OK" if item["success"] else "Ошибка")