Update app.py

app.py

@@ -1,432 +1,5 @@
-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_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
-
-# 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"
-]
-
-# 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."""
-    if parser == "pypdf":
-        loader = PyPDFLoader(file.name)
-        return loader.load_and_split()
-    elif parser == "llamaparse":
-        try:
-            documents = llama_parser.load_data(file.name)
-            return [Document(page_content=doc.text, metadata={"source": file.name}) for doc in documents]
-        except Exception as e:
-            print(f"Error using Llama Parse: {str(e)}")
-            print("Falling back to PyPDF parser")
-            loader = PyPDFLoader(file.name)
-            return loader.load_and_split()
-    else:
-        raise ValueError("Invalid parser specified. Use 'pypdf' or 'llamaparse'.")
-
-def get_embeddings():
-    return HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
-
-def update_vectors(files, parser):
-    global uploaded_documents
-    logging.info(f"Entering update_vectors with {len(files)} files and parser: {parser}")
-    
-    if not files:
-        logging.warning("No files provided for update_vectors")
-        return "Please upload at least one PDF file.", 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"
-        )
-    
-    embed = get_embeddings()
-    total_chunks = 0
-    
-    all_data = []
-    for file in files:
-        logging.info(f"Processing file: {file.name}")
-        try:
-            data = load_document(file, parser)
-            logging.info(f"Loaded {len(data)} chunks from {file.name}")
-            all_data.extend(data)
-            total_chunks += len(data)
-            if not any(doc["name"] == file.name for doc in uploaded_documents):
-                uploaded_documents.append({"name": file.name, "selected": True})
-                logging.info(f"Added new document to uploaded_documents: {file.name}")
-            else:
-                logging.info(f"Document already exists in uploaded_documents: {file.name}")
-        except Exception as e:
-            logging.error(f"Error processing file {file.name}: {str(e)}")
-    
-    logging.info(f"Total chunks processed: {total_chunks}")
-    
-    if os.path.exists("faiss_database"):
-        logging.info("Updating existing FAISS database")
-        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
-        database.add_documents(all_data)
-    else:
-        logging.info("Creating new FAISS database")
-        database = FAISS.from_documents(all_data, embed)
-    
-    database.save_local("faiss_database")
-    logging.info("FAISS database saved")
-    
-    return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}.", 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"
-    )
-
-def duckduckgo_search(query):
-    with DDGS() as ddgs:
-        results = ddgs.text(query, max_results=5)
-    return results
-
-class CitingSources(BaseModel):
-    sources: List[str] = Field(
-        ...,
-        description="List of sources to cite. Should be an URL of the source."
-    )
-
-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
-    }
-
-    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 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"""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."""
-
-    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=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, 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
-
-    retriever = database.as_retriever()
-    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)}")
-    
-    # Filter relevant_docs based on selected documents
-    filtered_docs = [doc for doc in relevant_docs if doc.metadata["source"] in selected_docs]
-    logging.info(f"Number of filtered documents: {len(filtered_docs)}")
-    
-    if not filtered_docs:
-        logging.warning(f"No relevant information found in the selected documents: {selected_docs}")
-        yield "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
-        return
-
-    for doc in filtered_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
-
-    context_str = "\n".join([doc.page_content for doc in filtered_docs])
-    logging.info(f"Total context length: {len(context_str)}")
-
-    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}'"""
-        
-        client = InferenceClient(model, token=huggingface_token)
-        
-        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=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
-        
-        logging.info("Finished generating response")
-
-def continue_response(last_response, context, query, model, temperature):
-    prompt = f"""Using the following context and partial response:
-
-Context:
-{context}
-
-Partial Response:
-{last_response}
-
-Continue the response to fully answer the query: '{query}'
-Make sure the continuation flows smoothly from the previous part."""
-
-    if model == "@cf/meta/llama-3.1-8b-instruct":
-        return get_response_from_cloudflare(prompt="", context=context, query=prompt, num_calls=1, temperature=temperature, search_type="pdf")
-    else:
-        client = InferenceClient(model, token=huggingface_token)
-        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:
-                yield message.choices[0].delta.content
-
-def chatbot_interface(message, history, use_web_search, model, temperature, num_calls, selected_docs):
-    if not message.strip():
-        return "", history
-
-    history = history + [(message, "")]
-
-    try:
-        last_response = ""
-        for response in respond(message, history, model, temperature, num_calls, use_web_search, selected_docs):
-            last_response = response
-            history[-1] = (message, response)
-            yield history
-
-        # Check if the response seems truncated
-        if not last_response.strip().endswith((".", "!", "?")):
-            history.append((None, "Response may be incomplete. Type 'continue' to generate more."))
-            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 continue_generation(history, use_web_search, model, temperature, selected_docs):
-    if not history:
-        return history, gr.Button.update(visible=False)
-
-    last_message = history[-1][0]
-    last_response = history[-1][1]
-
-    if use_web_search:
-        search_results = duckduckgo_search(last_message)
-        context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
-                            for result in search_results if 'body' in result)
-    else:
-        embed = get_embeddings()
-        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
-        retriever = database.as_retriever()
-        relevant_docs = retriever.get_relevant_documents(last_message)
-        filtered_docs = [doc for doc in relevant_docs if doc.metadata["source"] in selected_docs]
-        context = "\n".join([doc.page_content for doc in filtered_docs])
-
-    continuation = ""
-    for chunk in continue_response(last_response, context, last_message, model, temperature):
-        continuation += chunk
-        history[-1] = (last_message, last_response + continuation)
-        yield history, gr.Button.update(visible=True)
-
-    if not (last_response + continuation).strip().endswith((".", "!", "?")):
-        yield history, gr.Button.update(visible=True, text="Continue Generation")
-    else:
-        yield history, gr.Button.update(visible=False)
-
-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}")
-
-    # Check if the user wants to continue the previous response
-    if message.strip().lower() == "continue" and history:
-        last_message = history[-2][0]  # Get the last user message
-        last_response = history[-2][1]  # Get the last bot response
-        context = get_context(last_message, use_web_search, selected_docs)
-        for continuation in continue_response(last_response, context, last_message, model, temperature):
-            yield last_response + continuation
-    else:
-        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:
-                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."
-
-def get_context(message, use_web_search, selected_docs):
-    if use_web_search:
-        search_results = duckduckgo_search(message)
-        return "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
-                         for result in search_results if 'body' in result)
-    else:
-        embed = get_embeddings()
-        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
-        retriever = database.as_retriever()
-        relevant_docs = retriever.get_relevant_documents(message)
-        filtered_docs = [doc for doc in relevant_docs if doc.metadata["source"] in selected_docs]
-        return "\n".join([doc.page_content for doc in filtered_docs])
-
-
-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 = """
-/* Add your custom CSS here */
-"""
-
-uploaded_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"
-    )
-
-# Define the checkbox outside the demo block
-document_selector = gr.CheckboxGroup(label="Select documents to query")
-
-use_web_search = gr.Checkbox(label="Use Web Search", value=False)
-
 demo = gr.ChatInterface(
-    chatbot_interface,
+    respond,
     additional_inputs=[
         gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0]),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),
@@ -435,7 +8,7 @@ demo = gr.ChatInterface(
         document_selector  # Add the document selector to the chat interface
     ],
     title="AI-powered Web Search and PDF Chat Assistant",
-    description="Chat with your PDFs or use web search to answer questions. Type 'continue' to generate more if a response seems incomplete.",
+    description="Chat with your PDFs or use web search to answer questions.",
     theme=gr.themes.Soft(
         primary_hue="orange",
         secondary_hue="amber",
@@ -467,6 +40,7 @@ demo = gr.ChatInterface(
 # 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")
@@ -489,7 +63,6 @@ with demo:
     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.
-    8. If a response seems incomplete, type 'continue' to generate more.
     """
     )