assistant/main.py at main · prcore/assistant · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
import os
from random import choice
from string import ascii_letters, digits

import streamlit as st
from streamlit.delta_generator import DeltaGenerator

import chat
import index
import web


def init_session_state() -> None:
    if "human" not in st.session_state:
        st.session_state["human"] = []
    if "ai" not in st.session_state:
        st.session_state["ai"] = []
    if "input" not in st.session_state:
        st.session_state["input"] = ""
    if "input_disabled" not in st.session_state:
        st.session_state["input_disabled"] = False
    if "magic_mode" not in st.session_state:
        st.session_state["magic_mode"] = False


def check_input(container: DeltaGenerator, query: str) -> None:
    if not query:
        return

    print("-" * 24)
    print(f"Query from {web.get_ip()}: {query}")
    if query == os.getenv("MAGIC_WORD", "".join(choice(ascii_letters + digits) for _ in range(16))):
        st.session_state["magic_mode"] = True
        st.session_state["memory"] = chat.new_magic_memory()
        st.session_state["chain"] = chat.get_magic_chain()
        st.session_state["human"] = []
        st.session_state["ai"] = []
        st.session_state["input_disabled"] = False
        st.session_state["input"] = ""
        st.experimental_rerun()
        return

    st.session_state["input"] = ""
    st.session_state.human.append(query)
    web.display_conversation(container)

    output = chat.get_response(query)
    print(f"Response: {output}")
    st.session_state["input_disabled"] = False
    st.session_state.ai.append(output)
    web.display_conversation(container)
    st.experimental_rerun()


if __name__ == '__main__':
    web.set_web()
    init_session_state()
    index.create_index()
    index.load_vector_store()
    conversation_container = st.empty()
    web.display_conversation(conversation_container)
    user_input = chat.get_input()
    check_input(conversation_container, user_input)