Update app.py

app.py

@@ -1,16 +1,103 @@
+import os
+import json
+import re
 import gradio as gr
-from huggingface_hub import InferenceApi
-from duckduckgo_search import DDGS
 import requests
-import json
+from duckduckgo_search import DDGS
 from typing import List
 from pydantic import BaseModel, Field
-import os
+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
 
 # Environment variables and configurations
 huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
+llama_cloud_api_key = os.environ.get("LLAMA_CLOUD_API_KEY")
+
+# 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 = "pypdf") -> List[dict]:
+    """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 [{"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):
+    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, 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 generate_chunked_response(prompt, max_tokens=1000, max_chunks=5):
+    API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3"
+    headers = {"Authorization": f"Bearer {huggingface_token}"}
+    payload = {
+        "inputs": prompt,
+        "parameters": {
+            "max_new_tokens": max_tokens,
+            "temperature": 0.7,
+            "top_p": 0.95,
+            "top_k": 40,
+            "repetition_penalty": 1.1
+        }
+    }
+    
+    full_response = ""
+    for _ in range(max_chunks):
+        response = requests.post(API_URL, headers=headers, json=payload)
+        if response.status_code == 200:
+            result = response.json()
+            if isinstance(result, list) and len(result) > 0:
+                chunk = result[0].get('generated_text', '')
+                full_response += chunk
+                if chunk.endswith((".", "!", "?")):
+                    break
+            else:
+                break
+        else:
+            break
+    return full_response.strip()
 
-# Function to perform a DuckDuckGo search
 def duckduckgo_search(query):
     with DDGS() as ddgs:
         results = ddgs.text(query, max_results=5)
@@ -23,81 +110,106 @@ class CitingSources(BaseModel):
     )
 
 def get_response_with_search(query):
-    # Perform the web search
     search_results = duckduckgo_search(query)
-    
-    # Use the search results as context for the model
     context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
                         for result in search_results if 'body' in result)
     
-    # Prompt formatted for Mistral-7B-Instruct
     prompt = f"""<s>[INST] 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. [/INST]"""
     
-    # API endpoint for Mistral-7B-Instruct-v0.3
-    API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3"
+    generated_text = generate_chunked_response(prompt)
     
-    # Headers
-    headers = {"Authorization": f"Bearer {huggingface_token}"}
+    content_start = generated_text.find("[/INST]")
+    if content_start != -1:
+        generated_text = generated_text[content_start + 7:].strip()
     
-    # Payload
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            "max_new_tokens": 1000,
-            "temperature": 0.7,
-            "top_p": 0.95,
-            "top_k": 40,
-            "repetition_penalty": 1.1
-        }
-    }
+    parts = generated_text.split("Sources:", 1)
+    main_content = parts[0].strip()
+    sources = parts[1].strip() if len(parts) > 1 else ""
     
-    # Make the API call
-    response = requests.post(API_URL, headers=headers, json=payload)
-    
-    if response.status_code == 200:
-        result = response.json()
-        if isinstance(result, list) and len(result) > 0:
-            generated_text = result[0].get('generated_text', 'No text generated')
-            
-            # Remove the instruction part
-            content_start = generated_text.find("[/INST]")
-            if content_start != -1:
-                generated_text = generated_text[content_start + 7:].strip()
-            
-            # Split the response into main content and sources
-            parts = generated_text.split("Sources:", 1)
-            main_content = parts[0].strip()
-            sources = parts[1].strip() if len(parts) > 1 else ""
-            
-            return main_content, sources
-        else:
-            return f"Unexpected response format: {result}", ""
+    return main_content, sources
+
+def get_response_from_pdf(query):
+    embed = get_embeddings()
+    if os.path.exists("faiss_database"):
+        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
     else:
-        return f"Error: API returned status code {response.status_code}", ""
+        return "No documents available. Please upload PDF documents to answer questions.", ""
+
+    retriever = database.as_retriever()
+    relevant_docs = retriever.get_relevant_documents(query)
+    context_str = "\n".join([doc.page_content for doc in relevant_docs])
+
+    prompt = f"""<s>[INST] Using the following context from the PDF documents:
+{context_str}
+Write a detailed and complete response that answers the following user question: '{query}'
+After writing the response, please provide a list of sources used (document names) in your answer. [/INST]"""
 
-def chatbot_interface(message, history):
-    main_content, sources = get_response_with_search(message)
+    generated_text = generate_chunked_response(prompt)
+
+    content_start = generated_text.find("[/INST]")
+    if content_start != -1:
+        generated_text = generated_text[content_start + 7:].strip()
+
+    parts = generated_text.split("Sources:", 1)
+    main_content = parts[0].strip()
+    sources = parts[1].strip() if len(parts) > 1 else ""
+
+    return main_content, sources
+
+def chatbot_interface(message, history, use_web_search):
+    if use_web_search:
+        main_content, sources = get_response_with_search(message)
+    else:
+        main_content, sources = get_response_from_pdf(message)
+    
     formatted_response = f"{main_content}\n\nSources:\n{sources}"
     return formatted_response
 
-# Gradio chatbot interface
-iface = gr.ChatInterface(
-    fn=chatbot_interface,
-    title="AI-powered Web Search Assistant",
-    description="Ask questions, and I'll search the web and provide answers using the Mistral-7B-Instruct model.",
-    examples=[
-        ["What are the latest developments in AI?"],
-        ["Tell me about recent updates on GitHub"],
-        ["What are the best hotels in Galapagos, Ecuador?"],
-        ["Summarize recent advancements in Python programming"],
-    ],
-    retry_btn="Retry",
-    undo_btn="Undo",
-    clear_btn="Clear",
-)
+# 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 Document")
+    
+    update_output = gr.Textbox(label="Update Status")
+    update_button.click(update_vectors, inputs=[file_input, parser_dropdown], outputs=update_output)
+    
+    with gr.Row():
+        chatbot = gr.Chatbot(label="Conversation")
+        with gr.Column():
+            msg = gr.Textbox(label="Ask a question")
+            use_web_search = gr.Checkbox(label="Use Web Search", value=False)
+            submit = gr.Button("Submit")
+    
+    gr.Examples(
+        examples=[
+            ["What are the latest developments in AI?"],
+            ["Tell me about recent updates on GitHub"],
+            ["What are the best hotels in Galapagos, Ecuador?"],
+            ["Summarize recent advancements in Python programming"],
+        ],
+        inputs=msg,
+    )
+
+    submit.click(chatbot_interface, inputs=[msg, chatbot, use_web_search], outputs=[chatbot])
+    msg.submit(chatbot_interface, inputs=[msg, chatbot, use_web_search], outputs=[chatbot])
+
+    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. Ask questions in the textbox. 
+    4. Toggle "Use Web Search" to switch between PDF chat and web search.
+    5. Click "Submit" or press Enter to get a response.
+    """
+    )
 
 if __name__ == "__main__":
-    iface.launch(share=True)
+    demo.launch(share=True)