Update app.py

app.py

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Tuple
 import gradio as gr
 from langchain_openai import OpenAIEmbeddings
 from langchain_community.vectorstores import Chroma
@@ -16,7 +16,7 @@ openai.api_key = "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"
 def add_text(history, text: str):
     if not text:
         raise gr.Error("Enter text")
-    history = history + [(text, "")]
+    history.append((text, ""))
     return history
 
 class MyApp:
@@ -24,18 +24,15 @@ class MyApp:
         self.OPENAI_API_KEY: str = openai.api_key
         self.chain = None
         self.chat_history: list = []
-        self.N: int = 0
-        self.count: int = 0
         self.documents = None
         self.file_name = None
 
-    def __call__(self, file: str) -> Any:
-        if self.count == 0:
+    def __call__(self, file: str) -> ConversationalRetrievalChain:
+        if self.chain is None:
             self.chain = self.build_chain(file)
-            self.count += 1
         return self.chain
 
-    def process_file(self, file: str):
+    def process_file(self, file: str) -> Image:
         loader = PyMuPDFLoader(file.name)
         self.documents = loader.load()
         pattern = r"/([^/]+)$"
@@ -44,14 +41,13 @@ class MyApp:
             self.file_name = match.group(1)
         except:
             self.file_name = os.path.basename(file)
-        # Render the first page for display
         doc = fitz.open(file.name)
         page = doc[0]
         pix = page.get_pixmap(dpi=150)
         image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
         return image
 
-    def build_chain(self, file: str):
+    def build_chain(self, file: str) -> str:
         embeddings = OpenAIEmbeddings(openai_api_key=self.OPENAI_API_KEY)
         pdfsearch = Chroma.from_documents(
             self.documents,
@@ -72,23 +68,25 @@ def get_response(history, query, file):
     result = chain(
         {"question": query, "chat_history": app.chat_history}, return_only_outputs=True
     )
-    app.chat_history += [(query, result["answer"])]
-    app.N = list(result["source_documents"][0])[1][1]["page"]
-    highlighted_line = result["answer"]
+    app.chat_history.append((query, result["answer"]))
+    source_docs = result["source_documents"]
+    source_texts = []
+    for doc in source_docs:
+        source_texts.append(f"Page {doc.metadata['page'] + 1}: {doc.page_content}")
+    source_texts_str = "\n\n".join(source_texts)
     for char in result["answer"]:
         history[-1][-1] += char
-        yield history, "", f"Page: {app.N + 1}, Highlight: {highlighted_line}"
+        yield history, "", source_texts_str
 
-def render_file(file):
+def render_file(file) -> Image:
     doc = fitz.open(file.name)
-    page = doc[app.N]
+    page = doc[0]
     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):
+def purge_chat_and_render_first(file) -> Tuple[Image, list]:
     app.chat_history = []
-    app.count = 0
     doc = fitz.open(file.name)
     page = doc[0]
     pix = page.get_pixmap(dpi=150)
@@ -123,7 +121,7 @@ with gr.Blocks() as demo:
         )
         submit_btn = gr.Button("Submit", scale=1)
         refresh_btn = gr.Button("Refresh Chat", scale=1)
-        temperature_slider = gr.Slider(0, 1, value=0.0, label="Temperature")
+        source_texts_output = gr.Textbox(label="Source Texts", interactive=False)
 
     btn.upload(
         fn=purge_chat_and_render_first,
@@ -149,9 +147,7 @@ with gr.Blocks() as demo:
         outputs=[chatbot],
         queue=False,
     ).success(
-        fn=get_response, inputs=[chatbot, txt, btn], outputs=[chatbot, txt]
-    ).success(
-        fn=render_file, inputs=[btn], outputs=[show_img_processed]
+        fn=get_response, inputs=[chatbot, txt, btn], outputs=[chatbot, txt, source_texts_output]
     )
 
     refresh_btn.click(