Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import os
 import torch
-from langchain_community.vectorstores import FAISS, Chroma
+from langchain_community.vectorstores import Chroma
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.chains import ConversationalRetrievalChain
@@ -10,16 +10,16 @@ from langchain_community.llms import HuggingFaceEndpoint
 from langchain.memory import ConversationBufferMemory
 from langchain_community.retrievers import BM25Retriever
 from langchain.retrievers import EnsembleRetriever
-from langchain.retrievers.multi_query import MultiQueryRetriever
 
 # Environment variable for API token
-api_token = os.getenv("API_TOKEN")
+api_token = os.getenv("FirstToken")
+print(f"API Token loaded: {api_token[:5]}...")  # Debug: Show first 5 chars of token
 if not api_token:
     raise ValueError("Environment variable 'FirstToken' not set. Please set the Hugging Face API token.")
 
 # Available LLM models
 list_llm = [
-    "meta-llama/Meta-Llama-3-8B-Instruct",
+    "mistralai/Mixtral-8x7B-Instruct-v0.1",  # Publicly accessible
     "mistralai/Mistral-7B-Instruct-v0.2",
     "deepseek-ai/deepseek-llm-7b-chat"
 ]
@@ -55,11 +55,6 @@ def create_chromadb(splits, persist_directory="chroma_db"):
     )
     return chromadb
 
-def create_faissdb(splits):
-    """Create FAISS vector database from document splits."""
-    embeddings = HuggingFaceEmbeddings()
-    return FAISS.from_documents(splits, embeddings)
-
 # -----------------------------------------------------------------------------
 # Retrievers
 # -----------------------------------------------------------------------------
@@ -96,7 +91,11 @@ def initialize_database(list_file_obj, progress=gr.Progress()):
 # -----------------------------------------------------------------------------
 def initialize_llmchain(llm_model, temperature, max_tokens, top_k, retriever):
     """Initialize the language model chain with error handling."""
+    if retriever is None:
+        raise ValueError("Retriever is None. Please process documents first.")
+    
     try:
+        print(f"Initializing LLM: {llm_model} with token: {api_token[:5]}...")
         llm = HuggingFaceEndpoint(
             repo_id=llm_model,
             huggingfacehub_api_token=api_token,
@@ -127,6 +126,9 @@ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, retriever):
 # -----------------------------------------------------------------------------
 def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, retriever, progress=gr.Progress()):
     """Initialize the Language Model."""
+    if retriever is None:
+        return None, "Error: No database initialized. Please process documents first."
+    
     try:
         llm_name = list_llm[llm_option]
         print(f"Selected LLM model: {llm_name}")
@@ -150,7 +152,6 @@ def conversation(qa_chain, message, history, lang):
     if not qa_chain:
         return None, gr.update(value="Assistant not initialized"), history, "", 0, "", 0, "", 0
 
-    # Add language instruction
     lang_instruction = " (Responda em Português)" if lang == "pt" else " (Respond in English)"
     query = message + lang_instruction
     
@@ -159,13 +160,11 @@ def conversation(qa_chain, message, history, lang):
         response = qa_chain.invoke({"question": query, "chat_history": formatted_chat_history})
         answer = response["answer"].split("Helpful Answer:")[-1].strip() if "Helpful Answer:" in response["answer"] else response["answer"]
 
-        # Extract sources (handle cases where fewer than 3 documents are returned)
         sources = response["source_documents"]
         source_data = [("Unknown", 0)] * 3
         for i, doc in enumerate(sources[:3]):
             source_data[i] = (doc.page_content.strip(), doc.metadata["page"] + 1)
 
-        # Update history without the language instruction
         new_history = history + [(message, answer)]
         return (
             qa_chain, gr.update(value=""), new_history,
@@ -214,7 +213,7 @@ def demo():
                         slider_temperature = gr.Slider(0.01, 1.0, value=0.5, step=0.1, label="Analysis Precision")
                         slider_maxtokens = gr.Slider(128, 9192, value=4096, step=128, label="Response Length")
                         slider_topk = gr.Slider(1, 10, value=3, step=1, label="Analysis Diversity")
-                    qachain_btn = gr.Button("Initialize Assistant")
+                    qachain_btn = gr.Button("Initialize Assistant", interactive=False)  # Disabled by default
                     llm_progress = gr.Textbox(value="Not initialized", label="Assistant Status")
 
             with gr.Column(scale=2):
@@ -232,10 +231,36 @@ def demo():
 
         # Event Handlers
         language_btn.change(lambda x: "en" if x == "English" else "pt", inputs=language_btn, outputs=language)
-        db_btn.click(initialize_database, inputs=[document], outputs=[retriever, db_progress])
-        qachain_btn.click(initialize_LLM, inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, retriever], outputs=[qa_chain, llm_progress])
-        submit_btn.click(conversation, inputs=[qa_chain, msg, chatbot, language], outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page])
-        msg.submit(conversation, inputs=[qa_chain, msg, chatbot, language], outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page])
+
+        def enable_qachain_btn(retriever, status):
+            return gr.update(interactive=retriever is not None and "successfully" in status)
+
+        db_btn.click(
+            initialize_database,
+            inputs=[document],
+            outputs=[retriever, db_progress]
+        ).then(
+            enable_qachain_btn,
+            inputs=[retriever, db_progress],
+            outputs=[qachain_btn]
+        )
+
+        qachain_btn.click(
+            initialize_LLM,
+            inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, retriever],
+            outputs=[qa_chain, llm_progress]
+        )
+
+        submit_btn.click(
+            conversation,
+            inputs=[qa_chain, msg, chatbot, language],
+            outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page]
+        )
+        msg.submit(
+            conversation,
+            inputs=[qa_chain, msg, chatbot, language],
+            outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page]
+        )
 
     demo.launch(debug=True)