app.py

import os
import json
import re
import gradio as gr
import pandas as pd
import requests
import random
import urllib.parse
from tempfile import NamedTemporaryFile
from typing import List
from bs4 import BeautifulSoup
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyPDFLoader
from langchain_core.output_parsers import StrOutputParser
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.llms import HuggingFaceHub
from langchain_core.documents import Document  # Add this line

huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")

def load_document(file: NamedTemporaryFile) -> List[Document]:
    """Loads and splits the document into pages."""
    loader = PyPDFLoader(file.name)
    return loader.load_and_split()

def update_vectors(files):
    if not files:
        return "Please upload at least one PDF file."
    
    embed = get_embeddings()
    total_chunks = 0
    
    all_data = []
    for file in files:
        data = load_document(file)
        all_data.extend(data)
        total_chunks += len(data)
    
    if os.path.exists("faiss_database"):
        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
        database.add_documents(all_data)
    else:
        database = FAISS.from_documents(all_data, embed)
    
    database.save_local("faiss_database")
    
    return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files."

def get_embeddings():
    return HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")

def clear_cache():
    if os.path.exists("faiss_database"):
        os.remove("faiss_database")
        return "Cache cleared successfully."
    else:
        return "No cache to clear."

def get_model(temperature, top_p, repetition_penalty):
    return HuggingFaceHub(
        repo_id="mistralai/Mistral-7B-Instruct-v0.3",
        model_kwargs={
            "temperature": temperature,
            "top_p": top_p,
            "repetition_penalty": repetition_penalty,
            "max_length": 1000
        },
        huggingfacehub_api_token=huggingface_token
    )

def generate_chunked_response(model, prompt, max_tokens=1000, max_chunks=5):
    full_response = ""
    for i in range(max_chunks):
        chunk = model(prompt + full_response, max_new_tokens=max_tokens)
        chunk = chunk.strip()
        if chunk.endswith((".", "!", "?")):
            full_response += chunk
            break
        full_response += chunk
    return full_response.strip()

def extract_text_from_webpage(html):
    soup = BeautifulSoup(html, 'html.parser')
    for script in soup(["script", "style"]):
        script.extract()
    text = soup.get_text()
    lines = (line.strip() for line in text.splitlines())
    chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
    text = '\n'.join(chunk for chunk in chunks if chunk)
    return text

_useragent_list = [
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36",
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Edge/91.0.864.59 Safari/537.36",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Edge/91.0.864.59 Safari/537.36",
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Safari/537.36",
    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Safari/537.36",
]

def google_search(term, num_results=5, lang="en", timeout=5, safe="active", ssl_verify=None):
    escaped_term = urllib.parse.quote_plus(term)
    start = 0
    all_results = []
    max_chars_per_page = 8000

    print(f"Starting Google search for term: '{term}'")

    with requests.Session() as session:
        while start < num_results:
            try:
                user_agent = random.choice(_useragent_list)
                headers = {
                    'User-Agent': user_agent
                }
                resp = session.get(
                    url="https://www.google.com/search",
                    headers=headers,
                    params={
                        "q": term,
                        "num": num_results - start,
                        "hl": lang,
                        "start": start,
                        "safe": safe,
                    },
                    timeout=timeout,
                    verify=ssl_verify,
                )
                resp.raise_for_status()
                print(f"Successfully retrieved search results page (start={start})")
            except requests.exceptions.RequestException as e:
                print(f"Error retrieving search results: {e}")
                break

            soup = BeautifulSoup(resp.text, "html.parser")
            result_block = soup.find_all("div", attrs={"class": "g"})
            if not result_block:
                print("No results found on this page")
                break
            
            print(f"Found {len(result_block)} results on this page")
            for result in result_block:
                link = result.find("a", href=True)
                if link:
                    link = link["href"]
                    print(f"Processing link: {link}")
                    try:
                        webpage = session.get(link, headers=headers, timeout=timeout)
                        webpage.raise_for_status()
                        visible_text = extract_text_from_webpage(webpage.text)
                        if len(visible_text) > max_chars_per_page:
                            visible_text = visible_text[:max_chars_per_page] + "..."
                        all_results.append({"link": link, "text": visible_text})
                        print(f"Successfully extracted text from {link}")
                    except requests.exceptions.RequestException as e:
                        print(f"Error retrieving webpage content: {e}")
                        all_results.append({"link": link, "text": None})
                else:
                    print("No link found for this result")
                    all_results.append({"link": None, "text": None})
            start += len(result_block)

    print(f"Search completed. Total results: {len(all_results)}")
    
    if not all_results:
        print("No search results found. Returning a default message.")
        return [{"link": None, "text": "No information found in the web search results."}]

    return all_results

def rephrase_for_search(query, model):
    rephrase_prompt = PromptTemplate(
        input_variables=["query"],
        template="""
        Your task is to rephrase the given conversational query into a concise, search-engine-friendly format.
        Remove any conversational elements and focus on the core information need.
        Provide ONLY the rephrased query without any additional text or explanations.
        
        Conversational query: {query}
        
        Rephrased query:"""
    )
    
    chain = LLMChain(llm=model, prompt=rephrase_prompt)
    response = chain.run(query=query).strip()
    
    rephrased_query = response.replace("Rephrased query:", "").strip()
    
    if rephrased_query.lower() == query.lower() or len(rephrased_query) > len(query) * 1.5:
        common_words = set(['the', 'a', 'an', 'in', 'on', 'at', 'to', 'for', 'of', 'with', 'by', 'from', 'up', 'about', 'into', 'over', 'after'])
        keywords = [word.lower() for word in query.split() if word.lower() not in common_words]
        keywords = [word for word in keywords if word.isalnum()]
        return ' '.join(keywords)
    
    return rephrased_query

def ask_question(question, temperature, top_p, repetition_penalty, web_search):
    if not question:
        return "Please enter a question."

    model = get_model(temperature, top_p, repetition_penalty)
    embed = get_embeddings()

    if os.path.exists("faiss_database"):
        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
    else:
        database = None

    if web_search:
        original_query = question
        rephrased_query = rephrase_for_search(original_query, model)
        print(f"Original query: {original_query}")
        print(f"Rephrased query: {rephrased_query}")
        
        if rephrased_query == original_query:
            print("Warning: Query was not rephrased. Using original query for search.")
        
        search_results = google_search(rephrased_query)
        web_docs = [Document(page_content=result["text"], metadata={"source": result["link"]}) for result in search_results if result["text"]]
        
        if database is None:
            database = FAISS.from_documents(web_docs, embed)
        else:
            database.add_documents(web_docs)
        
        database.save_local("faiss_database")
        
        context_str = "\n".join([f"Source: {doc.metadata['source']}\nContent: {doc.page_content}" for doc in web_docs])
        
        prompt_template = """
        Answer the question based on the following web search results:
        Web Search Results:
        {context}
        Original Question: {original_question}
        Rephrased Search Query: {rephrased_query}
        If the web search results don't contain relevant information, state that the information is not available in the search results.
        Provide a concise and direct answer to the original question without mentioning the web search or these instructions.
        Do not include any source information in your answer.
        """
        prompt_val = ChatPromptTemplate.from_template(prompt_template)
        formatted_prompt = prompt_val.format(context=context_str, original_question=question, rephrased_query=rephrased_query)
    else:
        if database is None:
            return "No documents available. Please upload documents or enable web search to answer questions."

        retriever = database.as_retriever()
        relevant_docs = retriever.get_relevant_documents(question)
        context_str = "\n".join([doc.page_content for doc in relevant_docs])

        prompt_template = """
        Answer the question based on the following context:
        Context:
        {context}
        Current Question: {question}
        If the context doesn't contain relevant information, state that the information is not available.
        Provide a concise and direct answer to the question.
        Do not include any source information in your answer.
        """
        prompt_val = ChatPromptTemplate.from_template(prompt_template)
        formatted_prompt = prompt_val.format(context=context_str, question=question)

    full_response = generate_chunked_response(model, formatted_prompt)
    
    answer_patterns = [
        r"Provide a concise and direct answer to the question without mentioning the web search or these instructions:",
        r"Provide a concise and direct answer to the question:",
        r"Answer:",
        r"Provide a concise and direct answer to the original question without mentioning the web search or these instructions:"
    ]
    
    for pattern in answer_patterns:
        match = re.split(pattern, full_response, flags=re.IGNORECASE)
        if len(match) > 1:
            answer = match[-1].strip()
            break
    else:
        answer = full_response.strip()

    # Add sources section
    if web_search:
        sources = set(doc.metadata['source'] for doc in web_docs)
        sources_section = "\n\nSources:\n" + "\n".join(f"- {source}" for source in sources)
        answer += sources_section

    return answer

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Chat with your PDF documents and Web Search")
    
    with gr.Row():
        file_input = gr.Files(label="Upload your PDF documents", file_types=[".pdf"])
        update_button = gr.Button("Update Vector Store")
    
    update_output = gr.Textbox(label="Update Status")
    update_button.click(update_vectors, inputs=[file_input], outputs=update_output)
    
    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(label="Conversation")
            question_input = gr.Textbox(label="Ask a question about your documents or use web search")
            submit_button = gr.Button("Submit")
        with gr.Column(scale=1):
            temperature_slider = gr.Slider(label="Temperature", minimum=0.0, maximum=1.0, value=0.5, step=0.1)
            top_p_slider = gr.Slider(label="Top P", minimum=0.0, maximum=1.0, value=0.9, step=0.1)
            repetition_penalty_slider = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, value=1.0, step=0.1)
            web_search_checkbox = gr.Checkbox(label="Enable Web Search", value=False)

    def chat(question, history, temperature, top_p, repetition_penalty, web_search):
        answer = ask_question(question, temperature, top_p, repetition_penalty, web_search)
        history.append((question, answer))
        return "", history
    
    submit_button.click(chat, inputs=[question_input, chatbot, temperature_slider, top_p_slider, repetition_penalty_slider, web_search_checkbox], outputs=[question_input, chatbot])
    
    clear_button = gr.Button("Clear Cache")
    clear_output = gr.Textbox(label="Cache Status")
    clear_button.click(clear_cache, inputs=[], outputs=clear_output)

if __name__ == "__main__":
    demo.launch()