Synced repo using 'sync_with_huggingface' Github Action

.streamlit/config.toml

@@ -1,6 +1,2 @@
 [theme]
-primaryColor="#523EFD"
-backgroundColor="#FFFFFF"
-secondaryBackgroundColor="#D4CEFF"
-textColor="#262730"
-font="sans serif"
+base="dark"

.streamlit/secrets.example.toml

@@ -1,7 +1,8 @@
-MYSCALE_HOST = "msc-4a9e710a.us-east-1.aws.staging.myscale.cloud" # read-only database provided by MyScale
+MYSCALE_HOST = "msc-950b9f1f.us-east-1.aws.myscale.com" # read-only database provided by MyScale
 MYSCALE_PORT = 443
 MYSCALE_USER = "chatdata"
 MYSCALE_PASSWORD = "myscale_rocks"
+MYSCALE_ENABLE_HTTPS = true
 OPENAI_API_BASE = "https://api.openai.com/v1"
 OPENAI_API_KEY = "<your-openai-key>"
 UNSTRUCTURED_API = "<your-unstructured-io-api>" # optional if you don't upload documents

app.py

@@ -1,133 +1,86 @@
-import pandas as pd
-from os import environ
+import os
+import time
+
 import streamlit as st
 
-from callbacks.arxiv_callbacks import ChatDataSelfSearchCallBackHandler, \
-    ChatDataSelfAskCallBackHandler, ChatDataSQLSearchCallBackHandler, \
-    ChatDataSQLAskCallBackHandler
+from backend.constants.streamlit_keys import DATA_INITIALIZE_NOT_STATED, DATA_INITIALIZE_COMPLETED, \
+    DATA_INITIALIZE_STARTED
+from backend.constants.variables import DATA_INITIALIZE_STATUS, JUMP_QUERY_ASK, CHAINS_RETRIEVERS_MAPPING, \
+    TABLE_EMBEDDINGS_MAPPING, RETRIEVER_TOOLS, USER_NAME, GLOBAL_CONFIG, update_global_config
+from backend.construct.build_all import build_chains_and_retrievers, load_embedding_models, update_retriever_tools
+from backend.types.global_config import GlobalConfig
+from logger import logger
+from ui.chat_page import chat_page
+from ui.home import render_home
+from ui.retrievers import render_retrievers
 
-from chat import chat_page
-from login import login, back_to_main
-from lib.helper import build_tools, build_all, sel_map, display
 
+# warnings.filterwarnings("ignore", category=UserWarning)
 
-environ['OPENAI_API_BASE'] = st.secrets['OPENAI_API_BASE']
+def prepare_environment():
+    os.environ['TOKENIZERS_PARALLELISM'] = 'true'
+    os.environ["LANGCHAIN_TRACING_V2"] = "false"
+    # os.environ["LANGCHAIN_API_KEY"] = ""
+    os.environ["OPENAI_API_BASE"] = st.secrets['OPENAI_API_BASE']
+    os.environ["OPENAI_API_KEY"] = st.secrets['OPENAI_API_KEY']
+    os.environ["AUTH0_CLIENT_ID"] = st.secrets['AUTH0_CLIENT_ID']
+    os.environ["AUTH0_DOMAIN"] = st.secrets['AUTH0_DOMAIN']
+
+    update_global_config(GlobalConfig(
+        openai_api_base=st.secrets['OPENAI_API_BASE'],
+        openai_api_key=st.secrets['OPENAI_API_KEY'],
+        auth0_client_id=st.secrets['AUTH0_CLIENT_ID'],
+        auth0_domain=st.secrets['AUTH0_DOMAIN'],
+        myscale_user=st.secrets['MYSCALE_USER'],
+        myscale_password=st.secrets['MYSCALE_PASSWORD'],
+        myscale_host=st.secrets['MYSCALE_HOST'],
+        myscale_port=st.secrets['MYSCALE_PORT'],
+        query_model="gpt-3.5-turbo-0125",
+        chat_model="gpt-3.5-turbo-0125",
+        untrusted_api=st.secrets['UNSTRUCTURED_API'],
+        myscale_enable_https=st.secrets.get('MYSCALE_ENABLE_HTTPS', True),
+    ))
 
-st.set_page_config(page_title="ChatData",
-                   page_icon="https://myscale.com/favicon.ico")
-st.markdown(
-    f"""
-    <style>
-        .st-e4 {{
-            max-width: 500px
-        }}
-    </style>""",
-    unsafe_allow_html=True,
-)
-st.header("ChatData")
 
-if 'sel_map_obj' not in st.session_state or 'embeddings' not in st.session_state:
-    st.session_state["sel_map_obj"], st.session_state["embeddings"] = build_all()
-    st.session_state["tools"] = build_tools()
+# when refresh browser, all session keys will be cleaned.
+def initialize_session_state():
+    if DATA_INITIALIZE_STATUS not in st.session_state:
+        st.session_state[DATA_INITIALIZE_STATUS] = DATA_INITIALIZE_NOT_STATED
+        logger.info(f"Initialize session state key: {DATA_INITIALIZE_STATUS}")
+    if JUMP_QUERY_ASK not in st.session_state:
+        st.session_state[JUMP_QUERY_ASK] = False
+        logger.info(f"Initialize session state key: {JUMP_QUERY_ASK}")
 
-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
+def initialize_chat_data():
+    if st.session_state[DATA_INITIALIZE_STATUS] != DATA_INITIALIZE_COMPLETED:
+        start_time = time.time()
+        st.session_state[DATA_INITIALIZE_STATUS] = DATA_INITIALIZE_STARTED
+        st.session_state[TABLE_EMBEDDINGS_MAPPING] = load_embedding_models()
+        st.session_state[CHAINS_RETRIEVERS_MAPPING] = build_chains_and_retrievers()
+        st.session_state[RETRIEVER_TOOLS] = update_retriever_tools()
+        # mark data initialization finished.
+        st.session_state[DATA_INITIALIZE_STATUS] = DATA_INITIALIZE_COMPLETED
+        end_time = time.time()
+        logger.info(f"ChatData initialized finished in {round(end_time - start_time, 3)} seconds, "
+                    f"session state keys: {list(st.session_state.keys())}")
 
-            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
+st.set_page_config(
+    page_title="ChatData",
+    page_icon="https://myscale.com/favicon.ico",
+    initial_sidebar_state="expanded",
+    layout="wide",
+)
+
+prepare_environment()
+initialize_session_state()
+initialize_chat_data()
 
-            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 USER_NAME in st.session_state:
+    chat_page()
+else:
+    if st.session_state[JUMP_QUERY_ASK]:
+        render_retrievers()
+    else:
+        render_home()

backend/callbacks/arxiv_callbacks.py

@@ -0,0 +1,46 @@
+import json
+import textwrap
+from typing import Dict, Any, List
+
+from langchain.callbacks.streamlit.streamlit_callback_handler import (
+    LLMThought,
+    StreamlitCallbackHandler,
+)
+
+
+class LLMThoughtWithKnowledgeBase(LLMThought):
+    def on_tool_end(
+        self,
+        output: str,
+        color=None,
+        observation_prefix=None,
+        llm_prefix=None,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            self._container.markdown(
+                "\n\n".join(
+                    ["### Retrieved Documents:"]
+                    + [
+                        f"**{i+1}**: {textwrap.shorten(r['page_content'], width=80)}"
+                        for i, r in enumerate(json.loads(output))
+                    ]
+                )
+            )
+        except Exception as e:
+            super().on_tool_end(output, color, observation_prefix, llm_prefix, **kwargs)
+
+
+class ChatDataAgentCallBackHandler(StreamlitCallbackHandler):
+    def on_llm_start(
+        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+    ) -> None:
+        if self._current_thought is None:
+            self._current_thought = LLMThoughtWithKnowledgeBase(
+                parent_container=self._parent_container,
+                expanded=self._expand_new_thoughts,
+                collapse_on_complete=self._collapse_completed_thoughts,
+                labeler=self._thought_labeler,
+            )
+
+        self._current_thought.on_llm_start(serialized, prompts)

backend/callbacks/llm_thought_with_table.py

@@ -0,0 +1,36 @@
+from typing import Any, Dict, List
+
+import streamlit as st
+from langchain_core.outputs import LLMResult
+from streamlit.external.langchain import StreamlitCallbackHandler
+
+
+class ChatDataSelfQueryCallBack(StreamlitCallbackHandler):
+    def __init__(self):
+        super().__init__(st.container())
+        self._current_thought = None
+        self.progress_bar = st.progress(value=0.0, text="Executing ChatData SelfQuery CallBack...")
+
+    def on_llm_start(
+            self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+    ) -> None:
+        self.progress_bar.progress(value=0.35, text="Communicate with LLM...")
+        pass
+
+    def on_chain_end(self, outputs, **kwargs) -> None:
+        if len(kwargs['tags']) == 0:
+            self.progress_bar.progress(value=0.75, text="Searching in DB...")
+
+    def on_chain_start(self, serialized, inputs, **kwargs) -> None:
+
+        pass
+
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        st.markdown("### Generate filter by LLM \n"
+                    "> Here we get `query_constructor` results \n\n")
+
+        self.progress_bar.progress(value=0.5, text="Generate filter by LLM...")
+        for item in response.generations:
+            st.markdown(f"{item[0].text}")
+
+        pass

backend/callbacks/self_query_callbacks.py

@@ -0,0 +1,57 @@
+from typing import Dict, Any, List
+
+import streamlit as st
+from langchain.callbacks.streamlit.streamlit_callback_handler import (
+    StreamlitCallbackHandler,
+)
+from langchain.schema.output import LLMResult
+
+
+class CustomSelfQueryRetrieverCallBackHandler(StreamlitCallbackHandler):
+    def __init__(self):
+        super().__init__(st.container())
+        self._current_thought = None
+        self.progress_bar = st.progress(value=0.0, text="Executing ChatData SelfQuery...")
+
+    def on_llm_start(
+            self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
+    ) -> None:
+        self.progress_bar.progress(value=0.35, text="Communicate with LLM...")
+        pass
+
+    def on_chain_end(self, outputs, **kwargs) -> None:
+        if len(kwargs['tags']) == 0:
+            self.progress_bar.progress(value=0.75, text="Searching in DB...")
+        pass
+
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        st.markdown("### Generate filter by LLM \n"
+                    "> Here we get `query_constructor` results \n\n")
+        self.progress_bar.progress(value=0.5, text="Generate filter by LLM...")
+        for item in response.generations:
+            st.markdown(f"{item[0].text}")
+        pass
+
+
+class ChatDataSelfAskCallBackHandler(StreamlitCallbackHandler):
+    def __init__(self) -> None:
+        super().__init__(st.container())
+        self.progress_bar = st.progress(value=0.2, text="Executing ChatData SelfQuery Chain...")
+
+    def on_llm_start(self, serialized, prompts, **kwargs) -> None:
+        pass
+
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+
+        if len(kwargs['tags']) != 0:
+            self.progress_bar.progress(value=0.5, text="We got filter info from LLM...")
+            st.markdown("### Generate filter by LLM \n"
+                        "> Here we get `query_constructor` results \n\n")
+            for item in response.generations:
+                st.markdown(f"{item[0].text}")
+        pass
+
+    def on_chain_start(self, serialized, inputs, **kwargs) -> None:
+        cid = ".".join(serialized["id"])
+        if cid.endswith(".CustomStuffDocumentChain"):
+            self.progress_bar.progress(value=0.7, text="Asking LLM with related documents...")

backend/callbacks/vector_sql_callbacks.py

@@ -0,0 +1,53 @@
+import streamlit as st
+from langchain.callbacks.streamlit.streamlit_callback_handler import (
+    StreamlitCallbackHandler,
+)
+from langchain.schema.output import LLMResult
+from sql_formatter.core import format_sql
+
+
+class VectorSQLSearchDBCallBackHandler(StreamlitCallbackHandler):
+    def __init__(self) -> None:
+        self.progress_bar = st.progress(value=0.0, text="Writing SQL...")
+        self.status_bar = st.empty()
+        self.prog_value = 0
+        self.prog_interval = 0.2
+
+    def on_llm_start(self, serialized, prompts, **kwargs) -> None:
+        pass
+
+    def on_llm_end(
+            self,
+            response: LLMResult,
+            *args,
+            **kwargs,
+    ):
+        text = response.generations[0][0].text
+        if text.replace(" ", "").upper().startswith("SELECT"):
+            st.markdown("### Generated Vector Search SQL Statement \n"
+                        "> This sql statement is generated by LLM \n\n")
+            st.markdown(f"""```sql\n{format_sql(text, max_len=80)}\n```""")
+            self.prog_value += self.prog_interval
+            self.progress_bar.progress(
+                value=self.prog_value, text="Searching in DB...")
+
+    def on_chain_start(self, serialized, inputs, **kwargs) -> None:
+        cid = ".".join(serialized["id"])
+        self.prog_value += self.prog_interval
+        self.progress_bar.progress(
+            value=self.prog_value, text=f"Running Chain `{cid}`..."
+        )
+
+    def on_chain_end(self, outputs, **kwargs) -> None:
+        pass
+
+
+class VectorSQLSearchLLMCallBackHandler(VectorSQLSearchDBCallBackHandler):
+    def __init__(self, table: str) -> None:
+        self.progress_bar = st.progress(value=0.0, text="Writing SQL...")
+        self.status_bar = st.empty()
+        self.prog_value = 0
+        self.prog_interval = 0.1
+        self.table = table
+
+

backend/chains/retrieval_qa_with_sources.py

@@ -0,0 +1,70 @@
+import inspect
+from typing import Dict, Any, Optional, List
+
+from langchain.callbacks.manager import (
+    AsyncCallbackManagerForChainRun,
+    CallbackManagerForChainRun,
+)
+from langchain.chains.qa_with_sources.retrieval import RetrievalQAWithSourcesChain
+from langchain.docstore.document import Document
+
+from logger import logger
+
+
+class CustomRetrievalQAWithSourcesChain(RetrievalQAWithSourcesChain):
+    """QA with source chain for Chat ArXiv app with references
+
+    This chain will automatically assign reference number to the article,
+    Then parse it back to titles or anything else.
+    """
+
+    def _call(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[CallbackManagerForChainRun] = None,
+    ) -> Dict[str, str]:
+        logger.info(f"\033[91m\033[1m{self._chain_type}\033[0m")
+        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
+        accepts_run_manager = (
+            "run_manager" in inspect.signature(self._get_docs).parameters
+        )
+        if accepts_run_manager:
+            docs: List[Document] = self._get_docs(inputs, run_manager=_run_manager)
+        else:
+            docs: List[Document] = self._get_docs(inputs)  # type: ignore[call-arg]
+
+        answer = self.combine_documents_chain.run(
+            input_documents=docs, callbacks=_run_manager.get_child(), **inputs
+        )
+        # parse source with ref_id
+        sources = []
+        ref_cnt = 1
+        for d in docs:
+            ref_id = d.metadata['ref_id']
+            if f"Doc #{ref_id}" in answer:
+                answer = answer.replace(f"Doc #{ref_id}", f"#{ref_id}")
+            if f"#{ref_id}" in answer:
+                title = d.metadata['title'].replace('\n', '')
+                d.metadata['ref_id'] = ref_cnt
+                answer = answer.replace(f"#{ref_id}", f"{title} [{ref_cnt}]")
+                sources.append(d)
+                ref_cnt += 1
+
+        result: Dict[str, Any] = {
+            self.answer_key: answer,
+            self.sources_answer_key: sources,
+        }
+        if self.return_source_documents:
+            result["source_documents"] = docs
+        return result
+
+    async def _acall(
+        self,
+        inputs: Dict[str, Any],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
+    ) -> Dict[str, Any]:
+        raise NotImplementedError
+
+    @property
+    def _chain_type(self) -> str:
+        return "custom_retrieval_qa_with_sources_chain"

backend/chains/stuff_documents.py

@@ -0,0 +1,65 @@
+from typing import Any, List, Tuple
+
+from langchain.callbacks.manager import Callbacks
+from langchain.chains.combine_documents.stuff import StuffDocumentsChain
+from langchain.docstore.document import Document
+from langchain.schema.prompt_template import format_document
+
+
+class CustomStuffDocumentChain(StuffDocumentsChain):
+    """Combine arxiv documents with PDF reference number"""
+
+    def _get_inputs(self, docs: List[Document], **kwargs: Any) -> dict:
+        """Construct inputs from kwargs and docs.
+
+        Format and the join all the documents together into one input with name
+        `self.document_variable_name`. The pluck any additional variables
+        from **kwargs.
+
+        Args:
+            docs: List of documents to format and then join into single input
+            **kwargs: additional inputs to chain, will pluck any other required
+                arguments from here.
+
+        Returns:
+            dictionary of inputs to LLMChain
+        """
+        # Format each document according to the prompt
+        doc_strings = []
+        for doc_id, doc in enumerate(docs):
+            # add temp reference number in metadata
+            doc.metadata.update({'ref_id': doc_id})
+            doc.page_content = doc.page_content.replace('\n', ' ')
+            doc_strings.append(format_document(doc, self.document_prompt))
+        # Join the documents together to put them in the prompt.
+        inputs = {
+            k: v
+            for k, v in kwargs.items()
+            if k in self.llm_chain.prompt.input_variables
+        }
+        inputs[self.document_variable_name] = self.document_separator.join(
+            doc_strings)
+        return inputs
+
+    def combine_docs(
+            self, docs: List[Document], callbacks: Callbacks = None, **kwargs: Any
+    ) -> Tuple[str, dict]:
+        """Stuff all documents into one prompt and pass to LLM.
+
+        Args:
+            docs: List of documents to join together into one variable
+            callbacks: Optional callbacks to pass along
+            **kwargs: additional parameters to use to get inputs to LLMChain.
+
+        Returns:
+            The first element returned is the single string output. The second
+            element returned is a dictionary of other keys to return.
+        """
+        inputs = self._get_inputs(docs, **kwargs)
+        # Call predict on the LLM.
+        output = self.llm_chain.predict(callbacks=callbacks, **inputs)
+        return output, {}
+
+    @property
+    def _chain_type(self) -> str:
+        return "custom_stuff_document_chain"

backend/chat_bot/chat.py

@@ -0,0 +1,225 @@
+import time
+
+from os import environ
+from time import sleep
+import streamlit as st
+
+from backend.constants.prompts import DEFAULT_SYSTEM_PROMPT
+from backend.constants.streamlit_keys import CHAT_KNOWLEDGE_TABLE, CHAT_SESSION_MANAGER, \
+    CHAT_CURRENT_USER_SESSIONS, EL_SESSION_SELECTOR, USER_PRIVATE_FILES, \
+    EL_BUILD_KB_WITH_FILES, \
+    EL_PERSONAL_KB_NAME, EL_PERSONAL_KB_DESCRIPTION, \
+    USER_PERSONAL_KNOWLEDGE_BASES, AVAILABLE_RETRIEVAL_TOOLS, EL_PERSONAL_KB_NEEDS_REMOVE, \
+    EL_UPLOAD_FILES_STATUS, EL_SELECTED_KBS, EL_UPLOAD_FILES
+from backend.constants.variables import USER_INFO, USER_NAME, JUMP_QUERY_ASK, RETRIEVER_TOOLS
+from backend.construct.build_agents import build_agents
+from backend.chat_bot.session_manager import SessionManager
+from backend.callbacks.arxiv_callbacks import ChatDataAgentCallBackHandler
+
+from logger import logger
+
+environ["OPENAI_API_BASE"] = st.secrets["OPENAI_API_BASE"]
+
+TOOL_NAMES = {
+    "langchain_retriever_tool": "Self-querying retriever",
+    "vecsql_retriever_tool": "Vector SQL",
+}
+
+
+def on_chat_submit():
+    with st.session_state.next_round.container():
+        with st.chat_message("user"):
+            st.write(st.session_state.chat_input)
+        with st.chat_message("assistant"):
+            container = st.container()
+        st_callback = ChatDataAgentCallBackHandler(
+            container, collapse_completed_thoughts=False
+        )
+        ret = st.session_state.agent(
+            {"input": st.session_state.chat_input}, callbacks=[st_callback]
+        )
+        logger.info(f"ret:{ret}")
+
+
+def clear_history():
+    if "agent" in st.session_state:
+        st.session_state.agent.memory.clear()
+
+
+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]
+    if EL_SESSION_SELECTOR in st.session_state:
+        del st.session_state[EL_SESSION_SELECTOR]
+    if CHAT_CURRENT_USER_SESSIONS in st.session_state:
+        del st.session_state[CHAT_CURRENT_USER_SESSIONS]
+
+
+def refresh_sessions():
+    chat_session_manager: SessionManager = st.session_state[CHAT_SESSION_MANAGER]
+    current_user_name = st.session_state[USER_NAME]
+    current_user_sessions = chat_session_manager.list_sessions(current_user_name)
+
+    if not isinstance(current_user_sessions, dict) or not current_user_sessions:
+        # generate a default session for current user.
+        chat_session_manager.add_session(
+            user_id=current_user_name,
+            session_id=f"{current_user_name}?default",
+            system_prompt=DEFAULT_SYSTEM_PROMPT,
+        )
+        st.session_state[CHAT_CURRENT_USER_SESSIONS] = chat_session_manager.list_sessions(current_user_name)
+        current_user_sessions = st.session_state[CHAT_CURRENT_USER_SESSIONS]
+    else:
+        st.session_state[CHAT_CURRENT_USER_SESSIONS] = current_user_sessions
+
+    # load current user files.
+    st.session_state[USER_PRIVATE_FILES] = st.session_state[CHAT_KNOWLEDGE_TABLE].list_files(
+        current_user_name
+    )
+    # load current user private knowledge bases.
+    st.session_state[USER_PERSONAL_KNOWLEDGE_BASES] = \
+        st.session_state[CHAT_KNOWLEDGE_TABLE].list_private_knowledge_bases(current_user_name)
+    logger.info(f"current user name: {current_user_name}, "
+                f"user private knowledge bases: {st.session_state[USER_PERSONAL_KNOWLEDGE_BASES]}, "
+                f"user private files: {st.session_state[USER_PRIVATE_FILES]}")
+    st.session_state[AVAILABLE_RETRIEVAL_TOOLS] = {
+        # public retrieval tools
+        **st.session_state[RETRIEVER_TOOLS],
+        # private retrieval tools
+        **st.session_state[CHAT_KNOWLEDGE_TABLE].as_retrieval_tools(current_user_name),
+    }
+    # print(f"sel_session is {st.session_state.sel_session}, current_user_sessions is {current_user_sessions}")
+    print(f"current_user_sessions is {current_user_sessions}")
+    st.session_state[EL_SESSION_SELECTOR] = current_user_sessions[0]
+
+
+# process for session add and delete.
+def on_session_change_submit():
+    if "session_manager" in st.session_state and "session_editor" in st.session_state:
+        try:
+            for elem in st.session_state.session_editor["added_rows"]:
+                if len(elem) > 0 and "system_prompt" in elem and "session_id" in elem:
+                    if elem["session_id"] != "" and "?" not in elem["session_id"]:
+                        st.session_state.session_manager.add_session(
+                            user_id=st.session_state.user_name,
+                            session_id=f"{st.session_state.user_name}?{elem['session_id']}",
+                            system_prompt=elem["system_prompt"],
+                        )
+                    else:
+                        st.toast("`session_id` shouldn't be neither empty nor contain char `?`.", icon="❌")
+                        raise KeyError(
+                            "`session_id` shouldn't be neither empty nor contain char `?`."
+                        )
+                else:
+                    st.toast("`You should fill both `session_id` and `system_prompt` to add a column!", icon="❌")
+                    raise KeyError(
+                        "You should fill both `session_id` and `system_prompt` to add a column!"
+                    )
+            for elem in st.session_state.session_editor["deleted_rows"]:
+                user_name = st.session_state[USER_NAME]
+                session_id = st.session_state[CHAT_CURRENT_USER_SESSIONS][elem]['session_id']
+                user_with_session_id = f"{user_name}?{session_id}"
+                st.session_state.session_manager.remove_session(session_id=user_with_session_id)
+                st.toast(f"session `{user_with_session_id}` removed.", icon="✅")
+
+            refresh_sessions()
+        except Exception as e:
+            sleep(2)
+            st.error(f"{type(e)}: {str(e)}")
+        finally:
+            st.session_state.session_editor["added_rows"] = []
+            st.session_state.session_editor["deleted_rows"] = []
+        refresh_agent()
+
+
+def create_private_knowledge_base_as_tool():
+    current_user_name = st.session_state[USER_NAME]
+
+    if (
+            EL_PERSONAL_KB_NAME in st.session_state
+            and EL_PERSONAL_KB_DESCRIPTION in st.session_state
+            and EL_BUILD_KB_WITH_FILES in st.session_state
+            and len(st.session_state[EL_PERSONAL_KB_NAME]) > 0
+            and len(st.session_state[EL_PERSONAL_KB_DESCRIPTION]) > 0
+            and len(st.session_state[EL_BUILD_KB_WITH_FILES]) > 0
+    ):
+        st.session_state[CHAT_KNOWLEDGE_TABLE].create_private_knowledge_base(
+            user_id=current_user_name,
+            tool_name=st.session_state[EL_PERSONAL_KB_NAME],
+            tool_description=st.session_state[EL_PERSONAL_KB_DESCRIPTION],
+            files=[f["file_name"] for f in st.session_state[EL_BUILD_KB_WITH_FILES]],
+        )
+        refresh_sessions()
+    else:
+        st.session_state[EL_UPLOAD_FILES_STATUS].error(
+            "You should fill all fields to build up a tool!"
+        )
+        sleep(2)
+
+
+def remove_private_knowledge_bases():
+    if EL_PERSONAL_KB_NEEDS_REMOVE in st.session_state and st.session_state[EL_PERSONAL_KB_NEEDS_REMOVE]:
+        private_knowledge_bases_needs_remove = st.session_state[EL_PERSONAL_KB_NEEDS_REMOVE]
+        private_knowledge_base_names = [item["tool_name"] for item in private_knowledge_bases_needs_remove]
+        # remove these private knowledge bases.
+        st.session_state[CHAT_KNOWLEDGE_TABLE].remove_private_knowledge_bases(
+            user_id=st.session_state[USER_NAME],
+            private_knowledge_bases=private_knowledge_base_names
+        )
+        refresh_sessions()
+    else:
+        st.session_state[EL_UPLOAD_FILES_STATUS].error(
+            "You should specify at least one private knowledge base to delete!"
+        )
+        time.sleep(2)
+
+
+def refresh_agent():
+    with st.spinner("Initializing session..."):
+        user_name = st.session_state[USER_NAME]
+        session_id = st.session_state[EL_SESSION_SELECTOR]['session_id']
+        user_with_session_id = f"{user_name}?{session_id}"
+
+        if EL_SELECTED_KBS in st.session_state:
+            selected_knowledge_bases = st.session_state[EL_SELECTED_KBS]
+        else:
+            selected_knowledge_bases = ["Wikipedia + Vector SQL"]
+
+        logger.info(f"selected_knowledge_bases: {selected_knowledge_bases}")
+        if EL_SESSION_SELECTOR in st.session_state:
+            system_prompt = st.session_state[EL_SESSION_SELECTOR]["system_prompt"]
+        else:
+            system_prompt = DEFAULT_SYSTEM_PROMPT
+
+        st.session_state["agent"] = build_agents(
+            session_id=user_with_session_id,
+            tool_names=selected_knowledge_bases,
+            system_prompt=system_prompt
+        )
+
+
+def add_file():
+    user_name = st.session_state[USER_NAME]
+    if EL_UPLOAD_FILES not in st.session_state or len(st.session_state[EL_UPLOAD_FILES]) == 0:
+        st.session_state[EL_UPLOAD_FILES_STATUS].error("Please upload files!", icon="⚠️")
+        sleep(2)
+        return
+    try:
+        st.session_state[EL_UPLOAD_FILES_STATUS].info("Uploading...")
+        st.session_state[CHAT_KNOWLEDGE_TABLE].add_by_file(
+            user_id=user_name,
+            files=st.session_state[EL_UPLOAD_FILES]
+        )
+        refresh_sessions()
+    except ValueError as e:
+        st.session_state[EL_UPLOAD_FILES_STATUS].error("Failed to upload! " + str(e))
+        sleep(2)
+
+
+def clear_files():
+    st.session_state[CHAT_KNOWLEDGE_TABLE].clear(user_id=st.session_state[USER_NAME])
+    refresh_sessions()

backend/chat_bot/json_decoder.py

@@ -0,0 +1,24 @@
+import json
+import datetime
+
+
+class CustomJSONEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, datetime.datetime):
+            return datetime.datetime.isoformat(obj)
+        return json.JSONEncoder.default(self, obj)
+
+
+class CustomJSONDecoder(json.JSONDecoder):
+    def __init__(self, *args, **kwargs):
+        json.JSONDecoder.__init__(
+            self, object_hook=self.object_hook, *args, **kwargs)
+
+    def object_hook(self, source):
+        for k, v in source.items():
+            if isinstance(v, str):
+                try:
+                    source[k] = datetime.datetime.fromisoformat(str(v))
+                except:
+                    pass
+        return source

backend/chat_bot/message_converter.py

@@ -0,0 +1,67 @@
+import hashlib
+import json
+import time
+from typing import Any
+
+from langchain.memory.chat_message_histories.sql import DefaultMessageConverter
+from langchain.schema import BaseMessage, HumanMessage, AIMessage, SystemMessage, ChatMessage, FunctionMessage
+from langchain.schema.messages import ToolMessage
+from sqlalchemy.orm import declarative_base
+
+from backend.chat_bot.tools import create_message_history_table
+
+
+def _message_from_dict(message: dict) -> BaseMessage:
+    _type = message["type"]
+    if _type == "human":
+        return HumanMessage(**message["data"])
+    elif _type == "ai":
+        return AIMessage(**message["data"])
+    elif _type == "system":
+        return SystemMessage(**message["data"])
+    elif _type == "chat":
+        return ChatMessage(**message["data"])
+    elif _type == "function":
+        return FunctionMessage(**message["data"])
+    elif _type == "tool":
+        return ToolMessage(**message["data"])
+    elif _type == "AIMessageChunk":
+        message["data"]["type"] = "ai"
+        return AIMessage(**message["data"])
+    else:
+        raise ValueError(f"Got unexpected message type: {_type}")
+
+
+class DefaultClickhouseMessageConverter(DefaultMessageConverter):
+    """The default message converter for SQLChatMessageHistory."""
+
+    def __init__(self, table_name: str):
+        super().__init__(table_name)
+        self.model_class = create_message_history_table(table_name, declarative_base())
+
+    def to_sql_model(self, message: BaseMessage, session_id: str) -> Any:
+        time_stamp = time.time()
+        msg_id = hashlib.sha256(
+            f"{session_id}_{message}_{time_stamp}".encode('utf-8')).hexdigest()
+        user_id, _ = session_id.split("?")
+        return self.model_class(
+            id=time_stamp,
+            msg_id=msg_id,
+            user_id=user_id,
+            session_id=session_id,
+            type=message.type,
+            addtionals=json.dumps(message.additional_kwargs),
+            message=json.dumps({
+                "type": message.type,
+                "additional_kwargs": {"timestamp": time_stamp},
+                "data": message.dict()})
+        )
+
+    def from_sql_model(self, sql_message: Any) -> BaseMessage:
+        msg_dump = json.loads(sql_message.message)
+        msg = _message_from_dict(msg_dump)
+        msg.additional_kwargs = msg_dump["additional_kwargs"]
+        return msg
+
+    def get_sql_model_class(self) -> Any:
+        return self.model_class

backend/chat_bot/private_knowledge_base.py

@@ -0,0 +1,167 @@
+import hashlib
+from datetime import datetime
+from typing import List, Optional
+
+import pandas as pd
+from clickhouse_connect import get_client
+from langchain.schema.embeddings import Embeddings
+from langchain.vectorstores.myscale import MyScaleWithoutJSON, MyScaleSettings
+from streamlit.runtime.uploaded_file_manager import UploadedFile
+
+from backend.chat_bot.tools import parse_files, extract_embedding
+from backend.construct.build_retriever_tool import create_retriever_tool
+from logger import logger
+
+
+class ChatBotKnowledgeTable:
+    def __init__(self, host, port, username, password,
+                 embedding: Embeddings, parser_api_key: str, db="chat",
+                 kb_table="private_kb", tool_table="private_tool") -> None:
+        super().__init__()
+        personal_files_schema_ = f"""
+            CREATE TABLE IF NOT EXISTS {db}.{kb_table}(
+                entity_id String,
+                file_name String,
+                text String,
+                user_id String,
+                created_by DateTime,
+                vector Array(Float32),
+                CONSTRAINT cons_vec_len CHECK length(vector) = 768,
+                VECTOR INDEX vidx vector TYPE MSTG('metric_type=Cosine')
+            ) ENGINE = ReplacingMergeTree ORDER BY entity_id
+        """
+
+        # `tool_name` represent private knowledge database name.
+        private_knowledge_base_schema_ = f"""
+            CREATE TABLE IF NOT EXISTS {db}.{tool_table}(
+                tool_id String,
+                tool_name String,
+                file_names Array(String),
+                user_id String,
+                created_by DateTime,
+                tool_description String
+            ) ENGINE = ReplacingMergeTree ORDER BY tool_id
+        """
+        self.personal_files_table = kb_table
+        self.private_knowledge_base_table = tool_table
+        config = MyScaleSettings(
+            host=host,
+            port=port,
+            username=username,
+            password=password,
+            database=db,
+            table=kb_table,
+        )
+        self.client = get_client(
+            host=config.host,
+            port=config.port,
+            username=config.username,
+            password=config.password,
+        )
+        self.client.command("SET allow_experimental_object_type=1")
+        self.client.command(personal_files_schema_)
+        self.client.command(private_knowledge_base_schema_)
+        self.parser_api_key = parser_api_key
+        self.vector_store = MyScaleWithoutJSON(
+            embedding=embedding,
+            config=config,
+            must_have_cols=["file_name", "text", "created_by"],
+        )
+
+    # List all files with given `user_id`
+    def list_files(self, user_id: str):
+        query = f"""
+        SELECT DISTINCT file_name, COUNT(entity_id) AS num_paragraph, 
+            arrayMax(arrayMap(x->length(x), groupArray(text))) AS max_chars
+        FROM {self.vector_store.config.database}.{self.personal_files_table}
+        WHERE user_id = '{user_id}' GROUP BY file_name
+        """
+        return [r for r in self.vector_store.client.query(query).named_results()]
+
+    # Parse and embedding files
+    def add_by_file(self, user_id, files: List[UploadedFile]):
+        data = parse_files(self.parser_api_key, user_id, files)
+        data = extract_embedding(self.vector_store.embeddings, data)
+        self.vector_store.client.insert_df(
+            table=self.personal_files_table,
+            df=pd.DataFrame(data),
+            database=self.vector_store.config.database,
+        )
+
+    # Remove all files and private_knowledge_bases with given `user_id`
+    def clear(self, user_id: str):
+        self.vector_store.client.command(
+            f"DELETE FROM {self.vector_store.config.database}.{self.personal_files_table} "
+            f"WHERE user_id='{user_id}'"
+        )
+        query = f"""DELETE FROM {self.vector_store.config.database}.{self.private_knowledge_base_table} 
+                    WHERE user_id  = '{user_id}'"""
+        self.vector_store.client.command(query)
+
+    def create_private_knowledge_base(
+            self, user_id: str, tool_name: str, tool_description: str, files: Optional[List[str]] = None
+    ):
+        self.vector_store.client.insert_df(
+            self.private_knowledge_base_table,
+            pd.DataFrame(
+                [
+                    {
+                        "tool_id": hashlib.sha256(
+                            (user_id + tool_name).encode("utf-8")
+                        ).hexdigest(),
+                        "tool_name": tool_name,  # tool_name represent user's private knowledge base.
+                        "file_names": files,
+                        "user_id": user_id,
+                        "created_by": datetime.now(),
+                        "tool_description": tool_description,
+                    }
+                ]
+            ),
+            database=self.vector_store.config.database,
+        )
+
+    # Show all private knowledge bases with given `user_id`
+    def list_private_knowledge_bases(self, user_id: str, private_knowledge_base=None):
+        extended_where = f"AND tool_name = '{private_knowledge_base}'" if private_knowledge_base else ""
+        query = f"""
+        SELECT tool_name, tool_description, length(file_names) 
+        FROM {self.vector_store.config.database}.{self.private_knowledge_base_table}
+        WHERE user_id = '{user_id}' {extended_where}
+        """
+        return [r for r in self.vector_store.client.query(query).named_results()]
+
+    def remove_private_knowledge_bases(self, user_id: str, private_knowledge_bases: List[str]):
+        unique_list = list(set(private_knowledge_bases))
+        unique_list = ",".join([f"'{t}'" for t in unique_list])
+        query = f"""DELETE FROM {self.vector_store.config.database}.{self.private_knowledge_base_table}
+                    WHERE user_id  = '{user_id}' AND tool_name IN [{unique_list}]"""
+        self.vector_store.client.command(query)
+
+    def as_retrieval_tools(self, user_id, tool_name=None):
+        logger.info(f"")
+        private_knowledge_bases = self.list_private_knowledge_bases(user_id=user_id, private_knowledge_base=tool_name)
+        retrievers = {}
+        for private_kb in private_knowledge_bases:
+            file_names_sql = f"""
+            SELECT arrayJoin(file_names) FROM (
+                SELECT file_names 
+                FROM chat.private_tool
+                WHERE user_id = '{user_id}' AND tool_name = '{private_kb["tool_name"]}'
+            )
+            """
+            logger.info(f"user_id is {user_id}, file_names_sql is {file_names_sql}")
+            res = self.client.query(file_names_sql)
+            file_names = []
+            for line in res.result_rows:
+                file_names.append(line[0])
+            file_names = ', '.join(f"'{item}'" for item in file_names)
+            logger.info(f"user_id is {user_id}, file_names is {file_names}")
+            retrievers[private_kb["tool_name"]] = create_retriever_tool(
+                self.vector_store.as_retriever(
+                    search_kwargs={"where_str": f"user_id='{user_id}' AND file_name IN ({file_names})"},
+                ),
+                tool_name=private_kb["tool_name"],
+                description=private_kb["tool_description"],
+            )
+        return retrievers
+

backend/chat_bot/session_manager.py

@@ -0,0 +1,96 @@
+import json
+
+from backend.chat_bot.tools import create_session_table, create_message_history_table
+from backend.constants.variables import GLOBAL_CONFIG
+
+try:
+    from sqlalchemy.orm import declarative_base
+except ImportError:
+    from sqlalchemy.ext.declarative import declarative_base
+from datetime import datetime
+from sqlalchemy import orm, create_engine
+from logger import logger
+
+
+def get_sessions(engine, model_class, user_id):
+    with orm.sessionmaker(engine)() as session:
+        result = (
+            session.query(model_class)
+            .where(
+                model_class.session_id == user_id
+            )
+            .order_by(model_class.create_by.desc())
+        )
+    return json.loads(result)
+
+
+class SessionManager:
+    def __init__(
+            self,
+            session_state,
+            host,
+            port,
+            username,
+            password,
+            db='chat',
+            session_table='sessions',
+            msg_table='chat_memory'
+    ) -> None:
+        if GLOBAL_CONFIG.myscale_enable_https == False:
+            conn_str = f'clickhouse://{username}:{password}@{host}:{port}/{db}?protocol=http'
+        else:
+            conn_str = f'clickhouse://{username}:{password}@{host}:{port}/{db}?protocol=https'
+        self.engine = create_engine(conn_str, echo=False)
+        self.session_model_class = create_session_table(
+            session_table, declarative_base())
+        self.session_model_class.metadata.create_all(self.engine)
+        self.msg_model_class = create_message_history_table(msg_table, declarative_base())
+        self.msg_model_class.metadata.create_all(self.engine)
+        self.session_orm = orm.sessionmaker(self.engine)
+        self.session_state = session_state
+
+    def list_sessions(self, user_id: str):
+        with self.session_orm() as session:
+            result = (
+                session.query(self.session_model_class)
+                .where(
+                    self.session_model_class.user_id == user_id
+                )
+                .order_by(self.session_model_class.create_by.desc())
+            )
+            sessions = []
+            for r in result:
+                sessions.append({
+                    "session_id": r.session_id.split("?")[-1],
+                    "system_prompt": r.system_prompt,
+                })
+            return sessions
+
+    # Update sys_prompt with given session_id
+    def modify_system_prompt(self, session_id, sys_prompt):
+        with self.session_orm() as session:
+            obj = session.query(self.session_model_class).where(
+                self.session_model_class.session_id == session_id).first()
+            if obj:
+                obj.system_prompt = sys_prompt
+                session.commit()
+            else:
+                logger.warning(f"Session {session_id} not found")
+
+    # Add a session(session_id, sys_prompt)
+    def add_session(self, user_id: str, session_id: str, system_prompt: str, **kwargs):
+        with self.session_orm() as session:
+            elem = self.session_model_class(
+                user_id=user_id, session_id=session_id, system_prompt=system_prompt,
+                create_by=datetime.now(), additionals=json.dumps(kwargs)
+            )
+            session.add(elem)
+            session.commit()
+
+    # Remove a session and related chat history.
+    def remove_session(self, session_id: str):
+        with self.session_orm() as session:
+            # remove session
+            session.query(self.session_model_class).where(self.session_model_class.session_id == session_id).delete()
+            # remove related chat history.
+            session.query(self.msg_model_class).where(self.msg_model_class.session_id == session_id).delete()

backend/chat_bot/tools.py

@@ -0,0 +1,100 @@
+import hashlib
+from datetime import datetime
+from multiprocessing.pool import ThreadPool
+from typing import List
+
+import requests
+from clickhouse_sqlalchemy import types, engines
+from langchain.schema.embeddings import Embeddings
+from sqlalchemy import Column, Text
+from streamlit.runtime.uploaded_file_manager import UploadedFile
+
+
+def parse_files(api_key, user_id, files: List[UploadedFile]):
+    def parse_file(file: UploadedFile):
+        headers = {
+            "accept": "application/json",
+            "unstructured-api-key": api_key,
+        }
+        data = {"strategy": "auto", "ocr_languages": ["eng"]}
+        file_hash = hashlib.sha256(file.read()).hexdigest()
+        file_data = {"files": (file.name, file.getvalue(), file.type)}
+        response = requests.post(
+            url="https://api.unstructured.io/general/v0/general",
+            headers=headers,
+            data=data,
+            files=file_data
+        )
+        json_response = response.json()
+        if response.status_code != 200:
+            raise ValueError(str(json_response))
+        texts = [
+            {
+                "text": t["text"],
+                "file_name": t["metadata"]["filename"],
+                "entity_id": hashlib.sha256(
+                    (file_hash + t["text"]).encode()
+                ).hexdigest(),
+                "user_id": user_id,
+                "created_by": datetime.now(),
+            }
+            for t in json_response
+            if t["type"] == "NarrativeText" and len(t["text"].split(" ")) > 10
+        ]
+        return texts
+
+    with ThreadPool(8) as p:
+        rows = []
+        for r in p.imap_unordered(parse_file, files):
+            rows.extend(r)
+        return rows
+
+
+def extract_embedding(embeddings: Embeddings, texts):
+    if len(texts) > 0:
+        embeddings = embeddings.embed_documents(
+            [t["text"] for _, t in enumerate(texts)])
+        for i, _ in enumerate(texts):
+            texts[i]["vector"] = embeddings[i]
+        return texts
+    raise ValueError("No texts extracted!")
+
+
+def create_message_history_table(table_name: str, base_class):
+    class Message(base_class):
+        __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)
+        # should be additions, formal developer mistake spell it.
+        addtionals = Column(Text)
+        message = Column(Text)
+        __table_args__ = (
+            engines.MergeTree(
+                partition_by='session_id',
+                order_by=('id', 'msg_id')
+            ),
+            {'comment': 'Store Chat History'}
+        )
+
+    return Message
+
+
+def create_session_table(table_name: str, DynamicBase):
+    class Session(DynamicBase):
+        __tablename__ = table_name
+        user_id = Column(Text)
+        session_id = Column(Text, primary_key=True)
+        system_prompt = Column(Text)
+        # represent create time.
+        create_by = Column(types.DateTime)
+        # should be additions, formal developer mistake spell it.
+        additionals = Column(Text)
+        __table_args__ = (
+            engines.MergeTree(order_by=session_id),
+            {'comment': 'Store Session and Prompts'}
+        )
+
+    return Session

backend/constants/myscale_tables.py

@@ -0,0 +1,128 @@
+from typing import Dict, List
+import streamlit as st
+from langchain.chains.query_constructor.schema import AttributeInfo
+from langchain_community.embeddings import SentenceTransformerEmbeddings, HuggingFaceInstructEmbeddings
+from langchain.prompts import PromptTemplate
+
+from backend.types.table_config import TableConfig
+
+
+def hint_arxiv():
+    st.markdown("Here we provide some query samples.")
+    st.markdown("- If you want to search papers with filters")
+    st.markdown("1. ```What is a Bayesian network? Please use articles published later than Feb 2018 and with more "
+                "than 2 categories and whose title like `computer` and must have `cs.CV` in its category. ```")
+    st.markdown("2. ```What is a Bayesian network? Please use articles published later than Feb 2018```")
+    st.markdown("- If you want to ask questions based on arxiv papers stored in MyScaleDB")
+    st.markdown("1. ```Did Geoffrey Hinton wrote paper about Capsule Neural Networks?```")
+    st.markdown("2. ```Introduce some applications of GANs published around 2019.```")
+    st.markdown("3. ```请根据 2019 年左右的文章介绍一下 GAN 的应用都有哪些```")
+
+
+def hint_sql_arxiv():
+    st.markdown('''```sql
+CREATE TABLE default.ChatArXiv (
+    `abstract` String, 
+    `id` String, 
+    `vector` Array(Float32), 
+    `metadata` Object('JSON'), 
+    `pubdate` DateTime,
+    `title` String,
+    `categories` Array(String),
+    `authors` Array(String), 
+    `comment` String,
+    `primary_category` String,
+    VECTOR INDEX vec_idx vector TYPE MSTG('fp16_storage=1', 'metric_type=Cosine', 'disk_mode=3'), 
+    CONSTRAINT vec_len CHECK length(vector) = 768) 
+ENGINE = ReplacingMergeTree ORDER BY id
+```''')
+
+
+def hint_wiki():
+    st.markdown("Here we provide some query samples.")
+    st.markdown("1. ```Which company did Elon Musk found?```")
+    st.markdown("2. ```What is Iron Gwazi?```")
+    st.markdown("3. ```苹果的发源地是哪里？```")
+    st.markdown("4. ```What is a Ring in mathematics?```")
+    st.markdown("5. ```The producer of Rick and Morty.```")
+    st.markdown("6. ```How low is the temperature on Pluto?```")
+
+
+def hint_sql_wiki():
+    st.markdown('''```sql
+CREATE TABLE wiki.Wikipedia (
+    `id` String, 
+    `title` String, 
+    `text` String, 
+    `url` String, 
+    `wiki_id` UInt64, 
+    `views` Float32, 
+    `paragraph_id` UInt64, 
+    `langs` UInt32, 
+    `emb` Array(Float32), 
+    VECTOR INDEX vec_idx emb TYPE MSTG('fp16_storage=1', 'metric_type=Cosine', 'disk_mode=3'), 
+    CONSTRAINT emb_len CHECK length(emb) = 768) 
+ENGINE = ReplacingMergeTree ORDER BY id
+```''')
+
+
+MYSCALE_TABLES: Dict[str, TableConfig] = {
+    'Wikipedia': TableConfig(
+        database="wiki",
+        table="Wikipedia",
+        table_contents="Snapshort from Wikipedia for 2022. All in English.",
+        hint=hint_wiki,
+        hint_sql=hint_sql_wiki,
+        # doc_prompt 对 qa source chain 有用
+        doc_prompt=PromptTemplate(
+            input_variables=["page_content", "url", "title", "ref_id", "views"],
+            template="Title for Doc #{ref_id}: {title}\n\tviews: {views}\n\tcontent: {page_content}\nSOURCE: {url}"
+        ),
+        metadata_col_attributes=[
+            AttributeInfo(name="title", description="title of the wikipedia page", type="string"),
+            AttributeInfo(name="text", description="paragraph from this wiki page", type="string"),
+            AttributeInfo(name="views", description="number of views", type="float")
+        ],
+        must_have_col_names=['id', 'title', 'url', 'text', 'views'],
+        vector_col_name="emb",
+        text_col_name="text",
+        metadata_col_name="metadata",
+        emb_model=lambda: SentenceTransformerEmbeddings(
+            model_name='sentence-transformers/paraphrase-multilingual-mpnet-base-v2'
+        ),
+        tool_desc=("search_among_wikipedia", "Searches among Wikipedia and returns related wiki pages")
+    ),
+    'ArXiv Papers': TableConfig(
+        database="default",
+        table="ChatArXiv",
+        table_contents="Snapshort from Wikipedia for 2022. All in English.",
+        hint=hint_arxiv,
+        hint_sql=hint_sql_arxiv,
+        doc_prompt=PromptTemplate(
+            input_variables=["page_content", "id", "title", "ref_id", "authors", "pubdate", "categories"],
+            template="Title for Doc #{ref_id}: {title}\n\tAbstract: {page_content}\n\tAuthors: {authors}\n\t"
+                     "Date of Publication: {pubdate}\n\tCategories: {categories}\nSOURCE: {id}"
+        ),
+        metadata_col_attributes=[
+            AttributeInfo(name="pubdate", description="The year the paper is published", type="timestamp"),
+            AttributeInfo(name="authors", description="List of author names", type="list[string]"),
+            AttributeInfo(name="title", description="Title of the paper", type="string"),
+            AttributeInfo(name="categories", description="arxiv categories to this paper", type="list[string]"),
+            AttributeInfo(name="length(categories)", description="length of arxiv categories to this paper", type="int")
+        ],
+        must_have_col_names=['title', 'id', 'categories', 'abstract', 'authors', 'pubdate'],
+        vector_col_name="vector",
+        text_col_name="abstract",
+        metadata_col_name="metadata",
+        emb_model=lambda: HuggingFaceInstructEmbeddings(
+            model_name='hkunlp/instructor-xl',
+            embed_instruction="Represent the question for retrieving supporting scientific papers: "
+        ),
+        tool_desc=(
+            "search_among_scientific_papers",
+            "Searches among scientific papers from ArXiv and returns research papers"
+        )
+    )
+}
+
+ALL_TABLE_NAME: List[str] = [config.table for config in MYSCALE_TABLES.values()]

backend/constants/prompts.py

@@ -0,0 +1,128 @@
+from langchain.prompts import ChatPromptTemplate, \
+    SystemMessagePromptTemplate, HumanMessagePromptTemplate
+
+DEFAULT_SYSTEM_PROMPT = (
+    "Do your best to answer the questions. "
+    "Feel free to use any tools available to look up "
+    "relevant information. Please keep all details in query "
+    "when calling search functions."
+)
+
+COMBINE_PROMPT_TEMPLATE = (
+        "You are a helpful document assistant. "
+        "Your task is to provide information and answer any questions related to documents given below. "
+        "You should use the sections, title and abstract of the selected documents as your source of information "
+        "and try to provide concise and accurate answers to any questions asked by the user. "
+        "If you are unable to find relevant information in the given sections, "
+        "you will need to let the user know that the source does not contain relevant information but still try to "
+        "provide an answer based on your general knowledge. You must refer to the corresponding section name and page "
+        "that you refer to when answering. "
+        "The following is the related information about the document that will help you answer users' questions, "
+        "you MUST answer it using question's language:\n\n {summaries} "
+        "Now you should answer user's question. Remember you must use `Doc #` to refer papers:\n\n"
+)
+
+COMBINE_PROMPT = ChatPromptTemplate.from_strings(
+    string_messages=[(SystemMessagePromptTemplate, COMBINE_PROMPT_TEMPLATE),
+                     (HumanMessagePromptTemplate, '{question}')])
+
+MYSCALE_PROMPT = """
+You are a MyScale expert. Given an input question, first create a syntactically correct MyScale query to run, then look at the results of the query and return the answer to the input question.
+MyScale queries has a vector distance function called `DISTANCE(column, array)` to compute relevance to the user's question and sort the feature array column by the relevance. 
+When the query is asking for {top_k} closest row, you have to use this distance function to calculate distance to entity's array on vector column and order by the distance to retrieve relevant rows.
+
+*NOTICE*: `DISTANCE(column, array)` only accept an array column as its first argument and a `NeuralArray(entity)` as its second argument. You also need a user defined function called `NeuralArray(entity)` to retrieve the entity's array. 
+
+Unless the user specifies in the question a specific number of examples to obtain, query for at most {top_k} results using the LIMIT clause as per MyScale. You should only order according to the distance function.
+Never query for all columns from a table. You must query only the columns that are needed to answer the question. Wrap each column name in double quotes (") to denote them as delimited identifiers.
+Pay attention to use only the column names you can see in the tables below. Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.
+Pay attention to use today() function to get the current date, if the question involves "today". `ORDER BY` clause should always be after `WHERE` clause. DO NOT add semicolon to the end of SQL. Pay attention to the comment in table schema.
+Pay attention to the data type when using functions. Always use `AND` to connect conditions in `WHERE` and never use comma.
+Make sure you never write an isolated `WHERE` keyword and never use undesired condition to conrtain the query.
+
+Use the following format:
+
+======== table info ========
+<some table infos>
+
+Question: "Question here"
+SQLQuery: "SQL Query to run"
+
+
+Here are some examples:
+
+======== table info ========
+CREATE TABLE "ChatPaper" (
+	abstract String, 
+	id String, 
+	vector Array(Float32), 
+) ENGINE = ReplicatedReplacingMergeTree()
+ ORDER BY id
+ PRIMARY KEY id
+ 
+Question: What is Feartue Pyramid Network?
+SQLQuery: SELECT ChatPaper.abstract, ChatPaper.id FROM ChatPaper ORDER BY DISTANCE(vector, NeuralArray(PaperRank contribution)) LIMIT {top_k}
+
+
+======== table info ========
+CREATE TABLE "ChatPaper" (
+	abstract String, 
+	id String, 
+	vector Array(Float32), 
+	categories Array(String), 
+	pubdate DateTime, 
+	title String, 
+	authors Array(String), 
+	primary_category String
+) ENGINE = ReplicatedReplacingMergeTree()
+ ORDER BY id
+ PRIMARY KEY id
+ 
+Question: What is PaperRank? What is the contribution of those works? Use paper with more than 2 categories.
+SQLQuery: SELECT ChatPaper.title, ChatPaper.id, ChatPaper.authors FROM ChatPaper WHERE length(categories) > 2 ORDER BY DISTANCE(vector, NeuralArray(PaperRank contribution)) LIMIT {top_k}
+
+
+======== table info ========
+CREATE TABLE "ChatArXiv" (
+	primary_category String
+	categories Array(String), 
+	pubdate DateTime, 
+	abstract String, 
+	title String, 
+	paper_id String, 
+	vector Array(Float32), 
+	authors Array(String), 
+) ENGINE = MergeTree()
+ ORDER BY paper_id
+ PRIMARY KEY paper_id
+ 
+Question: Did Geoffrey Hinton wrote about Capsule Neural Networks? Please use articles published later than 2021.
+SQLQuery: SELECT ChatArXiv.title, ChatArXiv.paper_id, ChatArXiv.authors FROM ChatArXiv WHERE has(authors, 'Geoffrey Hinton') AND pubdate > parseDateTimeBestEffort('2021-01-01') ORDER BY DISTANCE(vector, NeuralArray(Capsule Neural Networks)) LIMIT {top_k}
+
+
+======== table info ========
+CREATE TABLE "PaperDatabase" (
+	abstract String, 
+	categories Array(String), 
+	vector Array(Float32), 
+	pubdate DateTime, 
+	id String, 
+	comments String,
+	title String, 
+	authors Array(String), 
+	primary_category String
+) ENGINE = MergeTree()
+ ORDER BY id
+ PRIMARY KEY id
+ 
+Question: Find papers whose abstract has Mutual Information in it.
+SQLQuery: SELECT PaperDatabase.title, PaperDatabase.id FROM PaperDatabase WHERE abstract ILIKE '%Mutual Information%' ORDER BY DISTANCE(vector, NeuralArray(Mutual Information)) LIMIT {top_k}
+
+ 
+Let's begin:
+
+======== table info ========
+{table_info}
+
+Question: {input}
+SQLQuery: """

backend/constants/streamlit_keys.py

@@ -0,0 +1,35 @@
+DATA_INITIALIZE_NOT_STATED = "data_initialize_not_started"
+DATA_INITIALIZE_STARTED = "data_initialize_started"
+DATA_INITIALIZE_COMPLETED = "data_initialize_completed"
+
+
+CHAT_SESSION = "sel_sess"
+CHAT_KNOWLEDGE_TABLE = "private_kb"
+
+CHAT_SESSION_MANAGER = "session_manager"
+CHAT_CURRENT_USER_SESSIONS = "current_sessions"
+
+EL_SESSION_SELECTOR = "el_session_selector"
+
+# all personal knowledge bases under a specific user.
+USER_PERSONAL_KNOWLEDGE_BASES = "user_tools"
+# all personal files under a specific user.
+USER_PRIVATE_FILES = "user_files"
+# public and personal knowledge bases.
+AVAILABLE_RETRIEVAL_TOOLS = "tools_with_users"
+
+EL_PERSONAL_KB_NEEDS_REMOVE = "el_personal_kb_needs_remove"
+
+# files needs upload
+EL_UPLOAD_FILES = "el_upload_files"
+EL_UPLOAD_FILES_STATUS = "el_upload_files_status"
+
+# use these files to build private knowledge base
+EL_BUILD_KB_WITH_FILES = "el_build_kb_with_files"
+# build a personal kb, given name.
+EL_PERSONAL_KB_NAME = "el_personal_kb_name"
+# build a personal kb, given description.
+EL_PERSONAL_KB_DESCRIPTION = "el_personal_kb_description"
+
+# knowledge bases selected by user.
+EL_SELECTED_KBS = "el_selected_kbs"

backend/constants/variables.py

@@ -0,0 +1,58 @@
+from backend.types.global_config import GlobalConfig
+
+# ***** str variables ***** #
+EMBEDDING_MODEL_PREFIX = "embedding_model"
+CHAINS_RETRIEVERS_MAPPING = "sel_map_obj"
+LANGCHAIN_RETRIEVER = "langchain_retriever"
+VECTOR_SQL_RETRIEVER = "vecsql_retriever"
+TABLE_EMBEDDINGS_MAPPING = "embeddings"
+RETRIEVER_TOOLS = "tools"
+DATA_INITIALIZE_STATUS = "data_initialized"
+UI_INITIALIZED = "ui_initialized"
+JUMP_QUERY_ASK = "jump_query_ask"
+USER_NAME = "user_name"
+USER_INFO = "user_info"
+
+DIVIDER_HTML = """
+    <div style="
+        height: 4px;
+        background: linear-gradient(to right, red, orange, yellow, green, blue, indigo, violet);
+        margin-top: 20px;
+        margin-bottom: 20px;
+    "></div>
+"""
+
+DIVIDER_THIN_HTML = """
+    <div style="
+        height: 2px;
+        background: linear-gradient(to right, blue, darkslateblue, indigo, violet);
+        margin-top: 20px;
+        margin-bottom: 20px;
+    "></div>
+"""
+
+
+class RetrieverButtons:
+    vector_sql_query_from_db = "vector_sql_query_from_db"
+    vector_sql_query_with_llm = "vector_sql_query_with_llm"
+    self_query_from_db = "self_query_from_db"
+    self_query_with_llm = "self_query_with_llm"
+
+
+GLOBAL_CONFIG = GlobalConfig()
+
+
+def update_global_config(new_config: GlobalConfig):
+    global GLOBAL_CONFIG
+    GLOBAL_CONFIG.openai_api_base = new_config.openai_api_base
+    GLOBAL_CONFIG.openai_api_key = new_config.openai_api_key
+    GLOBAL_CONFIG.auth0_client_id = new_config.auth0_client_id
+    GLOBAL_CONFIG.auth0_domain = new_config.auth0_domain
+    GLOBAL_CONFIG.myscale_user = new_config.myscale_user
+    GLOBAL_CONFIG.myscale_password = new_config.myscale_password
+    GLOBAL_CONFIG.myscale_host = new_config.myscale_host
+    GLOBAL_CONFIG.myscale_port = new_config.myscale_port
+    GLOBAL_CONFIG.query_model = new_config.query_model
+    GLOBAL_CONFIG.chat_model = new_config.chat_model
+    GLOBAL_CONFIG.untrusted_api = new_config.untrusted_api
+    GLOBAL_CONFIG.myscale_enable_https = new_config.myscale_enable_https

backend/construct/build_agents.py

@@ -0,0 +1,82 @@
+import os
+from typing import Sequence, List
+
+import streamlit as st
+from langchain.agents import AgentExecutor
+from langchain.schema.language_model import BaseLanguageModel
+from langchain.tools import BaseTool
+
+from backend.chat_bot.message_converter import DefaultClickhouseMessageConverter
+from backend.constants.prompts import DEFAULT_SYSTEM_PROMPT
+from backend.constants.streamlit_keys import AVAILABLE_RETRIEVAL_TOOLS
+from backend.constants.variables import GLOBAL_CONFIG, RETRIEVER_TOOLS
+from logger import logger
+
+try:
+    from sqlalchemy.orm import declarative_base
+except ImportError:
+    from sqlalchemy.ext.declarative import declarative_base
+from langchain.chat_models import ChatOpenAI
+from langchain.prompts.chat import MessagesPlaceholder
+from langchain.agents.openai_functions_agent.agent_token_buffer_memory import AgentTokenBufferMemory
+from langchain.agents.openai_functions_agent.base import OpenAIFunctionsAgent
+from langchain.schema.messages import SystemMessage
+from langchain.memory import SQLChatMessageHistory
+
+
+def create_agent_executor(
+        agent_name: str,
+        session_id: str,
+        llm: BaseLanguageModel,
+        tools: Sequence[BaseTool],
+        system_prompt: str,
+        **kwargs
+) -> AgentExecutor:
+    agent_name = agent_name.replace(" ", "_")
+    conn_str = f'clickhouse://{os.environ["MYSCALE_USER"]}:{os.environ["MYSCALE_PASSWORD"]}@{os.environ["MYSCALE_HOST"]}:{os.environ["MYSCALE_PORT"]}'
+    chat_memory = SQLChatMessageHistory(
+        session_id,
+        connection_string=f'{conn_str}/chat?protocol=http' if GLOBAL_CONFIG.myscale_enable_https == False else f'{conn_str}/chat?protocol=https',
+        custom_message_converter=DefaultClickhouseMessageConverter(agent_name))
+    memory = AgentTokenBufferMemory(llm=llm, chat_memory=chat_memory)
+
+    prompt = OpenAIFunctionsAgent.create_prompt(
+        system_message=SystemMessage(content=system_prompt),
+        extra_prompt_messages=[MessagesPlaceholder(variable_name="history")],
+    )
+    agent = OpenAIFunctionsAgent(llm=llm, tools=tools, prompt=prompt)
+    return AgentExecutor(
+        agent=agent,
+        tools=tools,
+        memory=memory,
+        verbose=True,
+        return_intermediate_steps=True,
+        **kwargs
+    )
+
+
+def build_agents(
+        session_id: str,
+        tool_names: List[str],
+        model: str = "gpt-3.5-turbo-0125",
+        temperature: float = 0.6,
+        system_prompt: str = DEFAULT_SYSTEM_PROMPT
+):
+    chat_llm = ChatOpenAI(
+        model_name=model,
+        temperature=temperature,
+        base_url=GLOBAL_CONFIG.openai_api_base,
+        api_key=GLOBAL_CONFIG.openai_api_key,
+        streaming=True
+    )
+    tools = st.session_state.get(AVAILABLE_RETRIEVAL_TOOLS, st.session_state.get(RETRIEVER_TOOLS))
+    selected_tools = [tools[k] for k in tool_names]
+    logger.info(f"create agent, use tools: {selected_tools}")
+    agent = create_agent_executor(
+        agent_name="chat_memory",
+        session_id=session_id,
+        llm=chat_llm,
+        tools=selected_tools,
+        system_prompt=system_prompt
+    )
+    return agent

backend/construct/build_all.py

@@ -0,0 +1,95 @@
+from logger import logger
+from typing import Dict, Any, Union
+
+import streamlit as st
+
+from backend.constants.myscale_tables import MYSCALE_TABLES
+from backend.constants.variables import CHAINS_RETRIEVERS_MAPPING
+from backend.construct.build_chains import build_retrieval_qa_with_sources_chain
+from backend.construct.build_retriever_tool import create_retriever_tool
+from backend.construct.build_retrievers import build_self_query_retriever, build_vector_sql_db_chain_retriever
+from backend.types.chains_and_retrievers import ChainsAndRetrievers, MetadataColumn
+
+from langchain_community.embeddings import HuggingFaceEmbeddings, HuggingFaceInstructEmbeddings, \
+    SentenceTransformerEmbeddings
+
+
+@st.cache_resource
+def load_embedding_model_for_table(table_name: str) -> \
+        Union[SentenceTransformerEmbeddings, HuggingFaceInstructEmbeddings]:
+    with st.spinner(f"Loading embedding models for [{table_name}] ..."):
+        embeddings = MYSCALE_TABLES[table_name].emb_model()
+    return embeddings
+
+
+@st.cache_resource
+def load_embedding_models() -> Dict[str, Union[HuggingFaceEmbeddings, HuggingFaceInstructEmbeddings]]:
+    embedding_models = {}
+    for table in MYSCALE_TABLES:
+        embedding_models[table] = load_embedding_model_for_table(table)
+    return embedding_models
+
+
+@st.cache_resource
+def update_retriever_tools():
+    retrievers_tools = {}
+    for table in MYSCALE_TABLES:
+        logger.info(f"Updating retriever tools [<retriever>, <sql_retriever>] for table {table}")
+        retrievers_tools.update(
+            {
+                f"{table} + Self Querying": create_retriever_tool(
+                    st.session_state[CHAINS_RETRIEVERS_MAPPING][table]["retriever"],
+                    *MYSCALE_TABLES[table].tool_desc
+                ),
+                f"{table} + Vector SQL": create_retriever_tool(
+                    st.session_state[CHAINS_RETRIEVERS_MAPPING][table]["sql_retriever"],
+                    *MYSCALE_TABLES[table].tool_desc
+                ),
+            })
+    return retrievers_tools
+
+
+@st.cache_resource
+def build_chains_retriever_for_table(table_name: str) -> ChainsAndRetrievers:
+    metadata_col_attributes = MYSCALE_TABLES[table_name].metadata_col_attributes
+
+    self_query_retriever = build_self_query_retriever(table_name)
+    self_query_chain = build_retrieval_qa_with_sources_chain(
+        table_name=table_name,
+        retriever=self_query_retriever,
+        chain_name="Self Query Retriever"
+    )
+
+    vector_sql_retriever = build_vector_sql_db_chain_retriever(table_name)
+    vector_sql_chain = build_retrieval_qa_with_sources_chain(
+        table_name=table_name,
+        retriever=vector_sql_retriever,
+        chain_name="Vector SQL DB Retriever"
+    )
+
+    metadata_columns = [
+        MetadataColumn(
+            name=attribute.name,
+            desc=attribute.description,
+            type=attribute.type
+        )
+        for attribute in metadata_col_attributes
+    ]
+    return ChainsAndRetrievers(
+        metadata_columns=metadata_columns,
+        # for self query
+        retriever=self_query_retriever,
+        chain=self_query_chain,
+        # for vector sql
+        sql_retriever=vector_sql_retriever,
+        sql_chain=vector_sql_chain
+    )
+
+
+@st.cache_resource
+def build_chains_and_retrievers() -> Dict[str, Dict[str, Any]]:
+    chains_and_retrievers = {}
+    for table in MYSCALE_TABLES:
+        logger.info(f"Building chains, retrievers for table {table}")
+        chains_and_retrievers[table] = build_chains_retriever_for_table(table).to_dict()
+    return chains_and_retrievers

backend/construct/build_chains.py

@@ -0,0 +1,39 @@
+from langchain.chains import LLMChain
+from langchain.chat_models import ChatOpenAI
+from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
+from langchain.schema import BaseRetriever
+import streamlit as st
+
+from backend.chains.retrieval_qa_with_sources import CustomRetrievalQAWithSourcesChain
+from backend.chains.stuff_documents import CustomStuffDocumentChain
+from backend.constants.myscale_tables import MYSCALE_TABLES
+from backend.constants.prompts import COMBINE_PROMPT
+from backend.constants.variables import GLOBAL_CONFIG
+
+
+def build_retrieval_qa_with_sources_chain(
+        table_name: str,
+        retriever: BaseRetriever,
+        chain_name: str = "<chain_name>"
+) -> CustomRetrievalQAWithSourcesChain:
+    with st.spinner(f'Building QA source chain named `{chain_name}` for MyScaleDB/{table_name} ...'):
+        # Assign ref_id for documents
+        custom_stuff_document_chain = CustomStuffDocumentChain(
+            llm_chain=LLMChain(
+                prompt=COMBINE_PROMPT,
+                llm=ChatOpenAI(
+                    model_name=GLOBAL_CONFIG.chat_model,
+                    openai_api_key=GLOBAL_CONFIG.openai_api_key,
+                    temperature=0.6
+                ),
+            ),
+            document_prompt=MYSCALE_TABLES[table_name].doc_prompt,
+            document_variable_name="summaries",
+        )
+        chain = CustomRetrievalQAWithSourcesChain(
+            retriever=retriever,
+            combine_documents_chain=custom_stuff_document_chain,
+            return_source_documents=True,
+            max_tokens_limit=12000,
+        )
+    return chain

backend/construct/build_chat_bot.py

@@ -0,0 +1,36 @@
+from backend.chat_bot.private_knowledge_base import ChatBotKnowledgeTable
+from backend.constants.streamlit_keys import CHAT_KNOWLEDGE_TABLE, CHAT_SESSION, CHAT_SESSION_MANAGER
+import streamlit as st
+
+from backend.constants.variables import GLOBAL_CONFIG, TABLE_EMBEDDINGS_MAPPING
+from backend.constants.prompts import DEFAULT_SYSTEM_PROMPT
+from backend.chat_bot.session_manager import SessionManager
+
+
+def build_chat_knowledge_table():
+    if CHAT_KNOWLEDGE_TABLE not in st.session_state:
+        st.session_state[CHAT_KNOWLEDGE_TABLE] = ChatBotKnowledgeTable(
+            host=GLOBAL_CONFIG.myscale_host,
+            port=GLOBAL_CONFIG.myscale_port,
+            username=GLOBAL_CONFIG.myscale_user,
+            password=GLOBAL_CONFIG.myscale_password,
+            # embedding=st.session_state[TABLE_EMBEDDINGS_MAPPING]["Wikipedia"],
+            embedding=st.session_state[TABLE_EMBEDDINGS_MAPPING]["ArXiv Papers"],
+            parser_api_key=GLOBAL_CONFIG.untrusted_api,
+        )
+
+
+def initialize_session_manager():
+    if CHAT_SESSION not in st.session_state:
+        st.session_state[CHAT_SESSION] = {
+            "session_id": "default",
+            "system_prompt": DEFAULT_SYSTEM_PROMPT,
+        }
+    if CHAT_SESSION_MANAGER not in st.session_state:
+        st.session_state[CHAT_SESSION_MANAGER] = SessionManager(
+            st.session_state,
+            host=GLOBAL_CONFIG.myscale_host,
+            port=GLOBAL_CONFIG.myscale_port,
+            username=GLOBAL_CONFIG.myscale_user,
+            password=GLOBAL_CONFIG.myscale_password,
+        )

backend/construct/build_retriever_tool.py

@@ -0,0 +1,45 @@
+import json
+from typing import List
+
+from langchain.pydantic_v1 import BaseModel, Field
+from langchain.schema import BaseRetriever, Document
+from langchain.tools import Tool
+
+from backend.chat_bot.json_decoder import CustomJSONEncoder
+
+
+class RetrieverInput(BaseModel):
+    query: str = Field(description="query to look up in retriever")
+
+
+def create_retriever_tool(
+        retriever: BaseRetriever,
+        tool_name: str,
+        description: str
+) -> Tool:
+    """Create a tool to do retrieval of documents.
+
+    Args:
+        retriever: The retriever to use for the retrieval
+        tool_name: The name for the tool. This will be passed to the language model,
+            so should be unique and somewhat descriptive.
+        description: The description for the tool. This will be passed to the language
+            model, so should be descriptive.
+
+    Returns:
+        Tool class to pass to an agent
+    """
+    def wrap(func):
+        def wrapped_retrieve(*args, **kwargs):
+            docs: List[Document] = func(*args, **kwargs)
+            return json.dumps([d.dict() for d in docs], cls=CustomJSONEncoder)
+
+        return wrapped_retrieve
+
+    return Tool(
+        name=tool_name,
+        description=description,
+        func=wrap(retriever.get_relevant_documents),
+        coroutine=retriever.aget_relevant_documents,
+        args_schema=RetrieverInput,
+    )

backend/construct/build_retrievers.py

@@ -0,0 +1,120 @@
+import streamlit as st
+from langchain.chat_models import ChatOpenAI
+from langchain.prompts.prompt import PromptTemplate
+from langchain.retrievers.self_query.base import SelfQueryRetriever
+from langchain.retrievers.self_query.myscale import MyScaleTranslator
+from langchain.utilities.sql_database import SQLDatabase
+from langchain.vectorstores import MyScaleSettings
+from langchain_experimental.retrievers.vector_sql_database import VectorSQLDatabaseChainRetriever
+from langchain_experimental.sql.vector_sql import VectorSQLDatabaseChain
+from sqlalchemy import create_engine, MetaData
+
+from backend.constants.myscale_tables import MYSCALE_TABLES
+from backend.constants.prompts import MYSCALE_PROMPT
+from backend.constants.variables import TABLE_EMBEDDINGS_MAPPING, GLOBAL_CONFIG
+from backend.retrievers.vector_sql_output_parser import VectorSQLRetrieveOutputParser
+from backend.vector_store.myscale_without_metadata import MyScaleWithoutMetadataJson
+from logger import logger
+
+
+@st.cache_resource
+def build_self_query_retriever(table_name: str) -> SelfQueryRetriever:
+    with st.spinner(f"Building VectorStore for MyScaleDB/{table_name} ..."):
+        myscale_connection = {
+            "host": GLOBAL_CONFIG.myscale_host,
+            "port": GLOBAL_CONFIG.myscale_port,
+            "username": GLOBAL_CONFIG.myscale_user,
+            "password": GLOBAL_CONFIG.myscale_password,
+        }
+        myscale_settings = MyScaleSettings(
+            **myscale_connection,
+            database=MYSCALE_TABLES[table_name].database,
+            table=MYSCALE_TABLES[table_name].table,
+            column_map={
+                "id": "id",
+                "text": MYSCALE_TABLES[table_name].text_col_name,
+                "vector": MYSCALE_TABLES[table_name].vector_col_name,
+                # TODO refine MyScaleDB metadata in langchain.
+                "metadata": MYSCALE_TABLES[table_name].metadata_col_name
+            }
+        )
+        myscale_vector_store = MyScaleWithoutMetadataJson(
+            embedding=st.session_state[TABLE_EMBEDDINGS_MAPPING][table_name],
+            config=myscale_settings,
+            must_have_cols=MYSCALE_TABLES[table_name].must_have_col_names
+        )
+
+    with st.spinner(f"Building SelfQueryRetriever for MyScaleDB/{table_name} ..."):
+        retriever: SelfQueryRetriever = SelfQueryRetriever.from_llm(
+            llm=ChatOpenAI(
+                model_name=GLOBAL_CONFIG.query_model,
+                base_url=GLOBAL_CONFIG.openai_api_base,
+                api_key=GLOBAL_CONFIG.openai_api_key,
+                temperature=0
+            ),
+            vectorstore=myscale_vector_store,
+            document_contents=MYSCALE_TABLES[table_name].table_contents,
+            metadata_field_info=MYSCALE_TABLES[table_name].metadata_col_attributes,
+            use_original_query=False,
+            structured_query_translator=MyScaleTranslator()
+        )
+    return retriever
+
+
+@st.cache_resource
+def build_vector_sql_db_chain_retriever(table_name: str) -> VectorSQLDatabaseChainRetriever:
+    """Get a group of relative docs from MyScaleDB"""
+    with st.spinner(f'Building Vector SQL Database Retriever for MyScaleDB/{table_name}...'):
+        if GLOBAL_CONFIG.myscale_enable_https == False:
+            engine = create_engine(
+                f'clickhouse://{GLOBAL_CONFIG.myscale_user}:{GLOBAL_CONFIG.myscale_password}@'
+                f'{GLOBAL_CONFIG.myscale_host}:{GLOBAL_CONFIG.myscale_port}'
+                f'/{MYSCALE_TABLES[table_name].database}?protocol=http'
+            )
+        else:
+            engine = create_engine(
+                f'clickhouse://{GLOBAL_CONFIG.myscale_user}:{GLOBAL_CONFIG.myscale_password}@'
+                f'{GLOBAL_CONFIG.myscale_host}:{GLOBAL_CONFIG.myscale_port}'
+                f'/{MYSCALE_TABLES[table_name].database}?protocol=https'
+            )
+        metadata = MetaData(bind=engine)
+        logger.info(f"{table_name} metadata is : {metadata}")
+        prompt = PromptTemplate(
+            input_variables=["input", "table_info", "top_k"],
+            template=MYSCALE_PROMPT,
+        )
+        # Custom `out_put_parser` rewrite search SQL, make it's possible to query custom column.
+        output_parser = VectorSQLRetrieveOutputParser.from_embeddings(
+            model=st.session_state[TABLE_EMBEDDINGS_MAPPING][table_name],
+            # rewrite columns needs be searched.
+            must_have_columns=MYSCALE_TABLES[table_name].must_have_col_names
+        )
+
+        # `db_chain` will generate a SQL
+        vector_sql_db_chain: VectorSQLDatabaseChain = VectorSQLDatabaseChain.from_llm(
+            llm=ChatOpenAI(
+                model_name=GLOBAL_CONFIG.query_model,
+                base_url=GLOBAL_CONFIG.openai_api_base,
+                api_key=GLOBAL_CONFIG.openai_api_key,
+                temperature=0
+            ),
+            prompt=prompt,
+            top_k=10,
+            return_direct=True,
+            db=SQLDatabase(
+                engine,
+                None,
+                metadata,
+                include_tables=[MYSCALE_TABLES[table_name].table],
+                max_string_length=1024
+            ),
+            sql_cmd_parser=output_parser,  # TODO needs update `langchain`, fix return type.
+            native_format=True
+        )
+
+        # `retriever` can search a group of documents with `db_chain`
+        vector_sql_db_chain_retriever = VectorSQLDatabaseChainRetriever(
+            sql_db_chain=vector_sql_db_chain,
+            page_content_key=MYSCALE_TABLES[table_name].text_col_name
+        )
+    return vector_sql_db_chain_retriever

backend/retrievers/self_query.py

@@ -0,0 +1,89 @@
+from typing import List
+
+import pandas as pd
+import streamlit as st
+from langchain.retrievers import SelfQueryRetriever
+from langchain_core.documents import Document
+from langchain_core.runnables import RunnableConfig
+
+from backend.chains.retrieval_qa_with_sources import CustomRetrievalQAWithSourcesChain
+from backend.constants.myscale_tables import MYSCALE_TABLES
+from backend.constants.variables import CHAINS_RETRIEVERS_MAPPING, DIVIDER_HTML, RetrieverButtons
+from backend.callbacks.self_query_callbacks import ChatDataSelfAskCallBackHandler, CustomSelfQueryRetrieverCallBackHandler
+from ui.utils import display
+from logger import logger
+
+
+def process_self_query(selected_table, query_type):
+    place_holder = st.empty()
+    logger.info(
+        f"button-1: {RetrieverButtons.self_query_from_db}, "
+        f"button-2: {RetrieverButtons.self_query_with_llm}, "
+        f"content: {st.session_state.query_self}"
+    )
+    with place_holder.expander('🪵 Chat Log', expanded=True):
+        try:
+            if query_type == RetrieverButtons.self_query_from_db:
+                callback = CustomSelfQueryRetrieverCallBackHandler()
+                retriever: SelfQueryRetriever = \
+                    st.session_state[CHAINS_RETRIEVERS_MAPPING][selected_table]["retriever"]
+                config: RunnableConfig = {"callbacks": [callback]}
+
+                relevant_docs = retriever.invoke(
+                    input=st.session_state.query_self,
+                    config=config
+                )
+
+                callback.progress_bar.progress(
+                    value=1.0, text="[Question -> LLM -> Query filter -> MyScaleDB -> Results] Done!✅")
+
+                st.markdown(f"### Self Query Results from `{selected_table}` \n"
+                            f"> Here we get documents from MyScaleDB by `SelfQueryRetriever` \n\n")
+                display(
+                    dataframe=pd.DataFrame(
+                        [{**d.metadata, 'abstract': d.page_content} for d in relevant_docs]
+                    ),
+                    columns_=MYSCALE_TABLES[selected_table].must_have_col_names
+                )
+            elif query_type == RetrieverButtons.self_query_with_llm:
+                # callback = CustomSelfQueryRetrieverCallBackHandler()
+                callback = ChatDataSelfAskCallBackHandler()
+                chain: CustomRetrievalQAWithSourcesChain = \
+                    st.session_state[CHAINS_RETRIEVERS_MAPPING][selected_table]["chain"]
+                chain_results = chain(st.session_state.query_self, callbacks=[callback])
+                callback.progress_bar.progress(
+                    value=1.0,
+                    text="[Question -> LLM -> Query filter -> MyScaleDB -> Related Results -> LLM -> LLM Answer] Done!✅"
+                )
+
+                documents_reference: List[Document] = chain_results["source_documents"]
+                st.markdown(f"### SelfQueryRetriever Results from `{selected_table}` \n"
+                            f"> Here we get documents from MyScaleDB by `SelfQueryRetriever` \n\n")
+                display(
+                    pd.DataFrame(
+                        [{**d.metadata, 'abstract': d.page_content} for d in documents_reference]
+                    )
+                )
+                st.markdown(
+                    f"### Answer from LLM \n"
+                    f"> The response of the LLM when given the `SelfQueryRetriever` results. \n\n"
+                )
+                st.write(chain_results['answer'])
+                st.markdown(
+                    f"### References from `{selected_table}`\n"
+                    f"> Here shows that which documents used by LLM \n\n"
+                )
+                if len(chain_results['sources']) == 0:
+                    st.write("No documents is used by LLM.")
+                else:
+                    display(
+                        dataframe=pd.DataFrame(
+                            [{**d.metadata, 'abstract': d.page_content} for d in chain_results['sources']]
+                        ),
+                        columns_=['ref_id'] + MYSCALE_TABLES[selected_table].must_have_col_names,
+                        index='ref_id'
+                    )
+            st.markdown(DIVIDER_HTML, unsafe_allow_html=True)
+        except Exception as e:
+            st.write('Oops 😵 Something bad happened...')
+            raise e

backend/retrievers/vector_sql_output_parser.py

@@ -0,0 +1,23 @@
+from typing import Dict, Any, List
+
+from langchain_experimental.sql.vector_sql import VectorSQLOutputParser
+
+
+class VectorSQLRetrieveOutputParser(VectorSQLOutputParser):
+    """Based on VectorSQLOutputParser
+    It also modify the SQL to get all columns
+    """
+    must_have_columns: List[str]
+
+    @property
+    def _type(self) -> str:
+        return "vector_sql_retrieve_custom"
+
+    def parse(self, text: str) -> Dict[str, Any]:
+        text = text.strip()
+        start = text.upper().find("SELECT")
+        if start >= 0:
+            end = text.upper().find("FROM")
+            text = text.replace(
+                text[start + len("SELECT") + 1: end - 1], ", ".join(self.must_have_columns))
+        return super().parse(text)

backend/retrievers/vector_sql_query.py

@@ -0,0 +1,95 @@
+from typing import List
+
+import pandas as pd
+import streamlit as st
+from langchain.schema import Document
+from langchain_experimental.retrievers.vector_sql_database import VectorSQLDatabaseChainRetriever
+
+from backend.chains.retrieval_qa_with_sources import CustomRetrievalQAWithSourcesChain
+from backend.constants.myscale_tables import MYSCALE_TABLES
+from backend.constants.variables import CHAINS_RETRIEVERS_MAPPING, DIVIDER_HTML, RetrieverButtons
+from backend.callbacks.vector_sql_callbacks import VectorSQLSearchDBCallBackHandler, VectorSQLSearchLLMCallBackHandler
+from ui.utils import display
+from logger import logger
+
+
+def process_sql_query(selected_table: str, query_type: str):
+    place_holder = st.empty()
+    logger.info(
+        f"button-1: {st.session_state[RetrieverButtons.vector_sql_query_from_db]}, "
+        f"button-2: {st.session_state[RetrieverButtons.vector_sql_query_with_llm]}, "
+        f"table: {selected_table}, "
+        f"content: {st.session_state.query_sql}"
+    )
+    with place_holder.expander('🪵 Query Log', expanded=True):
+        try:
+            if query_type == RetrieverButtons.vector_sql_query_from_db:
+                callback = VectorSQLSearchDBCallBackHandler()
+                vector_sql_retriever: VectorSQLDatabaseChainRetriever = \
+                    st.session_state[CHAINS_RETRIEVERS_MAPPING][selected_table]["sql_retriever"]
+                relevant_docs: List[Document] = vector_sql_retriever.get_relevant_documents(
+                    query=st.session_state.query_sql,
+                    callbacks=[callback]
+                )
+
+                callback.progress_bar.progress(
+                    value=1.0,
+                    text="[Question -> LLM -> SQL Statement -> MyScaleDB -> Results] Done! ✅"
+                )
+
+                st.markdown(f"### Vector Search Results from `{selected_table}` \n"
+                            f"> Here we get documents from MyScaleDB with given sql statement \n\n")
+                display(
+                    pd.DataFrame(
+                        [{**d.metadata, 'abstract': d.page_content} for d in relevant_docs]
+                    )
+                )
+            elif query_type == RetrieverButtons.vector_sql_query_with_llm:
+                callback = VectorSQLSearchLLMCallBackHandler(table=selected_table)
+                vector_sql_chain: CustomRetrievalQAWithSourcesChain = \
+                    st.session_state[CHAINS_RETRIEVERS_MAPPING][selected_table]["sql_chain"]
+                chain_results = vector_sql_chain(
+                    inputs=st.session_state.query_sql,
+                    callbacks=[callback]
+                )
+
+                callback.progress_bar.progress(
+                    value=1.0,
+                    text="[Question -> LLM -> SQL Statement -> MyScaleDB -> "
+                         "(Question,Results) -> LLM -> Results] Done! ✅"
+                )
+
+                documents_reference: List[Document] = chain_results["source_documents"]
+                st.markdown(f"### Vector Search Results from `{selected_table}` \n"
+                            f"> Here we get documents from MyScaleDB with given sql statement \n\n")
+                display(
+                    pd.DataFrame(
+                        [{**d.metadata, 'abstract': d.page_content} for d in documents_reference]
+                    )
+                )
+                st.markdown(
+                    f"### Answer from LLM \n"
+                    f"> The response of the LLM when given the vector search results. \n\n"
+                )
+                st.write(chain_results['answer'])
+                st.markdown(
+                    f"### References from `{selected_table}`\n"
+                    f"> Here shows that which documents used by LLM \n\n"
+                )
+                if len(chain_results['sources']) == 0:
+                    st.write("No documents is used by LLM.")
+                else:
+                    display(
+                        dataframe=pd.DataFrame(
+                            [{**d.metadata, 'abstract': d.page_content} for d in chain_results['sources']]
+                        ),
+                        columns_=['ref_id'] + MYSCALE_TABLES[selected_table].must_have_col_names,
+                        index='ref_id'
+                    )
+            else:
+                raise NotImplementedError(f"Unsupported query type: {query_type}")
+            st.markdown(DIVIDER_HTML, unsafe_allow_html=True)
+        except Exception as e:
+            st.write('Oops 😵 Something bad happened...')
+            raise e
+

backend/types/chains_and_retrievers.py

@@ -0,0 +1,34 @@
+from typing import Dict
+from dataclasses import dataclass
+from typing import List, Any
+from langchain.retrievers import SelfQueryRetriever
+from langchain_experimental.retrievers.vector_sql_database import VectorSQLDatabaseChainRetriever
+
+from backend.chains.retrieval_qa_with_sources import CustomRetrievalQAWithSourcesChain
+
+
+@dataclass
+class MetadataColumn:
+    name: str
+    desc: str
+    type: str
+
+
+@dataclass
+class ChainsAndRetrievers:
+    metadata_columns: List[MetadataColumn]
+    retriever: SelfQueryRetriever
+    chain: CustomRetrievalQAWithSourcesChain
+    sql_retriever: VectorSQLDatabaseChainRetriever
+    sql_chain: CustomRetrievalQAWithSourcesChain
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "metadata_columns": self.metadata_columns,
+            "retriever": self.retriever,
+            "chain": self.chain,
+            "sql_retriever": self.sql_retriever,
+            "sql_chain": self.sql_chain
+        }
+
+

backend/types/global_config.py

@@ -0,0 +1,22 @@
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class GlobalConfig:
+    openai_api_base: Optional[str] = ""
+    openai_api_key: Optional[str] = ""
+
+    auth0_client_id: Optional[str] = ""
+    auth0_domain: Optional[str] = ""
+
+    myscale_user: Optional[str] = ""
+    myscale_password: Optional[str] = ""
+    myscale_host: Optional[str] = ""
+    myscale_port: Optional[int] = 443
+
+    query_model: Optional[str] = ""
+    chat_model: Optional[str] = ""
+
+    untrusted_api: Optional[str] = ""
+    myscale_enable_https: Optional[bool] = True

backend/types/table_config.py

@@ -0,0 +1,25 @@
+from typing import Callable
+from langchain.chains.query_constructor.schema import AttributeInfo
+from langchain.prompts import PromptTemplate
+from dataclasses import dataclass
+from typing import List
+
+
+@dataclass
+class TableConfig:
+    database: str
+    table: str
+    table_contents: str
+    # column names
+    must_have_col_names: List[str]
+    vector_col_name: str
+    text_col_name: str
+    metadata_col_name: str
+    # hint for UI
+    hint: Callable
+    hint_sql: Callable
+    # for langchain
+    doc_prompt: PromptTemplate
+    metadata_col_attributes: List[AttributeInfo]
+    emb_model: Callable
+    tool_desc: tuple

backend/vector_store/myscale_without_metadata.py

@@ -0,0 +1,52 @@
+from typing import Any, Optional, List
+
+from langchain.docstore.document import Document
+from langchain.embeddings.base import Embeddings
+from langchain.vectorstores.myscale import MyScale, MyScaleSettings
+
+from logger import logger
+
+
+class MyScaleWithoutMetadataJson(MyScale):
+    def __init__(self, embedding: Embeddings, config: Optional[MyScaleSettings] = None, must_have_cols: List[str] = [],
+                 **kwargs: Any) -> None:
+        try:
+            super().__init__(embedding, config, **kwargs)
+        except Exception as e:
+            logger.error(e)
+        self.must_have_cols: List[str] = must_have_cols
+
+    def _build_qstr(
+            self, q_emb: List[float], topk: int, where_str: Optional[str] = None
+    ) -> str:
+        q_emb_str = ",".join(map(str, q_emb))
+        if where_str:
+            where_str = f"PREWHERE {where_str}"
+        else:
+            where_str = ""
+
+        q_str = f"""
+            SELECT {self.config.column_map['text']}, dist, {','.join(self.must_have_cols)}
+            FROM {self.config.database}.{self.config.table}
+            {where_str}
+            ORDER BY distance({self.config.column_map['vector']}, [{q_emb_str}]) 
+                AS dist {self.dist_order}
+            LIMIT {topk}
+            """
+        return q_str
+
+    def similarity_search_by_vector(self, embedding: List[float], k: int = 4, where_str: Optional[str] = None,
+                                    **kwargs: Any) -> List[Document]:
+        q_str = self._build_qstr(embedding, k, where_str)
+        try:
+            return [
+                Document(
+                    page_content=r[self.config.column_map["text"]],
+                    metadata={k: r[k] for k in self.must_have_cols},
+                )
+                for r in self.client.query(q_str).named_results()
+            ]
+        except Exception as e:
+            logger.error(
+                f"\033[91m\033[1m{type(e)}\033[0m \033[95m{str(e)}\033[0m")
+            return []

logger.py

@@ -0,0 +1,18 @@
+import logging
+
+
+def setup_logger():
+    logger_ = logging.getLogger('chat-data')
+    logger_.setLevel(logging.INFO)
+    if not logger_.handlers:
+        console_handler = logging.StreamHandler()
+        console_handler.setLevel(logging.INFO)
+        formatter = logging.Formatter(
+            '%(asctime)s - %(filename)s - %(funcName)s - %(levelname)s - %(message)s - [Thread ID: %(thread)d]'
+        )
+        console_handler.setFormatter(formatter)
+        logger_.addHandler(console_handler)
+    return logger_
+
+
+logger = setup_logger()

requirements.txt

@@ -1,15 +1,17 @@
-langchain @ git+https://github.com/myscale/langchain.git@preview#egg=langchain&subdirectory=libs/langchain
-langchain-experimental @ git+https://github.com/myscale/langchain.git@preview#egg=langchain-experimental&subdirectory=libs/experimental
-# https://github.com/PromtEngineer/localGPT/issues/722
-sentence_transformers==2.2.2
+langchain==0.2.1
+langchain-community==0.2.1
+langchain-core==0.2.1
+langchain-experimental==0.0.59
+langchain-openai==0.1.7
+sentence-transformers==2.2.2
 InstructorEmbedding
 pandas
-sentence_transformers
-streamlit==1.25
+streamlit
+streamlit-extras
 streamlit-auth0-component
 altair==4.2.2
 clickhouse-connect
-openai==0.28
+openai==1.35.3
 lark
 tiktoken
 sql-formatter

ui/chat_page.py

@@ -0,0 +1,196 @@
+import datetime
+import json
+
+import pandas as pd
+import streamlit as st
+from langchain_core.messages import HumanMessage, FunctionMessage
+from streamlit.delta_generator import DeltaGenerator
+
+from backend.chat_bot.json_decoder import CustomJSONDecoder
+from backend.constants.streamlit_keys import CHAT_CURRENT_USER_SESSIONS, EL_SESSION_SELECTOR, \
+    EL_UPLOAD_FILES_STATUS, USER_PRIVATE_FILES, EL_BUILD_KB_WITH_FILES, \
+    EL_PERSONAL_KB_NAME, EL_PERSONAL_KB_DESCRIPTION, \
+    USER_PERSONAL_KNOWLEDGE_BASES, AVAILABLE_RETRIEVAL_TOOLS, EL_PERSONAL_KB_NEEDS_REMOVE, \
+    CHAT_KNOWLEDGE_TABLE, EL_UPLOAD_FILES, EL_SELECTED_KBS
+from backend.constants.variables import DIVIDER_HTML, USER_NAME, RETRIEVER_TOOLS
+from backend.construct.build_chat_bot import build_chat_knowledge_table, initialize_session_manager
+from backend.chat_bot.chat import refresh_sessions, on_session_change_submit, refresh_agent, \
+    create_private_knowledge_base_as_tool, \
+    remove_private_knowledge_bases, add_file, clear_files, clear_history, back_to_main, on_chat_submit
+
+
+def render_session_manager():
+    with st.expander("🤖 Session Management"):
+        if CHAT_CURRENT_USER_SESSIONS not in st.session_state:
+            refresh_sessions()
+        st.markdown("Here you can update `session_id` and `system_prompt`")
+        st.markdown("- Click empty row to add a new item")
+        st.markdown("- If needs to delete an item, just click it and press `DEL` key")
+        st.markdown("- Don't forget to submit your change.")
+
+        st.data_editor(
+            data=st.session_state[CHAT_CURRENT_USER_SESSIONS],
+            num_rows="dynamic",
+            key="session_editor",
+            use_container_width=True,
+        )
+        st.button("⏫ Submit", on_click=on_session_change_submit, type="primary")
+
+
+def render_session_selection():
+    with st.expander("✅ Session Selection", expanded=True):
+        st.selectbox(
+            "Choose a `session` to chat",
+            options=st.session_state[CHAT_CURRENT_USER_SESSIONS],
+            index=None,
+            key=EL_SESSION_SELECTOR,
+            format_func=lambda x: x["session_id"],
+            on_change=refresh_agent,
+        )
+
+
+def render_files_manager():
+    with st.expander("📃 **Upload your personal files**", expanded=False):
+        st.markdown("- Files will be parsed by [Unstructured API](https://unstructured.io/api-key).")
+        st.markdown("- All files will be converted into vectors and stored in [MyScaleDB](https://myscale.com/).")
+        st.file_uploader(label="⏫ **Upload files**", key=EL_UPLOAD_FILES, accept_multiple_files=True)
+        # st.markdown("### Uploaded Files")
+        st.dataframe(
+            data=st.session_state[CHAT_KNOWLEDGE_TABLE].list_files(st.session_state[USER_NAME]),
+            use_container_width=True,
+        )
+        st.session_state[EL_UPLOAD_FILES_STATUS] = st.empty()
+        col_1, col_2 = st.columns(2)
+        with col_1:
+            st.button(label="Upload files", on_click=add_file)
+        with col_2:
+            st.button(label="Clear all files and tools", on_click=clear_files)
+
+
+def _render_create_personal_knowledge_bases(div: DeltaGenerator):
+    with div:
+        st.markdown("- If you haven't upload your personal files, please upload them first.")
+        st.markdown("- Select some **files** to build your `personal knowledge base`.")
+        st.markdown("- Once the your `personal knowledge base` is built, "
+                    "it will answer your questions using information from your personal **files**.")
+        st.multiselect(
+            label="⚡️Select some files to build a **personal knowledge base**",
+            options=st.session_state[USER_PRIVATE_FILES],
+            placeholder="You should upload some files first",
+            key=EL_BUILD_KB_WITH_FILES,
+            format_func=lambda x: x["file_name"],
+        )
+        st.text_input(
+            label="⚡️Personal knowledge base name",
+            value="get_relevant_documents",
+            key=EL_PERSONAL_KB_NAME
+        )
+        st.text_input(
+            label="⚡️Personal knowledge base description",
+            value="Searches from some personal files.",
+            key=EL_PERSONAL_KB_DESCRIPTION,
+        )
+        st.button(
+            label="Build 🔧",
+            on_click=create_private_knowledge_base_as_tool
+        )
+
+
+def _render_remove_personal_knowledge_bases(div: DeltaGenerator):
+    with div:
+        st.markdown("> Here is all your personal knowledge bases.")
+        if USER_PERSONAL_KNOWLEDGE_BASES in st.session_state and len(st.session_state[USER_PERSONAL_KNOWLEDGE_BASES]) > 0:
+            st.dataframe(st.session_state[USER_PERSONAL_KNOWLEDGE_BASES])
+        else:
+            st.warning("You don't have any personal knowledge bases, please create a new one.")
+        st.multiselect(
+            label="Choose a personal knowledge base to delete",
+            placeholder="Choose a personal knowledge base to delete",
+            options=st.session_state[USER_PERSONAL_KNOWLEDGE_BASES],
+            format_func=lambda x: x["tool_name"],
+            key=EL_PERSONAL_KB_NEEDS_REMOVE,
+        )
+        st.button("Delete", on_click=remove_private_knowledge_bases, type="primary")
+
+
+def render_personal_tools_build():
+    with st.expander("🔨 **Build your personal knowledge base**", expanded=True):
+        create_new_kb, kb_manager = st.tabs(["Create personal knowledge base", "Personal knowledge base management"])
+        _render_create_personal_knowledge_bases(create_new_kb)
+        _render_remove_personal_knowledge_bases(kb_manager)
+
+
+def render_knowledge_base_selector():
+    with st.expander("🙋 **Select some knowledge bases to query**", expanded=True):
+        st.markdown("- Knowledge bases come in two types: `public` and `private`.")
+        st.markdown("- All users can access our `public` knowledge bases.")
+        st.markdown("- Only you can access your `personal` knowledge bases.")
+        options = st.session_state[RETRIEVER_TOOLS].keys()
+        if AVAILABLE_RETRIEVAL_TOOLS in st.session_state:
+            options = st.session_state[AVAILABLE_RETRIEVAL_TOOLS]
+        st.multiselect(
+            label="Select some knowledge base tool",
+            placeholder="Please select some knowledge bases to query",
+            options=options,
+            default=["Wikipedia + Self Querying"],
+            key=EL_SELECTED_KBS,
+            on_change=refresh_agent,
+        )
+
+
+def chat_page():
+    # initialize resources
+    build_chat_knowledge_table()
+    initialize_session_manager()
+
+    # render sidebar
+    with st.sidebar:
+        left, middle, right = st.columns([1, 1, 2])
+        with left:
+            st.button(label="↩️ Log Out", help="log out and back to main page", on_click=back_to_main)
+        with right:
+            st.markdown(f"👤 `{st.session_state[USER_NAME]}`")
+        st.markdown(DIVIDER_HTML, unsafe_allow_html=True)
+        render_session_manager()
+        render_session_selection()
+        render_files_manager()
+        render_personal_tools_build()
+        render_knowledge_base_selector()
+
+    # render chat history
+    if "agent" not in st.session_state:
+        refresh_agent()
+    for msg in st.session_state.agent.memory.chat_memory.messages:
+        speaker = "user" if isinstance(msg, HumanMessage) else "assistant"
+        if isinstance(msg, FunctionMessage):
+            with st.chat_message(name="from knowledge base", avatar="📚"):
+                st.write(
+                    f"*{datetime.datetime.fromtimestamp(msg.additional_kwargs['timestamp']).isoformat()}*"
+                )
+                st.write("Retrieved from knowledge base:")
+                try:
+                    st.dataframe(
+                        pd.DataFrame.from_records(
+                            json.loads(msg.content, cls=CustomJSONDecoder)
+                        ),
+                        use_container_width=True,
+                    )
+                except Exception as e:
+                    st.warning(e)
+                    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.session_state["next_round"] = st.empty()
+    from streamlit import _bottom
+    with _bottom:
+        col1, col2 = st.columns([1, 16])
+        with col1:
+            st.button("🗑️", help="Clean chat history", on_click=clear_history, type="secondary")
+        with col2:
+            st.chat_input("Input Message", on_submit=on_chat_submit, key="chat_input")

ui/home.py

@@ -0,0 +1,156 @@
+import base64
+
+from streamlit_extras.add_vertical_space import add_vertical_space
+from streamlit_extras.card import card
+from streamlit_extras.colored_header import colored_header
+from streamlit_extras.mention import mention
+from streamlit_extras.tags import tagger_component
+
+from logger import logger
+import os
+
+import streamlit as st
+from auth0_component import login_button
+
+from backend.constants.variables import JUMP_QUERY_ASK, USER_INFO, USER_NAME, DIVIDER_HTML, DIVIDER_THIN_HTML
+from streamlit_extras.let_it_rain import rain
+
+
+def render_home():
+    render_home_header()
+    # st.divider()
+    # st.markdown(DIVIDER_THIN_HTML, unsafe_allow_html=True)
+    add_vertical_space(5)
+    render_home_content()
+    # st.divider()
+    st.markdown(DIVIDER_THIN_HTML, unsafe_allow_html=True)
+    render_home_footer()
+
+
+def render_home_header():
+    logger.info("render home header")
+    st.header("ChatData - Your Intelligent Assistant")
+    st.markdown(DIVIDER_THIN_HTML, unsafe_allow_html=True)
+    st.markdown("> [ChatData](https://github.com/myscale/ChatData) \
+                     is developed by [MyScale](https://myscale.com/), \
+                     it's an integration of [LangChain](https://www.langchain.com/) \
+                     and [MyScaleDB](https://github.com/myscale/myscaledb)")
+
+    tagger_component(
+        "Keywords:",
+        ["MyScaleDB", "LangChain", "VectorSearch", "ChatBot", "GPT", "arxiv", "wikipedia", "Personal Knowledge Base 📚"],
+        color_name=["darkslateblue", "green", "orange", "darkslategrey", "red", "crimson", "darkcyan", "darkgrey"],
+    )
+    text, col1, col2, col3, _ = st.columns([1, 1, 1, 1, 4])
+    with text:
+        st.markdown("Related:")
+    with col1.container():
+        mention(
+            label="streamlit",
+            icon="streamlit",
+            url="https://streamlit.io/",
+            write=True
+        )
+    with col2.container():
+        mention(
+            label="langchain",
+            icon="🦜🔗",
+            url="https://www.langchain.com/",
+            write=True
+        )
+    with col3.container():
+        mention(
+            label="streamlit-extras",
+            icon="🪢",
+            url="https://github.com/arnaudmiribel/streamlit-extras",
+            write=True
+        )
+
+
+def _render_self_query_chain_content():
+    col1, col2 = st.columns([1, 1], gap='large')
+    with col1.container():
+        st.image(image='../assets/home_page_background_1.png',
+                 caption=None,
+                 width=None,
+                 use_column_width=True,
+                 clamp=False,
+                 channels="RGB",
+                 output_format="PNG")
+    with col2.container():
+        st.header("VectorSearch & SelfQuery with Sources")
+        st.info("In this sample, you will learn how **LangChain** integrates with **MyScaleDB**.")
+        st.markdown("""This example demonstrates two methods for integrating MyScale into LangChain: [Vector SQL](https://api.python.langchain.com/en/latest/sql/langchain_experimental.sql.vector_sql.VectorSQLDatabaseChain.html) and [Self-querying retriever](https://python.langchain.com/v0.2/docs/integrations/retrievers/self_query/myscale_self_query/). For each method, you can choose one of the following options:
+
+1. `Retrieve from MyScaleDB ➡️` - The LLM (GPT) converts user queries into SQL statements with vector search, executes these searches in MyScaleDB, and retrieves relevant content.
+   
+2. `Retrieve and answer with LLM ➡️` - After retrieving relevant content from MyScaleDB, the user query along with the retrieved content is sent to the LLM (GPT), which then provides a comprehensive answer.""")
+        add_vertical_space(3)
+        _, middle, _ = st.columns([2, 1, 2], gap='small')
+        with middle.container():
+            st.session_state[JUMP_QUERY_ASK] = st.button("Try sample", use_container_width=False, type="secondary")
+
+
+def _render_chat_bot_content():
+    col1, col2 = st.columns(2, gap='large')
+    with col1.container():
+        st.image(image='../assets/home_page_background_2.png',
+                 caption=None,
+                 width=None,
+                 use_column_width=True,
+                 clamp=False,
+                 channels="RGB",
+                 output_format="PNG")
+    with col2.container():
+        st.header("Chat Bot")
+        st.info("Now you can try our chatbot, this chatbot is built with MyScale and LangChain.")
+        st.markdown("- You need to log in. We use `user_name` to identify each customer.")
+        st.markdown("- You can upload your own PDF files and build your own knowledge base. \
+                     (This is just a sample application. Please do not upload important or confidential files.)")
+        st.markdown("- A default session will be assigned as your initial chat session. \
+                     You can create and switch to other sessions to jump between different chat conversations.")
+        add_vertical_space(1)
+        _, middle, _ = st.columns([1, 2, 1], gap='small')
+        with middle.container():
+            if USER_NAME not in st.session_state:
+                login_button(clientId=os.environ["AUTH0_CLIENT_ID"],
+                             domain=os.environ["AUTH0_DOMAIN"],
+                             key="auth0")
+                # if user_info:
+                #     user_name = user_info.get("nickname", "default") + "_" + user_info.get("email", "null")
+                #     st.session_state[USER_NAME] = user_name
+                #     print(user_info)
+
+
+def render_home_content():
+    logger.info("render home content")
+    _render_self_query_chain_content()
+    add_vertical_space(3)
+    _render_chat_bot_content()
+
+
+def render_home_footer():
+    logger.info("render home footer")
+    st.write(
+        "Please 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.write("Our [privacy policy](https://myscale.com/privacy/), [terms of service](https://myscale.com/terms/)")
+
+    # st.write(
+    #     "Recommended to use the standalone version of Chat-Data, "
+    #     "available [here](https://myscale-chatdata.hf.space/)."
+    # )
+
+    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.rerun()
+    if st.session_state.jump_query_ask:
+        st.rerun()

ui/retrievers.py

@@ -0,0 +1,97 @@
+import streamlit as st
+from streamlit_extras.add_vertical_space import add_vertical_space
+
+from backend.constants.myscale_tables import MYSCALE_TABLES
+from backend.constants.variables import CHAINS_RETRIEVERS_MAPPING, RetrieverButtons
+from backend.retrievers.self_query import process_self_query
+from backend.retrievers.vector_sql_query import process_sql_query
+from backend.constants.variables import JUMP_QUERY_ASK, USER_NAME, USER_INFO
+
+
+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]
+
+
+def _render_table_selector() -> str:
+    col1, col2 = st.columns(2)
+    with col1:
+        selected_table = st.selectbox(
+            label='Each public knowledge base is stored in a MyScaleDB table, which is read-only.',
+            options=MYSCALE_TABLES.keys(),
+        )
+        MYSCALE_TABLES[selected_table].hint()
+    with col2:
+        add_vertical_space(1)
+        st.info(f"Here is your selected public knowledge base schema in MyScaleDB",
+                icon='📚')
+        MYSCALE_TABLES[selected_table].hint_sql()
+
+    return selected_table
+
+
+def render_retrievers():
+    st.button("⬅️ Back", key="back_sql", on_click=back_to_main)
+    st.subheader('Please choose a public knowledge base to search.')
+    selected_table = _render_table_selector()
+
+    tab_sql, tab_self_query = st.tabs(
+        tabs=['Vector SQL', 'Self-querying Retriever']
+    )
+
+    with tab_sql:
+        render_tab_sql(selected_table)
+
+    with tab_self_query:
+        render_tab_self_query(selected_table)
+
+
+def render_tab_sql(selected_table: str):
+    st.warning(
+        "When you input a query with filtering conditions, you need to ensure that your filters are applied only to "
+        "the metadata we provide. This table allows filters to be established on the following metadata fields:",
+        icon="⚠️")
+    st.dataframe(st.session_state[CHAINS_RETRIEVERS_MAPPING][selected_table]["metadata_columns"])
+
+    cols = st.columns([8, 3, 3, 2])
+    cols[0].text_input("Input your question:", key='query_sql')
+    with cols[1].container():
+        add_vertical_space(2)
+        st.button("Retrieve from MyScaleDB ➡️", key=RetrieverButtons.vector_sql_query_from_db)
+    with cols[2].container():
+        add_vertical_space(2)
+        st.button("Retrieve and answer with LLM ➡️", key=RetrieverButtons.vector_sql_query_with_llm)
+
+    if st.session_state[RetrieverButtons.vector_sql_query_from_db]:
+        process_sql_query(selected_table, RetrieverButtons.vector_sql_query_from_db)
+
+    if st.session_state[RetrieverButtons.vector_sql_query_with_llm]:
+        process_sql_query(selected_table, RetrieverButtons.vector_sql_query_with_llm)
+
+
+def render_tab_self_query(selected_table):
+    st.warning(
+        "When you input a query with filtering conditions, you need to ensure that your filters are applied only to "
+        "the metadata we provide. This table allows filters to be established on the following metadata fields:",
+        icon="⚠️")
+    st.dataframe(st.session_state[CHAINS_RETRIEVERS_MAPPING][selected_table]["metadata_columns"])
+
+    cols = st.columns([8, 3, 3, 2])
+    cols[0].text_input("Input your question:", key='query_self')
+
+    with cols[1].container():
+        add_vertical_space(2)
+        st.button("Retrieve from MyScaleDB ➡️", key='search_self')
+    with cols[2].container():
+        add_vertical_space(2)
+        st.button("Retrieve and answer with LLM ➡️", key='ask_self')
+
+    if st.session_state.search_self:
+        process_self_query(selected_table, RetrieverButtons.self_query_from_db)
+
+    if st.session_state.ask_self:
+        process_self_query(selected_table, RetrieverButtons.self_query_with_llm)

ui/utils.py

@@ -0,0 +1,18 @@
+import streamlit as st
+
+
+def display(dataframe, columns_=None, index=None):
+    if len(dataframe) > 0:
+        if index:
+            dataframe.set_index(index)
+        if columns_:
+            st.dataframe(dataframe[columns_])
+        else:
+            st.dataframe(dataframe)
+    else:
+        st.write(
+            "Sorry 😵 we didn't find any articles related to your query.\n\n"
+            "Maybe the LLM is too naughty that does not follow our instruction... \n\n"
+            "Please try again and use verbs that may match the datatype.",
+            unsafe_allow_html=True
+        )