Update app.py

app.py

@@ -17,10 +17,9 @@ from huggingface_hub import InferenceClient
 import inspect
 import logging
 import shutil
-import asyncio
-from concurrent.futures import ThreadPoolExecutor
+import numpy as np
+import soundfile as sf
 
-executor = ThreadPoolExecutor()
 
 # Set up basic configuration for logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -31,7 +30,6 @@ llama_cloud_api_key = os.environ.get("LLAMA_CLOUD_API_KEY")
 ACCOUNT_ID = os.environ.get("CLOUDFARE_ACCOUNT_ID")
 API_TOKEN = os.environ.get("CLOUDFLARE_AUTH_TOKEN")
 API_BASE_URL = "https://api.cloudflare.com/client/v4/accounts/a17f03e0f049ccae0c15cdcf3b9737ce/ai/run/"
-whisper_api = InferenceClient("openai/whisper-large-v2", token=huggingface_token)
 
 print(f"ACCOUNT_ID: {ACCOUNT_ID}")
 print(f"CLOUDFLARE_AUTH_TOKEN: {API_TOKEN[:5]}..." if API_TOKEN else "Not set")
@@ -401,82 +399,84 @@ def summarize_web_results(query: str, search_results: List[Dict[str, str]], conv
         return f"An error occurred during summarization: {str(e)}"
 
 # Modify the existing respond function to handle both PDF and web search
-# Modify your existing respond function to handle both PDF and web search
-async def respond(message, history, model, temperature, num_calls, use_web_search, selected_docs, progress=gr.Progress()):
-    logger.info(f"User Query: {message}")
-    logger.info(f"Model Used: {model}")
-    logger.info(f"Selected Documents: {selected_docs}")
-    logger.info(f"Use Web Search: {use_web_search}")
-
-    response = ""
-
-    try:
-        if use_web_search:
-            progress(0, desc="Starting web search")
-            original_query = message
-            rephrased_query = rephrase_query(message, conversation_manager)
-            logger.info(f"Original query: {original_query}")
-            logger.info(f"Rephrased query: {rephrased_query}")
-
-            final_summary = ""
-            for i in range(num_calls):
-                progress((i + 1) * 33, desc=f"Performing web search {i+1}/{num_calls}")
-                search_results = get_web_search_results(rephrased_query)
-                if not search_results:
-                    final_summary += f"No search results found for the query: {rephrased_query}\n\n"
-                elif "error" in search_results[0]:
-                    final_summary += search_results[0]["error"] + "\n\n"
-                else:
-                    summary = summarize_web_results(rephrased_query, search_results, conversation_manager)
-                    final_summary += summary + "\n\n"
-
-            if final_summary:
-                conversation_manager.add_interaction(original_query, final_summary)
-                response = final_summary
+def respond(message, history, model, temperature, num_calls, use_web_search, selected_docs, audio_input):
+    if audio_input:
+        message = transcribe_audio(audio_input)
+        logging.info(f"Transcribed audio: {message}")
+    
+    if not message.strip():
+        return "Please provide a text or audio query.", history
+
+    logging.info(f"User Query: {message}")
+    logging.info(f"Model Used: {model}")
+    logging.info(f"Selected Documents: {selected_docs}")
+    logging.info(f"Use Web Search: {use_web_search}")
+
+    if use_web_search:
+        original_query = message
+        rephrased_query = rephrase_query(message, conversation_manager)
+        logging.info(f"Original query: {original_query}")
+        logging.info(f"Rephrased query: {rephrased_query}")
+
+        final_summary = ""
+        for _ in range(num_calls):
+            search_results = get_web_search_results(rephrased_query)
+            if not search_results:
+                final_summary += f"No search results found for the query: {rephrased_query}\n\n"
+            elif "error" in search_results[0]:
+                final_summary += search_results[0]["error"] + "\n\n"
             else:
-                response = "Unable to generate a response. Please try a different query."
+                summary = summarize_web_results(rephrased_query, search_results, conversation_manager)
+                final_summary += summary + "\n\n"
 
+        if final_summary:
+            conversation_manager.add_interaction(original_query, final_summary)
+            yield final_summary
         else:
-            # Existing PDF search logic
-            progress(0, desc="Starting PDF search")
+            yield "Unable to generate a response. Please try a different query."
+    else:
+        # Existing PDF search logic
+        try:
             embed = get_embeddings()
             if os.path.exists("faiss_database"):
                 database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
                 retriever = database.as_retriever(search_kwargs={"k": 20})
                 
-                progress(33, desc="Retrieving relevant documents")
                 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:
-                    response = "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
-                else:
-                    context_str = "\n".join([doc.page_content for doc in relevant_docs])
-                    logger.info(f"Context length: {len(context_str)}")
-
-                    progress(66, desc="Generating response")
-                    if model.startswith("duckduckgo/"):
-                        # Use DuckDuckGo chat with context
-                        for partial_response in get_response_from_duckduckgo(message, model, context_str, num_calls, temperature):
-                            response += partial_response
-                    elif model == "@cf/meta/llama-3.1-8b-instruct":
-                        # 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"):
-                            response += partial_response
-                    else:
-                        # Use Hugging Face API
-                        for partial_response in get_response_from_pdf(message, model, selected_docs, num_calls=num_calls, temperature=temperature):
-                            response += partial_response
+                    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])
+                logging.info(f"Context length: {len(context_str)}")
             else:
-                response = "No documents available. Please upload PDF documents to answer questions."
-
-    except Exception as e:
-        logger.error(f"Error in respond function: {str(e)}")
-        response = f"An error occurred: {str(e)}. Please try again or select a different model."
-
-    progress(100, desc="Response generation complete")
-    return response
-
+                context_str = "No documents available."
+                yield "No documents available. Please upload PDF documents to answer questions."
+                return
+            
+            if model.startswith("duckduckgo/"):
+                # Use DuckDuckGo chat with context
+                for partial_response in get_response_from_duckduckgo(message, model, context_str, num_calls, temperature):
+                    yield partial_response
+            elif model == "@cf/meta/llama-3.1-8b-instruct":
+                # 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"):
+                    yield partial_response
+            else:
+                # Use Hugging Face API
+                for partial_response in get_response_from_pdf(message, model, selected_docs, num_calls=num_calls, temperature=temperature):
+                    yield partial_response
+        except Exception as e:
+            logging.error(f"Error with {model}: {str(e)}")
+            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, 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)
 
 def get_response_from_cloudflare(prompt, context, query, num_calls=3, temperature=0.2, search_type="pdf"):
@@ -615,33 +615,23 @@ Write a detailed and complete response that answers the following user question:
         
         logging.info("Finished generating response")
 
-# Modified transcribe function
-async def transcribe_with_timeout(audio_file, progress=gr.Progress()):
-    if audio_file is None:
-        logger.warning("No audio file provided")
-        return "No audio file provided"
+def transcribe_audio(audio_file):
+    client = InferenceClient("openai/whisper-large-v3", token=huggingface_token)
     
-    try:
-        logger.info(f"Starting transcription for file: {audio_file}")
-        progress(0, desc="Starting transcription")
-        
-        with open(audio_file, "rb") as f:
-            audio_data = f.read()
-        
-        logger.info("Audio file read successfully")
-        progress(50, desc="Audio loaded, sending to API")
-        
-        response = await asyncio.wait_for(asyncio.to_thread(whisper_api, audio_data), timeout=30)
-        
-        logger.info("Transcription complete")
-        progress(100, desc="Transcription complete")
-        return response["text"]
-    except asyncio.TimeoutError:
-        logger.error("Transcription timed out")
-        return "Transcription timed out. Please try again with a shorter audio clip."
-    except Exception as e:
-        logger.exception(f"Error during transcription: {str(e)}")
-        return f"Error during transcription: {str(e)}"
+    # Load the audio file
+    audio, sample_rate = sf.read(audio_file)
+    
+    # Ensure audio is mono
+    if len(audio.shape) > 1:
+        audio = audio.mean(axis=1)
+    
+    # Normalize audio
+    audio = (audio / np.max(np.abs(audio))) * 32767
+    audio = audio.astype(np.int16)
+    
+    # Transcribe
+    result = client.automatic_speech_recognition(audio, sampling_rate=sample_rate)
+    return result["text"]
 
 def vote(data: gr.LikeData):
     if data.liked:
@@ -692,39 +682,73 @@ use_web_search = gr.Checkbox(label="Use Web Search", value=False)
 
 custom_placeholder = "Ask a question (Note: You can toggle between Web Search and PDF Chat in Additional Inputs below)"
 
-def update_textbox(transcription):
-    return gr.Textbox.update(value=transcription)
+# Update the demo interface
+# Update the Gradio interface
+demo = gr.Interface(
+    fn=respond,
+    inputs=[
+        gr.Textbox(placeholder=custom_placeholder, container=False, scale=7),
+        gr.State([]),  # for history
+        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"),
+        gr.Checkbox(label="Use Web Search", value=True),
+        gr.CheckboxGroup(label="Select documents to query"),
+        gr.Audio(source="microphone", type="filepath")
+    ],
+    outputs=[gr.Chatbot()],
+    title="AI-powered PDF Chat and Web Search Assistant",
+    description="Chat with your PDFs or use web search to answer questions. You can type or speak your query.",
+    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?"],
+        ["What's the latest news about artificial intelligence?"]
+    ],
+    cache_examples=False,
+    analytics_enabled=False,
+)
 
-# Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# AI-powered PDF Chat and Web Search Assistant with Speech Input")
-    
+# Add file upload functionality
+# Add file upload functionality
+with demo:
+    gr.Markdown("## Chat Interface")
     with gr.Row():
-        with gr.Column(scale=1):
-            audio_input = gr.Audio(sources=["microphone"], type="filepath", label="Speak your query")
-            transcribe_button = gr.Button("Transcribe")
-        
-        with gr.Column(scale=2):
-            chatbot = gr.Chatbot(
-                show_copy_button=True,
-                bubble=True,
-                height=400,
-                value=initial_conversation()
-            )
-            query_textbox = gr.Textbox(
-                placeholder="Ask a question about the uploaded PDFs or any topic", 
-                container=False, 
-                scale=7
-            )
-            submit_button = gr.Button("Submit")
-    
-    with gr.Accordion("⚙️ Parameters", open=False):
-        model = gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[3])
-        temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature")
-        num_calls = gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls")
-        use_web_search = gr.Checkbox(label="Use Web Search", value=True)
-        document_selector = gr.CheckboxGroup(label="Select documents to query")
-    
+        text_input = gr.Textbox(
+            placeholder=custom_placeholder,
+            container=False,
+            scale=7
+        )
+        audio_input = gr.Audio(source="microphone", type="filepath")
+        submit_button = gr.Button("Submit")
+
+    chat_output = gr.Chatbot(
+        show_copy_button=True,
+        likeable=True,
+        layout="bubble",
+        height=400,
+        value=initial_conversation()
+    )
+with demo:
     gr.Markdown("## Upload and Manage PDF Documents")
     with gr.Row():
         file_input = gr.Files(label="Upload your PDF documents", file_types=[".pdf"])
@@ -735,36 +759,39 @@ with gr.Blocks() as demo:
     update_output = gr.Textbox(label="Update Status")
     delete_button = gr.Button("Delete Selected Documents")
     
-    # Connect components
-    transcribe_button.click(
-        transcribe_with_timeout, 
-        inputs=[audio_input], 
-        outputs=[query_textbox]
-    )
-    
-    submit_button.click(
-        respond,
-        inputs=[query_textbox, chatbot, model, temperature, num_calls, use_web_search, document_selector],
-        outputs=[chatbot]
-    )
-    
+    # Update both the output text and the document selector
     update_button.click(
         update_vectors, 
         inputs=[file_input, parser_dropdown], 
-        outputs=[update_output, document_selector]
+        outputs=[update_output, demo.additional_inputs[-1]]  # Use the CheckboxGroup from additional_inputs
     )
     
+    # Add the refresh button functionality
     refresh_button.click(
         refresh_documents, 
         inputs=[], 
-        outputs=[document_selector]
+        outputs=[demo.additional_inputs[-1]]  # Use the CheckboxGroup from additional_inputs
     )
     
+    # Add the delete button functionality
     delete_button.click(
         delete_documents,
-        inputs=[document_selector],
-        outputs=[update_output, document_selector]
+        inputs=[demo.additional_inputs[-1]],  # Use the CheckboxGroup from additional_inputs
+        outputs=[update_output, demo.additional_inputs[-1]]
+    )
+
+    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. 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.
+    """
     )
 
 if __name__ == "__main__":
-    demo.launch(debug=True, show_error=True)
+    demo.launch(share=True)