update session model

app.py

@@ -10,13 +10,22 @@ from callbacks.arxiv_callbacks import ChatDataSelfSearchCallBackHandler, \
 
 from chat import chat_page
 from login import login, back_to_main
+from helper import build_tools, build_agents, build_all, sel_map, display
 
 
-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.markdown(
+    f"""
+    <style>
+        .st-e4 {{
+            max-width: 500px
+        }}
+    </style>""",
+    unsafe_allow_html=True,
+)
 st.header("ChatData")
 
 if 'retriever' not in st.session_state:

chat.py

@@ -1,20 +1,37 @@
 import pandas as pd
 from os import environ
+from time import sleep
 import datetime
 import streamlit as st
+from lib.sessions import SessionManager
 from langchain.schema import HumanMessage, FunctionMessage
 
-from helper import build_agents
+from helper import (
+    build_agents,
+    MYSCALE_HOST,
+    MYSCALE_PASSWORD,
+    MYSCALE_PORT,
+    MYSCALE_USER,
+    DEFAULT_SYSTEM_PROMPT,
+)
 from login import back_to_main
 
-environ['OPENAI_API_BASE'] = st.secrets['OPENAI_API_BASE']
+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():
-    ret = st.session_state.agents[st.session_state.sel][st.session_state.ret_type]({"input": st.session_state.chat_input})
+    ret = st.session_state.agent({"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()
+    if "agent" in st.session_state:
+        st.session_state.agent.memory.clear()
 
 
 def back_to_main():
@@ -25,29 +42,156 @@ def back_to_main():
     if "jump_query_ask" in st.session_state:
         del st.session_state.jump_query_ask
 
+
+def on_session_change_submit():
+    if "session_manager" in st.session_state and "session_editor" in st.session_state:
+        print(st.session_state.session_editor)
+        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:
+                        raise KeyError(
+                            "`session_id` should NOT be neither empty nor contain question marks."
+                        )
+                else:
+                    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"]:
+                st.session_state.session_manager.remove_session(
+                    session_id=f"{st.session_state.user_name}?{st.session_state.current_sessions[elem]['session_id']}",
+                )
+            refresh_sessions()
+            if len(st.session_state.session_editor["deleted_rows"]) > 0:
+                try:
+                    dfl_indx = [
+                        x["session_id"] for x in st.session_state.current_sessions
+                    ].index("default")
+                except ValueError:
+                    dfl_indx = 0
+                st.session_state.sel_sess = st.session_state.current_sessions[dfl_indx]
+        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 build_session_manager():
+    return SessionManager(
+        host=MYSCALE_HOST,
+        port=MYSCALE_PORT,
+        username=MYSCALE_USER,
+        password=MYSCALE_PASSWORD,
+    )
+
+
+def refresh_sessions():
+    st.session_state[
+        "current_sessions"
+    ] = st.session_state.session_manager.list_sessions(st.session_state.user_name)
+    if type(st.session_state.current_sessions) is not dict and len(st.session_state.current_sessions) <= 0:
+        st.session_state.session_manager.add_session(
+            st.session_state.user_name,
+            f"{st.session_state.user_name}?default",
+            DEFAULT_SYSTEM_PROMPT,
+        )
+        st.session_state[
+            "current_sessions"
+        ] = st.session_state.session_manager.list_sessions(st.session_state.user_name)
+
+
+def refresh_agent():
+    with st.spinner("Initializing session..."):
+        print(
+            f"??? Changed to ",
+            f"{st.session_state.user_name}?{st.session_state.sel_sess['session_id']}",
+        )
+        st.session_state["agent"] = build_agents(
+            f"{st.session_state.user_name}?{st.session_state.sel_sess['session_id']}",
+            ["LangChain Self Query Retriever For Wikipedia"]
+            if "selected_tools" not in st.session_state
+            else st.session_state.selected_tools,
+            system_prompt=DEFAULT_SYSTEM_PROMPT
+            if "sel_sess" not in st.session_state
+            else st.session_state.sel_sess["system_prompt"],
+        )
+        st.session_state["session_manager"] = build_session_manager()
+
+
 def chat_page():
-    st.session_state["agents"] = build_agents(f"{st.session_state.user_name}?default")
+    if "sel_sess" not in st.session_state:
+        st.session_state["sel_sess"] = {
+            "session_id": "default",
+            "system_prompt": DEFAULT_SYSTEM_PROMPT,
+        }
     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")
+        with st.expander("Session Management"):
+            refresh_sessions()
+            st.data_editor(
+                st.session_state.current_sessions,
+                num_rows="dynamic",
+                key="session_editor",
+                use_container_width=True,
+            )
+            st.button("Submit Change!", on_click=on_session_change_submit)
+        with st.expander("Session Selection", expanded=True):
+            try:
+                dfl_indx = [
+                    x["session_id"] for x in st.session_state.current_sessions
+                ].index("default")
+            except ValueError:
+                dfl_indx = 0
+            st.selectbox(
+                "Choose a session be chat:",
+                options=st.session_state.current_sessions,
+                index=dfl_indx,
+                key="sel_sess",
+                format_func=lambda x: x["session_id"],
+                on_change=refresh_agent,
+            )
+            print(st.session_state.sel_sess)
+        with st.expander("Tool Settings", expanded=True):
+            st.multiselect(
+                "Knowledge Base",
+                st.session_state.tools.keys(),
+                default=["LangChain Self Query Retriever For Wikipedia"],
+                key="selected_tools",
+                on_change=refresh_agent,
+            )
         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:
+    if 'agent' not in st.session_state:
+        refresh_agent()
+    print("!!! ", st.session_state.agent.memory.chat_memory.session_id)
+    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("Knowledge Base", avatar="📖"):
-                print(type(msg.content))
-                st.write(f"*{datetime.datetime.fromtimestamp(msg.additional_kwargs['timestamp']).isoformat()}*")
+                st.write(
+                    f"*{datetime.datetime.fromtimestamp(msg.additional_kwargs['timestamp']).isoformat()}*"
+                )
                 st.write("Retrieved from knowledge base:")
                 try:
-                    st.dataframe(pd.DataFrame.from_records(map(dict, eval(msg.content))))
+                    st.dataframe(
+                        pd.DataFrame.from_records(map(dict, eval(msg.content)))
+                    )
                 except:
                     st.write(msg.content)
         else:
             if len(msg.content) > 0:
                 with st.chat_message(speaker):
                     print(type(msg), msg.dict())
-                    st.write(f"*{datetime.datetime.fromtimestamp(msg.additional_kwargs['timestamp']).isoformat()}*")
+                    st.write(
+                        f"*{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")
-    

helper.py

@@ -68,6 +68,12 @@ MYSCALE_PORT = st.secrets['MYSCALE_PORT']
 COMBINE_PROMPT = ChatPromptTemplate.from_strings(
     string_messages=[(SystemMessagePromptTemplate, combine_prompt_template),
                     (HumanMessagePromptTemplate, '{question}')])
+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."
+)
 
 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"
@@ -415,7 +421,7 @@ class DefaultClickhouseMessageConverter(DefaultMessageConverter):
         return self.model_class
 
 
-def create_agent_executor(name, session_id, llm, tools, **kwargs):
+def create_agent_executor(name, session_id, llm, tools, system_prompt, **kwargs):
     name = name.replace(" ", "_")
     conn_str = f'clickhouse://{MYSCALE_USER}:{MYSCALE_PASSWORD}@{MYSCALE_HOST}:{MYSCALE_PORT}'
     chat_memory = SQLChatMessageHistory(
@@ -425,12 +431,7 @@ def create_agent_executor(name, session_id, llm, tools, **kwargs):
     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."
-        )
+        content=system_prompt
     )
     prompt = OpenAIFunctionsAgent.create_prompt(
         system_message=_system_message,
@@ -463,38 +464,23 @@ def build_tools():
             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"],),
-        }
+        sel_map_obj.update({
+            f"LangChain Self Query Retriever For {k}": create_retriever_tool(st.session_state.sel_map_obj[k]["retriever"], *sel_map[k]["tool_desc"],),
+            f"Vector SQL Retriever For {k}": 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(session_id):
-    chat_llm = ChatOpenAI(model_name=chat_model_name, temperature=0.6, openai_api_base=OPENAI_API_BASE, openai_api_key=OPENAI_API_KEY)
-    agents = {}
-    cnt = 0
-    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",
-                session_id,
-                chat_llm,
-                tools=tools,
-                )
-            cnt += 1/6
-            p.progress(cnt, f"Building with Knowledge Base {k} via Retriever {n}...")
-    p.empty()
-    return agents
+def build_agents(session_id, tool_names, chat_model_name=chat_model_name, temperature=0.6, system_prompt=DEFAULT_SYSTEM_PROMPT):
+    chat_llm = ChatOpenAI(model_name=chat_model_name, temperature=temperature, openai_api_base=OPENAI_API_BASE, openai_api_key=OPENAI_API_KEY)
+    tools = [st.session_state.tools[k] for k in tool_names]
+    agent = create_agent_executor(
+        "chat_memory",
+        session_id,
+        chat_llm,
+        tools=tools,
+        system_prompt=system_prompt
+        )
+    return agent
 
 
 def display(dataframe, columns_=None, index=None):

lib/schemas.py

@@ -0,0 +1,52 @@
+from sqlalchemy import Column, Text
+from clickhouse_sqlalchemy import types, engines
+
+
+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)
+        user_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
+
+
+def create_session_table(table_name, DynamicBase):  # type: ignore
+    # Model decleared inside a function to have a dynamic table name
+    class Session(DynamicBase):
+        __tablename__ = table_name
+        user_id = Column(Text)
+        session_id = Column(Text, primary_key=True)
+        system_prompt = Column(Text)
+        create_by = Column(types.DateTime)
+        additionals = Column(Text)
+        __table_args__ = (
+            engines.ReplacingMergeTree(
+                order_by=('session_id')),
+            {'comment': 'Store Session and Prompts'}
+        )
+    return Session

lib/sessions.py

@@ -0,0 +1,68 @@
+import json
+try:
+    from sqlalchemy.orm import declarative_base
+except ImportError:
+    from sqlalchemy.ext.declarative import declarative_base
+from datetime import datetime
+from sqlalchemy import Column, Text, orm, create_engine
+from clickhouse_sqlalchemy import types, engines
+from .schemas import create_message_model, create_session_table
+
+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, host, port, username, password, db='chat', sess_table='sessions', msg_table='chat_memory') -> None:
+        conn_str = f'clickhouse://{username}:{password}@{host}:{port}/{db}?protocol=https'
+        self.engine = create_engine(conn_str, echo=False)
+        self.sess_model_class = create_session_table(sess_table, declarative_base())
+        self.sess_model_class.metadata.create_all(self.engine)
+        self.msg_model_class = create_message_model(msg_table, declarative_base())
+        self.msg_model_class.metadata.create_all(self.engine)
+        self.Session = orm.sessionmaker(self.engine)
+
+    def list_sessions(self, user_id):
+        with self.Session() as session:
+            result = (
+                session.query(self.sess_model_class)
+                .where(
+                    self.sess_model_class.user_id == user_id
+                )
+                .order_by(self.sess_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
+    
+    def modify_system_prompt(self, session_id, sys_prompt):
+        with self.Session() as session:
+            session.update(self.sess_model_class).where(self.sess_model_class==session_id).value(system_prompt=sys_prompt)
+            session.commit()
+    
+    def add_session(self, user_id, session_id, system_prompt, **kwargs):
+        with self.Session() as session:
+            elem = self.sess_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()
+    
+    def remove_session(self, session_id):
+        with self.Session() as session:
+            session.query(self.sess_model_class).where(self.sess_model_class.session_id==session_id).delete()
+            session.query(self.msg_model_class).where(self.msg_model_class.session_id==session_id).delete()
+
+