frontend/ui.py - master - igormayer/langgraph_sql_agent_web

frontend/ui.py
177 строк7 KB
Igor Mayer
fixes for running in docker
08 фев 2026, 21:16
08 фев 2026, 21:16988621e
import sys
import os
from dotenv import load_dotenv
import json
import streamlit as st
import requests
import httpx
import sqlglot
from sqlglot.errors import ParseError
import uuid

st.set_page_config(page_title="Langgraph SQL Agent Web", layout="wide")

if "thread_id" not in st.session_state:
    st.session_state.thread_id = str(uuid.uuid4())

BASE_URL = "http://api:8000"

# load env params
load_dotenv() 
model_name = os.getenv("MODEL_NAME")
model_host_type = os.getenv("MODEL_HOST_TYPE")
model_host_url = os.getenv("MODEL_HOST_URL")
model_api_key = os.getenv("MODEL_API_KEY")

db_dialect = os.getenv("DB_DIALECT", "clickhouse")
db_user = os.getenv("DB_USER", "default")
db_password = os.getenv("DB_PASSWORD", "learn_password")
db_host = os.getenv("DB_HOST", "localhost")
db_port = os.getenv("DB_PORT", "8123")
db_name = os.getenv("DB_NAME", "e_commerce_analytics")

db_uri = f"{db_dialect}://{db_user}:{db_password}@{db_host}:{db_port}/{db_name}"

if not model_name or not model_host_url:
    print("ERROR: Model not loaded!")
    sys.exit(1)

def format_sql_to_markdown(sql_query):
    if "DB:Exception" in sql_query:
        return f"```sql\n{sql_query}\n```"
    try:    
        formatted_sql = sqlglot.transpile(
            sql_query, 
            read="clickhouse",  # Исходный диалект
            write="clickhouse", # Целевой диалект, 
            pretty=True)[0]
        return f"```sql\n{formatted_sql}\n```"
    except Exception:
        # Если возникла ошибка — возвращаем исходный текст без изменений
        return f"```sql\n{sql_query}\n```"

# Инициализация истории сообщений
if "messages" not in st.session_state:
    st.session_state.messages = []

if st.sidebar.button("Новая сессия", use_container_width=True):
    st.session_state.thread_id = str(uuid.uuid4())
    st.session_state.messages = []
    st.rerun()

# Выведем ID в сайдбаре для отладки (по желанию)
st.sidebar.caption(f"Сессия: {st.session_state.thread_id}")

# Правый блок настроек
with st.sidebar:

    model_host_options = ["OLLAMA", "OPENAI API"]
    
    model_host_type = st.selectbox("Тип сервера модели", model_host_options, index=model_host_options.index(model_host_type))
    model_host_url = st.text_input("Адрес сервера модели", value=model_host_url)
    model_name = st.text_input("Название модели", value=model_name)
    model_api_key = st.text_input("API Key", value=model_api_key)

    st.divider()
    
    db_dialect = st.selectbox("Тип БД", ["clickhouse", "postgresql"])
    db_host = st.text_input("Сервер БД", value=db_host)
    db_port = st.text_input("Порт БД", value=db_port)
    db_name = st.text_input("Название БД", value=db_name)
    db_user = st.text_input("Логин БД", value=db_user)
    db_password = st.text_input("Пароль БД", type="password", value=db_password) 

    if st.button("Сохранить настройки"):
        settings_payload = {
            "db_type": db_dialect, "db_server": db_host, "db_name": db_name,
            "db_user": db_user, "db_pass": db_password, "ollama_url": model_host_url,
            "model_name": model_name
        }
        try:
            res = requests.post(f"{BASE_URL}/settings", json=settings_payload)
            if res.status_code == 200:
                st.success("Настройки применены")
        except Exception as e:
            st.error(f"Ошибка связи: {e}")

# Область диалога
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Строка ввода промта 
if prompt := st.chat_input("Введите ваш запрос..."):
    # Добавляем сообщение пользователя в интерфейс и состояние
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Список для сбора полного финального ответа (если нужно сохранить в историю)
    collected_responses = []

    settings = {
        'db_dialect': db_dialect,
        'db_host': db_host,
        'db_port': db_port,
        'db_name': db_name,
        'db_user': db_user,
        'db_password': db_password,
        'model_name': model_name,
        'model_host_url': model_host_url,
        'model_host_type': model_host_type,
        'model_api_key': model_api_key,
        'thread_id': st.session_state.thread_id
    }

    payload = {
        "prompt": prompt,
        "settings": settings
    }

    with httpx.stream("POST", f"{BASE_URL}/chat", json=payload, timeout=None) as r:
        for line in r.iter_lines():
            if not line.startswith("data: "):
                continue
            
            print(line)
            raw_data = json.loads(line[6:])
            payload = raw_data.get("data")
            msg_type = raw_data.get("type")

            # Системные статусы и SQL (выводим в отдельный блок ассистента)
            if msg_type == "status":
                msg = payload.get("message", "")
                sql = payload.get("sql", "")
                
                with st.chat_message("assistant"):
                    if msg:
                        st.info(msg) # Используем info для визуального отличия шага
                    if sql:
                        st.code(sql, language="sql")
                
                # Добавляем в историю (опционально)
                if msg: collected_responses.append(msg)
                if sql: collected_responses.append(f"```sql\n{sql}\n```")

            # Б. Обновления от узлов (каждый этап — отдельный блок)
            elif msg_type == "update":
                content_to_display = ""
                
                if "discovery" in payload:
                    content_to_display = "🔍 **Для запроса подходят следующие таблицы:**\n " + ", ".join(payload["discovery"].get("tables", ""))
                elif "schema" in payload:
                    content_to_display = "📋 **Схема данных:**\n" + format_sql_to_markdown(payload["schema"].get("schema", ""))
                elif "generator" in payload:
                    content_to_display = "⚙️ **Сгенерированный запрос:**\n" + format_sql_to_markdown(payload["generator"].get("sql_query", ""))
                elif "summarizer" in payload:
                    content_to_display = "✅ **Итоговый ответ:**\n" + payload["summarizer"].get("final_answer", "")

                if content_to_display:
                    with st.chat_message("assistant"):
                        st.markdown(content_to_display)
                    collected_responses.append(content_to_display)

    # Сохраняем всё накопленное взаимодействие в историю сессии
    if collected_responses:
        full_history_str = "\n\n".join(collected_responses)
        st.session_state.messages.append({"role": "assistant", "content": full_history_str})