Update app.py

app.py

@@ -1,48 +1,21 @@
 import os
 import json
-import re
 import gradio as gr
-import requests
 from duckduckgo_search import DDGS
 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
 
-
-# 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"
-]
-
 # Initialize LlamaParse
 llama_parser = LlamaParse(
-    api_key=llama_cloud_api_key,
+    api_key=os.environ.get("LLAMA_CLOUD_API_KEY"),
     result_type="markdown",
     num_workers=4,
     verbose=True,
@@ -69,7 +42,6 @@ def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[
 def get_embeddings():
     return HuggingFaceEmbeddings(model_name="avsolatorio/GIST-Embedding-v0")
 
-# Add this at the beginning of your script, after imports
 DOCUMENTS_FILE = "uploaded_documents.json"
 
 def load_documents():
@@ -82,7 +54,6 @@ def save_documents(documents):
     with open(DOCUMENTS_FILE, "w") as f:
         json.dump(documents, f)
 
-# Replace the global uploaded_documents with this
 uploaded_documents = load_documents()
 
 # Modify the update_vectors function
@@ -181,142 +152,37 @@ def delete_documents(selected_docs):
     
     return f"Deleted documents: {', '.join(deleted_docs)}", display_documents()
 
-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 refresh_documents():
+    global uploaded_documents
+    uploaded_documents = load_documents()
+    return display_documents()
 
-class CitingSources(BaseModel):
-    sources: List[str] = Field(
-        ...,
-        description="List of sources to cite. Should be an URL of the source."
+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"
     )
-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)
 
-from duckduckgo_search import DDGS
+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!")
+    ]
 
-def respond(message, history, use_web_search, model, temperature, num_calls, selected_docs):
+def respond(message, history, use_web_search, selected_docs, model):
     logging.info(f"User Query: {message}")
     logging.info(f"Search Type: {'Web Search' if use_web_search else 'PDF Search'}")
     logging.info(f"Selected Documents: {selected_docs}")
+    logging.info(f"Model Used: {model}")
 
     try:
         if use_web_search:
-            for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
+            for main_content, sources in get_response_with_search(message):
                 response = f"{main_content}\n\n{sources}"
                 yield response
         else:
@@ -324,10 +190,11 @@ def respond(message, history, use_web_search, model, temperature, num_calls, sel
             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
@@ -337,231 +204,90 @@ def respond(message, history, use_web_search, model, temperature, num_calls, sel
                 context_str = "No documents available."
                 yield "No documents available. Please upload PDF documents to answer questions."
                 return
-            
-            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}'"""
-
-            try:
-                response = DDGS().chat(prompt, model="llama-3-70b")
-                yield response
-            except Exception as e:
-                logging.error(f"Error with DuckDuckGo chat API: {str(e)}")
-                logging.info("Falling back to Hugging Face API")
-                yield from get_response_from_pdf(message, model, selected_docs, num_calls=num_calls, temperature=temperature)
 
+            for partial_response in get_response_from_duckduckgo(message, context_str, model):
+                yield partial_response
     except Exception as e:
-        logging.error(f"Error: {str(e)}")
-        yield f"An error occurred: {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"):
-    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
-    }
+        logging.error(f"Error in responding: {str(e)}")
+        yield f"An error occurred: {str(e)}. Please try again later."
 
-    full_response = ""
-    for i in range(num_calls):
+def get_response_from_duckduckgo(query, context, model):
+    with DDGS() as ddgs:
         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."
+            response = ddgs.chat(f"Using the following context:\n{context}\nRespond to the following query: {query}", model=model, timeout=30)
+            yield response
         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."
+            logging.error(f"Error in getting response from Duckduckgo: {str(e)}")
+            yield f"An error occurred: {str(e)}. 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 chatbot_interface(message, history, use_web_search, model, selected_docs):
+    if not message.strip():
+        return "", history
 
-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']}" for result in search_results])
-    
-    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."""
+    history = history + [(message, "")]
 
     try:
-        response = DDGS().chat(prompt, model="llama-3-70b")
-        yield response, ""
+        for response in respond(message, history, use_web_search, selected_docs, model):
+            history[-1] = (message, response)
+            yield history
+    except gr.CancelledError:
+        yield history
     except Exception as e:
-        logging.error(f"Error with DuckDuckGo chat API: {str(e)}")
-        logging.info("Falling back to Hugging Face API")
-        yield from get_response_from_huggingface(prompt, model, num_calls, temperature)
-
-def get_response_from_huggingface(prompt, model, num_calls=3, temperature=0.2):
-    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")
-        yield "No documents available. Please upload PDF documents to answer questions."
-        return
+        logging.error(f"Unexpected error in chatbot_interface: {str(e)}")
+        history[-1] = (message, f"An unexpected error occurred: {str(e)}")
+        yield history
 
-    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)}")
+def get_response_with_search(query):
+    search_results = duckduckgo_search(query)
+    web_search_database = create_web_search_vectors(search_results)
 
-    if not filtered_docs:
-        logging.warning(f"No documents found for the selected sources: {selected_docs}")
-        yield "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
+    if not web_search_database:
+        yield "No web search results available. Please try again.", ""
         return
 
-    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}")
+    retriever = web_search_database.as_retriever(search_kwargs={"k": 5})
     relevant_docs = retriever.get_relevant_documents(query)
-    logging.info(f"Number of relevant documents retrieved: {len(relevant_docs)}")
 
-    context_str = "\n".join([doc.page_content for doc in relevant_docs])
-    logging.info(f"Total context length: {len(context_str)}")
+    context = "\n".join([doc.page_content for doc in relevant_docs])
 
-    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}'"""
+    for partial_response in get_response_from_duckduckgo(query, context, "gpt-4o-mini"):  # Use the default model for web search
+        yield partial_response, ""  # Yield streaming response without sources
 
-    try:
-        response = DDGS().chat(prompt, model="llama-3-70b")
-        yield response
-    except Exception as e:
-        logging.error(f"Error with DuckDuckGo chat API: {str(e)}")
-        logging.info("Falling back to Hugging Face API")
-        yield from get_response_from_huggingface(prompt, model, num_calls, temperature)
-        
-        logging.info("Finished generating response")
+def create_web_search_vectors(search_results):
+    embed = get_embeddings()
 
-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}")
-
-css = """
-/* Fine-tune chatbox size */
-.chatbot-container {
-    height: 600px !important;
-    width: 100% !important;
-}
-.chatbot-container > div {
-    height: 100%;
-    width: 100%;
-}
-"""
-
-uploaded_documents = []
+    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']}))
 
-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"
-    )
+    return FAISS.from_documents(documents, embed)
 
-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()
 
-# Define the checkbox outside the demo block
-document_selector = gr.CheckboxGroup(label="Select documents to query")
+DUCKDUCKGO_CHAT_MODELS = [
+    "gpt-4o-mini",
+    "claude-3-haiku",
+    "llama-3.1-70b",
+    "mixtral-8x7b"
+]
 
+document_selector = gr.CheckboxGroup(label="Select documents to query")
 use_web_search = gr.Checkbox(label="Use Web Search", value=True)
+model_selector = gr.Dropdown(choices=DUCKDUCKGO_CHAT_MODELS, label="Select Duckduckgo Chat Model", value="gpt-4o-mini")
 
 custom_placeholder = "Ask a question (Note: You can toggle between Web Search and PDF Chat in Additional Inputs below)"
 
 demo = gr.ChatInterface(
-    respond,
+    chatbot_interface,
     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
+        document_selector,
+        model_selector
     ],
     title="AI-powered Web Search and PDF Chat Assistant",
     description="Chat with your PDFs or use web search to answer questions. Toggle between Web Search and PDF Chat in Additional Inputs below.",
@@ -592,16 +318,15 @@ 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,
-    value=initial_conversation()
-)
+    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")
 
@@ -610,21 +335,18 @@ with demo:
         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])
@@ -635,12 +357,12 @@ with demo:
     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.
+    4. Select the Duckduckgo Chat Model you want to use.
+    5. Ask questions in the chat interface. 
+    6. Toggle "Use Web Search" to switch between PDF chat and web search.
     7. Use the provided examples or ask your own questions.
     """
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=True)