updated app.py to use classes

app.py

@@ -12,18 +12,16 @@ import os
 import fitz
 from PIL import Image
 
-COUNT,N = 0,0
-chat_history = []
-chain = ''
-enable_box = gr.Textbox.update(value=None,placeholder= 'Upload your OpenAI API key',interactive=True)
-disable_box = gr.Textbox.update(value = 'OpenAI API key is Set',interactive=False)
-#os.environ['OPENAI_API_KEY'] = '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'
-
-def set_apikey(api_key):
-    os.environ['OPENAI_API_KEY'] = api_key
-    return disable_box
-def enable_api_box():
-    return enable_box
+from functools import lru_cache
+import json
+import uuid
+import re
+import chromadb
+from chromadb.config import Settings
+
+
+
+
 
 def add_text(history, text):
     if not text:
@@ -31,59 +29,75 @@ def add_text(history, text):
     history = history + [(text,'')] 
     return history
 
-def process_file(file):
-    if 'OPENAI_API_KEY' not in os.environ:
-        raise gr.Error('Upload your OpenAI API key')
-
-    loader = PyPDFLoader(file.name)
-    documents = loader.load()
-
-    embeddings = OpenAIEmbeddings()
+class my_app:
+    def __init__(self, OPENAI_API_KEY: None ) -> None:
+        self.OPENAI_API_KEY = OPENAI_API_KEY
+        self.chain = None
+        self.chat_history = []
+        self.N = 0
+        self.enable_box = gr.Textbox.update(value=None,placeholder= 'Upload your OpenAI API key',interactive=True)
+        self.disable_box = gr.Textbox.update(value = 'OpenAI API key is Set',interactive=False)
+ 
+    def set_apikey(self,api_key):
+        self.OPENAI_API_KEY = api_key
+        return self.disable_box
     
-    pdfsearch = Chroma.from_documents(documents, embeddings, metadatas=[{"source": f"{i}-pl"} for i in range(len(documents))])
-
-    chain = ConversationalRetrievalChain.from_llm(ChatOpenAI(temperature=0.3), 
-                                                  retriever=pdfsearch.as_retriever(search_kwargs={"k": 1}),
-                                                  return_source_documents=True,)
-    return chain
+    def enable_api_box(self):
+        return self.enable_box
+    
+    def chroma_client(self):
+        #create a chroma client
+        client = chromadb.Client()
+        #create a collecyion
+        collection = client.get_or_create_collection(name="my-collection")
 
-def generate_response(history, query, btn):
-    global COUNT, N, chat_history, chain
+        return client
     
-    if not btn:
-        raise gr.Error(message='Upload a PDF')
-    if COUNT == 0:
-            chain = process_file(btn)
-            COUNT += 1
+    def process_file(self,file):
+
+        loader = PyPDFLoader(file.name)
+        documents = loader.load()  
+        pattern = r"/([^/]+)$"
+        match = re.search(pattern, file.name)
+        file_name = match.group(1)
+        return documents, file_name
     
-    result = chain({"question": query, 'chat_history':chat_history},return_only_outputs=True)
-    chat_history += [(query, result["answer"])]
-    N = list(result['source_documents'][0])[1][1]['page']
-
-    for char in result['answer']:
-       history[-1][-1] += char
-       yield history,''
-
+    def build_chain(self, file):
+        documents, file_name = self.process_file(file)
+        #Load embeddings model
+        embeddings = OpenAIEmbeddings(openai_api_key=self.OPENAI_API_KEY) 
+        pdfsearch = Chroma.from_documents(documents, embeddings, collection_name= file_name,)
 
-def render_first(btn):
+        self.chain = ConversationalRetrievalChain.from_llm(ChatOpenAI(temperature=0.0, openai_api_key=self.OPENAI_API_KEY), 
+                                                  retriever=pdfsearch.as_retriever(search_kwargs={"k": 1}),
+                                                  return_source_documents=True,)
     
-    doc = fitz.open(btn.name)
-    page = doc[0]
-
-    #Render the page as a PNG image with a resolution of 300 DPI
-    pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))
-    image = Image.frombytes('RGB', [pix.width, pix.height], pix.samples)
-    return image
-
-def render_file(file):
-    global N
-    doc = fitz.open(file.name)
-    page = doc[N]
-    #Render the page as a PNG image with a resolution of 300 DPI
-    pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))
-    image = Image.frombytes('RGB', [pix.width, pix.height], pix.samples)
-    return image
-
+    def get_response(self, history, query, file):
+
+        if not file:
+            raise gr.Error(message='Upload a PDF')
+        
+        if not self.chain:
+            self.build_chain(file)
+
+        result = self.chain({"question": query, 'chat_history':self.chat_history},return_only_outputs=True)
+        self.chat_history += [(query, result["answer"])]
+        self.N = list(result['source_documents'][0])[1][1]['page']
+
+        for char in result['answer']:
+           history[-1][-1] += char
+           yield history,''
+
+    def render_file(self,file):
+                
+        doc = fitz.open(file.name)
+        page = doc[self.N]
+        #Render the page as a PNG image with a resolution of 300 DPI
+        pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))
+        image = Image.frombytes('RGB', [pix.width, pix.height], pix.samples)
+        return image
+        
+app = my_app()
 
 with gr.Blocks() as demo:
 
@@ -97,27 +111,30 @@ with gr.Blocks() as demo:
             chatbot = gr.Chatbot(value=[], elem_id='chatbot').style(height=650)
             show_img = gr.Image(label='Upload PDF', tool='select' ).style(height=680)
     with gr.Row():
-        with gr.Column(scale=0.70):
+        with gr.Column(scale=0.60):
             txt = gr.Textbox(
                         show_label=False,
                         placeholder="Enter text and press enter",
                     ).style(container=False)
-        with gr.Column(scale=0.15):
+        with gr.Column(scale=0.20):
             submit_btn = gr.Button('submit')
-        with gr.Column(scale=0.15):
+        with gr.Column(scale=0.20):
             btn = gr.UploadButton("📁 upload a PDF", file_types=[".pdf"]).style()
+        
+    
+
+    api_key.submit(fn=app.set_apikey, inputs=[api_key], outputs=[api_key])
+    change_api_key.click(fn= app.enable_api_box,outputs=[api_key])
+    btn.upload(fn=app.render_file, inputs=[btn], outputs=[show_img],)
     
-    api_key.submit(fn=set_apikey, inputs=[api_key], outputs=[api_key])
-    change_api_key.click(fn= enable_api_box,outputs=[api_key])
-    btn.upload(fn=render_first, inputs=[btn], outputs=[show_img],)
+    submit_btn.click(fn=add_text, inputs=[chatbot,txt], outputs=[chatbot, ], queue=False).success(fn=app.get_response,inputs = [chatbot, txt, btn],
+                                    outputs = [chatbot,txt]).success(fn=app.render_file,inputs = [btn], outputs=[show_img])
+
     
-    submit_btn.click(fn=add_text, inputs=[chatbot,txt], outputs=[chatbot, ], queue=False).success(fn=generate_response,inputs = [chatbot, txt, btn],
-                                    outputs = [chatbot,txt]).success(fn=render_file,inputs = [btn], outputs=[show_img])
+demo.queue()
+demo.launch()  
+
+
 
-    txt.submit(fn=add_text, inputs=[chatbot,txt], outputs=[chatbot, ], queue=False).success(fn=generate_response,inputs = [chatbot, txt, btn],
-                                    outputs = [chatbot,txt]).success(fn=render_file,inputs = [btn], outputs=[show_img])
 
 
-demo.queue()
-if __name__ == "__main__":
-    demo.launch()