Update main.py

main.py

@@ -1,81 +1,152 @@
-# main.py
 import os
 import streamlit as st
-import anthropic
-
+import logging
+from requests.exceptions import JSONDecodeError
 from langchain_community.embeddings import HuggingFaceInferenceAPIEmbeddings
-from langchain_community.embeddings import HuggingFaceBgeEmbeddings
 from langchain_community.vectorstores import SupabaseVectorStore
 from langchain_community.llms import HuggingFaceEndpoint
-from langchain_community.vectorstores import SupabaseVectorStore
-
 from langchain.chains import ConversationalRetrievalChain
 from langchain.memory import ConversationBufferMemory
-
 from supabase import Client, create_client
 from streamlit.logger import get_logger
-from stats import get_usage, add_usage
 
-supabase_url = st.secrets.SUPABASE_URL
-supabase_key = st.secrets.SUPABASE_KEY
-openai_api_key = st.secrets.openai_api_key
-anthropic_api_key = st.secrets.anthropic_api_key
-hf_api_key = st.secrets.hf_api_key
-username = st.secrets.username
+# Configure logging
+logger = get_logger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# Load secrets
+supabase_url = st.secrets["SUPABASE_URL"]
+supabase_key = st.secrets["SUPABASE_KEY"]
+hf_api_key = st.secrets["hf_api_key"]
+username = st.secrets["username"]
 
+# Initialize Supabase client
 supabase: Client = create_client(supabase_url, supabase_key)
-logger = get_logger(__name__)
 
-embeddings = HuggingFaceInferenceAPIEmbeddings(
+# Custom HuggingFaceInferenceAPIEmbeddings to handle JSONDecodeError
+class CustomHuggingFaceInferenceAPIEmbeddings(HuggingFaceInferenceAPIEmbeddings):
+    def embed_query(self, text: str):
+        try:
+            response = self.client.post(
+                json={"inputs": text, "options": {"use_cache": False}},
+                task="feature-extraction",
+            )
+            if response.status_code != 200:
+                logger.error(f"API request failed with status {response.status_code}: {response.text}")
+                return [0.0] * 384  # Return zero vector of expected dimension
+            try:
+                embeddings = response.json()
+                if not isinstance(embeddings, list) or not embeddings:
+                    logger.error(f"Invalid embeddings response: {embeddings}")
+                    return [0.0] * 384
+                return embeddings[0]
+            except JSONDecodeError as e:
+                logger.error(f"JSON decode error: {str(e)}, response: {response.text}")
+                return [0.0] * 384
+        except Exception as e:
+            logger.error(f"Error embedding query: {str(e)}")
+            return [0.0] * 384
+
+    def embed_documents(self, texts):
+        try:
+            response = self.client.post(
+                json={"inputs": texts, "options": {"use_cache": False}},
+                task="feature-extraction",
+            )
+            if response.status_code != 200:
+                logger.error(f"API request failed with status {response.status_code}: {response.text}")
+                return [[0.0] * 384 for _ in texts]
+            try:
+                embeddings = response.json()
+                if not isinstance(embeddings, list) or not embeddings:
+                    logger.error(f"Invalid embeddings response: {embeddings}")
+                    return [[0.0] * 384 for _ in texts]
+                return [emb[0] for emb in embeddings]
+            except JSONDecodeError as e:
+                logger.error(f"JSON decode error: {str(e)}, response: {response.text}")
+                return [[0.0] * 384 for _ in texts]
+        except Exception as e:
+            logger.error(f"Error embedding documents: {str(e)}")
+            return [[0.0] * 384 for _ in texts]
+
+# Initialize embeddings
+embeddings = CustomHuggingFaceInferenceAPIEmbeddings(
     api_key=hf_api_key,
     model_name="BAAI/bge-large-en-v1.5",
-    api_url="https://router.huggingface.co/hf-inference/pipeline/feature-extraction/",
 )
 
-if 'chat_history' not in st.session_state:
-    st.session_state['chat_history'] = []
+# Initialize session state
+if "chat_history" not in st.session_state:
+    st.session_state["chat_history"] = []
 
-vector_store = SupabaseVectorStore(supabase, embeddings, query_name='match_documents', table_name="documents")
-memory = ConversationBufferMemory(memory_key="chat_history", input_key='question', output_key='answer', return_messages=True)
+# Initialize vector store and memory
+vector_store = SupabaseVectorStore(
+    client=supabase,
+    embedding=embeddings,
+    query_name="match_documents",
+    table_name="documents",
+)
+memory = ConversationBufferMemory(
+    memory_key="chat_history",
+    input_key="question",
+    output_key="answer",
+    return_messages=True,
+)
 
+# Model configuration
 model = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-
 temperature = 0.1
 max_tokens = 500
-stats = str(get_usage(supabase)) 
+
+# Mock stats function (replace with your actual implementation)
+def get_usage(supabase):
+    return 100  # Replace with actual logic
+
+def add_usage(supabase, action, prompt, metadata):
+    pass  # Replace with actual logic
+
+stats = str(get_usage(supabase))
 
 def response_generator(query):
-    qa = None
-    add_usage(supabase, "chat", "prompt" + query, {"model": model, "temperature": temperature})
-    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" : temperature,
-                    "max_new_tokens" : max_tokens,
-                    # "repetition_penalty" : 1.1,
-                    "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": username}}), memory=memory, verbose=True, return_source_documents=True)
-    
-    # Generate model's response 
-    model_response = qa({"question": query})
-    logger.info('Result: %s', model_response["answer"])
-    sources = model_response["source_documents"]
-    logger.info('Sources: %s', model_response["source_documents"])
-
-    if len(sources) > 0:
-        response = model_response["answer"]
-    else:
-        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]."
-    
-    return response
-    
-# Set the theme
+    try:
+        add_usage(supabase, "chat", f"prompt: {query}", {"model": model, "temperature": temperature})
+        logger.info("Using HF model %s", model)
+        
+        endpoint_url = f"https://api-inference.huggingface.co/models/{model}"
+        model_kwargs = {
+            "temperature": temperature,
+            "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,
+        )
+        qa = ConversationalRetrievalChain.from_llm(
+            llm=hf,
+            retriever=vector_store.as_retriever(search_kwargs={"score_threshold": 0.6, "k": 4, "filter": {"user": username}}),
+            memory=memory,
+            verbose=True,
+            return_source_documents=True,
+        )
+        
+        # Use invoke instead of deprecated __call__
+        model_response = qa.invoke({"question": query})
+        logger.info("Result: %s", model_response["answer"])
+        sources = model_response["source_documents"]
+        logger.info("Sources: %s", sources)
+
+        if sources:
+            return model_response["answer"]
+        else:
+            return "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]."
+    except Exception as e:
+        logger.error(f"Error generating response: {str(e)}")
+        return "An error occurred while processing your request. Please try again later."
+
+# Streamlit UI
 st.set_page_config(
     page_title="Securade.ai - Safety Copilot",
     page_icon="https://securade.ai/favicon.ico",
@@ -83,55 +154,33 @@ st.set_page_config(
     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]"
-    }
+        "Get Help": "https://securade.ai",
+        "Report a Bug": "mailto:[email protected]",
+    },
 )
 
 st.title("👷‍♂️ Safety Copilot 🦺")
+st.markdown(
+    "Chat with your personal safety assistant about any health & safety related queries. "
+    "[[blog](https://securade.ai/blog/how-securade-ai-safety-copilot-transforms-worker-safety.html)|"
+    "[paper](https://securade.ai/assets/pdfs/Securade.ai-Safety-Copilot-Whitepaper.pdf)]"
+)
+st.markdown(f"_{stats} queries answered!_")
 
-st.markdown("Chat with your personal safety assistant about any health & safety related queries. [[blog](https://securade.ai/blog/how-securade-ai-safety-copilot-transforms-worker-safety.html)|[paper](https://securade.ai/assets/pdfs/Securade.ai-Safety-Copilot-Whitepaper.pdf)]")
-# st.markdown("Up-to-date with latest OSH regulations for Singapore, Indonesia, Malaysia & other parts of Asia.")
-st.markdown("_"+ stats + " queries answered!_")
-
-if 'chat_history' not in st.session_state:
-    st.session_state['chat_history'] = []
-    
-# Display chat messages from history on app rerun
+# Display chat history
 for message in st.session_state.chat_history:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
-        
-# Accept user input
-if prompt := st.chat_input("Ask a question"): 
-    # print(prompt)
-    # Add user message to chat history
+
+# Handle user input
+if prompt := st.chat_input("Ask a question"):
     st.session_state.chat_history.append({"role": "user", "content": prompt})
-    # Display user message in chat message container
     with st.chat_message("user"):
         st.markdown(prompt)
     
-    with st.spinner('Safety briefing in progress...'):
+    with st.spinner("Safety briefing in progress..."):
         response = response_generator(prompt)
     
-    # Display assistant response in chat message container
     with st.chat_message("assistant"):
         st.markdown(response)
-    # Add assistant response to chat history
-    # print(response)
-    st.session_state.chat_history.append({"role": "assistant", "content": response})
-
-# query = 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")
-# with columns[1]:
-#     clear_history = st.button("Clear History", type='secondary')
-    
-# st.markdown("---\n\n")
-
-# if clear_history:
-#     # Clear memory in Langchain
-#     memory.clear()
-#     st.session_state['chat_history'] = []
-#     st.experimental_rerun()
+    st.session_state.chat_history.append({"role": "assistant", "content": response})