RustX/CSV-ChatBot

1import os2import streamlit as st3from dotenv import load_dotenv4from io import BytesIO5from io import StringIO6import sys7import re8from langchain.agents import create_csv_agent9from langchain.chat_models import ChatOpenAI10from modules.history import ChatHistory11from modules.layout import Layout12from modules.utils import Utilities13from modules.sidebar import Sidebar14 15#To be able to update the changes made to modules in localhost,16#you can press the "r" key on the localhost page to refresh and reflect the changes made to the module files.17def reload_module(module_name):18    import importlib19    import sys20    if module_name in sys.modules:21        importlib.reload(sys.modules[module_name])22    return sys.modules[module_name]23 24history_module = reload_module('modules.history')25layout_module = reload_module('modules.layout')26utils_module = reload_module('modules.utils')27sidebar_module = reload_module('modules.sidebar')28 29ChatHistory = history_module.ChatHistory30Layout = layout_module.Layout31Utilities = utils_module.Utilities32Sidebar = sidebar_module.Sidebar33 34def init():35    load_dotenv()36    st.set_page_config(layout="wide", page_icon="💬", page_title="ChatBot-CSV")37 38 39def main():40 41    init()42    layout, sidebar, utils = Layout(), Sidebar(), Utilities()43    layout.show_header()44    user_api_key = utils.load_api_key()45 46    if not user_api_key:47        layout.show_api_key_missing()48    else:49        os.environ["OPENAI_API_KEY"] = user_api_key50        uploaded_file = utils.handle_upload()51 52        if uploaded_file:53            history = ChatHistory()54            sidebar.show_options()55 56            uploaded_file_content = BytesIO(uploaded_file.getvalue())57 58            try:59                chatbot = utils.setup_chatbot(60                    uploaded_file, st.session_state["model"], st.session_state["temperature"]61                )62                st.session_state["chatbot"] = chatbot63 64                if st.session_state["ready"]:65                    response_container, prompt_container = st.container(), st.container()66 67                    with prompt_container:68                        is_ready, user_input = layout.prompt_form()69 70                        history.initialize(uploaded_file)71                        if st.session_state["reset_chat"]:72                            history.reset(uploaded_file)73 74                        if is_ready:75                            history.append("user", user_input)76                            output = st.session_state["chatbot"].conversational_chat(user_input)77                            history.append("assistant", output)78 79                    history.generate_messages(response_container)80 81                    if st.session_state["show_csv_agent"]:82                        query = st.text_input(label="Use CSV agent for precise information about the structure of your csv file / csv 파일 구조에 대한 정확한 정보를 얻으려면 CSV 에이전트를 사용하십시오", placeholder="ex : how many rows in my file ? / 예: 내 파일에 몇 개의 행이 있습니까?")83                        if query != "":84 85                            old_stdout = sys.stdout86                            sys.stdout = captured_output = StringIO()87                            agent = create_csv_agent(ChatOpenAI(temperature=0), uploaded_file_content, verbose=True, max_iterations=8)88 89                            result = agent.run(query)90 91                            sys.stdout = old_stdout92                            thoughts = captured_output.getvalue()93 94                            cleaned_thoughts = re.sub(r'\x1b\[[0-9;]*[a-zA-Z]', '', thoughts)95                            cleaned_thoughts = re.sub(r'\[1m>', '', cleaned_thoughts)96 97                            with st.expander("Show agent's thoughts"):98                                st.write(cleaned_thoughts)99                            100                            st.write(result)101 102            except Exception as e:103                st.error(f"Error: {str(e)}")104 105    sidebar.about()106 107 108if __name__ == "__main__":109    main()