add parse and private knowledge base

app.py

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

chat.py

@@ -4,6 +4,7 @@ from time import sleep
 import datetime
 import streamlit as st
 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
@@ -15,6 +16,7 @@ from helper import (
     MYSCALE_PORT,
     MYSCALE_USER,
     DEFAULT_SYSTEM_PROMPT,
+    UNSTRUCTURED_API,
 )
 from login import back_to_main
 
@@ -49,6 +51,10 @@ def back_to_main():
         del st.session_state.user_name
     if "jump_query_ask" in st.session_state:
         del st.session_state.jump_query_ask
+    if "sel_sess" in st.session_state:
+        del st.session_state.sel_sess
+    if "current_sessions" in st.session_state:
+        del st.session_state.current_sessions
 
 
 def on_session_change_submit():
@@ -87,6 +93,7 @@ def on_session_change_submit():
 
 def build_session_manager():
     return SessionManager(
+        st.session_state,
         host=MYSCALE_HOST,
         port=MYSCALE_PORT,
         username=MYSCALE_USER,
@@ -130,7 +137,23 @@ def refresh_agent():
             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 add_file():
+    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)
+    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)
 
 
 def chat_page():
@@ -139,7 +162,17 @@ def chat_page():
             "session_id": "default",
             "system_prompt": DEFAULT_SYSTEM_PROMPT,
         }
-    st.session_state["session_manager"] = build_session_manager()
+    if "private_kb" not in st.session_state:
+        st.session_state["private_kb"] = PrivateKnowledgeBase(
+            host=MYSCALE_HOST,
+            port=MYSCALE_PORT,
+            username=MYSCALE_USER,
+            password=MYSCALE_PASSWORD,
+            embedding=st.session_state.embeddings['Wikipedia'],
+            parser_api_key=UNSTRUCTURED_API,
+        )
+    if "session_manager" not in st.session_state:
+        st.session_state["session_manager"] = build_session_manager()
     with st.sidebar:
         with st.expander("Session Management"):
             if "current_sessions" not in st.session_state:
@@ -179,13 +212,34 @@ def chat_page():
         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=["Wikipedia + Self Querying"],
-                key="selected_tools",
-                on_change=refresh_agent,
-            )
+            st.session_state["tool_status"] = st.empty()
+            tab_kb, tab_file, tab_build = st.tabs(["Knowledge Bases", "File Upload", "KB Builder"])
+            with tab_kb:
+                st.multiselect(
+                    "Select a Knowledge Base Tool",
+                    st.session_state.tools.keys(),
+                    default=["Wikipedia + Self Querying"],
+                    key="selected_tools",
+                    on_change=refresh_agent,
+                )
+            with tab_file:
+                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))
+                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 Chat History", on_click=clear_history)
         st.button("Logout", on_click=back_to_main)
     if 'agent' not in st.session_state:

helper.py

@@ -2,7 +2,7 @@
 import json
 import time
 import hashlib
-from typing import Dict, Any, List
+from typing import Dict, Any, List, Tuple
 import re
 import pandas as pd
 from os import environ
@@ -67,6 +67,7 @@ MYSCALE_USER = st.secrets['MYSCALE_USER']
 MYSCALE_PASSWORD = st.secrets['MYSCALE_PASSWORD']
 MYSCALE_HOST = st.secrets['MYSCALE_HOST']
 MYSCALE_PORT = st.secrets['MYSCALE_PORT']
+UNSTRUCTURED_API = st.secrets['UNSTRUCTURED_API']
 
 COMBINE_PROMPT = ChatPromptTemplate.from_strings(
     string_messages=[(SystemMessagePromptTemplate, combine_prompt_template),
@@ -348,17 +349,19 @@ def build_qa_chain(_sel: str, retriever: BaseRetriever, name: str="Self-query")
     return chain
 
 @st.cache_resource
-def build_all() -> Dict[str, Any]:
+def build_all() -> Tuple[Dict[str, Any], Dict[str, Any]]:
     """build all resources
 
     :return: sel_map_obj
     :rtype: Dict[str, Any]
     """
     sel_map_obj = {}
+    embeddings = {}
     for k in sel_map:
-        st.session_state[f'emb_model_{k}'] = build_embedding_model(k)
+        embeddings[k] = build_embedding_model(k)
+        st.session_state[f'emb_model_{k}'] = embeddings[k]
         sel_map_obj[k] = build_chains_retrievers(k)
-    return sel_map_obj
+    return sel_map_obj, embeddings
 
 def create_message_model(table_name, DynamicBase):  # type: ignore
     """

lib/private_kb.py

@@ -0,0 +1,138 @@
+import pandas as pd
+import hashlib
+import requests
+from typing import List
+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
+
+parser_url = "https://api.unstructured.io/general/v0/general"
+
+
+def parse_files(api_key, user_id, files: List[UploadedFile], collection="default"):
+    def parse_file(file: UploadedFile):
+        headers = {
+            "accept": "application/json",
+            "unstructured-api-key": api_key,
+        }
+        data = {"strategy": "auto", "ocr_languages": ["eng"]}
+        file_hash = hashlib.sha256(file.read()).hexdigest()
+        file_data = {"files": (file.name, file.getvalue(), file.type)}
+        response = requests.post(
+            parser_url, headers=headers, data=data, files=file_data
+        )
+        json_response = response.json()
+        if response.status_code != 200:
+            raise ValueError(str(json_response))
+        texts = [
+            {
+                "text": t["text"],
+                "file_name": t["metadata"]["filename"],
+                "entity_id": hashlib.sha256((file_hash + t["text"]).encode()).hexdigest(),
+                "user_id": user_id,
+                "collection_id": collection,
+                "created_by": datetime.now(),
+            }
+            for t in json_response
+            if t["type"] == "NarrativeText" and len(t["text"].split(" ")) > 10
+        ]
+        return texts
+
+    with ThreadPool(8) as p:
+        rows = []
+        for r in map(parse_file, files):
+            rows.extend(r)
+        return rows
+
+
+def extract_embedding(embeddings: Embeddings, texts):
+    if len(texts) > 0:
+        embs = embeddings.embed_documents([t["text"] for _, t in enumerate(texts)])
+        for i, _ in enumerate(texts):
+            texts[i]["vector"] = embs[i]
+        return texts
+    raise ValueError("No texts extracted!")
+
+
+class PrivateKnowledgeBase:
+    def __init__(
+        self,
+        host,
+        port,
+        username,
+        password,
+        embedding: Embeddings,
+        parser_api_key,
+        db="chat",
+        kb_table="private_kb",
+    ) -> None:
+        super().__init__()
+        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
+        """
+        config = MyScaleSettings(
+            host=host,
+            port=port,
+            username=username,
+            password=password,
+            database=db,
+            table=kb_table,
+        )
+        client = get_client(
+            host=config.host,
+            port=config.port,
+            username=config.username,
+            password=config.password,
+        )
+        client.command("SET allow_experimental_object_type=1")
+        client.command(schema_)
+        self.parser_api_key = parser_api_key
+        self.vstore = MyScaleWithoutJSON(
+            embedding=embedding,
+            config=config,
+            must_have_cols=["file_name", "text", "create_by"],
+        )
+        self.retriever = self.vstore.as_retriever()
+
+    def list_files(self, user_id):
+        query = f"""
+        SELECT DISTINCT file_name FROM {self.vstore.config.database}.{self.vstore.config.table}
+        WHERE user_id = '{user_id}'
+        """
+        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
+    ):
+        data = parse_files(self.parser_api_key, user_id, files, collection=collection)
+        data = extract_embedding(self.vstore.embeddings, data)
+        self.vstore.client.insert_df(
+            self.vstore.config.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"WHERE user_id='{user_id}'"
+        )
+
+    def _get_relevant_documents(self, query, *args, **kwargs):
+        return self.retriever._get_relevant_documents(query, *args, **kwargs)
+
+    async def _aget_relevant_documents(self, *args, **kwargs):
+        return self.retriever._aget_relevant_documents(*args, **kwargs)

lib/sessions.py

@@ -3,10 +3,12 @@ try:
     from sqlalchemy.orm import declarative_base
 except ImportError:
     from sqlalchemy.ext.declarative import declarative_base
+from langchain.schema import BaseChatMessageHistory
 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:
@@ -20,7 +22,8 @@ def get_sessions(engine, model_class, user_id):
     return json.loads(result)
 
 class SessionManager:
-    def __init__(self, host, port, username, password, db='chat', sess_table='sessions', msg_table='chat_memory') -> None:
+    def __init__(self, session_state, 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())
@@ -28,6 +31,7 @@ class SessionManager:
         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)
+        self.session_state = session_state
 
     def list_sessions(self, user_id):
         with self.Session() as session:
@@ -63,6 +67,10 @@ class SessionManager:
     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()
+            # session.query(self.msg_model_class).where(self.msg_model_class.session_id==session_id).delete()
+        if "agent" in self.session_state:
+            self.session_state.agent.memory.chat_memory.clear()
+        if "file_analyzer" in self.session_state:
+            self.session_state.file_analyzer.clear_files()