Update app.py

app.py

@@ -33,7 +33,8 @@ print(f"CLOUDFLARE_AUTH_TOKEN: {API_TOKEN[:5]}..." if API_TOKEN else "Not set")
 MODELS = [
     "mistralai/Mistral-7B-Instruct-v0.3",
     "mistralai/Mixtral-8x7B-Instruct-v0.1",
-    "@cf/meta/llama-3.1-8b-instruct"
+    "@cf/meta/llama-3.1-8b-instruct",
+    "mistralai/Mistral-Nemo-Instruct-2407"
 ]
 
 # Initialize LlamaParse
@@ -63,32 +64,60 @@ def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[
         raise ValueError("Invalid parser specified. Use 'pypdf' or 'llamaparse'.")
 
 def get_embeddings():
-    return HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+    return HuggingFaceEmbeddings(model_name="sentence-transformers/stsb-roberta-large")
 
 def update_vectors(files, parser):
+    global uploaded_documents
+    logging.info(f"Entering update_vectors with {len(files)} files and parser: {parser}")
+    
     if not files:
-        return "Please upload at least one PDF file."
+        logging.warning("No files provided for update_vectors")
+        return "Please upload at least one PDF file.", gr.CheckboxGroup(
+            choices=[doc["name"] for doc in uploaded_documents],
+            value=[doc["name"] for doc in uploaded_documents if doc["selected"]],
+            label="Select documents to query"
+        )
     
     embed = get_embeddings()
     total_chunks = 0
     
     all_data = []
     for file in files:
-        data = load_document(file, parser)
-        all_data.extend(data)
-        total_chunks += len(data)
+        logging.info(f"Processing file: {file.name}")
+        try:
+            data = load_document(file, parser)
+            logging.info(f"Loaded {len(data)} chunks from {file.name}")
+            all_data.extend(data)
+            total_chunks += len(data)
+            # Append new documents instead of replacing
+            if not any(doc["name"] == file.name for doc in uploaded_documents):
+                uploaded_documents.append({"name": file.name, "selected": True})
+                logging.info(f"Added new document to uploaded_documents: {file.name}")
+            else:
+                logging.info(f"Document already exists in uploaded_documents: {file.name}")
+        except Exception as e:
+            logging.error(f"Error processing file {file.name}: {str(e)}")
+    
+    logging.info(f"Total chunks processed: {total_chunks}")
     
     if os.path.exists("faiss_database"):
+        logging.info("Updating existing FAISS database")
         database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
         database.add_documents(all_data)
     else:
+        logging.info("Creating new FAISS database")
         database = FAISS.from_documents(all_data, embed)
     
     database.save_local("faiss_database")
+    logging.info("FAISS database saved")
     
-    return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}."
+    return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}.", gr.CheckboxGroup(
+        choices=[doc["name"] for doc in uploaded_documents],
+        value=[doc["name"] for doc in uploaded_documents if doc["selected"]],
+        label="Select documents to query"
+    )
 
-def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=3, temperature=0.2, should_stop=False):
+def generate_chunked_response(prompt, model, max_tokens=10000, num_calls=3, temperature=0.2, should_stop=False):
     print(f"Starting generate_chunked_response with {num_calls} calls")
     full_response = ""
     messages = [{"role": "user", "content": prompt}]
@@ -214,27 +243,39 @@ def retry_last_response(history, use_web_search, model, temperature, num_calls):
     
     return chatbot_interface(last_user_msg, history, use_web_search, model, temperature, num_calls)
 
-def respond(message, history, model, temperature, num_calls, use_web_search):
+def respond(message, history, model, temperature, num_calls, use_web_search, selected_docs):
     logging.info(f"User Query: {message}")
     logging.info(f"Model Used: {model}")
     logging.info(f"Search Type: {'Web Search' if use_web_search else 'PDF Search'}")
 
+    logging.info(f"Selected Documents: {selected_docs}")
+
     try:
         if use_web_search:
             for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
                 response = f"{main_content}\n\n{sources}"
                 first_line = response.split('\n')[0] if response else ''
-                logging.info(f"Generated Response (first line): {first_line}")
+#                logging.info(f"Generated Response (first line): {first_line}")
                 yield response
         else:
             embed = get_embeddings()
             if os.path.exists("faiss_database"):
                 database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
                 retriever = database.as_retriever()
-                relevant_docs = retriever.get_relevant_documents(message)
+                
+                # Filter relevant documents based on user selection
+                all_relevant_docs = retriever.get_relevant_documents(message)
+                relevant_docs = [doc for doc in all_relevant_docs if doc.metadata["source"] in selected_docs]
+                
+                if not relevant_docs:
+                    yield "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
+                    return
+
                 context_str = "\n".join([doc.page_content for doc in relevant_docs])
             else:
                 context_str = "No documents available."
+                yield "No documents available. Please upload PDF documents to answer questions."
+                return
             
             if model == "@cf/meta/llama-3.1-8b-instruct":
                 # Use Cloudflare API
@@ -244,7 +285,7 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
                     yield partial_response
             else:
                 # Use Hugging Face API
-                for partial_response in get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature):
+                for partial_response in get_response_from_pdf(message, model, selected_docs, num_calls=num_calls, temperature=temperature):
                     first_line = partial_response.split('\n')[0] if partial_response else ''
                     logging.info(f"Generated Response (first line): {first_line}")
                     yield partial_response
@@ -253,7 +294,7 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
         if "microsoft/Phi-3-mini-4k-instruct" in model:
             logging.info("Falling back to Mistral model due to Phi-3 error")
             fallback_model = "mistralai/Mistral-7B-Instruct-v0.3"
-            yield from respond(message, history, fallback_model, temperature, num_calls, use_web_search)
+            yield from respond(message, history, fallback_model, temperature, num_calls, use_web_search, selected_docs)
         else:
             yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
 
@@ -284,7 +325,8 @@ After writing the document, please provide a list of sources used in your respon
     payload = {
         "messages": inputs,
         "stream": True,
-        "temperature": temperature
+        "temperature": temperature,
+        "max_tokens": 32000
     }
 
     full_response = ""
@@ -335,7 +377,7 @@ After writing the document, please provide a list of sources used in your respon
         for i in range(num_calls):
             for message in client.chat_completion(
                 messages=[{"role": "user", "content": prompt}],
-                max_tokens=1000,
+                max_tokens=10000,
                 temperature=temperature,
                 stream=True,
             ):
@@ -344,23 +386,46 @@ After writing the document, please provide a list of sources used in your respon
                     main_content += chunk
                     yield main_content, ""  # Yield partial main content without sources
 
-def get_response_from_pdf(query, model, num_calls=3, temperature=0.2):
+def get_response_from_pdf(query, model, selected_docs, num_calls=3, temperature=0.2):
+    logging.info(f"Entering get_response_from_pdf with query: {query}, model: {model}, selected_docs: {selected_docs}")
+    
     embed = get_embeddings()
     if os.path.exists("faiss_database"):
+        logging.info("Loading FAISS database")
         database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
     else:
+        logging.warning("No FAISS database found")
         yield "No documents available. Please upload PDF documents to answer questions."
         return
 
     retriever = database.as_retriever()
+    logging.info(f"Retrieving relevant documents for query: {query}")
     relevant_docs = retriever.get_relevant_documents(query)
-    context_str = "\n".join([doc.page_content for doc in relevant_docs])
+    logging.info(f"Number of relevant documents retrieved: {len(relevant_docs)}")
+    
+    # Filter relevant_docs based on selected documents
+    filtered_docs = [doc for doc in relevant_docs if doc.metadata["source"] in selected_docs]
+    logging.info(f"Number of filtered documents: {len(filtered_docs)}")
+    
+    if not filtered_docs:
+        logging.warning(f"No relevant information found in the selected documents: {selected_docs}")
+        yield "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
+        return
+
+    for doc in filtered_docs:
+        logging.info(f"Document source: {doc.metadata['source']}")
+        logging.info(f"Document content preview: {doc.page_content[:100]}...")  # Log first 100 characters of each document
+
+    context_str = "\n".join([doc.page_content for doc in filtered_docs])
+    logging.info(f"Total context length: {len(context_str)}")
 
     if model == "@cf/meta/llama-3.1-8b-instruct":
+        logging.info("Using Cloudflare API")
         # Use Cloudflare API with the retrieved context
         for response in get_response_from_cloudflare(prompt="", context=context_str, query=query, num_calls=num_calls, temperature=temperature, search_type="pdf"):
             yield response
     else:
+        logging.info("Using Hugging Face API")
         # Use Hugging Face API
         prompt = f"""Using the following context from the PDF documents:
 {context_str}
@@ -370,9 +435,10 @@ Write a detailed and complete response that answers the following user question:
         
         response = ""
         for i in range(num_calls):
+            logging.info(f"API call {i+1}/{num_calls}")
             for message in client.chat_completion(
                 messages=[{"role": "user", "content": prompt}],
-                max_tokens=1000,
+                max_tokens=10000,
                 temperature=temperature,
                 stream=True,
             ):
@@ -380,6 +446,8 @@ Write a detailed and complete response that answers the following user question:
                     chunk = message.choices[0].delta.content
                     response += chunk
                     yield response  # Yield partial response
+        
+        logging.info("Finished generating response")
 
 def vote(data: gr.LikeData):
     if data.liked:
@@ -388,22 +456,44 @@ def vote(data: gr.LikeData):
         print(f"You downvoted this response: {data.value}")
 
 css = """
-/* Add your custom CSS here */
+/* Fine-tune chatbox size */
+.chatbot-container {
+    height: 600px !important;
+    width: 100% !important;
+}
+.chatbot-container > div {
+    height: 100%;
+    width: 100%;
+}
 """
 
+uploaded_documents = []
+
+def display_documents():
+    return gr.CheckboxGroup(
+        choices=[doc["name"] for doc in uploaded_documents],
+        value=[doc["name"] for doc in uploaded_documents if doc["selected"]],
+        label="Select documents to query"
+    )
+
 # Define the checkbox outside the demo block
-use_web_search = gr.Checkbox(label="Use Web Search", value=False)
+document_selector = gr.CheckboxGroup(label="Select documents to query")
+
+use_web_search = gr.Checkbox(label="Use Web Search", value=True)
+
+custom_placeholder = "Ask a question (Note: You can toggle between Web Search and PDF Chat in Additional Inputs below)"
 
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0]),
+        gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[3]),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),
         gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls"),
-        use_web_search  # Add this line to include the checkbox
+        use_web_search,
+        document_selector
     ],
     title="AI-powered Web Search and PDF Chat Assistant",
-    description="Chat with your PDFs or use web search to answer questions.",
+    description="Chat with your PDFs or use web search to answer questions. Toggle between Web Search and PDF Chat in Additional Inputs below.",
     theme=gr.themes.Soft(
         primary_hue="orange",
         secondary_hue="amber",
@@ -422,7 +512,6 @@ demo = gr.ChatInterface(
         color_accent_soft_dark="transparent",
         code_background_fill_dark="#140b0b"
     ),
-    
     css=css,
     examples=[
         ["Tell me about the contents of the uploaded PDFs."],
@@ -431,6 +520,13 @@ demo = gr.ChatInterface(
     ],
     cache_examples=False,
     analytics_enabled=False,
+    textbox=gr.Textbox(placeholder=custom_placeholder, container=False, scale=7),
+    chatbot = gr.Chatbot(  
+    show_copy_button=True,
+    likeable=True,
+    layout="bubble",
+    height=400,
+)
 )
 
 # Add file upload functionality
@@ -443,18 +539,22 @@ with demo:
         update_button = gr.Button("Upload Document")
     
     update_output = gr.Textbox(label="Update Status")
-    update_button.click(update_vectors, inputs=[file_input, parser_dropdown], outputs=update_output)
-
     
+    # Update both the output text and the document selector
+    update_button.click(update_vectors, 
+                        inputs=[file_input, parser_dropdown], 
+                        outputs=[update_output, document_selector])
+
     gr.Markdown(
     """
     ## How to use
     1. Upload PDF documents using the file input at the top.
     2. Select the PDF parser (pypdf or llamaparse) and click "Upload Document" to update the vector store.
-    3. Ask questions in the chat interface. 
-    4. Toggle "Use Web Search" to switch between PDF chat and web search.
-    5. Adjust Temperature and Number of API Calls to fine-tune the response generation.
-    6. Use the provided examples or ask your own questions.
+    3. Select the documents you want to query using the checkboxes.
+    4. Ask questions in the chat interface. 
+    5. Toggle "Use Web Search" to switch between PDF chat and web search.
+    6. Adjust Temperature and Number of API Calls to fine-tune the response generation.
+    7. Use the provided examples or ask your own questions.
     """
     )