add knowledge base management

app.py

@@ -10,7 +10,7 @@ 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 lib.helper import build_tools, build_all, sel_map, display
 
 
 

callbacks/arxiv_callbacks.py

@@ -3,70 +3,79 @@ 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 LLMThought, 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:
         self.progress_bar = st.progress(value=0.0, text="Working...")
         self.tokens_stream = ""
-    
+
     def on_llm_start(self, serialized, prompts, **kwargs) -> None:
         pass
-        
+
     def on_text(self, text: str, **kwargs) -> None:
         self.progress_bar.progress(value=0.2, text="Asking LLM...")
-        
+
     def on_chain_end(self, outputs, **kwargs) -> None:
-        self.progress_bar.progress(value=0.6, text='Searching in DB...')
-        if 'repr' in outputs:
-            st.markdown('### Generated Filter')
+        self.progress_bar.progress(value=0.6, text="Searching in DB...")
+        if "repr" in outputs:
+            st.markdown("### Generated Filter")
             st.markdown(f"```python\n{outputs['repr']}\n```", unsafe_allow_html=True)
-    
+
     def on_chain_start(self, serialized, inputs, **kwargs) -> None:
         pass
 
+
 class ChatDataSelfAskCallBackHandler(StreamlitCallbackHandler):
     def __init__(self) -> None:
-        self.progress_bar = st.progress(value=0.0, text='Searching DB...')
+        self.progress_bar = st.progress(value=0.0, text="Searching DB...")
         self.status_bar = st.empty()
         self.prog_value = 0.0
         self.prog_map = {
-            'langchain.chains.qa_with_sources.retrieval.RetrievalQAWithSourcesChain': 0.2,
-            'langchain.chains.combine_documents.map_reduce.MapReduceDocumentsChain': 0.4,
-            'langchain.chains.combine_documents.stuff.StuffDocumentsChain': 0.8
+            "langchain.chains.qa_with_sources.retrieval.RetrievalQAWithSourcesChain": 0.2,
+            "langchain.chains.combine_documents.map_reduce.MapReduceDocumentsChain": 0.4,
+            "langchain.chains.combine_documents.stuff.StuffDocumentsChain": 0.8,
         }
 
     def on_llm_start(self, serialized, prompts, **kwargs) -> None:
         pass
-        
+
     def on_text(self, text: str, **kwargs) -> None:
         pass
-        
+
     def on_chain_start(self, serialized, inputs, **kwargs) -> None:
-        cid = '.'.join(serialized['id']) 
-        if cid != 'langchain.chains.llm.LLMChain':
-            self.progress_bar.progress(value=self.prog_map[cid], text=f'Running Chain `{cid}`...')
+        cid = ".".join(serialized["id"])
+        if cid != "langchain.chains.llm.LLMChain":
+            self.progress_bar.progress(
+                value=self.prog_map[cid], text=f"Running Chain `{cid}`..."
+            )
             self.prog_value = self.prog_map[cid]
         else:
             self.prog_value += 0.1
-            self.progress_bar.progress(value=self.prog_value, text=f'Running Chain `{cid}`...')
+            self.progress_bar.progress(
+                value=self.prog_value, text=f"Running Chain `{cid}`..."
+            )
 
     def on_chain_end(self, outputs, **kwargs) -> None:
         pass
-    
+
 
 class ChatDataSQLSearchCallBackHandler(StreamlitCallbackHandler):
     def __init__(self) -> None:
-        self.progress_bar = st.progress(value=0.0, text='Writing SQL...')
+        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,
@@ -74,41 +83,56 @@ class ChatDataSQLSearchCallBackHandler(StreamlitCallbackHandler):
         **kwargs,
     ):
         text = response.generations[0][0].text
-        if text.replace(' ', '').upper().startswith('SELECT'):
-            st.write('We generated Vector SQL for you:')
-            st.markdown(f'''```sql\n{format_sql(text, max_len=80)}\n```''')
+        if text.replace(" ", "").upper().startswith("SELECT"):
+            st.write("We generated Vector SQL for you:")
+            st.markdown(f"""```sql\n{format_sql(text, max_len=80)}\n```""")
             print(f"Vector SQL: {text}")
             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']) 
+        cid = ".".join(serialized["id"])
         self.prog_value += self.prog_interval
-        self.progress_bar.progress(value=self.prog_value, text=f'Running Chain `{cid}`...')
-   
+        self.progress_bar.progress(
+            value=self.prog_value, text=f"Running Chain `{cid}`..."
+        )
+
     def on_chain_end(self, outputs, **kwargs) -> None:
         pass
-    
+
+
 class ChatDataSQLAskCallBackHandler(ChatDataSQLSearchCallBackHandler):
     def __init__(self) -> None:
-        self.progress_bar = st.progress(value=0.0, text='Writing SQL...')
+        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
-        
-        
+
+
 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:
+    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))]))
+            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:
@@ -120,4 +144,4 @@ class ChatDataAgentCallBackHandler(StreamlitCallbackHandler):
                 labeler=self._thought_labeler,
             )
 
-        self._current_thought.on_llm_start(serialized, prompts)
+        self._current_thought.on_llm_start(serialized, prompts)

chat.py

@@ -1,3 +1,4 @@
+import json
 import pandas as pd
 from os import environ
 from time import sleep
@@ -7,9 +8,12 @@ from lib.sessions import SessionManager
 from lib.private_kb import PrivateKnowledgeBase
 from langchain.schema import HumanMessage, FunctionMessage
 from callbacks.arxiv_callbacks import ChatDataAgentCallBackHandler
-from langchain.callbacks.streamlit.streamlit_callback_handler import StreamlitCallbackHandler
+from langchain.callbacks.streamlit.streamlit_callback_handler import (
+    StreamlitCallbackHandler,
+)
+from lib.json_conv import CustomJSONDecoder
 
-from helper import (
+from lib.helper import (
     build_agents,
     MYSCALE_HOST,
     MYSCALE_PASSWORD,
@@ -30,12 +34,16 @@ TOOL_NAMES = {
 
 def on_chat_submit():
     with st.session_state.next_round.container():
-        with st.chat_message('user'):
+        with st.chat_message("user"):
             st.write(st.session_state.chat_input)
-        with st.chat_message('assistant'):
+        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])
+        st_callback = ChatDataAgentCallBackHandler(
+            container, collapse_completed_thoughts=False
+        )
+        ret = st.session_state.agent(
+            {"input": st.session_state.chat_input}, callbacks=[st_callback]
+        )
         print(ret)
 
 
@@ -105,7 +113,10 @@ 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:
+    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",
@@ -114,14 +125,64 @@ def refresh_sessions():
         st.session_state[
             "current_sessions"
         ] = st.session_state.session_manager.list_sessions(st.session_state.user_name)
-    
+    st.session_state["user_files"] = st.session_state.private_kb.list_files(
+        st.session_state.user_name
+    )
+    st.session_state["user_tools"] = st.session_state.private_kb.list_tools(
+        st.session_state.user_name
+    )
+    st.session_state["tools_with_users"] = {
+        **st.session_state.tools,
+        **st.session_state.private_kb.as_tools(st.session_state.user_name),
+    }
     try:
-        dfl_indx = [x["session_id"] for x in st.session_state.current_sessions].index("default" if "" not in st.session_state else st.session_state.sel_session["session_id"])
+        dfl_indx = [x["session_id"] for x in st.session_state.current_sessions].index(
+            "default"
+            if "" not in st.session_state
+            else st.session_state.sel_session["session_id"]
+        )
     except ValueError:
         dfl_indx = 0
     st.session_state.sel_sess = st.session_state.current_sessions[dfl_indx]
 
 
+def build_kb_as_tool():
+    if (
+        "b_tool_name" in st.session_state
+        and "b_tool_desc" in st.session_state
+        and "b_tool_files" in st.session_state
+        and len(st.session_state.b_tool_name) > 0
+        and len(st.session_state.b_tool_desc) > 0
+        and len(st.session_state.b_tool_files) > 0
+    ):
+        st.session_state.private_kb.create_tool(
+            st.session_state.user_name,
+            st.session_state.b_tool_name,
+            st.session_state.b_tool_desc,
+            [f["file_name"] for f in st.session_state.b_tool_files],
+        )
+        refresh_sessions()
+    else:
+        st.session_state.tool_status.error(
+            "You should fill all fields to build up a tool!"
+        )
+        sleep(2)
+
+
+def remove_kb():
+    if "r_tool_names" in st.session_state and len(st.session_state.r_tool_names) > 0:
+        st.session_state.private_kb.remove_tools(
+            st.session_state.user_name,
+            [f["tool_name"] for f in st.session_state.r_tool_names],
+        )
+        refresh_sessions()
+    else:
+        st.session_state.tool_status.error(
+            "You should specify at least one tool to delete!"
+        )
+        sleep(2)
+
+
 def refresh_agent():
     with st.spinner("Initializing session..."):
         print(
@@ -138,22 +199,29 @@ def refresh_agent():
             else st.session_state.sel_sess["system_prompt"],
         )
 
+
 def add_file():
-    if 'uploaded_files' not in st.session_state or len(st.session_state.uploaded_files) == 0:
+    if (
+        "uploaded_files" not in st.session_state
+        or len(st.session_state.uploaded_files) == 0
+    ):
         st.session_state.tool_status.error("Please upload files!", icon="⚠️")
         sleep(2)
         return
     try:
         st.session_state.tool_status.info("Uploading...")
-        print([(f.name, f.type) for f in st.session_state.uploaded_files])
-        st.session_state.private_kb.add_by_file(st.session_state.user_name,
-                                                st.session_state.uploaded_files)
+        st.session_state.private_kb.add_by_file(
+            st.session_state.user_name, st.session_state.uploaded_files
+        )
+        refresh_sessions()
     except ValueError as e:
         st.session_state.tool_status.error("Failed to upload! " + str(e))
         sleep(2)
-    
+
+
 def clear_files():
     st.session_state.private_kb.clear(st.session_state.user_name)
+    refresh_sessions()
 
 
 def chat_page():
@@ -168,7 +236,7 @@ def chat_page():
             port=MYSCALE_PORT,
             username=MYSCALE_USER,
             password=MYSCALE_PASSWORD,
-            embedding=st.session_state.embeddings['Wikipedia'],
+            embedding=st.session_state.embeddings["Wikipedia"],
             parser_api_key=UNSTRUCTURED_API,
         )
     if "session_manager" not in st.session_state:
@@ -177,12 +245,21 @@ def chat_page():
         with st.expander("Session Management"):
             if "current_sessions" not in st.session_state:
                 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.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",
@@ -191,12 +268,18 @@ 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="❤️")
+            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
-                ].index("default" if "" not in st.session_state else st.session_state.sel_session["session_id"])
+                ].index(
+                    "default"
+                    if "" not in st.session_state
+                    else st.session_state.sel_session["session_id"]
+                )
             except Exception as e:
                 print("*** ", str(e))
                 dfl_indx = 0
@@ -210,39 +293,84 @@ 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.info(
+                "We provides you several knowledge base tools for you. We are building more tools!",
+                icon="🔧",
+            )
             st.session_state["tool_status"] = st.empty()
-            tab_kb, tab_file, tab_build = st.tabs(["Knowledge Bases", "File Upload", "KB Builder"])
+            tab_kb, tab_file = st.tabs(
+                [
+                    "Knowledge Bases",
+                    "File Upload",
+                ]
+            )
             with tab_kb:
+                st.markdown("#### Build You Own Knowledge")
+                st.multiselect(
+                    "Select Files to Build up",
+                    st.session_state.user_files,
+                    placeholder="You should upload files first",
+                    key="b_tool_files",
+                    format_func=lambda x: x["file_name"],
+                )
+                st.text_input("Tool Name", "get_relevant_documents", key="b_tool_name")
+                st.text_input(
+                    "Tool Description",
+                    "Searches among user's private files and returns related documents",
+                    key="b_tool_desc",
+                )
+                st.button("Build!", on_click=build_kb_as_tool)
+                st.markdown("### Knowledge Base Selection")
+                if (
+                    "user_tools" in st.session_state
+                    and len(st.session_state.user_tools) > 0
+                ):
+                    st.markdown("***User Created Knowledge Bases***")
+                    st.dataframe(st.session_state.user_tools)
                 st.multiselect(
                     "Select a Knowledge Base Tool",
-                    st.session_state.tools.keys(),
+                    st.session_state.tools.keys()
+                    if "tools_with_users" not in st.session_state
+                    else st.session_state.tools_with_users,
                     default=["Wikipedia + Self Querying"],
                     key="selected_tools",
                     on_change=refresh_agent,
                 )
+                st.markdown("### Delete Knowledge Base")
+                st.multiselect(
+                    "Choose Knowledge Base to Remove",
+                    st.session_state.user_tools,
+                    format_func=lambda x: x["tool_name"],
+                    key="r_tool_names",
+                )
+                st.button("Delete", on_click=remove_kb)
             with tab_file:
-                st.file_uploader("Upload files", key="uploaded_files", accept_multiple_files=True)
+                st.info(
+                    (
+                        "We adopted [Unstructured API](https://unstructured.io/api-key) "
+                        "here and we only store the processed texts from your documents. "
+                        "For privacy concerns, please refer to "
+                        "[our policy issue](https://myscale.com/privacy/)."
+                    ),
+                    icon="📃",
+                )
+                st.file_uploader(
+                    "Upload files", key="uploaded_files", accept_multiple_files=True
+                )
                 st.markdown("### Uploaded Files")
-                st.dataframe(st.session_state.private_kb.list_files(st.session_state.user_name))
+                st.dataframe(
+                    st.session_state.private_kb.list_files(st.session_state.user_name),
+                    use_container_width=True,
+                )
                 col_1, col_2 = st.columns(2)
                 with col_1:
                     st.button("Add Files", on_click=add_file)
                 with col_2:
-                    st.button("Clear Files", on_click=clear_files)
-            # with tab_build:
-            #     st.text_input("Give this knowledge base a description:")
-            #     col_3, col_4 = st.columns(2)
-            #     with col_3:
-            #         st.button("Build Your KB!")
-            #     with col_4:
-            #         st.button("Delete Your KB")
-                
-            
+                    st.button("Clear Files and All Tools", on_click=clear_files)
+
         st.button("Clear Chat History", on_click=clear_history)
         st.button("Logout", on_click=back_to_main)
-    if 'agent' not in st.session_state:
+    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:
@@ -255,7 +383,10 @@ def chat_page():
                 st.write("Retrieved from knowledge base:")
                 try:
                     st.dataframe(
-                        pd.DataFrame.from_records(map(dict, eval(msg.content)))
+                        pd.DataFrame.from_records(
+                            json.loads(msg.content, cls=CustomJSONDecoder)
+                        ),
+                        use_container_width=True,
                     )
                 except:
                     st.write(msg.content)

helper.py → lib/helper.py

@@ -49,10 +49,12 @@ 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 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
+from .json_conv import CustomJSONEncoder
+
 environ['TOKENIZERS_PARALLELISM'] = 'true'
 environ['OPENAI_API_BASE'] = st.secrets['OPENAI_API_BASE']
 
@@ -495,7 +497,7 @@ def create_retriever_tool(
     def wrap(func):
         def wrapped_retrieve(*args, **kwargs):
             docs: List[Document] = func(*args, **kwargs)
-            return json.dumps([d.dict() for d in docs])
+            return json.dumps([d.dict() for d in docs], cls=CustomJSONEncoder)
         return wrapped_retrieve
     
     return Tool(
@@ -533,12 +535,13 @@ def build_agents(session_id, tool_names, chat_model_name=chat_model_name, temper
     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]
+    tools = st.session_state.tools if "tools_with_users" not in st.session_state else st.session_state.tools_with_users
+    sel_tools = [tools[k] for k in tool_names]
     agent = create_agent_executor(
         "chat_memory",
         session_id,
         chat_llm,
-        tools=tools,
+        tools=sel_tools,
         system_prompt=system_prompt
         )
     return agent

lib/json_conv.py

@@ -0,0 +1,21 @@
+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

lib/private_kb.py

@@ -1,18 +1,19 @@
 import pandas as pd
 import hashlib
 import requests
-from typing import List
+from typing import List, Optional
 from datetime import datetime
 from langchain.schema.embeddings import Embeddings
 from streamlit.runtime.uploaded_file_manager import UploadedFile
 from clickhouse_connect import get_client
 from multiprocessing.pool import ThreadPool
 from langchain.vectorstores.myscale import MyScaleWithoutJSON, MyScaleSettings
+from .helper import create_retriever_tool
 
 parser_url = "https://api.unstructured.io/general/v0/general"
 
 
-def parse_files(api_key, user_id, files: List[UploadedFile], collection="default"):
+def parse_files(api_key, user_id, files: List[UploadedFile]):
     def parse_file(file: UploadedFile):
         headers = {
             "accept": "application/json",
@@ -31,9 +32,10 @@ def parse_files(api_key, user_id, files: List[UploadedFile], collection="default
             {
                 "text": t["text"],
                 "file_name": t["metadata"]["filename"],
-                "entity_id": hashlib.sha256((file_hash + t["text"]).encode()).hexdigest(),
+                "entity_id": hashlib.sha256(
+                    (file_hash + t["text"]).encode()
+                ).hexdigest(),
                 "user_id": user_id,
-                "collection_id": collection,
                 "created_by": datetime.now(),
             }
             for t in json_response
@@ -43,7 +45,7 @@ def parse_files(api_key, user_id, files: List[UploadedFile], collection="default
 
     with ThreadPool(8) as p:
         rows = []
-        for r in map(parse_file, files):
+        for r in p.imap_unordered(parse_file, files):
             rows.extend(r)
         return rows
 
@@ -68,21 +70,33 @@ class PrivateKnowledgeBase:
         parser_api_key,
         db="chat",
         kb_table="private_kb",
+        tool_table="private_tool",
     ) -> None:
         super().__init__()
-        schema_ = f"""
+        kb_schema_ = f"""
             CREATE TABLE IF NOT EXISTS {db}.{kb_table}(
                 entity_id String,
                 file_name String,
                 text String,
                 user_id String,
-                collection_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_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.kb_table = kb_table
+        self.tool_table = tool_table
         config = MyScaleSettings(
             host=host,
             port=port,
@@ -98,41 +112,101 @@ class PrivateKnowledgeBase:
             password=config.password,
         )
         client.command("SET allow_experimental_object_type=1")
-        client.command(schema_)
+        client.command(kb_schema_)
+        client.command(tool_schema_)
         self.parser_api_key = parser_api_key
         self.vstore = MyScaleWithoutJSON(
             embedding=embedding,
             config=config,
-            must_have_cols=["file_name", "text", "create_by"],
+            must_have_cols=["file_name", "text", "created_by"],
         )
-        self.retriever = self.vstore.as_retriever()
 
-    def list_files(self, user_id):
+    def list_files(self, user_id, tool_name=None):
         query = f"""
-        SELECT DISTINCT file_name FROM {self.vstore.config.database}.{self.vstore.config.table}
-        WHERE user_id = '{user_id}'
+        SELECT DISTINCT file_name, COUNT(entity_id) AS num_paragraph, 
+            arrayMax(arrayMap(x->length(x), groupArray(text))) AS max_chars
+        FROM {self.vstore.config.database}.{self.kb_table}
+        WHERE user_id = '{user_id}' GROUP BY file_name
         """
         return [r for r in self.vstore.client.query(query).named_results()]
 
     def add_by_file(
-        self, user_id, files: List[UploadedFile], collection="default", **kwargs
+        self, user_id, files: List[UploadedFile], **kwargs
     ):
-        data = parse_files(self.parser_api_key, user_id, files, collection=collection)
+        data = parse_files(self.parser_api_key, user_id, files)
         data = extract_embedding(self.vstore.embeddings, data)
         self.vstore.client.insert_df(
-            self.vstore.config.table,
+            self.kb_table,
             pd.DataFrame(data),
             database=self.vstore.config.database,
         )
 
     def clear(self, user_id):
         self.vstore.client.command(
-            f"DELETE FROM {self.vstore.config.database}.{self.vstore.config.table} "
+            f"DELETE FROM {self.vstore.config.database}.{self.kb_table} "
             f"WHERE user_id='{user_id}'"
         )
+        query = f"""DELETE FROM {self.vstore.config.database}.{self.tool_table} 
+                    WHERE user_id  = '{user_id}'"""
+        self.vstore.client.command(query)
 
-    def _get_relevant_documents(self, query, *args, **kwargs):
-        return self.retriever._get_relevant_documents(query, *args, **kwargs)
+    def create_tool(
+        self, user_id, tool_name, tool_description, files: Optional[List[str]] = None
+    ):
+        self.vstore.client.insert_df(
+            self.tool_table,
+            pd.DataFrame(
+                [
+                    {
+                        "tool_id": hashlib.sha256(
+                            (user_id + tool_name).encode("utf-8")
+                        ).hexdigest(),
+                        "tool_name": tool_name,
+                        "file_names": files,
+                        "user_id": user_id,
+                        "created_by": datetime.now(),
+                        "tool_description": tool_description,
+                    }
+                ]
+            ),
+            database=self.vstore.config.database,
+        )
 
-    async def _aget_relevant_documents(self, *args, **kwargs):
-        return self.retriever._aget_relevant_documents(*args, **kwargs)
+    def list_tools(self, user_id, tool_name=None):
+        extended_where = f"AND tool_name = '{tool_name}'" if tool_name else ""
+        query = f"""
+        SELECT tool_name, tool_description, length(file_names) 
+        FROM {self.vstore.config.database}.{self.tool_table}
+        WHERE user_id = '{user_id}' {extended_where}
+        """
+        return [r for r in self.vstore.client.query(query).named_results()]
+
+    def remove_tools(self, user_id, tool_names):
+        tool_names = ",".join([f"'{t}'" for t in tool_names])
+        query = f"""DELETE FROM {self.vstore.config.database}.{self.tool_table}
+                    WHERE user_id  = '{user_id}' AND tool_name IN [{tool_names}]"""
+        self.vstore.client.command(query)
+
+    def as_tools(self, user_id, tool_name=None):
+        tools = self.list_tools(user_id=user_id, tool_name=tool_name)
+        retrievers = {
+            t["tool_name"]: create_retriever_tool(
+                self.vstore.as_retriever(
+                    search_kwargs={
+                        "where_str": (
+                            f"user_id='{user_id}' "
+                            f"""AND file_name IN (
+                                SELECT arrayJoin(file_names) FROM (
+                                    SELECT file_names 
+                                    FROM {self.vstore.config.database}.{self.tool_table}
+                                    WHERE user_id = '{user_id}' AND tool_name = '{t['tool_name']}')
+                        )"""
+                        )
+                    },
+                ),
+                name=t["tool_name"],
+                description=t["tool_description"],
+            )
+            for t in tools
+        }
+        return retrievers

lib/sessions.py

@@ -8,7 +8,6 @@ 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
-from .private_kb import PrivateKnowledgeBase
 
 def get_sessions(engine, model_class, user_id):
     with orm.sessionmaker(engine)() as session: