add stats

main.py

@@ -7,6 +7,10 @@ from question import chat_with_doc
 from langchain.embeddings import HuggingFaceInferenceAPIEmbeddings
 from langchain.vectorstores import SupabaseVectorStore
 from supabase import Client, create_client
+from stats import add_usage
+from langchain.llms import HuggingFaceEndpoint
+from langchain.chains import ConversationalRetrievalChain
+from langchain.memory import ConversationBufferMemory
 
 supabase_url = st.secrets.SUPABASE_URL
 supabase_key = st.secrets.SUPABASE_KEY
@@ -19,7 +23,6 @@ username = st.secrets.username
 
 # embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key)
 
-
 embeddings = HuggingFaceInferenceAPIEmbeddings(
     api_key=hf_api_key,
     model_name="BAAI/bge-large-en-v1.5"
@@ -36,38 +39,71 @@ if anthropic_api_key:
     models += ["claude-v1", "claude-v1.3",
                "claude-instant-v1-100k", "claude-instant-v1.1-100k"]
 
-# Set the theme
-st.set_page_config(
-    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]"
-    }
-)
+if 'question' in st.query_params:
+    query = st.query_params['question']
+    model = "meta-llama/Llama-2-70b-chat-hf"
+    temp = 0.1
+    max_tokens = 500
+    add_usage(supabase, "api", "prompt" + query, {"model": model, "temperature": temp})
+    # print(st.session_state['max_tokens'])
+    endpoint_url = ("https://api-inference.huggingface.co/models/"+ model)
+    model_kwargs = {"temperature" : temp,
+                    "max_new_tokens" : 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
+    )
+    memory = ConversationBufferMemory(memory_key="chat_history", input_key='question', output_key='answer', return_messages=True)
+    qa = ConversationalRetrievalChain.from_llm(hf, retriever=vector_store.as_retriever(search_kwargs={"score_threshold": 0.8, "k": 4,"filter": {"user": username}}), memory=memory, return_source_documents=True)
+    model_response = qa({"question": query})
+    # print( model_response["answer"])
+    sources = model_response["source_documents"]
+    # print(sources)
+    if len(sources) > 0:
+        json = {"response": model_response["answer"]}
+        st.code(json, language="json")
+    else:
+        json = {"response": "I am sorry, I do not have enough information to provide an answer. If there is a public source of data that you would like to add, please email [email protected]."}
+        st.code(json, language="json")
+    memory.clear()
+else:
+    # Set the theme
+    st.set_page_config(
+        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("👷‍♂️ Safety Copilot 🦺")
 
-st.title("👷‍♂️ Safety Copilot 🦺")
-st.markdown("Chat with your personal assistant about health and safety information.")
+    st.markdown("Chat with your personal safety assistant about any health & safety related queries.")
+    st.markdown("Up-to-date with latest OSH regulations for Singapore, Indonesia, Malaysia & other parts of Asia.")
 
-st.markdown("---\n\n")
+    st.markdown("---\n\n")
 
-# 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
+    # 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)
+    chat_with_doc(st.session_state['model'], vector_store, stats_db=supabase)
 
-st.markdown("---\n\n")
+    st.markdown("---\n\n")

question.py

@@ -7,7 +7,7 @@ from langchain.llms import OpenAI
 from langchain.llms import HuggingFaceEndpoint
 from langchain.chat_models import ChatAnthropic
 from langchain.vectorstores import SupabaseVectorStore
-from stats import add_usage
+from stats import add_usage, get_usage
 
 memory = ConversationBufferMemory(memory_key="chat_history", input_key='question', output_key='answer', return_messages=True)
 openai_api_key = st.secrets.openai_api_key
@@ -15,13 +15,13 @@ anthropic_api_key = st.secrets.anthropic_api_key
 hf_api_key = st.secrets.hf_api_key
 logger = get_logger(__name__)
 
-
 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")
+      
+    stats = str(get_usage(stats_db)) 
+    question = st.text_area("## Ask a question (" + stats + " queries answered so far)", max_chars=500)
     columns = st.columns(2)
     with columns[0]:
         button = st.button("Ask")
@@ -62,16 +62,21 @@ def chat_with_doc(model, vector_store: SupabaseVectorStore, stats_db):
                 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)
+            qa = ConversationalRetrievalChain.from_llm(hf, retriever=vector_store.as_retriever(search_kwargs={"score_threshold": 0.8, "k": 4,"filter": {"user": st.session_state["username"]}}), 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"])
-
-        st.session_state['chat_history'].append(("Safety Copilot", model_response["answer"]))
+        sources = model_response["source_documents"]
         logger.info('Sources: %s', model_response["source_documents"])
+
+        if len(sources) > 0:
+            st.session_state['chat_history'].append(("Safety Copilot", model_response["answer"]))
+        else:
+            st.session_state['chat_history'].append(("Safety Copilot", "I am sorry, I do not have enough information to provide an answer. If there is a public source of data that you would like to add, please email [email protected]."))
+        
         
     # Display chat history
     st.empty()

requirements.txt

@@ -3,7 +3,7 @@ Markdown==3.4.3
 openai==0.27.6
 pdf2image==1.16.3
 pypdf==3.8.1
-streamlit==1.22.0
+streamlit==1.30.0
 StrEnum==0.4.10
 supabase==1.0.3
 tiktoken==0.4.0

stats.py

@@ -29,3 +29,8 @@ def add_usage(supabase, type, details, metadata):
         "details": details,
         "metadata": metadata
     }).execute()
+    
+def get_usage(supabase):
+    # Returns the number of rows in the stats table for the last 24 hours
+    response = supabase.table("stats").select("id", count="exact").execute()
+    return response.count