update co pilot

*.tar.* filter=lfs diff=lfs merge=lfs -text

main.py

@@ -3,15 +3,10 @@ import os
 import tempfile
 
 import streamlit as st
-from files import file_uploader, url_uploader
 from question import chat_with_doc
-from brain import brain
 from langchain.embeddings import HuggingFaceInferenceAPIEmbeddings
 from langchain.vectorstores import SupabaseVectorStore
 from supabase import Client, create_client
-from explorer import view_document
-from stats import get_usage_today
-from st_login_form import login_form
 
 supabase_url = st.secrets.SUPABASE_URL
 supabase_key = st.secrets.SUPABASE_KEY
@@ -20,6 +15,7 @@ anthropic_api_key = st.secrets.anthropic_api_key
 hf_api_key = st.secrets.hf_api_key
 supabase: Client = create_client(supabase_url, supabase_key)
 self_hosted = st.secrets.self_hosted
+username = st.secrets.username
 
 # embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key)
 
@@ -42,103 +38,36 @@ if anthropic_api_key:
 
 # Set the theme
 st.set_page_config(
-    page_title="meraKB",
-    layout="wide",
-    initial_sidebar_state="expanded",
+    page_title="Securade.ai - Safety Copilot",
+    page_icon="https://securade.ai/favicon.ico",
+    layout="centered",
+    initial_sidebar_state="collapsed",
+    menu_items={
+        "About": "# Securade.ai Safety Copilot v0.1\n [https://securade.ai](https://securade.ai)",
+        "Get Help" : "https://securade.ai",
+        "Report a Bug": "mailto:[email protected]"
+    }
 )
 
-
-st.title("🧠 meraKB - Your digital brain 🧠")
-st.markdown("Store your knowledge in a vector store and chat with it.")
-if self_hosted == "false":
-    st.markdown('**📢 Note: In the public demo, access to functionality is restricted. You can only use the GPT-3.5-turbo model and upload files up to 1Mb. To use more models and upload larger files, consider self-hosting meraKB.**')
+st.title("👷‍♂️ Safety Copilot 🦺")
+st.markdown("Chat with your personal assistant about health and safety information.")
 
 st.markdown("---\n\n")
 
-st.session_state["overused"] = False
-if self_hosted == "false":
-    usage = get_usage_today(supabase)
-    if usage > st.secrets.usage_limit:
-        st.markdown(
-            f"<span style='color:red'>You have used {usage} tokens today, which is more than your daily limit of {st.secrets.usage_limit} tokens. Please come back later or consider self-hosting.</span>", unsafe_allow_html=True)
-        st.session_state["overused"] = True
-    else:
-        st.markdown(f"<span style='color:blue'>Usage today: {usage} tokens out of {st.secrets.usage_limit}</span>", unsafe_allow_html=True)
-    st.write("---")
-
-client = login_form()
-
-if st.session_state["authenticated"]:
-    if st.session_state["username"]:
-        st.success(f"Welcome {st.session_state['username']}")
-    else:
-        st.session_state["username"] = 'guest'
-        st.success("Welcome guest")
-
-    # Initialize session state variables
-    if 'model' not in st.session_state:
-        st.session_state['model'] = "meta-llama/Llama-2-70b-chat-hf"
-    if 'temperature' not in st.session_state:
-        st.session_state['temperature'] = 0.1
-    if 'chunk_size' not in st.session_state:
-        st.session_state['chunk_size'] = 500
-    if 'chunk_overlap' not in st.session_state:
-        st.session_state['chunk_overlap'] = 0
-    if 'max_tokens' not in st.session_state:
-        st.session_state['max_tokens'] = 500
-
-    # Create a radio button for user to choose between adding knowledge or asking a question
-    user_choice = st.radio(
-        "Choose an action", ('Add Knowledge', 'Chat with your Brain', 'Forget', "Explore"))
-
-    st.markdown("---\n\n")
-
-    if user_choice == 'Add Knowledge':
-        # Display chunk size and overlap selection only when adding knowledge
-        st.sidebar.title("Configuration")
-        st.sidebar.markdown(
-            "Choose your chunk size and overlap for adding knowledge.")
-        st.session_state['chunk_size'] = st.sidebar.slider(
-            "Select Chunk Size", 100, 1000, st.session_state['chunk_size'], 50)
-        st.session_state['chunk_overlap'] = st.sidebar.slider(
-            "Select Chunk Overlap", 0, 100, st.session_state['chunk_overlap'], 10)
-        
-        # Create two columns for the file uploader and URL uploader
-        col1, col2 = st.columns(2)
-        
-        with col1:
-            file_uploader(supabase, vector_store)
-        with col2:
-            url_uploader(supabase, vector_store)
-    elif user_choice == 'Chat with your Brain':
-        # Display model and temperature selection only when asking questions
-        st.sidebar.title("Configuration")
-        st.sidebar.markdown(
-            "Choose your model and temperature for asking questions.")
-        if self_hosted != "false":
-            st.session_state['model'] = st.sidebar.selectbox(
-            "Select Model", models, index=(models).index(st.session_state['model']))
-        else:
-            st.sidebar.write("**Model**: gpt-3.5-turbo")
-            st.sidebar.write("**Self Host to unlock more models such as claude-v1 and GPT4**")
-            st.session_state['model'] = "gpt-3.5-turbo"
-        st.session_state['temperature'] = st.sidebar.slider(
-            "Select Temperature", 0.1, 1.0, st.session_state['temperature'], 0.1)
-        if st.secrets.self_hosted != "false":
-            st.session_state['max_tokens'] = st.sidebar.slider(
-                "Select Max Tokens", 500, 4000, st.session_state['max_tokens'], 500)
-        else:
-            st.session_state['max_tokens'] = 500
-        
-        chat_with_doc(st.session_state['model'], vector_store, stats_db=supabase)
-    elif user_choice == 'Forget':
-        st.sidebar.title("Configuration")
-
-        brain(supabase)
-    elif user_choice == 'Explore':
-        st.sidebar.title("Configuration")
-        view_document(supabase)
-
-    st.markdown("---\n\n")
-else:
-    st.error("Not authenticated")
+# Initialize session state variables
+if 'model' not in st.session_state:
+    st.session_state['model'] = "meta-llama/Llama-2-70b-chat-hf"
+if 'temperature' not in st.session_state:
+    st.session_state['temperature'] = 0.1
+if 'chunk_size' not in st.session_state:
+    st.session_state['chunk_size'] = 500
+if 'chunk_overlap' not in st.session_state:
+    st.session_state['chunk_overlap'] = 0
+if 'max_tokens' not in st.session_state:
+    st.session_state['max_tokens'] = 500
+if 'username' not in st.session_state:
+    st.session_state['username'] = username
+
+chat_with_doc(st.session_state['model'], vector_store, stats_db=supabase)
+
+st.markdown("---\n\n")

question.py

@@ -16,30 +16,19 @@ hf_api_key = st.secrets.hf_api_key
 logger = get_logger(__name__)
 
 
-def count_tokens(question, model):
-    count = f'Words: {len(question.split())}'
-    if model.startswith("claude"):
-        count += f' | Tokens: {anthropic.count_tokens(question)}'
-    return count
-
-
 def chat_with_doc(model, vector_store: SupabaseVectorStore, stats_db):
     
     if 'chat_history' not in st.session_state:
         st.session_state['chat_history'] = []
         
-    
-    
     question = st.text_area("## Ask a question")
-    columns = st.columns(3)
+    columns = st.columns(2)
     with columns[0]:
         button = st.button("Ask")
     with columns[1]:
-        count_button = st.button("Count Tokens", type='secondary')
-    with columns[2]:
         clear_history = st.button("Clear History", type='secondary')
-    
-    
+        
+    st.markdown("---\n\n")
     
     if clear_history:
         # Clear memory in Langchain
@@ -49,48 +38,43 @@ def chat_with_doc(model, vector_store: SupabaseVectorStore, stats_db):
 
     if button:
         qa = None
-        if not st.session_state["overused"]:
-            add_usage(stats_db, "chat", "prompt" + question, {"model": model, "temperature": st.session_state['temperature']})
-            if model.startswith("gpt"):
-                logger.info('Using OpenAI model %s', model)
-                qa = ConversationalRetrievalChain.from_llm(
-                    OpenAI(
-                        model_name=st.session_state['model'], openai_api_key=openai_api_key, temperature=st.session_state['temperature'], max_tokens=st.session_state['max_tokens']), vector_store.as_retriever(), memory=memory, verbose=True)
-            elif anthropic_api_key and model.startswith("claude"):
-                logger.info('Using Anthropics model %s', model)
-                qa = ConversationalRetrievalChain.from_llm(
-                    ChatAnthropic(
-                        model=st.session_state['model'], anthropic_api_key=anthropic_api_key, temperature=st.session_state['temperature'], max_tokens_to_sample=st.session_state['max_tokens']), vector_store.as_retriever(), memory=memory, verbose=True, max_tokens_limit=102400)
-            elif hf_api_key:
-                logger.info('Using HF model %s', model)
-                # print(st.session_state['max_tokens'])
-                endpoint_url = ("https://api-inference.huggingface.co/models/"+ model)
-                model_kwargs = {"temperature" : st.session_state['temperature'],
-                                "max_new_tokens" : st.session_state['max_tokens'],
-                                "return_full_text" : False}
-                hf = HuggingFaceEndpoint(
-                    endpoint_url=endpoint_url,
-                    task="text-generation",
-                    huggingfacehub_api_token=hf_api_key,
-                    model_kwargs=model_kwargs
-                )
-                qa = ConversationalRetrievalChain.from_llm(hf, retriever=vector_store.as_retriever(), memory=memory, verbose=True, return_source_documents=True)
-            
-            st.session_state['chat_history'].append(("You", question))
-
-            # Generate model's response and add it to chat history
-            model_response = qa({"question": question})
-            logger.info('Result: %s', model_response["answer"])
+        add_usage(stats_db, "chat", "prompt" + question, {"model": model, "temperature": st.session_state['temperature']})
+        if model.startswith("gpt"):
+            logger.info('Using OpenAI model %s', model)
+            qa = ConversationalRetrievalChain.from_llm(
+                OpenAI(
+                    model_name=st.session_state['model'], openai_api_key=openai_api_key, temperature=st.session_state['temperature'], max_tokens=st.session_state['max_tokens']), vector_store.as_retriever(), memory=memory, verbose=True)
+        elif anthropic_api_key and model.startswith("claude"):
+            logger.info('Using Anthropics model %s', model)
+            qa = ConversationalRetrievalChain.from_llm(
+                ChatAnthropic(
+                    model=st.session_state['model'], anthropic_api_key=anthropic_api_key, temperature=st.session_state['temperature'], max_tokens_to_sample=st.session_state['max_tokens']), vector_store.as_retriever(), memory=memory, verbose=True, max_tokens_limit=102400)
+        elif hf_api_key:
+            logger.info('Using HF model %s', model)
+            # print(st.session_state['max_tokens'])
+            endpoint_url = ("https://api-inference.huggingface.co/models/"+ model)
+            model_kwargs = {"temperature" : st.session_state['temperature'],
+                            "max_new_tokens" : st.session_state['max_tokens'],
+                            "return_full_text" : False}
+            hf = HuggingFaceEndpoint(
+                endpoint_url=endpoint_url,
+                task="text-generation",
+                huggingfacehub_api_token=hf_api_key,
+                model_kwargs=model_kwargs
+            )
+            qa = ConversationalRetrievalChain.from_llm(hf, retriever=vector_store.as_retriever(search_kwargs={"score_threshold": 0.6, "k": 4,"filter": {"user": st.session_state["username"]}}), memory=memory, verbose=True, return_source_documents=True)
+        
+        st.session_state['chat_history'].append(("You", question))
 
-            st.session_state['chat_history'].append(("meraKB", model_response["answer"]))
-            logger.info('Sources: %s', model_response["source_documents"])
+        # Generate model's response and add it to chat history
+        model_response = qa({"question": question})
+        logger.info('Result: %s', model_response["answer"])
 
-            # Display chat history
-            st.empty()
-            for speaker, text in st.session_state['chat_history']:
-                st.markdown(f"**{speaker}:** {text}")
-        else:
-            st.error("You have used all your free credits. Please try again later or self host.")
+        st.session_state['chat_history'].append(("Safety Copilot", model_response["answer"]))
+        logger.info('Sources: %s', model_response["source_documents"])
         
-    if count_button:
-        st.write(count_tokens(question, model))
+    # Display chat history
+    st.empty()
+    chat_history = st.session_state['chat_history']
+    for speaker, text in chat_history:
+        st.markdown(f"**{speaker}:** {text}")