File size: 5,224 Bytes
550e87b
3378b23
550e87b
be06203
 
316a678
3f940cd
4bcb630
6a2d3ac
 
c3b9b9a
 
4bcb630
be06203
83ed4d1
6a2d3ac
c3b9b9a
 
7b594ac
 
 
 
 
6a2d3ac
b8c8744
83ed4d1
 
 
 
 
 
 
 
 
6a2d3ac
 
 
 
 
 
4bcb630
c3b9b9a
4bcb630
c3b9b9a
4bcb630
 
83ed4d1
c3b9b9a
 
4bcb630
c3b9b9a
6a2d3ac
 
 
c3b9b9a
6a2d3ac
c3b9b9a
f6c85ec
c89ea47
f6c85ec
c3b9b9a
c89ea47
 
 
 
 
6a2d3ac
 
b8c8744
6a2d3ac
 
 
c89ea47
6a2d3ac
bc46efe
b8c8744
 
 
bc46efe
6a2d3ac
316a678
6a2d3ac
b8c8744
6a2d3ac
 
 
7b594ac
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6a2d3ac
7b594ac
 
 
 
 
 
6a2d3ac
 
7b594ac
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6a2d3ac
7b594ac
 
3378b23
b8c8744
3378b23
 
 
 
 
 
 
 
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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
import os
import json

import streamlit as st
import streamlit.components.v1 as components
from chat_history import insert_chat_history
from css import load_css
from custom_pgvector import CustomPGVector
from langchain import OpenAI
from langchain.callbacks import get_openai_callback
from langchain.chains import ConversationalRetrievalChain
from langchain.chains.conversation.memory import ConversationBufferMemory
from langchain.embeddings import GPT4AllEmbeddings
from message import Message
import sqlalchemy

CONNECTION_STRING = "postgresql+psycopg2://localhost/sorbobot"

st.set_page_config(layout="wide")

st.title("Sorbobot - Le futur de la recherche scientifique interactive")

chat_column, doc_column = st.columns([2, 1])


def connect() -> sqlalchemy.engine.Connection:
    engine = sqlalchemy.create_engine(CONNECTION_STRING)
    conn = engine.connect()
    return conn


conn = connect()


def initialize_session_state():
    if "history" not in st.session_state:
        st.session_state.history = []
    if "token_count" not in st.session_state:
        st.session_state.token_count = 0
    if "conversation" not in st.session_state:
        embeddings = GPT4AllEmbeddings()

        db = CustomPGVector(
            embedding_function=embeddings,
            table_name="article",
            column_name="abstract_embedding",
            connection=conn,
        )

        retriever = db.as_retriever()

        llm = OpenAI(
            temperature=0,
            openai_api_key=os.environ["OPENAI_API_KEY"],
            model="text-davinci-003",
        )

        memory = ConversationBufferMemory(
            output_key="answer", memory_key="chat_history", return_messages=True
        )
        st.session_state.conversation = ConversationalRetrievalChain.from_llm(
            llm=llm,
            retriever=retriever,
            verbose=True,
            memory=memory,
            return_source_documents=True,
        )


def on_click_callback():
    with get_openai_callback() as cb:
        human_prompt = st.session_state.human_prompt
        llm_response = st.session_state.conversation(human_prompt)
        st.session_state.history.append(Message("human", human_prompt))
        st.session_state.history.append(
            Message(
                "ai", llm_response["answer"], documents=llm_response["source_documents"]
            )
        )
        st.session_state.token_count += cb.total_tokens
        insert_chat_history(conn, human_prompt, llm_response["answer"])


load_css()
initialize_session_state()

with chat_column:
    chat_placeholder = st.container()
    prompt_placeholder = st.form("chat-form")
    information_placeholder = st.empty()

    with chat_placeholder:
        for chat in st.session_state.history:
            div = f"""
                <div class="chat-row 
                    {'' if chat.origin == 'ai' else 'row-reverse'}">
                    <img class="chat-icon" src="./app/static/{
                        'ai_icon.png' if chat.origin == 'ai' 
                                    else 'user_icon.png'}"
                        width=32 height=32>
                    <div class="chat-bubble
                    {'ai-bubble' if chat.origin == 'ai' else 'human-bubble'}">
                        &#8203;{chat.message}
                    </div>
                </div>
            """
            st.markdown(div, unsafe_allow_html=True)

        for _ in range(3):
            st.markdown("")

    with prompt_placeholder:
        st.markdown("**Chat**")
        cols = st.columns((6, 1))
        cols[0].text_input(
            "Chat",
            value="Hello bot",
            label_visibility="collapsed",
            key="human_prompt",
        )
        cols[1].form_submit_button(
            "Submit",
            type="primary",
            on_click=on_click_callback,
        )

    information_placeholder.caption(
        f"""
    Used {st.session_state.token_count} tokens \n
    Debug Langchain conversation: 
    {st.session_state.conversation.memory.buffer}
    """
    )

    components.html(
        """
    <script>
    const streamlitDoc = window.parent.document;

    const buttons = Array.from(
        streamlitDoc.querySelectorAll('.stButton > button')
    );
    const submitButton = buttons.find(
        el => el.innerText === 'Submit'
    );

    streamlitDoc.addEventListener('keydown', function(e) {
        switch (e.key) {
            case 'Enter':
                submitButton.click();
                break;
        }
    });
    </script>
    """,
        height=0,
        width=0,
    )

with doc_column:
    if len(st.session_state.history) > 0:
        st.markdown("**Source documents**")
        for doc in st.session_state.history[-1].documents:
            doc_content = json.loads(doc.page_content)

            expander = st.expander(doc_content["title"])
            expander.markdown("**" + doc_content["doi"] + "**")
            expander.markdown(doc_content["abstract"])
            expander.markdown("**Authors** : " + doc_content["authors"])
            expander.markdown("**Keywords** : " + doc_content["keywords"])
            expander.markdown("**Distance** : " + str(doc_content["distance"]))