Update app.py

app.py

@@ -247,42 +247,23 @@ def respond(message, history, model, temperature, num_calls, use_web_search, sel
     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:
-            search_results = duckduckgo_search(message)
-            context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
-                                for result in search_results if 'body' in result)
-            
-            prompt = f"""Using the following context:
-            {context}
-            Write a detailed and complete research document that fulfills the following user request: '{message}'
-            After writing the document, please provide a list of sources used in your response."""
-
-            if model == "@cf/meta/llama-3.1-8b-instruct":
-                for response in get_response_from_cloudflare(prompt="", context=context, query=message, num_calls=num_calls, temperature=temperature, search_type="web"):
-                    yield response
-            else:
-                client = InferenceClient(model, token=huggingface_token)
-                response = ""
-                for i in range(num_calls):
-                    for msg in client.chat_completion(
-                        messages=[{"role": "user", "content": prompt}],
-                        max_tokens=10000,
-                        temperature=temperature,
-                        stream=True,
-                    ):
-                        if msg.choices and msg.choices[0].delta and msg.choices[0].delta.content:
-                            chunk = msg.choices[0].delta.content
-                            response += chunk
-                            yield response
+            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}")
+                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()
                 
+                # 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]
                 
@@ -292,34 +273,30 @@ def respond(message, history, model, temperature, num_calls, use_web_search, sel
 
                 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":
-                for response in get_response_from_cloudflare(prompt="", context=context_str, query=message, num_calls=num_calls, temperature=temperature, search_type="pdf"):
-                    yield response
+                # Use Cloudflare API
+                for partial_response in get_response_from_cloudflare(prompt="", context=context_str, query=message, num_calls=num_calls, temperature=temperature, search_type="pdf"):
+                    first_line = partial_response.split('\n')[0] if partial_response else ''
+                    logging.info(f"Generated Response (first line): {first_line}")
+                    yield partial_response
             else:
-                prompt = f"""Using the following context from the PDF documents:
-                {context_str}
-                Write a detailed and complete response that answers the following user question: '{message}'"""
-                
-                client = InferenceClient(model, token=huggingface_token)
-                response = ""
-                for i in range(num_calls):
-                    for msg in client.chat_completion(
-                        messages=[{"role": "user", "content": prompt}],
-                        max_tokens=10000,
-                        temperature=temperature,
-                        stream=True,
-                    ):
-                        if msg.choices and msg.choices[0].delta and msg.choices[0].delta.content:
-                            chunk = msg.choices[0].delta.content
-                            response += chunk
-                            yield response
-
+                # Use Hugging Face API
+                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
     except Exception as e:
         logging.error(f"Error with {model}: {str(e)}")
-        yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
+        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, selected_docs)
+        else:
+            yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
 
 logging.basicConfig(level=logging.DEBUG)
 
@@ -504,39 +481,57 @@ document_selector = gr.CheckboxGroup(label="Select documents to query")
 
 use_web_search = gr.Checkbox(label="Use Web Search", value=True)
 
-with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# AI-powered Web Search and PDF Chat Assistant")
-    gr.Markdown("Chat with your PDFs or use web search to answer questions")
-
-    use_web_search = gr.Checkbox(label="Use Web Search", value=True)
-    document_selector = gr.CheckboxGroup(label="Select documents to query")
-
-    chat_interface = gr.ChatInterface(
-        respond,
-        additional_inputs=[
-            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"),
-            document_selector  # Moved use_web_search out, kept document_selector here
-        ],
-        title="Chat Interface",
-        description="Ask questions about your PDFs or use web search",
-        examples=[
-            ["Tell me about the contents of the uploaded PDFs."],
-            ["What are the main topics discussed in the documents?"],
-            ["Can you summarize the key points from the PDFs?"]
-        ],
-        cache_examples=False,
-        analytics_enabled=False,
-    )
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        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,
+        document_selector  # Add the document selector to the chat interface
+    ],
+    title="AI-powered Web Search and PDF Chat Assistant",
+    description="Chat with your PDFs or use web search to answer questions (Please use toggle under Additional Inputs to swithc between PDF and Web Search, Default Value Web Search)",
+    theme=gr.themes.Soft(
+        primary_hue="orange",
+        secondary_hue="amber",
+        neutral_hue="gray",
+        font=[gr.themes.GoogleFont("Exo"), "ui-sans-serif", "system-ui", "sans-serif"]
+    ).set(
+        body_background_fill_dark="#0c0505",
+        block_background_fill_dark="#0c0505",
+        block_border_width="1px",
+        block_title_background_fill_dark="#1b0f0f",
+        input_background_fill_dark="#140b0b",
+        button_secondary_background_fill_dark="#140b0b",
+        border_color_accent_dark="#1b0f0f",
+        border_color_primary_dark="#1b0f0f",
+        background_fill_secondary_dark="#0c0505",
+        color_accent_soft_dark="transparent",
+        code_background_fill_dark="#140b0b"
+    ),
+    css=css,
+    examples=[
+        ["Tell me about the contents of the uploaded PDFs."],
+        ["What are the main topics discussed in the documents?"],
+        ["Can you summarize the key points from the PDFs?"]
+    ],
+    cache_examples=False,
+    analytics_enabled=False,
+)
+
+# Add file upload functionality
+with demo:
+    gr.Markdown("## Upload PDF Documents")
 
     with gr.Row():
         file_input = gr.Files(label="Upload your PDF documents", file_types=[".pdf"])
         parser_dropdown = gr.Dropdown(choices=["pypdf", "llamaparse"], label="Select PDF Parser", value="llamaparse")
         update_button = gr.Button("Upload Document")
-
+    
     update_output = gr.Textbox(label="Update Status")
-
+    
+    # Update both the output text and the document selector
     update_button.click(update_vectors, 
                         inputs=[file_input, parser_dropdown], 
                         outputs=[update_output, document_selector])
@@ -544,12 +539,12 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     gr.Markdown(
     """
     ## How to use
-    1. Use the "Use Web Search" checkbox at the top to switch between PDF chat and web search.
-    2. Upload PDF documents using the file input at the bottom.
-    3. Select the PDF parser (pypdf or llamaparse) and click "Upload Document" to update the vector store.
-    4. Select the documents you want to query using the checkboxes.
-    5. Adjust Temperature and Number of API Calls in the Additional Inputs to fine-tune the response generation.
-    6. Type your questions in the chat interface and press Enter to get responses.
+    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. 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.
     """
     )