update chat

README.md

@@ -5,7 +5,7 @@ colorFrom: pink
 colorTo: purple
 sdk: streamlit
 sdk_version: 1.20.0
-app_file: app.py
+app_file: chat.py
 pinned: true
 license: mit
 ---

app.py

@@ -14,308 +14,18 @@ from langchain.prompts import PromptTemplate, ChatPromptTemplate, \
 from langchain.prompts.prompt import PromptTemplate
 from langchain.chat_models import ChatOpenAI
 from langchain import OpenAI
-from langchain.chains.query_constructor.base import AttributeInfo, VirtualColumnName
-from langchain.retrievers.self_query.base import SelfQueryRetriever
-from langchain.retrievers.self_query.myscale import MyScaleTranslator
-from langchain.embeddings import HuggingFaceInstructEmbeddings, SentenceTransformerEmbeddings
-from langchain.vectorstores import MyScaleSettings
-from chains.arxiv_chains import MyScaleWithoutMetadataJson
 import re
 import pandas as pd
 from os import environ
 import streamlit as st
 import datetime
-environ['TOKENIZERS_PARALLELISM'] = 'true'
+from helper import build_all, sel_map, display
 environ['OPENAI_API_BASE'] = st.secrets['OPENAI_API_BASE']
 
-
 st.set_page_config(page_title="ChatData")
 
 st.header("ChatData")
 
-# query_model_name = "gpt-3.5-turbo-instruct"
-query_model_name = "text-davinci-003"
-chat_model_name = "gpt-3.5-turbo-16k"
-
-
-def hint_arxiv():
-    st.info("We provides you metadata columns below for query. Please choose a natural expression to describe filters on those columns.\n\n"
-            "For example: \n\n"
-            "*If you want to search papers with complex filters*:\n\n"
-            "- 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.\n\n"
-            "*If you want to ask questions based on papers in database*:\n\n"
-            "- What is PageRank?\n"
-            "- Did Geoffrey Hinton wrote paper about Capsule Neural Networks?\n"
-            "- Introduce some applications of GANs published around 2019.\n"
-            "- 请根据 2019 年左右的文章介绍一下 GAN 的应用都有哪些\n"
-            "- Veuillez présenter les applications du GAN sur la base des articles autour de 2019 ?\n"
-            "- Is it possible to synthesize room temperature super conductive material?")
-
-
-def hint_sql_arxiv():
-    st.info("You can retrieve papers with button `Query` or ask questions based on retrieved papers with button `Ask`.", icon='💡')
-    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.info("We provides you metadata columns below for query. Please choose a natural expression to describe filters on those columns.\n\n"
-            "For example: \n\n"
-            "- Which company did Elon Musk found?\n"
-            "- What is Iron Gwazi?\n"
-            "- What is a Ring in mathematics?\n"
-            "- 苹果的发源地是那里？\n")
-
-
-def hint_sql_wiki():
-    st.info("You can retrieve papers with button `Query` or ask questions based on retrieved papers with button `Ask`.", icon='💡')
-    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
-```''')
-
-
-sel_map = {
-    'Wikipedia': {
-        "database": "wiki",
-        "table": "Wikipedia",
-        "hint": hint_wiki,
-        "hint_sql": hint_sql_wiki,
-        "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_cols": [
-            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_cols": ['id', 'title', 'url', 'text', 'views'],
-        "vector_col": "emb",
-        "text_col": "text",
-        "metadata_col": "metadata",
-        "emb_model": lambda: SentenceTransformerEmbeddings(
-            model_name='sentence-transformers/paraphrase-multilingual-mpnet-base-v2',)
-    },
-    'ArXiv Papers': {
-        "database": "default",
-        "table": "ChatArXiv",
-        "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\tDate of Publication: {pubdate}\n\tCategories: {categories}\nSOURCE: {id}"),
-        "metadata_cols": [
-            AttributeInfo(
-                name=VirtualColumnName(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_cols": ['title', 'id', 'categories', 'abstract', 'authors', 'pubdate'],
-        "vector_col": "vector",
-        "text_col": "abstract",
-        "metadata_col": "metadata",
-        "emb_model": lambda: HuggingFaceInstructEmbeddings(
-            model_name='hkunlp/instructor-xl',
-            embed_instruction="Represent the question for retrieving supporting scientific papers: ")
-    }
-}
-
-
-def try_eval(x):
-    try:
-        return eval(x, {'datetime': datetime})
-    except:
-        return x
-
-
-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\nMaybe the LLM is too naughty that does not follow our instruction... \n\nPlease try again and use verbs that may match the datatype.", unsafe_allow_html=True)
-
-
-def build_embedding_model(_sel):
-    with st.spinner("Loading Model..."):
-        embeddings = sel_map[_sel]["emb_model"]()
-    return embeddings
-
-
-def build_retriever(_sel):
-    with st.spinner(f"Connecting DB for {_sel}..."):
-        myscale_connection = {
-            "host": st.secrets['MYSCALE_HOST'],
-            "port": st.secrets['MYSCALE_PORT'],
-            "username": st.secrets['MYSCALE_USER'],
-            "password": st.secrets['MYSCALE_PASSWORD'],
-        }
-
-        config = MyScaleSettings(**myscale_connection,
-                                 database=sel_map[_sel]["database"],
-                                 table=sel_map[_sel]["table"],
-                                 column_map={
-                                     "id": "id",
-                                     "text": sel_map[_sel]["text_col"],
-                                     "vector": sel_map[_sel]["vector_col"],
-                                     "metadata": sel_map[_sel]["metadata_col"]
-                                 })
-        doc_search = MyScaleWithoutMetadataJson(st.session_state[f"emb_model_{_sel}"], config, 
-                                                must_have_cols=sel_map[_sel]['must_have_cols'])
-
-    with st.spinner(f"Building Self Query Retriever for {_sel}..."):
-        metadata_field_info = sel_map[_sel]["metadata_cols"]
-        retriever = SelfQueryRetriever.from_llm(
-            OpenAI(model_name=query_model_name, openai_api_key=st.secrets['OPENAI_API_KEY'], temperature=0),
-            doc_search, "Scientific papers indexes with abstracts. All in English.", metadata_field_info,
-            use_original_query=False, structured_query_translator=MyScaleTranslator())
-
-    COMBINE_PROMPT = ChatPromptTemplate.from_strings(
-        string_messages=[(SystemMessagePromptTemplate, combine_prompt_template),
-                         (HumanMessagePromptTemplate, '{question}')])
-    OPENAI_API_KEY = st.secrets['OPENAI_API_KEY']
-
-    with st.spinner(f'Building QA Chain with Self-query for {_sel}...'):
-        chain = ArXivQAwithSourcesChain(
-            retriever=retriever,
-            combine_documents_chain=ArXivStuffDocumentChain(
-                llm_chain=LLMChain(
-                    prompt=COMBINE_PROMPT,
-                    llm=ChatOpenAI(model_name=chat_model_name,
-                                   openai_api_key=OPENAI_API_KEY, temperature=0.6),
-                ),
-                document_prompt=sel_map[_sel]["doc_prompt"],
-                document_variable_name="summaries",
-
-            ),
-            return_source_documents=True,
-            max_tokens_limit=12000,
-        )
-
-    with st.spinner(f'Building Vector SQL Database Retriever for {_sel}...'):
-        MYSCALE_USER = st.secrets['MYSCALE_USER']
-        MYSCALE_PASSWORD = st.secrets['MYSCALE_PASSWORD']
-        MYSCALE_HOST = st.secrets['MYSCALE_HOST']
-        MYSCALE_PORT = st.secrets['MYSCALE_PORT']
-        engine = create_engine(
-            f'clickhouse://{MYSCALE_USER}:{MYSCALE_PASSWORD}@{MYSCALE_HOST}:{MYSCALE_PORT}/{sel_map[_sel]["database"]}?protocol=https')
-        metadata = MetaData(bind=engine)
-        PROMPT = PromptTemplate(
-            input_variables=["input", "table_info", "top_k"],
-            template=_myscale_prompt,
-        )
-        output_parser = VectorSQLRetrieveCustomOutputParser.from_embeddings(
-            model=st.session_state[f'emb_model_{_sel}'], must_have_columns=sel_map[_sel]["must_have_cols"])
-        sql_query_chain = VectorSQLDatabaseChain.from_llm(
-            llm=OpenAI(model_name=query_model_name, openai_api_key=OPENAI_API_KEY, temperature=0),
-            prompt=PROMPT,
-            top_k=10,
-            return_direct=True,
-            db=SQLDatabase(engine, None, metadata, max_string_length=1024),
-            sql_cmd_parser=output_parser,
-            native_format=True
-        )
-        sql_retriever = VectorSQLDatabaseChainRetriever(
-            sql_db_chain=sql_query_chain, page_content_key=sel_map[_sel]["text_col"])
-
-    with st.spinner(f'Building QA Chain with Vector SQL for {_sel}...'):
-        sql_chain = ArXivQAwithSourcesChain(
-            retriever=sql_retriever,
-            combine_documents_chain=ArXivStuffDocumentChain(
-                llm_chain=LLMChain(
-                    prompt=COMBINE_PROMPT,
-                    llm=ChatOpenAI(model_name=chat_model_name,
-                                   openai_api_key=OPENAI_API_KEY, temperature=0.6),
-                ),
-                document_prompt=sel_map[_sel]["doc_prompt"],
-                document_variable_name="summaries",
-
-            ),
-            return_source_documents=True,
-            max_tokens_limit=12000,
-        )
-
-    return {
-        "metadata_columns": [{'name': m.name.name if type(m.name) is VirtualColumnName else m.name, 'desc': m.description, 'type': m.type} for m in metadata_field_info],
-        "retriever": retriever,
-        "chain": chain,
-        "sql_retriever": sql_retriever,
-        "sql_chain": sql_chain
-    }
-
-
-@st.cache_resource
-def build_all():
-    sel_map_obj = {}
-    for k in sel_map:
-        st.session_state[f'emb_model_{k}'] = build_embedding_model(k)
-        sel_map_obj[k] = build_retriever(k)
-    return sel_map_obj
-
-
 if 'retriever' not in st.session_state:
     st.session_state["sel_map_obj"] = build_all()
 

chat.py

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

helper.py

@@ -0,0 +1,506 @@
+
+import json
+import time
+import hashlib
+from typing import Dict, Any
+import re
+import pandas as pd
+from os import environ
+import streamlit as st
+import datetime
+
+from sqlalchemy import Column, Text, create_engine, MetaData
+from langchain.agents import AgentExecutor
+try:
+    from sqlalchemy.orm import declarative_base
+except ImportError:
+    from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+from clickhouse_sqlalchemy import (
+    Table, make_session, get_declarative_base, types, engines
+)
+from langchain_experimental.sql.vector_sql import VectorSQLDatabaseChain
+from langchain_experimental.retrievers.vector_sql_database import VectorSQLDatabaseChainRetriever
+from langchain.utilities.sql_database import SQLDatabase
+from langchain.chains import LLMChain
+from sqlalchemy import create_engine, MetaData
+from langchain.prompts import PromptTemplate, ChatPromptTemplate, \
+    SystemMessagePromptTemplate, HumanMessagePromptTemplate
+from langchain.prompts.prompt import PromptTemplate
+from langchain.chat_models import ChatOpenAI
+from langchain.schema import BaseRetriever
+from langchain import OpenAI
+from langchain.chains.query_constructor.base import AttributeInfo, VirtualColumnName
+from langchain.retrievers.self_query.base import SelfQueryRetriever
+from langchain.retrievers.self_query.myscale import MyScaleTranslator
+from langchain.embeddings import HuggingFaceInstructEmbeddings, SentenceTransformerEmbeddings
+from langchain.vectorstores import MyScaleSettings
+from chains.arxiv_chains import MyScaleWithoutMetadataJson
+from langchain.schema import Document
+from langchain.prompts.prompt import PromptTemplate
+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 import BaseMessage, HumanMessage, AIMessage, FunctionMessage, SystemMessage
+from langchain.memory import SQLChatMessageHistory
+from langchain.memory.chat_message_histories.sql import \
+    BaseMessageConverter, DefaultMessageConverter
+from langchain.schema.messages import BaseMessage, _message_to_dict, messages_from_dict
+from langchain.agents.agent_toolkits import create_retriever_tool
+from prompts.arxiv_prompt import combine_prompt_template, _myscale_prompt
+from chains.arxiv_chains import ArXivQAwithSourcesChain, ArXivStuffDocumentChain
+from chains.arxiv_chains import VectorSQLRetrieveCustomOutputParser
+environ['TOKENIZERS_PARALLELISM'] = 'true'
+environ['OPENAI_API_BASE'] = st.secrets['OPENAI_API_BASE']
+
+# query_model_name = "gpt-3.5-turbo-instruct"
+query_model_name = "text-davinci-003"
+chat_model_name = "gpt-3.5-turbo-16k"
+
+
+OPENAI_API_KEY = st.secrets['OPENAI_API_KEY']
+OPENAI_API_BASE = st.secrets['OPENAI_API_BASE']
+MYSCALE_USER = st.secrets['MYSCALE_USER']
+MYSCALE_PASSWORD = st.secrets['MYSCALE_PASSWORD']
+MYSCALE_HOST = st.secrets['MYSCALE_HOST']
+MYSCALE_PORT = st.secrets['MYSCALE_PORT']
+
+COMBINE_PROMPT = ChatPromptTemplate.from_strings(
+    string_messages=[(SystemMessagePromptTemplate, combine_prompt_template),
+                    (HumanMessagePromptTemplate, '{question}')])
+
+def hint_arxiv():
+    st.info("We provides you metadata columns below for query. Please choose a natural expression to describe filters on those columns.\n\n"
+            "For example: \n\n"
+            "*If you want to search papers with complex filters*:\n\n"
+            "- 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.\n\n"
+            "*If you want to ask questions based on papers in database*:\n\n"
+            "- What is PageRank?\n"
+            "- Did Geoffrey Hinton wrote paper about Capsule Neural Networks?\n"
+            "- Introduce some applications of GANs published around 2019.\n"
+            "- 请根据 2019 年左右的文章介绍一下 GAN 的应用都有哪些\n"
+            "- Veuillez présenter les applications du GAN sur la base des articles autour de 2019 ?\n"
+            "- Is it possible to synthesize room temperature super conductive material?")
+
+
+def hint_sql_arxiv():
+    st.info("You can retrieve papers with button `Query` or ask questions based on retrieved papers with button `Ask`.", icon='💡')
+    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.info("We provides you metadata columns below for query. Please choose a natural expression to describe filters on those columns.\n\n"
+            "For example: \n\n"
+            "- Which company did Elon Musk found?\n"
+            "- What is Iron Gwazi?\n"
+            "- What is a Ring in mathematics?\n"
+            "- 苹果的发源地是那里？\n")
+
+
+def hint_sql_wiki():
+    st.info("You can retrieve papers with button `Query` or ask questions based on retrieved papers with button `Ask`.", icon='💡')
+    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
+```''')
+
+
+sel_map = {
+    'Wikipedia': {
+        "database": "wiki",
+        "table": "Wikipedia",
+        "hint": hint_wiki,
+        "hint_sql": hint_sql_wiki,
+        "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_cols": [
+            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_cols": ['id', 'title', 'url', 'text', 'views'],
+        "vector_col": "emb",
+        "text_col": "text",
+        "metadata_col": "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': {
+        "database": "default",
+        "table": "ChatArXiv",
+        "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\tDate of Publication: {pubdate}\n\tCategories: {categories}\nSOURCE: {id}"),
+        "metadata_cols": [
+            AttributeInfo(
+                name=VirtualColumnName(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_cols": ['title', 'id', 'categories', 'abstract', 'authors', 'pubdate'],
+        "vector_col": "vector",
+        "text_col": "abstract",
+        "metadata_col": "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"),
+    }
+}
+
+def build_embedding_model(_sel):
+    """Build embedding model
+    """
+    with st.spinner("Loading Model..."):
+        embeddings = sel_map[_sel]["emb_model"]()
+    return embeddings
+
+
+def build_chains_retrievers(_sel: str) -> Dict[str, Any]:
+    """build chains and retrievers
+
+    :param _sel: selected knowledge base
+    :type _sel: str
+    :return: _description_
+    :rtype: Dict[str, Any]
+    """
+    metadata_field_info = sel_map[_sel]["metadata_cols"]
+    retriever = build_self_query(_sel)
+    chain = build_qa_chain(_sel, retriever, name="Self Query Retriever")
+    sql_retriever = build_vector_sql(_sel)
+    sql_chain = build_qa_chain(_sel, sql_retriever, name="Vector SQL")
+
+    return {
+        "metadata_columns": [{'name': m.name.name if type(m.name) is VirtualColumnName else m.name, 'desc': m.description, 'type': m.type} for m in metadata_field_info],
+        "retriever": retriever,
+        "chain": chain,
+        "sql_retriever": sql_retriever,
+        "sql_chain": sql_chain
+    }
+    
+def build_self_query(_sel: str) -> SelfQueryRetriever:
+    """Build self querying retriever
+
+    :param _sel: selected knowledge base
+    :type _sel: str
+    :return: retriever used by chains
+    :rtype: SelfQueryRetriever
+    """
+    with st.spinner(f"Connecting DB for {_sel}..."):
+        myscale_connection = {
+            "host": MYSCALE_HOST,
+            "port": MYSCALE_PORT,
+            "username": MYSCALE_USER,
+            "password": MYSCALE_PASSWORD,
+        }
+        config = MyScaleSettings(**myscale_connection,
+                                 database=sel_map[_sel]["database"],
+                                 table=sel_map[_sel]["table"],
+                                 column_map={
+                                     "id": "id",
+                                     "text": sel_map[_sel]["text_col"],
+                                     "vector": sel_map[_sel]["vector_col"],
+                                     "metadata": sel_map[_sel]["metadata_col"]
+                                 })
+        doc_search = MyScaleWithoutMetadataJson(st.session_state[f"emb_model_{_sel}"], config, 
+                                                must_have_cols=sel_map[_sel]['must_have_cols'])
+
+    with st.spinner(f"Building Self Query Retriever for {_sel}..."):
+        metadata_field_info = sel_map[_sel]["metadata_cols"]
+        retriever = SelfQueryRetriever.from_llm(
+            OpenAI(model_name=query_model_name, openai_api_key=OPENAI_API_KEY, temperature=0),
+            doc_search, "Scientific papers indexes with abstracts. All in English.", metadata_field_info,
+            use_original_query=False, structured_query_translator=MyScaleTranslator())
+    return retriever
+
+def build_vector_sql(_sel: str)->VectorSQLDatabaseChainRetriever:
+    """Build Vector SQL Database Retriever
+
+    :param _sel: selected knowledge base
+    :type _sel: str
+    :return: retriever used by chains
+    :rtype: VectorSQLDatabaseChainRetriever
+    """
+    with st.spinner(f'Building Vector SQL Database Retriever for {_sel}...'):
+        engine = create_engine(
+            f'clickhouse://{MYSCALE_USER}:{MYSCALE_PASSWORD}@{MYSCALE_HOST}:{MYSCALE_PORT}/{sel_map[_sel]["database"]}?protocol=https')
+        metadata = MetaData(bind=engine)
+        PROMPT = PromptTemplate(
+            input_variables=["input", "table_info", "top_k"],
+            template=_myscale_prompt,
+        )
+        output_parser = VectorSQLRetrieveCustomOutputParser.from_embeddings(
+            model=st.session_state[f'emb_model_{_sel}'], must_have_columns=sel_map[_sel]["must_have_cols"])
+        sql_query_chain = VectorSQLDatabaseChain.from_llm(
+            llm=OpenAI(model_name=query_model_name, openai_api_key=OPENAI_API_KEY, temperature=0),
+            prompt=PROMPT,
+            top_k=10,
+            return_direct=True,
+            db=SQLDatabase(engine, None, metadata, max_string_length=1024),
+            sql_cmd_parser=output_parser,
+            native_format=True
+        )
+        sql_retriever = VectorSQLDatabaseChainRetriever(
+            sql_db_chain=sql_query_chain, page_content_key=sel_map[_sel]["text_col"])
+    return sql_retriever
+    
+def build_qa_chain(_sel: str, retriever: BaseRetriever, name: str="Self-query") -> ArXivQAwithSourcesChain:
+    """_summary_
+
+    :param _sel: selected knowledge base
+    :type _sel: str
+    :param retriever: retriever used by chains
+    :type retriever: BaseRetriever
+    :param name: display name, defaults to "Self-query"
+    :type name: str, optional
+    :return: QA chain interacts with user
+    :rtype: ArXivQAwithSourcesChain
+    """
+    with st.spinner(f'Building QA Chain with {name} for {_sel}...'):
+        chain = ArXivQAwithSourcesChain(
+            retriever=retriever,
+            combine_documents_chain=ArXivStuffDocumentChain(
+                llm_chain=LLMChain(
+                    prompt=COMBINE_PROMPT,
+                    llm=ChatOpenAI(model_name=chat_model_name,
+                                   openai_api_key=OPENAI_API_KEY, temperature=0.6),
+                ),
+                document_prompt=sel_map[_sel]["doc_prompt"],
+                document_variable_name="summaries",
+
+            ),
+            return_source_documents=True,
+            max_tokens_limit=12000,
+        )
+    return chain
+
+@st.cache_resource
+def build_all() -> Dict[str, Any]:
+    """build all resources
+
+    :return: sel_map_obj
+    :rtype: Dict[str, Any]
+    """
+    sel_map_obj = {}
+    for k in sel_map:
+        st.session_state[f'emb_model_{k}'] = build_embedding_model(k)
+        sel_map_obj[k] = build_chains_retrievers(k)
+    return sel_map_obj
+
+def create_message_model(table_name, DynamicBase):  # type: ignore
+    """
+    Create a message model for a given table name.
+
+    Args:
+        table_name: The name of the table to use.
+        DynamicBase: The base class to use for the model.
+
+    Returns:
+        The model class.
+
+    """
+
+    # Model decleared inside a function to have a dynamic table name
+    class Message(DynamicBase):
+        __tablename__ = table_name
+        id = Column(types.Float64)
+        session_id = Column(Text)
+        msg_id = Column(Text, primary_key=True)
+        type = Column(Text)
+        addtionals = Column(Text)
+        message = Column(Text)
+        __table_args__ = (
+            engines.ReplacingMergeTree(
+                partition_by='session_id',
+                order_by=('id', 'msg_id')),
+            {'comment': 'Store Chat History'}
+        )
+
+    return Message
+
+class DefaultClickhouseMessageConverter(DefaultMessageConverter):
+    """The default message converter for SQLChatMessageHistory."""
+
+    def __init__(self, table_name: str):
+        self.model_class = create_message_model(table_name, declarative_base())
+
+    def to_sql_model(self, message: BaseMessage, session_id: str) -> Any:
+        tstamp = time.time()
+        msg_id = hashlib.sha256(f"{session_id}_{message}_{tstamp}".encode('utf-8')).hexdigest()
+        return self.model_class(
+            id=tstamp, 
+            msg_id=msg_id,
+            session_id=session_id, 
+            type=message.type,
+            addtionals=json.dumps(message.additional_kwargs),
+            message=json.dumps({
+                "type": message.type, 
+                "additional_kwargs": {"timestamp": tstamp},
+                "data": message.dict()})
+        )
+    def from_sql_model(self, sql_message: Any) -> BaseMessage:
+        msg_dump = json.loads(sql_message.message)
+        msg = messages_from_dict([msg_dump])[0]
+        msg.additional_kwargs = msg_dump["additional_kwargs"]
+        return msg
+    
+    def get_sql_model_class(self) -> Any:
+        return self.model_class
+
+
+def create_agent_executor(name, session_id, llm, tools, **kwargs):
+    name = name.replace(" ", "_")
+    conn_str = f'clickhouse://{MYSCALE_USER}:{MYSCALE_PASSWORD}@{MYSCALE_HOST}:{MYSCALE_PORT}'
+    chat_memory = SQLChatMessageHistory(
+        session_id,
+        connection_string=f'{conn_str}/chat?protocol=https',
+        custom_message_converter=DefaultClickhouseMessageConverter(name))
+    memory = AgentTokenBufferMemory(llm=llm, chat_memory=chat_memory)
+    
+    _system_message = SystemMessage(
+        content=(
+            "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."
+        )
+    )
+    prompt = OpenAIFunctionsAgent.create_prompt(
+        system_message=_system_message,
+        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
+    )
+    
+@st.cache_resource
+def build_tools():
+    """build all resources
+
+    :return: sel_map_obj
+    :rtype: Dict[str, Any]
+    """
+    sel_map_obj = {}
+    for k in sel_map:
+        if f'emb_model_{k}' not in st.session_state:
+            st.session_state[f'emb_model_{k}'] = build_embedding_model(k)
+        if "sel_map_obj" not in st.session_state:
+            st.session_state["sel_map_obj"] = {}
+        if k not in st.session_state.sel_map_obj:
+            st.session_state["sel_map_obj"][k] = {}
+        if "langchain_retriever" not in st.session_state.sel_map_obj[k] or "vecsql_retriever" not in st.session_state.sel_map_obj[k]:
+            st.session_state.sel_map_obj[k].update(build_chains_retrievers(k))
+        sel_map_obj[k] = {
+            "langchain_retriever_tool": create_retriever_tool(st.session_state.sel_map_obj[k]["retriever"], *sel_map[k]["tool_desc"],),
+            "vecsql_retriever_tool": create_retriever_tool(st.session_state.sel_map_obj[k]["sql_retriever"], *sel_map[k]["tool_desc"],),
+        }
+    return sel_map_obj
+
+@st.cache_resource(max_entries=1)
+def build_agents(username):
+    chat_llm = ChatOpenAI(model_name=chat_model_name, temperature=0.6, openai_api_base=OPENAI_API_BASE, openai_api_key=OPENAI_API_KEY)
+    agents = {}
+    cnt = 0
+    p = st.progress(0.0, "Building agents with different knowledge base...")
+    for k in [*sel_map.keys(), 'ArXiv + Wikipedia']:
+        for m, n in [("langchain_retriever_tool", "Self-querying retriever"), ("vecsql_retriever_tool", "Vector SQL")]:
+            if k == 'ArXiv + Wikipedia':
+                tools = [st.session_state.tools[k][m] for k in sel_map.keys()]
+            elif k == 'Null':
+                tools = []
+            else:
+                tools = [st.session_state.tools[k][m]]
+            if k not in agents:
+                agents[k] = {}
+            agents[k][n] = create_agent_executor(
+                "chat_memory",
+                username,
+                chat_llm,
+                tools=tools,
+                )
+            cnt += 1/6
+            p.progress(cnt, f"Building with Knowledge Base {k} via Retriever {n}...")
+    p.empty()
+    return agents
+
+
+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\nMaybe the LLM is too naughty that does not follow our instruction... \n\nPlease try again and use verbs that may match the datatype.", unsafe_allow_html=True)

requirements.txt

@@ -3,7 +3,8 @@ langchain-experimental @ git+https://github.com/myscale/langchain.git@preview#eg
 InstructorEmbedding
 pandas
 sentence_transformers
-streamlit==1.20
+streamlit==1.25
+streamlit-auth0-component
 altair==4.2.2
 clickhouse-connect
 openai