app.py

import os
import json
import re
import gradio as gr
import requests
import random
import urllib.parse
from tempfile import NamedTemporaryFile
from bs4 import BeautifulSoup
from typing import List
from pydantic import BaseModel, Field
from huggingface_hub import InferenceApi
from duckduckgo_search import DDGS
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.llms import HuggingFaceHub
from langchain_core.documents import Document
from sentence_transformers import SentenceTransformer
from llama_parse import LlamaParse
import logging

# Set up logging
logging.basicConfig(level=logging.DEBUG, 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")

# Initialize SentenceTransformer and LlamaParse
sentence_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
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 = "pypdf") -> List[Document]:
    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 update_vectors(files, parser):
    if not files:
        return "Please upload at least one PDF file."
    
    embed = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
    total_chunks = 0
    
    all_data = []
    for file in files:
        data = load_document(file, parser)
        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 using {parser}."

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 duckduckgo_search(query):
    logging.debug(f"Performing DuckDuckGo search for query: {query}")
    with DDGS() as ddgs:
        results = list(ddgs.text(query, max_results=5))
    logging.debug(f"Search returned {len(results)} results")
    return results

def get_response_with_search(query, temperature, top_p, repetition_penalty, use_pdf=False):
    logging.debug(f"Getting response for query: {query}")
    logging.debug(f"Parameters: temperature={temperature}, top_p={top_p}, repetition_penalty={repetition_penalty}, use_pdf={use_pdf}")
    
    model = get_model(temperature, top_p, repetition_penalty)
    embed = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")

    if use_pdf and os.path.exists("faiss_database"):
        logging.debug("Using PDF database for context")
        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
        retriever = database.as_retriever()
        relevant_docs = retriever.get_relevant_documents(query)
        context = "\n".join([f"Content: {doc.page_content}\nSource: {doc.metadata['source']}\n" for doc in relevant_docs])
    else:
        logging.debug("Using web search for context")
        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)
    
    logging.debug(f"Context generated. Length: {len(context)} characters")

    prompt = f"""<s>[INST] Using the following context:
{context}
Write a detailed and complete research document that fulfills the following user request: '{query}'
After the main content, provide a list of sources used in your response, prefixed with 'Sources:'. 
Do not include any part of these instructions in your response. [/INST]"""

    logging.debug("Sending prompt to model")
    response = model(prompt)
    logging.debug(f"Received response from model. Length: {len(response)} characters")
    
    main_content, sources = split_response(response)
    
    logging.debug(f"Split response. Main content length: {len(main_content)}, Sources length: {len(sources)}")
    return main_content, sources

def split_response(response):
    logging.debug("Splitting response")
    logging.debug(f"Original response: {response[:100]}...") # Log first 100 characters
    
    # Remove any remaining instruction text
    response = re.sub(r'\[/?INST\]', '', response)
    response = re.sub(r'~~.*?~~', '', response, flags=re.DOTALL)
    
    logging.debug(f"After removing instructions: {response[:100]}...") # Log first 100 characters
    
    # Split the response into main content and sources
    parts = response.split("Sources:", 1)
    main_content = parts[0].strip()
    sources = parts[1].strip() if len(parts) > 1 else ""
    
    logging.debug(f"Main content starts with: {main_content[:100]}...") # Log first 100 characters
    logging.debug(f"Sources: {sources[:100]}...") # Log first 100 characters
    
    return main_content, sources

def chatbot_interface(message, history, temperature, top_p, repetition_penalty, use_pdf):
    logging.debug(f"Chatbot interface called with message: {message}")
    main_content, sources = get_response_with_search(message, temperature, top_p, repetition_penalty, use_pdf)
    formatted_response = f"{main_content}\n\nSources:\n{sources}"
    logging.debug(f"Formatted response: {formatted_response[:100]}...") # Log first 100 characters
    return formatted_response

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# AI-powered Web Search and PDF Chat Assistant")
    
    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="pypdf")
        update_button = gr.Button("Upload PDF")
    
    update_output = gr.Textbox(label="Update Status")
    update_button.click(update_vectors, inputs=[file_input, parser_dropdown], outputs=update_output)
    
    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(label="Conversation")
            msg = gr.Textbox(label="Ask a question")
            submit_button = gr.Button("Submit")
        with gr.Column(scale=1):
            temperature = gr.Slider(label="Temperature", minimum=0.0, maximum=1.0, value=0.7, step=0.1)
            top_p = gr.Slider(label="Top P", minimum=0.0, maximum=1.0, value=0.95, step=0.05)
            repetition_penalty = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, value=1.1, step=0.1)
            use_pdf = gr.Checkbox(label="Use PDF Documents", value=False)

    def respond(message, chat_history, temperature, top_p, repetition_penalty, use_pdf):
        bot_message = chatbot_interface(message, chat_history, temperature, top_p, repetition_penalty, use_pdf)
        chat_history.append((message, bot_message))
        return "", chat_history

    submit_button.click(respond, inputs=[msg, chatbot, temperature, top_p, repetition_penalty, use_pdf], outputs=[msg, 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()