Update app.py

app.py

@@ -20,9 +20,8 @@ huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
 llama_cloud_api_key = os.environ.get("LLAMA_CLOUD_API_KEY")
 
 MODELS = [
-    "google/gemma-2-9b",
-    "mistralai/Mixtral-8x7B-Instruct-v0.1",
     "mistralai/Mistral-7B-Instruct-v0.3",
+    "mistralai/Mixtral-8x7B-Instruct-v0.1",
     "microsoft/Phi-3-mini-4k-instruct"
 ]
 
@@ -78,76 +77,53 @@ def update_vectors(files, parser):
     
     return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}."
 
-def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=3, temperature=0.2, stop_clicked=None):
+def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=3, temperature=0.2, should_stop=False):
     print(f"Starting generate_chunked_response with {num_calls} calls")
     client = InferenceClient(model, token=huggingface_token)
-    full_responses = []
+    full_response = ""
     messages = [{"role": "user", "content": prompt}]
     
     for i in range(num_calls):
         print(f"Starting API call {i+1}")
-        if (isinstance(stop_clicked, gr.State) and stop_clicked.value) or stop_clicked:
+        if should_stop:
             print("Stop clicked, breaking loop")
             break
         try:
-            response = ""
             for message in client.chat_completion(
                 messages=messages,
                 max_tokens=max_tokens,
                 temperature=temperature,
                 stream=True,
             ):
-                if (isinstance(stop_clicked, gr.State) and stop_clicked.value) or stop_clicked:
+                if should_stop:
                     print("Stop clicked during streaming, breaking")
                     break
                 if message.choices and message.choices[0].delta and message.choices[0].delta.content:
                     chunk = message.choices[0].delta.content
-                    response += chunk
-            print(f"API call {i+1} response: {response[:100]}...")
-            full_responses.append(response)
+                    full_response += chunk
+            print(f"API call {i+1} completed")
         except Exception as e:
             print(f"Error in generating response: {str(e)}")
     
-    # Combine responses and clean up
-    combined_response = " ".join(full_responses)
-    clean_response = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', combined_response, flags=re.DOTALL)
+    # Clean up the response
+    clean_response = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', full_response, flags=re.DOTALL)
     clean_response = clean_response.replace("Using the following context:", "").strip()
     clean_response = clean_response.replace("Using the following context from the PDF documents:", "").strip()
     
-    # Split the response into main content and sources
-    parts = re.split(r'\n\s*Sources:\s*\n', clean_response, flags=re.IGNORECASE, maxsplit=1)
-    main_content = parts[0].strip()
-    sources = parts[1].strip() if len(parts) > 1 else ""
-
-    # Process main content
-    paragraphs = main_content.split('\n\n')
+    # Remove duplicate paragraphs and sentences
+    paragraphs = clean_response.split('\n\n')
     unique_paragraphs = []
     for paragraph in paragraphs:
         if paragraph not in unique_paragraphs:
-            unique_sentences = []
             sentences = paragraph.split('. ')
+            unique_sentences = []
             for sentence in sentences:
                 if sentence not in unique_sentences:
                     unique_sentences.append(sentence)
             unique_paragraphs.append('. '.join(unique_sentences))
     
-    final_content = '\n\n'.join(unique_paragraphs)
+    final_response = '\n\n'.join(unique_paragraphs)
     
-    # Process sources
-    if sources:
-        source_lines = sources.split('\n')
-        unique_sources = []
-        for line in source_lines:
-            if line.strip() and line not in unique_sources:
-                unique_sources.append(line)
-        final_sources = '\n'.join(unique_sources)
-        final_response = f"{final_content}\n\nSources:\n{final_sources}"
-    else:
-        final_response = final_content
-
-    # Remove any content after the sources
-    final_response = re.sub(r'(Sources:.*?)(?:\n\n|\Z).*', r'\1', final_response, flags=re.DOTALL)
-
     print(f"Final clean response: {final_response[:100]}...")
     return final_response
 
@@ -161,104 +137,148 @@ class CitingSources(BaseModel):
         ...,
         description="List of sources to cite. Should be an URL of the source."
     )
+def chatbot_interface(message, history, use_web_search, model, temperature, num_calls):
+    if not message.strip():
+        return "", history
+
+    history = history + [(message, "")]
+
+    try:
+        if use_web_search:
+            for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
+                history[-1] = (message, f"{main_content}\n\n{sources}")
+                yield history
+        else:
+            for partial_response in get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature):
+                history[-1] = (message, partial_response)
+                yield history
+    except gr.CancelledError:
+        yield history
+
+def retry_last_response(history, use_web_search, model, temperature, num_calls):
+    if not history:
+        return history
+    
+    last_user_msg = history[-1][0]
+    history = history[:-1]  # Remove the last response
+    
+    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):
+    if use_web_search:
+        for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
+            yield f"{main_content}\n\n{sources}"
+    else:
+        for partial_response in get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature):
+            yield partial_response
 
-def get_response_with_search(query, model, num_calls=3, temperature=0.2, stop_clicked=None):
+def get_response_with_search(query, model, num_calls=3, temperature=0.2):
     search_results = duckduckgo_search(query)
     context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
                         for result in search_results if 'body' in result)
     
-    prompt = f"""<s>[INST] Using the following context:
+    prompt = f"""Using the following context:
 {context}
 Write a detailed and complete research document that fulfills the following user request: '{query}'
-After writing the document, please provide a list of sources used in your response. [/INST]"""
-    
-    generated_text = generate_chunked_response(prompt, model, num_calls=num_calls, temperature=temperature, stop_clicked=stop_clicked)
-   
-    # Clean the response
-    clean_text = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', generated_text, flags=re.DOTALL)
-    clean_text = clean_text.replace("Using the following context:", "").strip()
+After writing the document, please provide a list of sources used in your response."""
     
-    # Split the content and sources
-    parts = clean_text.split("Sources:", 1)
-    main_content = parts[0].strip()
-    sources = parts[1].strip() if len(parts) > 1 else ""
+    client = InferenceClient(model, token=huggingface_token)
     
-    return main_content, sources
-
-def get_response_from_pdf(query, model, num_calls=3, temperature=0.2, stop_clicked=None):
+    main_content = ""
+    for i in range(num_calls):
+        for message in client.chat_completion(
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=1000,
+            temperature=temperature,
+            stream=True,
+        ):
+            if message.choices and message.choices[0].delta and message.choices[0].delta.content:
+                chunk = message.choices[0].delta.content
+                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):
     embed = get_embeddings()
     if os.path.exists("faiss_database"):
         database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
     else:
-        return "No documents available. Please upload PDF documents to answer questions."
+        yield "No documents available. Please upload PDF documents to answer questions."
+        return
 
     retriever = database.as_retriever()
     relevant_docs = retriever.get_relevant_documents(query)
     context_str = "\n".join([doc.page_content for doc in relevant_docs])
 
-    prompt = f"""<s>[INST] Using the following context from the PDF documents:
+    prompt = f"""Using the following context from the PDF documents:
 {context_str}
-Write a detailed and complete response that answers the following user question: '{query}'
-Do not include a list of sources in your response. [/INST]"""
-
-    generated_text = generate_chunked_response(prompt, model, num_calls=num_calls, temperature=temperature, stop_clicked=stop_clicked)
-
-    # Clean the response
-    clean_text = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', generated_text, flags=re.DOTALL)
-    clean_text = clean_text.replace("Using the following context from the PDF documents:", "").strip()
-
-    return clean_text
+Write a detailed and complete response that answers the following user question: '{query}'"""
 
-def chatbot_interface(message, history, use_web_search, model, temperature):
-    if not message.strip():  # Check if the message is empty or just whitespace
-        return history
-
-    if use_web_search:
-        main_content, sources = get_response_with_search(message, model, temperature)
-        formatted_response = f"{main_content}\n\nSources:\n{sources}"
-    else:
-        response = get_response_from_pdf(message, model, temperature)
-        formatted_response = response
-
-    # Check if the last message in history is the same as the current message
-    if history and history[-1][0] == message:
-        # Replace the last response instead of adding a new one
-        history[-1] = (message, formatted_response)
+    client = InferenceClient(model, token=huggingface_token)
+    
+    response = ""
+    for i in range(num_calls):
+        for message in client.chat_completion(
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=1000,
+            temperature=temperature,
+            stream=True,
+        ):
+            if message.choices and message.choices[0].delta and message.choices[0].delta.content:
+                chunk = message.choices[0].delta.content
+                response += chunk
+                yield response  # Yield partial response
+
+def vote(data: gr.LikeData):
+    if data.liked:
+        print(f"You upvoted this response: {data.value}")
     else:
-        # Add the new message-response pair
-        history.append((message, formatted_response))
-
-    return history
-
-
-def clear_and_update_chat(message, history, use_web_search, model, temperature):
-    updated_history = chatbot_interface(message, history, use_web_search, model, temperature)
-    return "", updated_history  # Return empty string to clear the input
-
-def retry_last_response(history):
-    if history:
-        last_user_message = history[-1][0]
-        return last_user_message, history[:-1]
-    return "", history
-
-def undo_last_interaction(history):
-    if len(history) >= 1:
-        return history[:-1]
-    return history
-
-def clear_conversation():
-    return []
-
-def stop_generation():
-    global is_generating
-    is_generating = False
-
-with gr.Blocks() as demo:
-    is_generating = gr.State(False)
-    stop_clicked = gr.State(False)
+        print(f"You downvoted this response: {data.value}")
+
+css = """
+/* Add your custom CSS here */
+"""
+
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[1]),
+        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=False)
+    ],
+    title="AI-powered Web Search and PDF Chat Assistant",
+    description="Chat with your PDFs or use web search to answer questions.",
+    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,
+)
 
-    gr.Markdown("# AI-powered Web Search and PDF Chat Assistant")
-    
+# 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")
@@ -266,111 +286,18 @@ with gr.Blocks() as demo:
     
     update_output = gr.Textbox(label="Update Status")
     update_button.click(update_vectors, inputs=[file_input, parser_dropdown], outputs=update_output)
-    
-    chatbot = gr.Chatbot(label="Conversation")
-    msg = gr.Textbox(label="Ask a question")
-    use_web_search = gr.Checkbox(label="Use Web Search", value=False)
-    
-    with gr.Row():
-        model_dropdown = gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[1])
-        temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature")
-        num_calls_slider = gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls")
-    
-    with gr.Row():
-        submit_btn = gr.Button("Send")
-        stop_btn = gr.Button("Stop")
-        retry_btn = gr.Button("Retry")
-        undo_btn = gr.Button("Undo")
-        clear_btn = gr.Button("Clear")
-
-    def protected_generate_response(message, history, use_web_search, model, temperature, num_calls, is_generating, stop_clicked):
-        print("Starting protected_generate_response")
-        if is_generating:
-            print("Already generating, returning")
-            return message, history, is_generating, stop_clicked
-    
-        is_generating = True
-        
-        if isinstance(stop_clicked, gr.State):
-            stop_clicked.value = False
-        else:
-            stop_clicked = False
-    
-        try:
-            print(f"Generating response for: {message}")
-            if use_web_search:
-                print("Using web search")
-                main_content, sources = get_response_with_search(message, model, num_calls=num_calls, temperature=temperature, stop_clicked=stop_clicked)
-                formatted_response = f"{main_content}\n\nSources:\n{sources}"
-            else:
-                print("Using PDF search")
-                formatted_response = get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature, stop_clicked=stop_clicked)
-    
-            print(f"Generated response: {formatted_response[:100]}...")
-    
-        except Exception as e:
-            print(f"Error generating response: {str(e)}")
-            formatted_response = "I'm sorry, but I encountered an error while generating the response. Please try again."
-        
-        is_generating = False
-        print(f"Returning final response")
-        return "", history + [(message, formatted_response)], is_generating, stop_clicked
-
-    def on_submit(message, history, use_web_search, model, temperature, num_calls, is_generating, stop_clicked):
-        print(f"Submit button clicked with message: {message}")
-        _, new_history, new_is_generating, new_stop_clicked = protected_generate_response(
-            message, history, use_web_search, model, temperature, num_calls, is_generating, stop_clicked
-        )
-        print(f"New history has {len(new_history)} items")
-        return "", new_history, new_is_generating, new_stop_clicked
-
-    submit_btn.click(
-        on_submit,
-        inputs=[msg, chatbot, use_web_search, model_dropdown, temperature_slider, num_calls_slider, is_generating, stop_clicked],
-        outputs=[msg, chatbot, is_generating, stop_clicked],
-        show_progress=True
-        )
-    stop_btn.click(
-        lambda: True,
-        None,
-        stop_clicked
-    )
-
-    retry_btn.click(
-        retry_last_response,
-        inputs=[chatbot],
-        outputs=[msg, chatbot]
-    ).then(
-        on_submit,
-        inputs=[msg, chatbot, use_web_search, model_dropdown, temperature_slider, num_calls_slider, is_generating, stop_clicked],
-        outputs=[msg, chatbot, is_generating, stop_clicked]
-    )
-
-    undo_btn.click(undo_last_interaction, inputs=[chatbot], outputs=[chatbot])
-    clear_btn.click(clear_conversation, outputs=[chatbot])
-
-    gr.Examples(
-        examples=[
-            ["What are the latest developments in AI?"],
-            ["Tell me about recent updates on GitHub"],
-            ["What are the best hotels in Galapagos, Ecuador?"],
-            ["Summarize recent advancements in Python programming"],
-        ],
-        inputs=msg,
-    )
 
     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 textbox. 
+    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 sliders to fine-tune the response generation.
-    6. Click "Send" or press Enter to get a response.
-    7. Use "Retry" to regenerate the last response, "Undo" to remove the last interaction, and "Clear" to reset the conversation.
-    8. Click "Stop" during generation to halt the process.
+    5. Adjust Temperature and Number of API Calls to fine-tune the response generation.
+    6. Use the provided examples or ask your own questions.
     """
     )
+
 if __name__ == "__main__":
     demo.launch(share=True)