Merge pull request #14 from KaiwuX/main

linancn · Nov 23, 2023 · 3004cde · 3004cde
2 parents 4ed2d65 + bfc6197
commit 3004cde
Show file tree

Hide file tree

Showing 5 changed files with 102 additions and 77 deletions.
diff --git a/src/Chat.py b/src/Chat.py
@@ -105,13 +105,17 @@
                 search_knowledge_base = st.toggle(
                     ui.search_knowledge_base_checkbox_label, value=False
                 )
-                search_online = st.toggle(ui.search_internet_checkbox_label, value=True)
+                search_online = st.toggle(
+                    ui.search_internet_checkbox_label, value=False
+                )
                 search_wikipedia = st.toggle(
                     ui.search_wikipedia_checkbox_label, value=False
                 )
                 search_arxiv = st.toggle(ui.search_arxiv_checkbox_label, value=False)
 
-                search_docs = st.toggle(ui.search_docs_checkbox_label, value=False)
+                search_docs = st.toggle(
+                    ui.search_docs_checkbox_label, value=False, disabled=True
+                )
 
                 # search_knowledge_base = True
                 # search_online = st.toggle(ui.search_internet_checkbox_label, value=False)
@@ -285,44 +289,59 @@ def main():
         if "xata_history_refresh" not in st.session_state:
             user_query = st.chat_input(placeholder=ui.chat_human_placeholder)
             if user_query:
-                try:
-                    st.chat_message("user", avatar=ui.chat_user_avatar).markdown(
-                        user_query
-                    )
-                    st.session_state["messages"].append(
-                        {"role": "user", "content": user_query}
-                    )
-                    human_message = HumanMessage(
-                        content=user_query,
-                        additional_kwargs={"id": st.session_state["username"]},
-                    )
-                    st.session_state["xata_history"].add_message(human_message)
-
-                    # check text sensitivity
-                    answer = check_text_sensitivity(user_query)["answer"]
-                    if answer is not None:
-                        with st.chat_message("assistant", avatar=ui.chat_ai_avatar):
-                            st.markdown(answer)
-                            st.session_state["messages"].append(
-                                {
-                                    "role": "assistant",
-                                    "content": answer,
-                                }
-                            )
-                            ai_message = AIMessage(
-                                content=answer,
-                                additional_kwargs={"id": st.session_state["username"]},
-                            )
-                            st.session_state["xata_history"].add_message(ai_message)
+                st.chat_message("user", avatar=ui.chat_user_avatar).markdown(user_query)
+                st.session_state["messages"].append(
+                    {"role": "user", "content": user_query}
+                )
+                human_message = HumanMessage(
+                    content=user_query,
+                    additional_kwargs={"id": st.session_state["username"]},
+                )
+                st.session_state["xata_history"].add_message(human_message)
+
+                # check text sensitivity
+                answer = check_text_sensitivity(user_query)["answer"]
+                if answer is not None:
+                    with st.chat_message("assistant", avatar=ui.chat_ai_avatar):
+                        st.markdown(answer)
+                        st.session_state["messages"].append(
+                            {
+                                "role": "assistant",
+                                "content": answer,
+                            }
+                        )
+                        ai_message = AIMessage(
+                            content=answer,
+                            additional_kwargs={"id": st.session_state["username"]},
+                        )
+                        st.session_state["xata_history"].add_message(ai_message)
+                else:
+                    if len(st.session_state["messages"]) == 2:
+                        chat_history_recent = ""
+                        chat_history_recent_neat = ""
                     else:
+                        current_message = st.session_state["messages"][-8:][1:][:-1]
+                        for item in current_message:
+                            item.pop("avatar", None)
+
                         chat_history_response = chat_history_chain()(
-                            {"input": st.session_state["xata_history"].messages[-6:]},
+                            {"input": current_message},
                         )
 
                         chat_history_recent = chat_history_response["text"]
-
+                        chat_history_recent_neat = chat_history_recent
+
+                    if (
+                        search_knowledge_base
+                        or search_online
+                        or search_wikipedia
+                        or search_arxiv
+                        or search_docs
+                    ):
                         func_calling_response = func_calling_chain().run(
                             chat_history_recent
+                            + " Latest message: "
+                            + str(st.session_state["messages"][-1])
                         )
 
                         query = func_calling_response.get("query")
@@ -365,41 +384,42 @@ def main():
                         )
 
                         input = f"""You Must:
-    - Respond to "{user_query}" by using information from "{docs_response}" (if available) and your own knowledge to provide a logical, clear, and critically analyzed reply in the same language.
-    - Use the chat context from "{chat_history_recent}" (if available) to adjust the level of detail in your response.
-    - Employ bullet points selectively, where they add clarity or organization.
-    - Cite sources in-text using the Author-Date citation style where applicable.
-    - Provide a list of full references, with hyperlinks, at the end for only the sources mentioned in the text.
-
-    Avoid:
-    - Repeating information or including redundancies.
-    - Translating cited references into the query's language.
-    - Prefacing responses with any designation such as "AI:"."""
-
-                        with st.chat_message("assistant", avatar=ui.chat_ai_avatar):
-                            st_cb = StreamHandler(st.empty())
-                            response = main_chain()(
-                                {"input": input},
-                                callbacks=[st_cb],
-                            )
+- Respond to "{user_query}" by using information from "{docs_response}" (if available) and your own knowledge to provide a logical, clear, and critically analyzed reply in the same language.
+- Use the chat context from "{chat_history_recent}" (if available) to adjust the level of detail in your response.
+- Employ bullet points selectively, where they add clarity or organization.
+- Cite sources in-text using the Author-Date citation style where applicable.
+- Provide a list of full references, with hyperlinks, at the end for only the sources mentioned in the text.
+
+Avoid:
+- Repeating information or including redundancies.
+- Translating cited references into the query's language.
+- Prefacing responses with any designation such as "AI:"."""
+
+                    else:
+                        input = f"""Respond to "{user_query}". If "{chat_history_recent_neat}" is not empty, use it as chat context."""
 
-                            st.session_state["messages"].append(
-                                {
-                                    "role": "assistant",
-                                    "content": response["text"],
-                                }
-                            )
-                            ai_message = AIMessage(
-                                content=response["text"],
-                                additional_kwargs={"id": st.session_state["username"]},
-                            )
-                            st.session_state["xata_history"].add_message(ai_message)
-                    if len(st.session_state["messages"]) == 3:
-                        st.session_state["xata_history_refresh"] = True
-                        st.rerun()
-                except:
                     with st.chat_message("assistant", avatar=ui.chat_ai_avatar):
-                        st.markdown(ui.chat_error_message)
+                        st_cb = StreamHandler(st.empty())
+                        response = main_chain()(
+                            {"input": input},
+                            callbacks=[st_cb],
+                        )
+
+                        st.session_state["messages"].append(
+                            {
+                                "role": "assistant",
+                                "content": response["text"],
+                            }
+                        )
+                        ai_message = AIMessage(
+                            content=response["text"],
+                            additional_kwargs={"id": st.session_state["username"]},
+                        )
+                        st.session_state["xata_history"].add_message(ai_message)
+
+                if len(st.session_state["messages"]) == 3:
+                    st.session_state["xata_history_refresh"] = True
+                    st.rerun()
         else:
             user_query = st.chat_input(placeholder=ui.chat_human_placeholder)
             del st.session_state["xata_history_refresh"]

diff --git a/src/ui/tiangong-en.py b/src/ui/tiangong-en.py
@@ -127,4 +127,5 @@
     "sidebar_delete_button_label": "Delete Chat",
     "sidebar_newchat_label": "New Chat",
     "chat_error_message": "Oops, we are facing an extremely high traffic. Please try again later.",
+    "wix_login_wait": "Please wait..."
 }
diff --git a/src/ui_config.py b/src/ui_config.py
@@ -61,6 +61,7 @@ class UI:
     chat_ai_welcome: str
     chat_human_placeholder: str
     chat_error_message: str
+    wix_login_wait:str
 
 
 def create_ui_from_config():

diff --git a/src/utils.py b/src/utils.py
@@ -1248,7 +1248,7 @@ def chat_history_chain():
         verbose=langchain_verbose,
     )
 
-    template = """Return highly concise and well-organized chat history from: {input}"""
+    template = """Return highly concise and well-organized chat history from: {input}. Do not include any references."""
     prompt = PromptTemplate(
         input_variables=["input"],
         template=template,
@@ -1282,14 +1282,12 @@ def main_chain():
 
     llm_chat = ChatOpenAI(
         model=llm_model,
-        temperature=0,
+        temperature=0.2,
         streaming=True,
         verbose=langchain_verbose,
     )
 
-    template = """You MUST ONLY response to science-related quests.
-    DO NOT return any information on politics, ethnicity, gender, national sovereignty, or other sensitive topics.
-    {input}"""
+    template = """{input} DO NOT return any information on politics, ethnicity, gender, national sovereignty, or other sensitive topics."""
 
     prompt = PromptTemplate(
         input_variables=["input"],

diff --git a/src/wix_oauth.py b/src/wix_oauth.py
@@ -141,6 +141,7 @@ def check_wix_oauth() -> (bool, str, str):
                 submit = st.form_submit_button(
                     ui.wix_login_button_label, type="primary", use_container_width=True
                 )
+            st.info(ui.wix_login_wait)
             st.link_button(
                 label=ui.wix_signup_button_label,
                 url=ui.wix_signup_button_url,
@@ -176,15 +177,19 @@ def check_wix_oauth() -> (bool, str, str):
                     member_access_token = get_member_access_token(
                         code=st.session_state["wix_return_data"]
                     )
-                    subsription = get_subscription(
+                    subscription = get_subscription(
                         member_access_token=member_access_token
                     )
 
-                    auth = True
-
-                    placeholder.empty()
-
-                    return auth, username, subsription
+                    # Check if the subscription is not None (i.e., user has an active subscription)
+                    if subscription is not None:
+                        auth = True
+                        placeholder.empty()
+                        return auth, username, subscription
+                    else:
+                        with col_center:
+                            st.error("Login failed: No active subscription found.", icon="🚫")
+                        return False, None, None
             else:
                 with col_center:
                     st.error(ui.wix_login_error_text, icon=ui.wix_login_error_icon)