update chat memory schema

app.py

@@ -1,125 +1,126 @@
-from prompts.arxiv_prompt import combine_prompt_template, _myscale_prompt
-from callbacks.arxiv_callbacks import ChatDataSelfSearchCallBackHandler, \
-    ChatDataSelfAskCallBackHandler, ChatDataSQLSearchCallBackHandler, \
-    ChatDataSQLAskCallBackHandler
-from chains.arxiv_chains import ArXivQAwithSourcesChain, ArXivStuffDocumentChain
-from chains.arxiv_chains import VectorSQLRetrieveCustomOutputParser
-from langchain_experimental.sql.vector_sql import VectorSQLDatabaseChain
-from langchain_experimental.retrievers.vector_sql_database import VectorSQLDatabaseChainRetriever
-from langchain.utilities.sql_database import SQLDatabase
-from langchain.chains import LLMChain
-from sqlalchemy import create_engine, MetaData
-from langchain.prompts import PromptTemplate, ChatPromptTemplate, \
-    SystemMessagePromptTemplate, HumanMessagePromptTemplate
-from langchain.prompts.prompt import PromptTemplate
-from langchain.chat_models import ChatOpenAI
-from langchain import OpenAI
-import re
+import json
+import time
 import pandas as pd
 from os import environ
 import streamlit as st
-import datetime
-from helper import build_all, sel_map, display
-environ['OPENAI_API_BASE'] = st.secrets['OPENAI_API_BASE']
 
-st.set_page_config(page_title="ChatData")
+from callbacks.arxiv_callbacks import ChatDataSelfSearchCallBackHandler, \
+    ChatDataSelfAskCallBackHandler, ChatDataSQLSearchCallBackHandler, \
+    ChatDataSQLAskCallBackHandler
+
+from chat import chat_page
+from login import login, back_to_main
+
+
+from helper import build_tools, build_agents, build_all, sel_map, display
+
+environ['OPENAI_API_BASE'] = st.secrets['OPENAI_API_BASE']
 
+st.set_page_config(page_title="ChatData", page_icon="https://myscale.com/favicon.ico")
 st.header("ChatData")
 
 if 'retriever' not in st.session_state:
     st.session_state["sel_map_obj"] = build_all()
+    st.session_state["tools"] = build_tools()
 
-sel = st.selectbox('Choose the knowledge base you want to ask with:',
-                   options=['ArXiv Papers', 'Wikipedia'])
-sel_map[sel]['hint']()
-tab_sql, tab_self_query = st.tabs(['Vector SQL', 'Self-Query Retrievers'])
-with tab_sql:
-    sel_map[sel]['hint_sql']()
-    st.text_input("Ask a question:", key='query_sql')
-    cols = st.columns([1, 1, 7])
-    cols[0].button("Query", key='search_sql')
-    cols[1].button("Ask", key='ask_sql')
-    plc_hldr = st.empty()
-    if st.session_state.search_sql:
-        plc_hldr = st.empty()
-        print(st.session_state.query_sql)
-        with plc_hldr.expander('Query Log', expanded=True):
-            callback = ChatDataSQLSearchCallBackHandler()
-            try:
-                docs = st.session_state.sel_map_obj[sel]["sql_retriever"].get_relevant_documents(
-                    st.session_state.query_sql, callbacks=[callback])
-                callback.progress_bar.progress(value=1.0, text="Done!")
-                docs = pd.DataFrame(
-                    [{**d.metadata, 'abstract': d.page_content} for d in docs])
-                display(docs)
-            except Exception as e:
-                st.write('Oops 😵 Something bad happened...')
-                raise e
+if login():
+    if "user_name" in st.session_state:
+        chat_page()
+    elif "jump_query_ask" in st.session_state and st.session_state.jump_query_ask:
+        
+        sel = st.selectbox('Choose the knowledge base you want to ask with:',
+                        options=['ArXiv Papers', 'Wikipedia'])
+        sel_map[sel]['hint']()
+        tab_sql, tab_self_query = st.tabs(['Vector SQL', 'Self-Query Retrievers'])
+        with tab_sql:
+            sel_map[sel]['hint_sql']()
+            st.text_input("Ask a question:", key='query_sql')
+            cols = st.columns([1, 1, 1, 4])
+            cols[0].button("Query", key='search_sql')
+            cols[1].button("Ask", key='ask_sql')
+            cols[2].button("Back", key='back_sql', on_click=back_to_main)
+            plc_hldr = st.empty()
+            if st.session_state.search_sql:
+                plc_hldr = st.empty()
+                print(st.session_state.query_sql)
+                with plc_hldr.expander('Query Log', expanded=True):
+                    callback = ChatDataSQLSearchCallBackHandler()
+                    try:
+                        docs = st.session_state.sel_map_obj[sel]["sql_retriever"].get_relevant_documents(
+                            st.session_state.query_sql, callbacks=[callback])
+                        callback.progress_bar.progress(value=1.0, text="Done!")
+                        docs = pd.DataFrame(
+                            [{**d.metadata, 'abstract': d.page_content} for d in docs])
+                        display(docs)
+                    except Exception as e:
+                        st.write('Oops 😵 Something bad happened...')
+                        raise e
 
-    if st.session_state.ask_sql:
-        plc_hldr = st.empty()
-        print(st.session_state.query_sql)
-        with plc_hldr.expander('Chat Log', expanded=True):
-            callback = ChatDataSQLAskCallBackHandler()
-            try:
-                ret = st.session_state.sel_map_obj[sel]["sql_chain"](
-                    st.session_state.query_sql, callbacks=[callback])
-                callback.progress_bar.progress(value=1.0, text="Done!")
-                st.markdown(
-                    f"### Answer from LLM\n{ret['answer']}\n### References")
-                docs = ret['sources']
-                docs = pd.DataFrame(
-                    [{**d.metadata, 'abstract': d.page_content} for d in docs])
-                display(
-                    docs, ['ref_id'] + sel_map[sel]["must_have_cols"], index='ref_id')
-            except Exception as e:
-                st.write('Oops 😵 Something bad happened...')
-                raise e
+            if st.session_state.ask_sql:
+                plc_hldr = st.empty()
+                print(st.session_state.query_sql)
+                with plc_hldr.expander('Chat Log', expanded=True):
+                    callback = ChatDataSQLAskCallBackHandler()
+                    try:
+                        ret = st.session_state.sel_map_obj[sel]["sql_chain"](
+                            st.session_state.query_sql, callbacks=[callback])
+                        callback.progress_bar.progress(value=1.0, text="Done!")
+                        st.markdown(
+                            f"### Answer from LLM\n{ret['answer']}\n### References")
+                        docs = ret['sources']
+                        docs = pd.DataFrame(
+                            [{**d.metadata, 'abstract': d.page_content} for d in docs])
+                        display(
+                            docs, ['ref_id'] + sel_map[sel]["must_have_cols"], index='ref_id')
+                    except Exception as e:
+                        st.write('Oops 😵 Something bad happened...')
+                        raise e
 
 
-with tab_self_query:
-    st.info("You can retrieve papers with button `Query` or ask questions based on retrieved papers with button `Ask`.", icon='💡')
-    st.dataframe(st.session_state.sel_map_obj[sel]["metadata_columns"])
-    st.text_input("Ask a question:", key='query_self')
-    cols = st.columns([1, 1, 7])
-    cols[0].button("Query", key='search_self')
-    cols[1].button("Ask", key='ask_self')
-    plc_hldr = st.empty()
-    if st.session_state.search_self:
-        plc_hldr = st.empty()
-        print(st.session_state.query_self)
-        with plc_hldr.expander('Query Log', expanded=True):
-            call_back = None
-            callback = ChatDataSelfSearchCallBackHandler()
-            try:
-                docs = st.session_state.sel_map_obj[sel]["retriever"].get_relevant_documents(
-                    st.session_state.query_self, callbacks=[callback])
-                print(docs)
-                callback.progress_bar.progress(value=1.0, text="Done!")
-                docs = pd.DataFrame(
-                    [{**d.metadata, 'abstract': d.page_content} for d in docs])
-                display(docs, sel_map[sel]["must_have_cols"])
-            except Exception as e:
-                st.write('Oops 😵 Something bad happened...')
-                raise e
+        with tab_self_query:
+            st.info("You can retrieve papers with button `Query` or ask questions based on retrieved papers with button `Ask`.", icon='💡')
+            st.dataframe(st.session_state.sel_map_obj[sel]["metadata_columns"])
+            st.text_input("Ask a question:", key='query_self')
+            cols = st.columns([1, 1, 1, 4])
+            cols[0].button("Query", key='search_self')
+            cols[1].button("Ask", key='ask_self')
+            cols[2].button("Back", key='back_self', on_click=back_to_main)
+            plc_hldr = st.empty()
+            if st.session_state.search_self:
+                plc_hldr = st.empty()
+                print(st.session_state.query_self)
+                with plc_hldr.expander('Query Log', expanded=True):
+                    call_back = None
+                    callback = ChatDataSelfSearchCallBackHandler()
+                    try:
+                        docs = st.session_state.sel_map_obj[sel]["retriever"].get_relevant_documents(
+                            st.session_state.query_self, callbacks=[callback])
+                        print(docs)
+                        callback.progress_bar.progress(value=1.0, text="Done!")
+                        docs = pd.DataFrame(
+                            [{**d.metadata, 'abstract': d.page_content} for d in docs])
+                        display(docs, sel_map[sel]["must_have_cols"])
+                    except Exception as e:
+                        st.write('Oops 😵 Something bad happened...')
+                        raise e
 
-    if st.session_state.ask_self:
-        plc_hldr = st.empty()
-        print(st.session_state.query_self)
-        with plc_hldr.expander('Chat Log', expanded=True):
-            call_back = None
-            callback = ChatDataSelfAskCallBackHandler()
-            try:
-                ret = st.session_state.sel_map_obj[sel]["chain"](
-                    st.session_state.query_self, callbacks=[callback])
-                callback.progress_bar.progress(value=1.0, text="Done!")
-                st.markdown(
-                    f"### Answer from LLM\n{ret['answer']}\n### References")
-                docs = ret['sources']
-                docs = pd.DataFrame(
-                    [{**d.metadata, 'abstract': d.page_content} for d in docs])
-                display(
-                    docs, ['ref_id'] + sel_map[sel]["must_have_cols"], index='ref_id')
-            except Exception as e:
-                st.write('Oops 😵 Something bad happened...')
-                raise e
+            if st.session_state.ask_self:
+                plc_hldr = st.empty()
+                print(st.session_state.query_self)
+                with plc_hldr.expander('Chat Log', expanded=True):
+                    call_back = None
+                    callback = ChatDataSelfAskCallBackHandler()
+                    try:
+                        ret = st.session_state.sel_map_obj[sel]["chain"](
+                            st.session_state.query_self, callbacks=[callback])
+                        callback.progress_bar.progress(value=1.0, text="Done!")
+                        st.markdown(
+                            f"### Answer from LLM\n{ret['answer']}\n### References")
+                        docs = ret['sources']
+                        docs = pd.DataFrame(
+                            [{**d.metadata, 'abstract': d.page_content} for d in docs])
+                        display(
+                            docs, ['ref_id'] + sel_map[sel]["must_have_cols"], index='ref_id')
+                    except Exception as e:
+                        st.write('Oops 😵 Something bad happened...')
+                        raise e

callbacks/arxiv_callbacks.py

@@ -17,8 +17,9 @@ class ChatDataSelfSearchCallBackHandler(StreamlitCallbackHandler):
         
     def on_chain_end(self, outputs, **kwargs) -> None:
         self.progress_bar.progress(value=0.6, text='Searching in DB...')
-        st.markdown('### Generated Filter')
-        st.write(outputs['text'], unsafe_allow_html=True)
+        if 'repr' in outputs:
+            st.markdown('### Generated Filter')
+            st.markdown(f"```python\n{outputs['repr']}\n```", unsafe_allow_html=True)
     
     def on_chain_start(self, serialized, inputs, **kwargs) -> None:
         pass

chat.py

@@ -1,31 +1,14 @@
-import json
-import time
 import pandas as pd
 from os import environ
 import datetime
 import streamlit as st
-from langchain.schema import Document
+from langchain.schema import HumanMessage, FunctionMessage
 
-from callbacks.arxiv_callbacks import ChatDataSelfSearchCallBackHandler, \
-    ChatDataSelfAskCallBackHandler, ChatDataSQLSearchCallBackHandler, \
-    ChatDataSQLAskCallBackHandler
-
-from langchain.schema import BaseMessage, HumanMessage, AIMessage, FunctionMessage, SystemMessage
-from auth0_component import login_button
-
-
-from helper import build_tools, build_agents, build_all, sel_map, display
+from helper import build_agents
+from login import back_to_main
 
 environ['OPENAI_API_BASE'] = st.secrets['OPENAI_API_BASE']
 
-st.set_page_config(page_title="ChatData", page_icon="https://myscale.com/favicon.ico")
-st.header("ChatData")
-
-
-if 'retriever' not in st.session_state:
-    st.session_state["sel_map_obj"] = build_all()
-    st.session_state["tools"] = build_tools()
-
 def on_chat_submit():
     ret = st.session_state.agents[st.session_state.sel][st.session_state.ret_type]({"input": st.session_state.chat_input})
     print(ret)
@@ -33,45 +16,7 @@ def on_chat_submit():
 def clear_history():
     st.session_state.agents[st.session_state.sel][st.session_state.ret_type].memory.clear()
 
-AUTH0_CLIENT_ID = st.secrets['AUTH0_CLIENT_ID']
-AUTH0_DOMAIN = st.secrets['AUTH0_DOMAIN']
 
-def login():
-    if "user_name" in st.session_state or ("jump_query_ask" in st.session_state and st.session_state.jump_query_ask):
-        return True
-    st.subheader("🤗 Welcom to [MyScale](https://myscale.com)'s [ChatData](https://github.com/myscale/ChatData)! 🤗 ")
-    st.write("You can now chat with ArXiv and Wikipedia! 🌟\n")
-    st.write("Built purely with streamlit 👑 , LangChain 🦜🔗 and love ❤️ for AI!")
-    st.write("Follow us on [Twitter](https://x.com/myscaledb) and [Discord](https://discord.gg/D2qpkqc4Jq)!")
-    st.write("For more details, please refer to [our repository on GitHub](https://github.com/myscale/ChatData)!")
-    st.divider()
-    col1, col2 = st.columns(2, gap='large')
-    with col1.container():
-        st.write("Try out MyScale's Self-query and Vector SQL retrievers!")
-        st.write("In this demo, you will be able to see how those retrievers "
-                 "**digest** -> **translate** -> **retrieve** -> **answer** to your question!")
-        st.session_state["jump_query_ask"] = st.button("Query / Ask")
-    with col2.container():
-        # st.warning("To use chat, please jump to [https://myscale-chatdata.hf.space](https://myscale-chatdata.hf.space)")
-        st.write("Now with the power of LangChain's Conversantional Agents, we are able to build "
-                 "an RAG-enabled chatbot within one MyScale instance! ")
-        st.write("Log in to Chat with RAG!")
-        login_button(AUTH0_CLIENT_ID, AUTH0_DOMAIN, "auth0")
-    st.divider()
-    st.write("- [Privacy Policy](https://myscale.com/privacy/)\n"
-             "- [Terms of Sevice](https://myscale.com/terms/)")
-    if st.session_state.auth0 is not None:
-        st.session_state.user_info = dict(st.session_state.auth0) 
-        if 'email' in st.session_state.user_info:
-            email = st.session_state.user_info["email"]
-        else:
-            email = f"{st.session_state.user_info['nickname']}@{st.session_state.user_info['sub']}"
-        st.session_state["user_name"] = email
-        del st.session_state.auth0
-        st.experimental_rerun()
-    if st.session_state.jump_query_ask:
-        st.experimental_rerun()
-        
 def back_to_main():
     if "user_info" in st.session_state:
         del st.session_state.user_info
@@ -80,127 +25,29 @@ def back_to_main():
     if "jump_query_ask" in st.session_state:
         del st.session_state.jump_query_ask
 
-if login():
-    if "user_name" in st.session_state:
-        st.session_state["agents"] = build_agents(st.session_state.user_name)
-        with st.sidebar:
-            st.radio("Retriever Type", ["Self-querying retriever", "Vector SQL"], key="ret_type")
-            st.selectbox("Knowledge Base", ["ArXiv Papers", "Wikipedia", "ArXiv + Wikipedia"], key="sel")
-            st.button("Clear Chat History", on_click=clear_history)
-            st.button("Logout", on_click=back_to_main)
-        for msg in st.session_state.agents[st.session_state.sel][st.session_state.ret_type].memory.chat_memory.messages:
-            speaker = "user" if isinstance(msg, HumanMessage) else "assistant"
-            if isinstance(msg, FunctionMessage):
-                with st.chat_message("Knowledge Base", avatar="📖"):
-                    print(type(msg.content))
+def chat_page():
+    st.session_state["agents"] = build_agents(f"{st.session_state.user_name}?default")
+    with st.sidebar:
+        st.radio("Retriever Type", ["Self-querying retriever", "Vector SQL"], key="ret_type")
+        st.selectbox("Knowledge Base", ["ArXiv Papers", "Wikipedia", "ArXiv + Wikipedia"], key="sel")
+        st.button("Clear Chat History", on_click=clear_history)
+        st.button("Logout", on_click=back_to_main)
+    for msg in st.session_state.agents[st.session_state.sel][st.session_state.ret_type].memory.chat_memory.messages:
+        speaker = "user" if isinstance(msg, HumanMessage) else "assistant"
+        if isinstance(msg, FunctionMessage):
+            with st.chat_message("Knowledge Base", avatar="📖"):
+                print(type(msg.content))
+                st.write(f"*{datetime.datetime.fromtimestamp(msg.additional_kwargs['timestamp']).isoformat()}*")
+                st.write("Retrieved from knowledge base:")
+                try:
+                    st.dataframe(pd.DataFrame.from_records(map(dict, eval(msg.content))))
+                except:
+                    st.write(msg.content)
+        else:
+            if len(msg.content) > 0:
+                with st.chat_message(speaker):
+                    print(type(msg), msg.dict())
                     st.write(f"*{datetime.datetime.fromtimestamp(msg.additional_kwargs['timestamp']).isoformat()}*")
-                    st.write("Retrieved from knowledge base:")
-                    try:
-                        st.dataframe(pd.DataFrame.from_records(map(dict, eval(msg.content))))
-                    except:
-                        st.write(msg.content)
-            else:
-                if len(msg.content) > 0:
-                    with st.chat_message(speaker):
-                        print(type(msg), msg.dict())
-                        st.write(f"*{datetime.datetime.fromtimestamp(msg.additional_kwargs['timestamp']).isoformat()}*")
-                        st.write(f"{msg.content}")
-        st.chat_input("Input Message", on_submit=on_chat_submit, key="chat_input")
-    elif "jump_query_ask" in st.session_state and st.session_state.jump_query_ask:
-        
-        sel = st.selectbox('Choose the knowledge base you want to ask with:',
-                        options=['ArXiv Papers', 'Wikipedia'])
-        sel_map[sel]['hint']()
-        tab_sql, tab_self_query = st.tabs(['Vector SQL', 'Self-Query Retrievers'])
-        with tab_sql:
-            sel_map[sel]['hint_sql']()
-            st.text_input("Ask a question:", key='query_sql')
-            cols = st.columns([1, 1, 1, 4])
-            cols[0].button("Query", key='search_sql')
-            cols[1].button("Ask", key='ask_sql')
-            cols[2].button("Back", key='back_sql', on_click=back_to_main)
-            plc_hldr = st.empty()
-            if st.session_state.search_sql:
-                plc_hldr = st.empty()
-                print(st.session_state.query_sql)
-                with plc_hldr.expander('Query Log', expanded=True):
-                    callback = ChatDataSQLSearchCallBackHandler()
-                    try:
-                        docs = st.session_state.sel_map_obj[sel]["sql_retriever"].get_relevant_documents(
-                            st.session_state.query_sql, callbacks=[callback])
-                        callback.progress_bar.progress(value=1.0, text="Done!")
-                        docs = pd.DataFrame(
-                            [{**d.metadata, 'abstract': d.page_content} for d in docs])
-                        display(docs)
-                    except Exception as e:
-                        st.write('Oops 😵 Something bad happened...')
-                        raise e
-
-            if st.session_state.ask_sql:
-                plc_hldr = st.empty()
-                print(st.session_state.query_sql)
-                with plc_hldr.expander('Chat Log', expanded=True):
-                    callback = ChatDataSQLAskCallBackHandler()
-                    try:
-                        ret = st.session_state.sel_map_obj[sel]["sql_chain"](
-                            st.session_state.query_sql, callbacks=[callback])
-                        callback.progress_bar.progress(value=1.0, text="Done!")
-                        st.markdown(
-                            f"### Answer from LLM\n{ret['answer']}\n### References")
-                        docs = ret['sources']
-                        docs = pd.DataFrame(
-                            [{**d.metadata, 'abstract': d.page_content} for d in docs])
-                        display(
-                            docs, ['ref_id'] + sel_map[sel]["must_have_cols"], index='ref_id')
-                    except Exception as e:
-                        st.write('Oops 😵 Something bad happened...')
-                        raise e
-
-
-        with tab_self_query:
-            st.info("You can retrieve papers with button `Query` or ask questions based on retrieved papers with button `Ask`.", icon='💡')
-            st.dataframe(st.session_state.sel_map_obj[sel]["metadata_columns"])
-            st.text_input("Ask a question:", key='query_self')
-            cols = st.columns([1, 1, 1, 4])
-            cols[0].button("Query", key='search_self')
-            cols[1].button("Ask", key='ask_self')
-            cols[2].button("Back", key='back_self', on_click=back_to_main)
-            plc_hldr = st.empty()
-            if st.session_state.search_self:
-                plc_hldr = st.empty()
-                print(st.session_state.query_self)
-                with plc_hldr.expander('Query Log', expanded=True):
-                    call_back = None
-                    callback = ChatDataSelfSearchCallBackHandler()
-                    try:
-                        docs = st.session_state.sel_map_obj[sel]["retriever"].get_relevant_documents(
-                            st.session_state.query_self, callbacks=[callback])
-                        print(docs)
-                        callback.progress_bar.progress(value=1.0, text="Done!")
-                        docs = pd.DataFrame(
-                            [{**d.metadata, 'abstract': d.page_content} for d in docs])
-                        display(docs, sel_map[sel]["must_have_cols"])
-                    except Exception as e:
-                        st.write('Oops 😵 Something bad happened...')
-                        raise e
-
-            if st.session_state.ask_self:
-                plc_hldr = st.empty()
-                print(st.session_state.query_self)
-                with plc_hldr.expander('Chat Log', expanded=True):
-                    call_back = None
-                    callback = ChatDataSelfAskCallBackHandler()
-                    try:
-                        ret = st.session_state.sel_map_obj[sel]["chain"](
-                            st.session_state.query_self, callbacks=[callback])
-                        callback.progress_bar.progress(value=1.0, text="Done!")
-                        st.markdown(
-                            f"### Answer from LLM\n{ret['answer']}\n### References")
-                        docs = ret['sources']
-                        docs = pd.DataFrame(
-                            [{**d.metadata, 'abstract': d.page_content} for d in docs])
-                        display(
-                            docs, ['ref_id'] + sel_map[sel]["must_have_cols"], index='ref_id')
-                    except Exception as e:
-                        st.write('Oops 😵 Something bad happened...')
-                        raise e
+                    st.write(f"{msg.content}")
+    st.chat_input("Input Message", on_submit=on_chat_submit, key="chat_input")
+    

helper.py

@@ -369,6 +369,7 @@ def create_message_model(table_name, DynamicBase):  # type: ignore
         __tablename__ = table_name
         id = Column(types.Float64)
         session_id = Column(Text)
+        user_id = Column(Text)
         msg_id = Column(Text, primary_key=True)
         type = Column(Text)
         addtionals = Column(Text)
@@ -391,9 +392,11 @@ class DefaultClickhouseMessageConverter(DefaultMessageConverter):
     def to_sql_model(self, message: BaseMessage, session_id: str) -> Any:
         tstamp = time.time()
         msg_id = hashlib.sha256(f"{session_id}_{message}_{tstamp}".encode('utf-8')).hexdigest()
+        user_id, _ = session_id.split("?")
         return self.model_class(
             id=tstamp, 
             msg_id=msg_id,
+            user_id=user_id,
             session_id=session_id, 
             type=message.type,
             addtionals=json.dumps(message.additional_kwargs),
@@ -467,7 +470,7 @@ def build_tools():
     return sel_map_obj
 
 @st.cache_resource(max_entries=1)
-def build_agents(username):
+def build_agents(session_id):
     chat_llm = ChatOpenAI(model_name=chat_model_name, temperature=0.6, openai_api_base=OPENAI_API_BASE, openai_api_key=OPENAI_API_KEY)
     agents = {}
     cnt = 0
@@ -484,7 +487,7 @@ def build_agents(username):
                 agents[k] = {}
             agents[k][n] = create_agent_executor(
                 "chat_memory",
-                username,
+                session_id,
                 chat_llm,
                 tools=tools,
                 )

login.py

@@ -0,0 +1,53 @@
+import json
+import time
+import pandas as pd
+from os import environ
+import streamlit as st
+from auth0_component import login_button
+
+AUTH0_CLIENT_ID = st.secrets['AUTH0_CLIENT_ID']
+AUTH0_DOMAIN = st.secrets['AUTH0_DOMAIN']
+
+def login():
+    if "user_name" in st.session_state or ("jump_query_ask" in st.session_state and st.session_state.jump_query_ask):
+        return True
+    st.subheader("🤗 Welcom to [MyScale](https://myscale.com)'s [ChatData](https://github.com/myscale/ChatData)! 🤗 ")
+    st.write("You can now chat with ArXiv and Wikipedia! 🌟\n")
+    st.write("Built purely with streamlit 👑 , LangChain 🦜🔗 and love ❤️ for AI!")
+    st.write("Follow us on [Twitter](https://x.com/myscaledb) and [Discord](https://discord.gg/D2qpkqc4Jq)!")
+    st.write("For more details, please refer to [our repository on GitHub](https://github.com/myscale/ChatData)!")
+    st.divider()
+    col1, col2 = st.columns(2, gap='large')
+    with col1.container():
+        st.write("Try out MyScale's Self-query and Vector SQL retrievers!")
+        st.write("In this demo, you will be able to see how those retrievers "
+                 "**digest** -> **translate** -> **retrieve** -> **answer** to your question!")
+        st.session_state["jump_query_ask"] = st.button("Query / Ask")
+    with col2.container():
+        # st.warning("To use chat, please jump to [https://myscale-chatdata.hf.space](https://myscale-chatdata.hf.space)")
+        st.write("Now with the power of LangChain's Conversantional Agents, we are able to build "
+                 "an RAG-enabled chatbot within one MyScale instance! ")
+        st.write("Log in to Chat with RAG!")
+        login_button(AUTH0_CLIENT_ID, AUTH0_DOMAIN, "auth0")
+    st.divider()
+    st.write("- [Privacy Policy](https://myscale.com/privacy/)\n"
+             "- [Terms of Sevice](https://myscale.com/terms/)")
+    if st.session_state.auth0 is not None:
+        st.session_state.user_info = dict(st.session_state.auth0) 
+        if 'email' in st.session_state.user_info:
+            email = st.session_state.user_info["email"]
+        else:
+            email = f"{st.session_state.user_info['nickname']}@{st.session_state.user_info['sub']}"
+        st.session_state["user_name"] = email
+        del st.session_state.auth0
+        st.experimental_rerun()
+    if st.session_state.jump_query_ask:
+        st.experimental_rerun()
+
+def back_to_main():
+    if "user_info" in st.session_state:
+        del st.session_state.user_info
+    if "user_name" in st.session_state:
+        del st.session_state.user_name
+    if "jump_query_ask" in st.session_state:
+        del st.session_state.jump_query_ask