Update app.py

app.py

@@ -1,13 +1,57 @@
+import os
+import json
+import re
 import gradio as gr
+import requests
 from duckduckgo_search import DDGS
-from typing import List, Dict
-import os
+from typing import List
+from pydantic import BaseModel, Field
+from tempfile import NamedTemporaryFile
+from langchain_community.vectorstores import FAISS
+from langchain_core.vectorstores import VectorStore
+from langchain_core.documents import Document
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from llama_parse import LlamaParse
+from langchain_core.documents import Document
+from huggingface_hub import InferenceClient
+import inspect
 import logging
+import shutil
+
 
-logging.basicConfig(level=logging.INFO)
+# Set up basic configuration for logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # Environment variables and configurations
 huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
+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/"
+
+print(f"ACCOUNT_ID: {ACCOUNT_ID}")
+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",
+    "mistralai/Mistral-Nemo-Instruct-2407",
+    "gpt-4o-mini",
+    "claude-3-haiku",
+    "llama-3.1-70b",
+    "mixtral-8x7b"
+]
+
+# Initialize LlamaParse
+llama_parser = LlamaParse(
+    api_key=llama_cloud_api_key,
+    result_type="markdown",
+    num_workers=4,
+    verbose=True,
+    language="en",
+)
 
 def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[Document]:
     """Loads and splits the document into pages."""
@@ -141,30 +185,349 @@ def delete_documents(selected_docs):
     
     return f"Deleted documents: {', '.join(deleted_docs)}", display_documents()
 
-def get_response_from_pdf(query: str, selected_docs: List[str], num_calls: int = 3, temperature: float = 0.2) -> str:
-    logging.info(f"Entering get_response_from_pdf with query: {query}, selected_docs: {selected_docs}")
+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}]
+    
+    if model == "@cf/meta/llama-3.1-8b-instruct":
+        # Cloudflare API
+        for i in range(num_calls):
+            print(f"Starting Cloudflare API call {i+1}")
+            if should_stop:
+                print("Stop clicked, breaking loop")
+                break
+            try:
+                response = requests.post(
+                    f"https://api.cloudflare.com/client/v4/accounts/{ACCOUNT_ID}/ai/run/@cf/meta/llama-3.1-8b-instruct",
+                    headers={"Authorization": f"Bearer {API_TOKEN}"},
+                    json={
+                        "stream": true,
+                        "messages": [
+                            {"role": "system", "content": "You are a friendly assistant"},
+                            {"role": "user", "content": prompt}
+                        ],
+                        "max_tokens": max_tokens,
+                        "temperature": temperature
+                    },
+                    stream=true
+                )
+                
+                for line in response.iter_lines():
+                    if should_stop:
+                        print("Stop clicked during streaming, breaking")
+                        break
+                    if line:
+                        try:
+                            json_data = json.loads(line.decode('utf-8').split('data: ')[1])
+                            chunk = json_data['response']
+                            full_response += chunk
+                        except json.JSONDecodeError:
+                            continue
+                print(f"Cloudflare API call {i+1} completed")
+            except Exception as e:
+                print(f"Error in generating response from Cloudflare: {str(e)}")
+    else:
+        # Original Hugging Face API logic
+        client = InferenceClient(model, token=huggingface_token)
+        
+        for i in range(num_calls):
+            print(f"Starting Hugging Face API call {i+1}")
+            if should_stop:
+                print("Stop clicked, breaking loop")
+                break
+            try:
+                for message in client.chat_completion(
+                    messages=messages,
+                    max_tokens=max_tokens,
+                    temperature=temperature,
+                    stream=True,
+                ):
+                    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
+                        full_response += chunk
+                print(f"Hugging Face API call {i+1} completed")
+            except Exception as e:
+                print(f"Error in generating response from Hugging Face: {str(e)}")
+    
+    # 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()
+    
+    # Remove duplicate paragraphs and sentences
+    paragraphs = clean_response.split('\n\n')
+    unique_paragraphs = []
+    for paragraph in paragraphs:
+        if paragraph not in unique_paragraphs:
+            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_response = '\n\n'.join(unique_paragraphs)
+    
+    print(f"Final clean response: {final_response[:100]}...")
+    return final_response
+
+def duckduckgo_search(query):
+    with DDGS() as ddgs:
+        results = ddgs.text(query, max_results=5)
+    return results
+
+def duckduckgo_chat(keywords, model, timeout=30):
+    with DDGS() as ddgs:
+        response = ddgs.chat(keywords, model=model, timeout=timeout)
+    return response
+
+class CitingSources(BaseModel):
+    sources: List[str] = Field(
+        ...,
+        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:
+        for response in respond(message, history, model, temperature, num_calls, use_web_search):
+            history[-1] = (message, response)
+            yield history
+    except gr.CancelledError:
+        yield history
+    except Exception as e:
+        logging.error(f"Unexpected error in chatbot_interface: {str(e)}")
+        history[-1] = (message, f"An unexpected error occurred: {str(e)}")
+        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, 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}")
+                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(search_kwargs={"k": 20})
+                
+                # 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
+                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
+            elif model in ["gpt-4o-mini", "claude-3-haiku", "llama-3.1-70b", "mixtral-8x7b"]:
+                # Use DuckDuckGo Chat API
+                for partial_response in get_response_from_duckduckgo(message, model, 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
+            else:
+                # 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)}")
+        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)
+
+def get_response_from_duckduckgo(message, model, num_calls=3, temperature=0.2):
+    for i in range(num_calls):
+        try:
+            response = duckduckgo_chat(message, model=model, timeout=30)
+            yield response
+        except Exception as e:
+            logging.error(f"Error in generating response from DuckDuckGo: {str(e)}")
+            yield f"I apologize, but an error occurred: {str(e)}. Please try again later."
+
+def get_response_from_cloudflare(prompt, context, query, num_calls=3, temperature=0.2, search_type="pdf"):
+    headers = {
+        "Authorization": f"Bearer {API_TOKEN}",
+        "Content-Type": "application/json"
+    }
+    model = "@cf/meta/llama-3.1-8b-instruct"
+
+    if search_type == "pdf":
+        instruction = f"""Using the following context from the PDF documents:
+{context}
+Write a detailed and complete response that answers the following user question: '{query}'"""
+    else:  # web search
+        instruction = 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."""
+
+    inputs = [
+        {"role": "system", "content": instruction},
+        {"role": "user", "content": query}
+    ]
+
+    payload = {
+        "messages": inputs,
+        "stream": True,
+        "temperature": temperature,
+        "max_tokens": 32000
+    }
+
+    full_response = ""
+    for i in range(num_calls):
+        try:
+            with requests.post(f"{API_BASE_URL}{model}", headers=headers, json=payload, stream=True) as response:
+                if response.status_code == 200:
+                    for line in response.iter_lines():
+                        if line:
+                            try:
+                                json_response = json.loads(line.decode('utf-8').split('data: ')[1])
+                                if 'response' in json_response:
+                                    chunk = json_response['response']
+                                    full_response += chunk
+                                    yield full_response
+                            except (json.JSONDecodeError, IndexError) as e:
+                                logging.error(f"Error parsing streaming response: {str(e)}")
+                                continue
+                else:
+                    logging.error(f"HTTP Error: {response.status_code}, Response: {response.text}")
+                    yield f"I apologize, but I encountered an HTTP error: {response.status_code}. Please try again later."
+        except Exception as e:
+            logging.error(f"Error in generating response from Cloudflare: {str(e)}")
+            yield f"I apologize, but an error occurred: {str(e)}. Please try again later."
+    
+    if not full_response:
+        yield "I apologize, but I couldn't generate a response at this time. Please try again later."
+
+def create_web_search_vectors(search_results):
+    embed = get_embeddings()
+    
+    documents = []
+    for result in search_results:
+        if 'body' in result:
+            content = f"{result['title']}\n{result['body']}\nSource: {result['href']}"
+            documents.append(Document(page_content=content, metadata={"source": result['href']}))
+    
+    return FAISS.from_documents(documents, embed)
+
+def get_response_with_search(query, model, num_calls=3, temperature=0.2):
+    search_results = duckduckgo_search(query)
+    web_search_database = create_web_search_vectors(search_results)
+    
+    if not web_search_database:
+        yield "No web search results available. Please try again.", ""
+        return
+    
+    retriever = web_search_database.as_retriever(search_kwargs={"k": 5})
+    relevant_docs = retriever.get_relevant_documents(query)
+    
+    context = "\n".join([doc.page_content for doc in relevant_docs])
+    
+    prompt = f"""Using the following context from web search results:
+{context}
+You are an expert AI assistant, write a detailed and complete research document that fulfills the following user request: '{query}'
+Base your entire response strictly on the information retrieved from trusted sources. Importantly, only include information that is directly supported by the retrieved content. 
+If any part of the information cannot be verified from the given sources, clearly state that it could not be confirmed.
+After writing the document, please provide a list of sources used in your response."""
+
+    if model == "@cf/meta/llama-3.1-8b-instruct":
+        # Use Cloudflare API
+        for response in get_response_from_cloudflare(prompt="", context=context, query=query, num_calls=num_calls, temperature=temperature, search_type="web"):
+            yield response, ""  # Yield streaming response without sources
+    else:
+        # Use Hugging Face API
+        client = InferenceClient(model, token=huggingface_token)
+        
+        main_content = ""
+        for i in range(num_calls):
+            for message in client.chat_completion(
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=10000,
+                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, 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")
-        return "No documents available. Please upload PDF documents to answer questions."
+        yield "No documents available. Please upload PDF documents to answer questions."
+        return
 
     # Pre-filter the documents
-    filtered_docs = [doc for doc_id, doc in database.docstore._dict.items() if isinstance(doc, Document) and doc.metadata.get("source") in selected_docs]
+    filtered_docs = []
+    for doc_id, doc in database.docstore._dict.items():
+        if isinstance(doc, Document) and doc.metadata.get("source") in selected_docs:
+            filtered_docs.append(doc)
+    
     logging.info(f"Number of documents after pre-filtering: {len(filtered_docs)}")
+
     if not filtered_docs:
         logging.warning(f"No documents found for the selected sources: {selected_docs}")
-        return "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
+        yield "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
+        return
 
     # Create a new FAISS index with only the selected documents
     filtered_db = FAISS.from_documents(filtered_docs, embed)
+    
     retriever = filtered_db.as_retriever(search_kwargs={"k": 10})
     logging.info(f"Retrieving relevant documents for query: {query}")
     relevant_docs = retriever.get_relevant_documents(query)
     logging.info(f"Number of relevant documents retrieved: {len(relevant_docs)}")
+
     for doc in relevant_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
@@ -172,159 +535,172 @@ def get_response_from_pdf(query: str, selected_docs: List[str], num_calls: int =
     context_str = "\n".join([doc.page_content for doc in relevant_docs])
     logging.info(f"Total context length: {len(context_str)}")
 
-    prompt = f"""Using the following context from the PDF documents:
+    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}
 Write a detailed and complete response that answers the following user question: '{query}'"""
-
-    response = ""
-    for i in range(num_calls):
-        logging.info(f"API call {i+1}/{num_calls}")
-        response += DDGS().chat(prompt, model="llama-3.1-70b", max_tokens=1024, temperature=temperature)
-
-    logging.info("Finished generating response")
-    return response
-
-class ConversationManager:
-    def __init__(self):
-        self.history = []
-        self.current_context = None
-
-    def add_interaction(self, query, response):
-        self.history.append((query, response))
-        self.current_context = f"Previous query: {query}\nPrevious response summary: {response[:200]}..."
-
-    def get_context(self):
-        return self.current_context
-
-def get_web_search_results(query: str, max_results: int = 10) -> List[Dict[str, str]]:
-    try:
-        results = list(DDGS().text(query, max_results=max_results))
-        if not results:
-            print(f"No results found for query: {query}")
-        return results
-    except Exception as e:
-        print(f"An error occurred during web search: {str(e)}")
-        return [{"error": f"An error occurred during web search: {str(e)}"}]
-
-def rephrase_query(original_query: str, conversation_manager: ConversationManager) -> str:
-    context = conversation_manager.get_context()
-    if context:
-        prompt = f"""You are a highly intelligent conversational chatbot. Your task is to analyze the given context and new query, then decide whether to rephrase the query with or without incorporating the context. Follow these steps:
-
-        1. Determine if the new query is a continuation of the previous conversation or an entirely new topic.
-        2. If it's a continuation, rephrase the query by incorporating relevant information from the context to make it more specific and contextual.
-        3. If it's a new topic, rephrase the query to make it more appropriate for a web search, focusing on clarity and accuracy without using the previous context.
-        4. Provide ONLY the rephrased query without any additional explanation or reasoning.
         
-        Context: {context}
+        client = InferenceClient(model, token=huggingface_token)
         
-        New query: {original_query}
+        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=10000,
+                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
         
-        Rephrased query:"""
-        response = DDGS().chat(prompt, model="llama-3.1-70b")
-        # Extract only the rephrased query, removing any explanations
-        rephrased_query = response.split('\n')[0].strip()
-        return rephrased_query
-    return original_query
-
-def summarize_results(query: str, search_results: List[Dict[str, str]], conversation_manager: ConversationManager) -> str:
-    try:
-        context = conversation_manager.get_context()
-        search_context = "\n\n".join([f"Title: {result['title']}\nContent: {result['body']}" for result in search_results])
-
-        prompt = f"""You are a highly intelligent & expert analyst and your job is to skillfully articulate the web search results about '{query}' and considering the context: {context}, 
-        You have to create a comprehensive news summary FOCUSING on the context provided to you. 
-        Include key facts, relevant statistics, and expert opinions if available. 
-        Ensure the article is well-structured with an introduction, main body, and conclusion, IF NECESSARY. 
-        Address the query in the context of the ongoing conversation IF APPLICABLE.
-        Cite sources directly within the generated text and not at the end of the generated text, integrating URLs where appropriate to support the information provided:
+        logging.info("Finished generating response")
 
-        {search_context}
-
-        Article:"""
-
-        summary = DDGS().chat(prompt, model="llama-3-70b")
-        return summary
-    except Exception as e:
-        return f"An error occurred during summarization: {str(e)}"
+def vote(data: gr.LikeData):
+    if data.liked:
+        print(f"You upvoted this response: {data.value}")
+    else:
+        print(f"You downvoted this response: {data.value}")
 
-conversation_manager = ConversationManager()
+css = """
+/* Fine-tune chatbox size */
+.chatbot-container {
+    height: 600px !important;
+    width: 100% !important;
+}
+.chatbot-container > div {
+    height: 100%;
+    width: 100%;
+}
+"""
 
-def respond(message, chat_history, temperature, num_api_calls):
-    final_summary = ""
-    original_query = message
-    rephrased_query = rephrase_query(message, conversation_manager)
+uploaded_documents = []
 
-    logging.info(f"Original query: {original_query}")
-    logging.info(f"Rephrased query: {rephrased_query}")
+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 or delete"
+    )
 
-    for _ in range(num_api_calls):
-        search_results = get_web_search_results(rephrased_query)
+def initial_conversation():
+    return [
+        (None, "Welcome! I'm your AI assistant for web search and PDF analysis. Here's how you can use me:\n\n"
+                "1. Set the toggle for Web Search and PDF Search from the checkbox in Additional Inputs drop down window\n"
+                "2. Use web search to find information\n"
+                "3. Upload the documents and ask questions about uploaded PDF documents by selecting your respective document\n"
+                "4. For any queries feel free to reach out @[email protected] or discord - shreyas094\n\n"
+                "To get started, upload some PDFs or ask me a question!")
+    ]
+# Add this new function
+def refresh_documents():
+    global uploaded_documents
+    uploaded_documents = load_documents()
+    return display_documents()
 
-        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_results(rephrased_query, search_results, conversation_manager)
-            final_summary += summary + "\n\n"
+# Define the checkbox outside the demo block
+document_selector = gr.CheckboxGroup(label="Select documents to query")
 
-    if final_summary:
-        conversation_manager.add_interaction(original_query, final_summary)
-        return final_summary
-    else:
-        return "Unable to generate a response. Please try a different query."
+use_web_search = gr.Checkbox(label="Use Web Search", value=True)
 
-# The rest of your code (CSS, theme, and Gradio interface setup) remains the same
-css = """
-Your custom CSS here
-"""
-
-custom_placeholder = "Ask me anything about web content"
-
-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"
-)
+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[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.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls"),
+        use_web_search,
+        document_selector
     ],
     title="AI-powered Web Search and PDF Chat Assistant",
-    description="This AI-powered Web Search and PDF Chat Assistant combines real-time web search capabilities with advanced language processing.",
-    theme=theme,
+    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",
+        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=[
-        ["What is AI"],
-        ["Any recent news on US Banks"],
-        ["Who is Donald Trump"]
+        ["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,
     textbox=gr.Textbox(placeholder=custom_placeholder, container=False, scale=7),
-    chatbot=gr.Chatbot(
-        show_copy_button=True,
-        likeable=True,
-        layout="bubble",
-        height=400,
-    )
+    chatbot = gr.Chatbot(  
+    show_copy_button=True,
+    likeable=True,
+    layout="bubble",
+    height=400,
+    value=initial_conversation()
 )
+)
+
+# Add file upload functionality
+with demo:
+    gr.Markdown("## Upload and Manage 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")
+        refresh_button = gr.Button("Refresh Document List")
+    
+    update_output = gr.Textbox(label="Update Status")
+    delete_button = gr.Button("Delete Selected Documents")
+    
+    # Update both the output text and the document selector
+    update_button.click(update_vectors, 
+                        inputs=[file_input, parser_dropdown], 
+                        outputs=[update_output, document_selector])
+    
+    # Add the refresh button functionality
+    refresh_button.click(refresh_documents, 
+                         inputs=[], 
+                         outputs=[document_selector])
+    
+    # Add the delete button functionality
+    delete_button.click(delete_documents,
+                        inputs=[document_selector],
+                        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. 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.
+    """
+    )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch(share=True)