Update app.py

app.py

@@ -11,31 +11,23 @@ import os
 import re
 import openai
 
-openai.api_key = "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"
-
-# Load the saved PDF and prepare the chain
+# MyApp class to handle the processes
 class MyApp:
     def __init__(self) -> None:
-        self.OPENAI_API_KEY: str = openai.api_key
+        self.OPENAI_API_KEY: str = None  # Initialize with None
         self.chain = None
         self.chat_history: list = []
         self.documents = None
         self.file_name = None
 
-    def __call__(self, file: str) -> ConversationalRetrievalChain:
-        if self.chain is None:
-            self.chain = self.build_chain(file)
-        return self.chain
+    def set_api_key(self, api_key: str):
+        self.OPENAI_API_KEY = api_key
+        openai.api_key = api_key
 
     def process_file(self, file) -> Image.Image:
         loader = PyMuPDFLoader(file.name)
         self.documents = loader.load()
-        pattern = r"/([^/]+)$"
-        match = re.search(pattern, file.name)
-        try:
-            self.file_name = match.group(1)
-        except:
-            self.file_name = os.path.basename(file)
+        self.file_name = os.path.basename(file.name)
         doc = fitz.open(file.name)
         page = doc[0]
         pix = page.get_pixmap(dpi=150)
@@ -56,6 +48,14 @@ class MyApp:
         )
         return "Vector database built successfully!"
 
+# Function to add text to chat history
+def add_text(history: List[Tuple[str, str]], text: str) -> List[Tuple[str, str]]:
+    if not text:
+        raise gr.Error("Enter text")
+    history.append((text, ""))
+    return history
+
+# Function to get response from the model
 def get_response(history, query, file):
     if not file:
         raise gr.Error(message="Upload a PDF")
@@ -76,6 +76,7 @@ def get_response(history, query, file):
         app.chat_history.append((query, "I have no information about it. Feed me knowledge, please!"))
         return history, f"I have no information about it. Feed me knowledge, please! Error: {str(e)}"
 
+# Function to render file
 def render_file(file) -> Image.Image:
     doc = fitz.open(file.name)
     page = doc[0]
@@ -83,6 +84,7 @@ def render_file(file) -> Image.Image:
     image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
     return image
 
+# Function to purge chat and render first page of PDF
 def purge_chat_and_render_first(file) -> Tuple[Image.Image, list]:
     app.chat_history = []
     doc = fitz.open(file.name)
@@ -91,6 +93,7 @@ def purge_chat_and_render_first(file) -> Tuple[Image.Image, list]:
     image = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
     return image, []
 
+# Function to refresh chat
 def refresh_chat():
     app.chat_history = []
     return []
@@ -99,12 +102,29 @@ app = MyApp()
 
 # Pre-process the saved PDF file
 saved_file_path = "track_training.pdf"
-app.process_file(open(saved_file_path, 'rb'))
-app.build_chain(open(saved_file_path, 'rb'))
+with open(saved_file_path, 'rb') as saved_file:
+    app.process_file(saved_file)
+    app.build_chain(saved_file)
+
+# Function to set API key
+def set_api_key(api_key):
+    app.set_api_key(api_key)
+    return f"API Key set to {api_key[:4]}...{api_key[-4:]}"
 
+# Gradio interface
 with gr.Blocks() as demo:
     with gr.Tab("Inst RAG"):
         with gr.Column():
+            api_key_input = gr.Textbox(label="OpenAI API Key", type="password", placeholder="Enter your OpenAI API Key")
+            api_key_btn = gr.Button("Set API Key")
+            api_key_status = gr.Textbox(value="API Key status", interactive=False)
+
+            api_key_btn.click(
+                fn=set_api_key,
+                inputs=[api_key_input],
+                outputs=[api_key_status]
+            )
+
             with gr.Row():
                 btn = gr.UploadButton("📁 Upload a PDF", file_types=[".pdf"])
                 show_img = gr.Image(label="Uploaded PDF")