Update app.py

app.py

@@ -1,26 +1,13 @@
-import os
-import json
 import gradio as gr
 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.documents import Document
-from langchain_community.document_loaders import PyPDFLoader
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from llama_parse import LlamaParse
+from typing import List, Dict
+import os
 import logging
-import shutil
-
-# Initialize LlamaParse
-llama_parser = LlamaParse(
-    api_key=os.environ.get("LLAMA_CLOUD_API_KEY"),
-    result_type="markdown",
-    num_workers=4,
-    verbose=True,
-    language="en",
-)
+
+logging.basicConfig(level=logging.INFO)
+
+# Environment variables and configurations
+huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
 
 def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[Document]:
     """Loads and splits the document into pages."""
@@ -42,6 +29,7 @@ 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():
@@ -54,6 +42,7 @@ 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
@@ -152,180 +141,180 @@ def delete_documents(selected_docs):
     
     return f"Deleted documents: {', '.join(deleted_docs)}", display_documents()
 
-def refresh_documents():
-    global uploaded_documents
-    uploaded_documents = load_documents()
-    return display_documents()
-
-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 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 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 @desai.[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, 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):
-                response = f"{main_content}\n\n{sources}"
-                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})
+    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."
+
+    # 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]
+    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."
+
+    # 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
 
-                # 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]
+    context_str = "\n".join([doc.page_content for doc in relevant_docs])
+    logging.info(f"Total context length: {len(context_str)}")
 
-                if not relevant_docs:
-                    yield "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
-                    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: '{query}'"""
 
-                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
+    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)
 
-            for partial_response in get_response_from_duckduckgo(message, context_str, model):
-                yield partial_response
-    except Exception as e:
-        logging.error(f"Error in responding: {str(e)}")
-        yield f"An error occurred: {str(e)}. Please try again later."
+    logging.info("Finished generating response")
+    return response
 
-def get_response_from_duckduckgo(query, context, model):
-    with DDGS() as ddgs:
-        try:
-            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 getting response from Duckduckgo: {str(e)}")
-            yield f"An error occurred: {str(e)}. Please try again later."
+class ConversationManager:
+    def __init__(self):
+        self.history = []
+        self.current_context = None
 
-def chatbot_interface(message, history, use_web_search, model, selected_docs):
-    if not message.strip():
-        return "", history
+    def add_interaction(self, query, response):
+        self.history.append((query, response))
+        self.current_context = f"Previous query: {query}\nPrevious response summary: {response[:200]}..."
 
-    history = history + [(message, "")]
+    def get_context(self):
+        return self.current_context
 
+def get_web_search_results(query: str, max_results: int = 10) -> List[Dict[str, str]]:
     try:
-        for response in respond(message, history, use_web_search, selected_docs, model):
-            history[-1] = (message, response)
-            yield history
-    except gr.CancelledError:
-        yield history
+        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:
-        logging.error(f"Unexpected error in chatbot_interface: {str(e)}")
-        history[-1] = (message, f"An unexpected error occurred: {str(e)}")
-        yield history
-
-def get_response_with_search(query):
-    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
+        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}
+        
+        New query: {original_query}
+        
+        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])
 
-    retriever = web_search_database.as_retriever(search_kwargs={"k": 5})
-    relevant_docs = retriever.get_relevant_documents(query)
+        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:
 
-    context = "\n".join([doc.page_content for doc in relevant_docs])
+        {search_context}
 
-    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
+        Article:"""
 
-def create_web_search_vectors(search_results):
-    embed = get_embeddings()
+        summary = DDGS().chat(prompt, model="llama-3-70b")
+        return summary
+    except Exception as e:
+        return f"An error occurred during summarization: {str(e)}"
 
-    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']}))
+conversation_manager = ConversationManager()
 
-    return FAISS.from_documents(documents, embed)
+def respond(message, chat_history, temperature, num_api_calls):
+    final_summary = ""
+    original_query = message
+    rephrased_query = rephrase_query(message, conversation_manager)
 
-def refresh_documents():
-    global uploaded_documents
-    uploaded_documents = load_documents()
-    return display_documents()
+    logging.info(f"Original query: {original_query}")
+    logging.info(f"Rephrased query: {rephrased_query}")
 
-DUCKDUCKGO_CHAT_MODELS = [
-    "gpt-4o-mini",
-    "claude-3-haiku",
-    "llama-3.1-70b",
-    "mixtral-8x7b"
-]
+    for _ in range(num_api_calls):
+        search_results = get_web_search_results(rephrased_query)
 
-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")
+        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"
 
-custom_placeholder = "Ask a question (Note: You can toggle between Web Search and PDF Chat in Additional Inputs below)"
+    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."
 
+# The rest of your code (CSS, theme, and Gradio interface setup) remains the same
 css = """
-/* Fine-tune chatbox size */
-.chatbot-container {
-    height: 600px !important;
-    width: 100% !important;
-}
-.chatbot-container > div {
-    height: 100%;
-    width: 100%;
-}
+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"
+)
+
 demo = gr.ChatInterface(
-    chatbot_interface,
+    respond,
     additional_inputs=[
-        use_web_search,
-        document_selector,
-        model_selector
+        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")
     ],
     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.",
-    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"
-    ),
+    description="This AI-powered Web Search and PDF Chat Assistant combines real-time web search capabilities with advanced language processing.",
+    theme=theme,
     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 is AI"],
+        ["Any recent news on US Banks"],
+        ["Who is Donald Trump"]
     ],
     cache_examples=False,
     analytics_enabled=False,
@@ -335,46 +324,7 @@ demo = gr.ChatInterface(
         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"])
-        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_button.click(update_vectors, 
-                        inputs=[file_input, parser_dropdown], 
-                        outputs=[update_output, document_selector])
-
-    refresh_button.click(refresh_documents, 
-                         inputs=[], 
-                         outputs=[document_selector])
-
-    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. 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()