Update app.py

app.py

@@ -22,8 +22,13 @@ from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.llms import HuggingFaceHub
 from langchain_core.documents import Document
 from sentence_transformers import SentenceTransformer
+import nest_asyncio
+from llama_parse import LlamaParse
+
+nest_asyncio.apply()
 
 huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
+llama_cloud_api_key = os.environ.get("LLAMA_CLOUD_API_KEY")
 
 # Load SentenceTransformer model
 sentence_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
@@ -108,12 +113,28 @@ class EnhancedContextDrivenChatbot:
         
         return contextualized_question, topics, self.entity_tracker
         
-def load_document(file: NamedTemporaryFile) -> List[Document]:
+# 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[Document]:
     """Loads and splits the document into pages."""
-    loader = PyPDFLoader(file.name)
-    return loader.load_and_split()
+    if parser == "pypdf":
+        loader = PyPDFLoader(file.name)
+        return loader.load_and_split()
+    elif parser == "llamaparse":
+        documents = llama_parser.load_data(file.name)
+        # Convert LlamaParse output to langchain Document format
+        return [Document(page_content=doc.text, metadata={"source": file.name}) for doc in documents]
+    else:
+        raise ValueError("Invalid parser specified. Use 'pypdf' or 'llamaparse'.")
 
-def update_vectors(files):
+def update_vectors(files, parser):
     if not files:
         return "Please upload at least one PDF file."
     
@@ -122,7 +143,7 @@ def update_vectors(files):
     
     all_data = []
     for file in files:
-        data = load_document(file)
+        data = load_document(file, parser)
         all_data.extend(data)
         total_chunks += len(data)
     
@@ -134,7 +155,7 @@ def update_vectors(files):
     
     database.save_local("faiss_database")
     
-    return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files."
+    return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}."
 
 def get_embeddings():
     return HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
@@ -410,17 +431,17 @@ def ask_question(question, temperature, top_p, repetition_penalty, web_search, c
 
     return "An unexpected error occurred. Please try again later."
 
-# Gradio interface
 # Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Context-Driven Conversational Chatbot")
+    gr.Markdown("# Enhanced Context-Driven Conversational Chatbot")
     
     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], outputs=update_output)
+    update_button.click(update_vectors, inputs=[file_input, parser_dropdown], outputs=update_output)
     
     with gr.Row():
         with gr.Column(scale=2):
@@ -433,10 +454,10 @@ with gr.Blocks() as demo:
             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)
 
-    context_driven_chatbot = EnhancedContextDrivenChatbot()
+    enhanced_context_driven_chatbot = EnhancedContextDrivenChatbot()
 
     def chat(question, history, temperature, top_p, repetition_penalty, web_search):
-        answer = ask_question(question, temperature, top_p, repetition_penalty, web_search, context_driven_chatbot)
+        answer = ask_question(question, temperature, top_p, repetition_penalty, web_search, enhanced_context_driven_chatbot)
         history.append((question, answer))
         return "", history