scripts/smoke_local.py

"""Локальный smoke-тест работоспособности Ru2SQL.

Прогоняет ключевые слои без поднятия Streamlit/FastAPI как отдельных
процессов. Покрывает: импорты, demo-базу, vocabulary, prompt,
постобработку, guardrail, новый API через FastAPI TestClient и
опционально — короткий инференс реальной модели.

Запуск:
    python scripts/smoke_local.py
    python scripts/smoke_local.py --with-model    # медленно, грузит Qwen
"""

from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

GREEN = "\033[32m"
RED = "\033[31m"
YELLOW = "\033[33m"
RESET = "\033[0m"

results: list[tuple[str, str, str]] = []


def check(name: str):
    def decorator(fn):
        try:
            t0 = time.time()
            fn()
            dt = time.time() - t0
            results.append(("OK", name, f"{dt*1000:.0f} мс"))
        except Exception as e:
            results.append(("FAIL", name, f"{type(e).__name__}: {e}"))
        return fn
    return decorator


# ──────────────────────────────────────────────────────────────────────
# Слой 1 — импорты, конфиг
# ──────────────────────────────────────────────────────────────────────
@check("Импорты ядра")
def _():
    from src.config import settings
    from src.api.schemas import GenerateRequest, QueryRequest, SchemaRequest
    from src.business.vocabulary import BusinessVocabulary
    from src.data.prompt import build_chat_messages, BASE_SYSTEM_PROMPT
    from src.data.schema_provider import (
        SpiderSchemaProvider, ConnectionSchemaProvider, TableSchema,
    )
    from src.db.connector import DbConnector
    from src.db.executor import SqlExecutor
    from src.models.postprocess import postprocess, is_select_only, is_valid_sql
    assert settings.base_model_name.startswith("Qwen")


# ──────────────────────────────────────────────────────────────────────
# Слой 2 — демо-база и read-only guardrail
# ──────────────────────────────────────────────────────────────────────
@check("Демо-база sales.sqlite: подключение и схема")
def _():
    from src.db.connector import DbConnector
    db = ROOT / "data" / "demo" / "sales.sqlite"
    assert db.exists(), f"Не найден файл {db} — запусти data/demo/create_demo_db.py"
    conn = DbConnector(str(db))
    tables = conn.list_tables()
    assert set(tables) == {"customers", "managers", "products", "orders", "order_items"}


@check("Демо-база: реальный SELECT выручки")
def _():
    from src.db.executor import SqlExecutor
    ex = SqlExecutor(str(ROOT / "data" / "demo" / "sales.sqlite"))
    res = ex.run("SELECT SUM(amount) FROM orders WHERE status='paid'")
    assert res.success, res.error
    assert res.rows[0][0] > 0


@check("Read-only guardrail: DELETE отклоняется драйвером")
def _():
    from src.db.executor import SqlExecutor
    ex = SqlExecutor(str(ROOT / "data" / "demo" / "sales.sqlite"))
    res = ex.run("DELETE FROM orders")
    assert not res.success
    assert ex.run("SELECT COUNT(*) FROM orders").rows[0][0] == 120


# ──────────────────────────────────────────────────────────────────────
# Слой 3 — BusinessVocabulary из YAML
# ──────────────────────────────────────────────────────────────────────
@check("BusinessVocabulary из configs/example_vocabulary.yaml")
def _():
    from src.business.vocabulary import BusinessVocabulary
    vocab = BusinessVocabulary.from_yaml(ROOT / "configs" / "example_vocabulary.yaml")
    assert bool(vocab)
    assert "выручка" in vocab.terms
    assert "SUM(orders.amount)" in vocab.render_system_context()


# ──────────────────────────────────────────────────────────────────────
# Слой 4 — PromptBuilder с реальной схемой и словарём
# ──────────────────────────────────────────────────────────────────────
@check("PromptBuilder: vocabulary уходит в system, не в user")
def _():
    from src.business.vocabulary import BusinessVocabulary
    from src.data.prompt import build_chat_messages
    from src.db.connector import DbConnector
    vocab = BusinessVocabulary.from_yaml(ROOT / "configs" / "example_vocabulary.yaml")
    schema = DbConnector(str(ROOT / "data" / "demo" / "sales.sqlite")).render_schema()
    msgs = build_chat_messages(schema, "Какая выручка за 2026 год?", vocabulary=vocab)
    assert "SUM(orders.amount)" in msgs[0]["content"]
    assert "SUM(orders.amount)" not in msgs[1]["content"]


# ──────────────────────────────────────────────────────────────────────
# Слой 5 — постобработка и guardrail
# ──────────────────────────────────────────────────────────────────────
@check("Постобработка: markdown, префиксы, truncated, мусор")
def _():
    from src.models.postprocess import postprocess, is_select_only
    cases = [
        ("```sql\nSELECT 1;\n```", lambda s: s.upper().startswith("SELECT")),
        ("Ответ: SELECT name FROM customers;", lambda s: s.startswith("SELECT name")),
        ("SELECT * FROM orders WHERE", lambda s: s == ""),
        ("просто текст без SQL", lambda s: s == ""),
    ]
    for raw, ok in cases:
        assert ok(postprocess(raw)), f"Сбой на: {raw!r} → {postprocess(raw)!r}"
    assert is_select_only("SELECT 1") is True
    assert is_select_only("DROP TABLE x") is False
    assert is_select_only("DELETE FROM x") is False


# ──────────────────────────────────────────────────────────────────────
# Слой 6 — FastAPI через TestClient с замоканным InferenceEngine
# ──────────────────────────────────────────────────────────────────────
@check("FastAPI: /health, /schema, /query с подменённым engine")
def _():
    try:
        from fastapi.testclient import TestClient
    except ImportError:
        raise AssertionError("Поставь fastapi[all] или httpx — TestClient недоступен")

    from src.api import dependencies as deps
    from src.api.main import app
    from src.models.inference import GenerationResult

    class FakeEngine:
        loaded = True
        base_model_name = "Qwen/Qwen2.5-Coder-3B-Instruct"
        def generate(self, schema, question, vocabulary=None, **kw):
            # Эмулируем валидный SQL — pipeline должен пропустить через guardrail.
            sql = "SELECT SUM(amount) FROM orders WHERE status = 'paid'"
            return GenerationResult(sql=sql, raw_output=sql)

    app.dependency_overrides[deps.get_engine] = lambda: FakeEngine()
    try:
        with TestClient(app) as client:
            # /health
            r = client.get("/health")
            assert r.status_code == 200
            assert r.json()["model_loaded"] is True

            # /schema на реальной demo-базе
            db_path = ROOT / "data" / "demo" / "sales.sqlite"
            r = client.post("/schema", json={
                "connection_string": str(db_path),
                "include_samples": True,
            })
            assert r.status_code == 200, r.text
            tables = r.json()["tables"]
            assert {t["name"] for t in tables} == {
                "customers", "managers", "products", "orders", "order_items"
            }

            # /query на реальной demo-базе, FakeEngine отдаст валидный SELECT
            r = client.post("/query", json={
                "question": "Какая выручка за оплаченные заказы?",
                "connection_string": str(db_path),
                "execute": True,
                "vocabulary": {
                    "company": "Демо",
                    "terms": {"выручка": "SUM(amount) WHERE status='paid'"},
                },
            })
            assert r.status_code == 200, r.text
            body = r.json()
            assert body["is_valid_sql"] is True
            assert body["execution"] is not None
            assert body["execution"]["rows"][0][0] > 0

            # /query с DELETE — должен быть отбит guardrail'ом
            class DropEngine(FakeEngine):
                def generate(self, *a, **kw):
                    return GenerationResult(
                        sql="DELETE FROM orders WHERE id=1",
                        raw_output="DELETE FROM orders WHERE id=1",
                    )
            app.dependency_overrides[deps.get_engine] = lambda: DropEngine()
            r = client.post("/query", json={
                "question": "Удали заказ 1",
                "connection_string": str(db_path),
                "execute": True,
            })
            assert r.status_code == 200, r.text
            body = r.json()
            assert body["execution"] is None
            assert body["error"] and "гвардейл" in body["error"].lower()
    finally:
        app.dependency_overrides.clear()


# ──────────────────────────────────────────────────────────────────────
# Слой 7 — опциональный инференс реальной модели
# ──────────────────────────────────────────────────────────────────────
def run_model_smoke():
    @check("InferenceEngine: загрузка модели и одна генерация")
    def _():
        from src.business.vocabulary import BusinessVocabulary
        from src.db.connector import DbConnector
        from src.models.inference import InferenceEngine

        engine = InferenceEngine()
        engine.load()
        assert engine.loaded
        schema = DbConnector(str(ROOT / "data" / "demo" / "sales.sqlite")).render_schema()
        vocab = BusinessVocabulary.from_yaml(ROOT / "configs" / "example_vocabulary.yaml")
        res = engine.generate(schema, "Какая суммарная выручка за 2026 год?", vocab)
        assert res.sql and res.sql.upper().startswith("SELECT"), \
            f"Модель вернула: {res.raw_output!r}"


# ──────────────────────────────────────────────────────────────────────
# main
# ──────────────────────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--with-model", action="store_true",
        help="Дополнительно прогнать реальный инференс Qwen (медленно, ~30 сек на CPU)",
    )
    args = parser.parse_args()

    if args.with_model:
        run_model_smoke()

    print()
    print("=" * 64)
    print("Smoke-проверка Ru2SQL")
    print("=" * 64)
    for status, name, info in results:
        color = GREEN if status == "OK" else RED
        mark = "✓" if status == "OK" else "✗"
        print(f"  {color}{mark}{RESET} {name}  {YELLOW}[{info}]{RESET}")
    ok = sum(1 for s, _, _ in results if s == "OK")
    print("=" * 64)
    summary_color = GREEN if ok == len(results) else RED
    print(f"  {summary_color}{ok} / {len(results)} проверок пройдено{RESET}")
    print("=" * 64)
    if ok < len(results):
        print()
        print("Подсказка: запусти 'pytest -v' для подробных диагностик.")
    sys.exit(0 if ok == len(results) else 1)


if __name__ == "__main__":
    main()