improve chat experience

app.py

@@ -28,7 +28,7 @@ st.markdown(
 )
 st.header("ChatData")
 
-if 'retriever' not in st.session_state:
+if 'sel_map_obj' not in st.session_state:
     st.session_state["sel_map_obj"] = build_all()
     st.session_state["tools"] = build_tools()
 

callbacks/arxiv_callbacks.py

@@ -1,8 +1,11 @@
 import streamlit as st
-from typing import Dict, Any
+import json
+import textwrap
+from typing import Dict, Any, List
 from sql_formatter.core import format_sql
-from langchain.callbacks.streamlit.streamlit_callback_handler import StreamlitCallbackHandler
+from langchain.callbacks.streamlit.streamlit_callback_handler import LLMThought, StreamlitCallbackHandler
 from langchain.schema.output import LLMResult
+from streamlit.delta_generator import DeltaGenerator
 
 class ChatDataSelfSearchCallBackHandler(StreamlitCallbackHandler):
     def __init__(self) -> None:
@@ -91,4 +94,30 @@ class ChatDataSQLAskCallBackHandler(ChatDataSQLSearchCallBackHandler):
         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.prog_interval = 0.1
+        
+        
+class LLMThoughtWithKB(LLMThought):
+    def on_tool_end(self, output: str, color: str | None = None, observation_prefix: str | None = None, llm_prefix: str | None = 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 = LLMThoughtWithKB(
+                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)

chat.py

@@ -5,6 +5,8 @@ import datetime
 import streamlit as st
 from lib.sessions import SessionManager
 from langchain.schema import HumanMessage, FunctionMessage
+from callbacks.arxiv_callbacks import ChatDataAgentCallBackHandler
+from langchain.callbacks.streamlit.streamlit_callback_handler import StreamlitCallbackHandler
 
 from helper import (
     build_agents,
@@ -25,8 +27,14 @@ TOOL_NAMES = {
 
 
 def on_chat_submit():
-    ret = st.session_state.agent({"input": st.session_state.chat_input})
-    print(ret)
+    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])
+        print(ret)
 
 
 def clear_history():
@@ -136,6 +144,12 @@ def chat_page():
     with st.sidebar:
         with st.expander("Session Management"):
             refresh_sessions()
+            st.info("Here you can set up your session! \n\nYou can **change your prompt** here!", 
+                    icon="🤖")
+            st.info(("**Add columns by clicking the empty row**.\n"
+                     "And **delete columns by selecting rows with a press on `DEL` Key**"), 
+                    icon="💡")
+            st.info("Don't forget to **click `Submit Change` to save your change**!", icon="📒")
             st.data_editor(
                 st.session_state.current_sessions,
                 num_rows="dynamic",
@@ -144,6 +158,8 @@ def chat_page():
             )
             st.button("Submit Change!", on_click=on_session_change_submit)
         with st.expander("Session Selection", expanded=True):
+            st.info("Here you can select your session!", icon="🤖")
+            st.info("If no session is attach to your account, then we will add a default session to you!", icon="❤️")
             try:
                 dfl_indx = [
                     x["session_id"] for x in st.session_state.current_sessions
@@ -152,7 +168,7 @@ def chat_page():
                 print("*** ", str(e))
                 dfl_indx = 0
             st.selectbox(
-                "Choose a session be chat:",
+                "Choose a session to chat:",
                 options=st.session_state.current_sessions,
                 index=dfl_indx,
                 key="sel_sess",
@@ -161,10 +177,12 @@ def chat_page():
             )
             print(st.session_state.sel_sess)
         with st.expander("Tool Settings", expanded=True):
+            st.info("Here you can select your tools.", icon="🔧")
+            st.info("We provides you several knowledge base tools for you. We are building more tools!", icon="👷‍♂️")
             st.multiselect(
                 "Knowledge Base",
                 st.session_state.tools.keys(),
-                default=["LangChain Self Query Retriever For Wikipedia"],
+                default=["Wikipedia + Self Querying"],
                 key="selected_tools",
                 on_change=refresh_agent,
             )
@@ -195,4 +213,5 @@ def chat_page():
                         f"*{datetime.datetime.fromtimestamp(msg.additional_kwargs['timestamp']).isoformat()}*"
                     )
                     st.write(f"{msg.content}")
+    st.session_state["next_round"] = st.empty()
     st.chat_input("Input Message", on_submit=on_chat_submit, key="chat_input")

helper.py

@@ -2,12 +2,15 @@
 import json
 import time
 import hashlib
-from typing import Dict, Any
+from typing import Dict, Any, List
 import re
 import pandas as pd
 from os import environ
 import streamlit as st
 import datetime
+from langchain.schema import BaseRetriever
+from langchain.tools import Tool
+from langchain.pydantic_v1 import BaseModel, Field
 
 from sqlalchemy import Column, Text, create_engine, MetaData
 from langchain.agents import AgentExecutor
@@ -28,7 +31,7 @@ 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.schema import BaseRetriever, Document
 from langchain import OpenAI
 from langchain.chains.query_constructor.base import AttributeInfo, VirtualColumnName
 from langchain.retrievers.self_query.base import SelfQueryRetriever
@@ -36,12 +39,12 @@ 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.schema.messages import BaseMessage, HumanMessage, AIMessage, FunctionMessage,\
+    SystemMessage, ChatMessage, ToolMessage
 from langchain.memory import SQLChatMessageHistory
 from langchain.memory.chat_message_histories.sql import \
     BaseMessageConverter, DefaultMessageConverter
@@ -389,6 +392,26 @@ def create_message_model(table_name, DynamicBase):  # type: ignore
 
     return Message
 
+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."""
 
@@ -411,9 +434,10 @@ class DefaultClickhouseMessageConverter(DefaultMessageConverter):
                 "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 = _message_from_dict(msg_dump)
         msg.additional_kwargs = msg_dump["additional_kwargs"]
         return msg
     
@@ -447,6 +471,38 @@ def create_agent_executor(name, session_id, llm, tools, system_prompt, **kwargs)
         **kwargs
     )
     
+class RetrieverInput(BaseModel):
+    query: str = Field(description="query to look up in retriever")
+
+def create_retriever_tool(
+    retriever: BaseRetriever, name: str, description: str
+) -> Tool:
+    """Create a tool to do retrieval of documents.
+
+    Args:
+        retriever: The retriever to use for the retrieval
+        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])
+        return wrapped_retrieve
+    
+    return Tool(
+        name=name,
+        description=description,
+        func=wrap(retriever.get_relevant_documents),
+        coroutine=retriever.aget_relevant_documents,
+        args_schema=RetrieverInput,
+    )
+    
 @st.cache_resource
 def build_tools():
     """build all resources
@@ -465,13 +521,15 @@ def build_tools():
         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.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"],),
+            f"{k} + Self Querying": create_retriever_tool(st.session_state.sel_map_obj[k]["retriever"], *sel_map[k]["tool_desc"],),
+            f"{k} + Vector SQL": create_retriever_tool(st.session_state.sel_map_obj[k]["sql_retriever"], *sel_map[k]["tool_desc"],),
         })
     return sel_map_obj
 
 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)
+    chat_llm = ChatOpenAI(model_name=chat_model_name, temperature=temperature, 
+                          openai_api_base=OPENAI_API_BASE, openai_api_key=OPENAI_API_KEY, streaming=True,
+                          )
     tools = [st.session_state.tools[k] for k in tool_names]
     agent = create_agent_executor(
         "chat_memory",