Update app.py

app.py

@@ -1,22 +1,17 @@
-
 from typing import Any
 import gradio as gr
 from langchain_openai import OpenAIEmbeddings
 from langchain_community.vectorstores import Chroma
-
 from langchain.chains import ConversationalRetrievalChain
 from langchain_openai import ChatOpenAI
-
 from langchain_community.document_loaders import PyMuPDFLoader
-
 import fitz
 from PIL import Image
 import os
 import re
 import openai
 
-openai.api_key = "sk-proj-9Wyp6JPj6G0QDIENCob2T3BlbkFJR2y4KpNGyS9C1KamMm9S"
-
+openai.api_key = "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"
 
 def add_text(history, text: str):
     if not text:
@@ -24,7 +19,6 @@ def add_text(history, text: str):
     history = history + [(text, "")]
     return history
 
-
 class MyApp:
     def __init__(self) -> None:
         self.OPENAI_API_KEY: str = openai.api_key
@@ -48,12 +42,10 @@ class MyApp:
             file_name = match.group(1)
         except:
             file_name = os.path.basename(file)
-
         return documents, file_name
 
     def build_chain(self, file: str):
         documents, file_name = self.process_file(file)
-        # Load embeddings model
         embeddings = OpenAIEmbeddings(openai_api_key=self.OPENAI_API_KEY)
         pdfsearch = Chroma.from_documents(
             documents,
@@ -67,7 +59,6 @@ class MyApp:
         )
         return chain
 
-
 def get_response(history, query, file):
     if not file:
         raise gr.Error(message="Upload a PDF")
@@ -77,34 +68,30 @@ def get_response(history, query, file):
     )
     app.chat_history += [(query, result["answer"])]
     app.N = list(result["source_documents"][0])[1][1]["page"]
+    highlighted_line = result["answer"]
     for char in result["answer"]:
         history[-1][-1] += char
-        yield history, ""
-
+        yield history, "", f"Page: {app.N + 1}, Highlight: {highlighted_line}"
 
 def render_file(file):
     doc = fitz.open(file.name)
     page = doc[app.N]
-    # Render the page as a PNG image with a resolution of 150 DPI
     pix = page.get_pixmap(dpi=150)
     image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
     return image
 
-
 def purge_chat_and_render_first(file):
-    print("purge_chat_and_render_first")
-    # Purges the previous chat session so that the bot has no concept of previous documents
     app.chat_history = []
     app.count = 0
-
-    # Use PyMuPDF to render the first page of the uploaded document
     doc = fitz.open(file.name)
     page = doc[0]
-    # Render the page as a PNG image with a resolution of 150 DPI
     pix = page.get_pixmap(dpi=150)
     image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
     return image, []
 
+def refresh_chat():
+    app.chat_history = []
+    return []
 
 app = MyApp()
 
@@ -112,21 +99,18 @@ with gr.Blocks() as demo:
     with gr.Column():
         with gr.Row():
             with gr.Column(scale=2):
-                with gr.Row():
-                    chatbot = gr.Chatbot(value=[], elem_id="chatbot")
-                with gr.Row():
-                    txt = gr.Textbox(
-                        show_label=False,
-                        placeholder="Enter text and press submit",
-                        scale=2
-                    )
-                    submit_btn = gr.Button("Submit", scale=1)
+                chatbot = gr.Chatbot(value=[], elem_id="chatbot")
+                txt = gr.Textbox(
+                    show_label=False,
+                    placeholder="Enter text and press submit",
+                    scale=2
+                )
+                submit_btn = gr.Button("Submit", scale=1)
+                refresh_btn = gr.Button("Refresh Chat", scale=1)
 
             with gr.Column(scale=1):
-                with gr.Row():
-                    show_img = gr.Image(label="Upload PDF")
-                with gr.Row():
-                    btn = gr.UploadButton("📁 Upload a PDF", file_types=[".pdf"])
+                show_img = gr.Image(label="Upload PDF")
+                btn = gr.UploadButton("📁 Upload a PDF", file_types=[".pdf"])
 
     btn.upload(
         fn=purge_chat_and_render_first,
@@ -147,7 +131,11 @@ with gr.Blocks() as demo:
         fn=render_file, inputs=[btn], outputs=[show_img]
     )
 
+    refresh_btn.click(
+        fn=refresh_chat,
+        inputs=[],
+        outputs=[chatbot],
+    )
+
 demo.queue()
 demo.launch()
-
-