src/streamlit_app.py

import streamlit as st
from langchain_core.messages import HumanMessage, AIMessage
from get_classification import get_graph_class
from datetime import datetime


def message_to_dict(messages):
    result = []
    for message in messages:
        if isinstance(message, HumanMessage) or isinstance(message, AIMessage):
            print("message:", message.content)
        if message.content == "" or message.content is None:
            continue
        if isinstance(message, HumanMessage):
            result.append({"role": "user", "content": message.content})
        elif isinstance(message, AIMessage):
            result.append({"role": "assistant", "content": message.content})
    print("-" * 100)
    return result


def find_last_bot_message(messages):
    """Находит последнее сообщение бота"""
    for message in messages[::-1]:
        if isinstance(message, AIMessage) and len(message.content) > 0:
            return message.content
    return None


def display_chat_messages():
    """Отображает историю сообщений в чате"""
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])


def save_broken_case():
    messages_dict = st.session_state.messages
    result_str = ""
    for elem in messages_dict:
        role = elem["role"]
        content = elem["content"]
        result_str += f"{role}: {content}\n"
    
    current_datetime = datetime.now()
    formatted_datetime = current_datetime.strftime("%Y-%m-%d %H:%M:%S")

    # with open("/Users/admin/my_documents/retrieval_part/services/broken_cases.txt", "a") as file:
    #     file.write(f"{formatted_datetime}\n" + result_str + "\n" + "-" * 50 + "\n\n")


def handle_user_input():
    """Обрабатывает ввод пользователя и генерирует ответ бота"""
    if prompt := st.chat_input("Введите ваш вопрос"):
        st.session_state.messages.append({"role": "user", "content": prompt})
        if prompt.lower().startswith("log"):
            save_broken_case()
            st.session_state.messages = []
            display_chat_messages()
        else:
            with st.chat_message("user"):
                st.markdown(prompt)
            
            try:
                if "bot" not in st.session_state:
                    print("init!!!!")
                    st.session_state.bot = get_graph_class(prompt)
                st.session_state.bot.invoke(prompt)
            except Exception as e:
                st.error(f"Ошибка: {str(e)}")
            
            # Извлекаем последнее сообщение бота
            last_bot_message = find_last_bot_message(st.session_state.bot.messages)
            st.session_state.messages.append(
                {"role": "assistant", "content": last_bot_message}
            )
            with st.chat_message("assistant"):
                st.markdown(last_bot_message)


def clear_chat():
    """Очищает чат и пересоздает бота"""
    st.session_state.messages = []
    if "bot" in st.session_state:
        del st.session_state.bot


def main():
    """Основная функция приложения"""
    # Заголовок приложения
    st.title("Чат-бот технической поддержки")

    # Кнопка очистки чата
    if st.button("Clear"):
        clear_chat()

    if "messages" not in st.session_state:
        st.session_state.messages = []

    # Отображение чата и обработка ввода
    display_chat_messages()
    handle_user_input()


if __name__ == "__main__":
    main()