optimize code

app.py

@@ -2,36 +2,25 @@ from typing import List, Optional, Tuple
 from queue import Empty, Queue
 from threading import Thread
 from bot.web_scrapping.crawler_and_indexer import content_crawler_and_index
-from bot.web_scrapping.searchable_index import SearchableIndex
 from bot.utils.callbacks import QueueCallback
 from bot.utils.constanst import set_api_key
 from bot.utils.show_log import logger
+from bot.web_scrapping.default import *
 from langchain.chat_models import ChatOpenAI
 from langchain.prompts import HumanMessagePromptTemplate
 from langchain.schema import AIMessage, BaseMessage, HumanMessage, SystemMessage
 import gradio as gr
 
 set_api_key()
-MODELS_NAMES = ["gpt-3.5-turbo"]
-DEFAULT_TEMPERATURE = 0.7
-
-ChatHistory = List[str]
-
-default_system_prompt = 'Put your prompt here'
-default_system_format = 'txt'
 human_message_prompt_template = HumanMessagePromptTemplate.from_template("{text}")
 
 
-def learning_feedbacks():
-    return 'Training Completed'
-
-
-def bot_learning(urls, file_formats, chat_mode=False):
-    index = content_crawler_and_index(url=str(urls), file_format=file_formats)
+def bot_learning(urls, file_formats, llm, prompt, chat_mode=False):
+    index = content_crawler_and_index(url=str(urls), llm=llm, prompt=prompt, file_format=file_formats)
     if chat_mode:
         return index
     else:
-        fb = learning_feedbacks()
+        fb = 'Training Completed'
         return fb
 
 
@@ -54,10 +43,10 @@ def chat_start(
     job_done = object()
     messages.append(HumanMessage(content=f':{message}'))
     chatbot_messages.append((message, ""))
-    index = bot_learning(urls='NO_URL', file_formats='txt', chat_mode=True)
+    index = bot_learning(urls='NO_URL', file_formats='txt', llm=chat, prompt=message, chat_mode=True)
 
     def query_retrieval():
-        response = SearchableIndex.query(message, chat, index)
+        response = index.query()
         chatbot_message = AIMessage(content=response)
         messages.append(chatbot_message)
         queue.put(job_done)
@@ -105,90 +94,96 @@ def on_apply_settings_button_click(
     return chat, *on_clear_button_click(system_prompt)
 
 
-with gr.Blocks() as demo:
-    system_prompt = gr.State(default_system_prompt)
-    messages = gr.State([SystemMessage(content=default_system_prompt)])
-    chat = gr.State(None)
+def main():
+    with gr.Blocks() as demo:
+        system_prompt = gr.State(default_system_prompt)
+        messages = gr.State([SystemMessage(content=default_system_prompt)])
+        chat = gr.State(None)
 
-    with gr.Column(elem_id="col_container"):
-        gr.Markdown("# Welcome to OWN-GPT! 🤖")
-        gr.Markdown(
-            "Demo Chat Bot Platform"
-        )
-
-        chatbot = gr.Chatbot()
-        with gr.Column():
-            message = gr.Textbox(label="Type some message")
-            message.submit(
-                chat_start,
-                [chat, message, chatbot, messages],
-                [chat, message, chatbot, messages],
-                queue=True,
-            )
-            message_button = gr.Button("Submit", variant="primary")
-            message_button.click(
-                chat_start,
-                [chat, message, chatbot, messages],
-                [chat, message, chatbot, messages],
+        with gr.Column(elem_id="col_container"):
+            gr.Markdown("# Welcome to OWN-GPT! 🤖")
+            gr.Markdown(
+                "Demo Chat Bot Platform"
             )
-        with gr.Column():
-            learning_status = gr.Textbox(label='Training Status')
-            url = gr.Textbox(label="URL to Documents")
-            file_format = gr.Textbox(label="Set your file format:", placeholder='Example: pdf, txt')
-            url.submit(
-                bot_learning,
-                [url, file_format],
-                [learning_status]
-            )
-            training_button = gr.Button("Training", variant="primary")
-            training_button.click(
-                bot_learning,
-                [url, file_format],
-                [learning_status]
-            )
-        with gr.Row():
+
+            chatbot = gr.Chatbot()
             with gr.Column():
-                clear_button = gr.Button("Clear")
-                clear_button.click(
-                    on_clear_button_click,
-                    [system_prompt],
-                    [message, chatbot, messages],
-                    queue=False,
+                message = gr.Textbox(label="Type some message")
+                message.submit(
+                    chat_start,
+                    [chat, message, chatbot, messages],
+                    [chat, message, chatbot, messages],
+                    queue=True,
                 )
-            with gr.Accordion("Settings", open=False):
-                model_name = gr.Dropdown(
-                    choices=MODELS_NAMES, value=MODELS_NAMES[0], label="model"
+                message_button = gr.Button("Submit", variant="primary")
+                message_button.click(
+                    chat_start,
+                    [chat, message, chatbot, messages],
+                    [chat, message, chatbot, messages],
                 )
-                temperature = gr.Slider(
-                    minimum=0.0,
-                    maximum=1.0,
-                    value=0.7,
-                    step=0.1,
-                    label="temperature",
-                    interactive=True,
+            with gr.Column():
+                learning_status = gr.Textbox(label='Training Status')
+                url = gr.Textbox(label="URL to Documents")
+                file_format = gr.Textbox(label="Set your file format:", placeholder='Example: pdf, txt')
+                url.submit(
+                    bot_learning,
+                    [url, file_format, chat, message],
+                    [learning_status]
                 )
-                apply_settings_button = gr.Button("Apply")
-                apply_settings_button.click(
-                    on_apply_settings_button_click,
-                    [system_prompt, model_name, temperature],
-                    [chat, message, chatbot, messages],
+                training_button = gr.Button("Training", variant="primary")
+                training_button.click(
+                    bot_learning,
+                    [url, file_format, chat, message],
+                    [learning_status]
                 )
+            with gr.Row():
+                with gr.Column():
+                    clear_button = gr.Button("Clear")
+                    clear_button.click(
+                        on_clear_button_click,
+                        [system_prompt],
+                        [message, chatbot, messages],
+                        queue=False,
+                    )
+                with gr.Accordion("Settings", open=False):
+                    model_name = gr.Dropdown(
+                        choices=MODELS_NAMES, value=MODELS_NAMES[0], label="model"
+                    )
+                    temperature = gr.Slider(
+                        minimum=0.0,
+                        maximum=1.0,
+                        value=0.7,
+                        step=0.1,
+                        label="temperature",
+                        interactive=True,
+                    )
+                    apply_settings_button = gr.Button("Apply")
+                    apply_settings_button.click(
+                        on_apply_settings_button_click,
+                        [system_prompt, model_name, temperature],
+                        [chat, message, chatbot, messages],
+                    )
 
-        with gr.Column():
-            system_prompt_area = gr.TextArea(
-                default_system_prompt, lines=4, label="prompt", interactive=True
-            )
-            system_prompt_area.input(
-                system_prompt_handler,
-                inputs=[system_prompt_area],
-                outputs=[system_prompt],
+            with gr.Column():
+                system_prompt_area = gr.TextArea(
+                    default_system_prompt, lines=4, label="prompt", interactive=True
+                )
+                system_prompt_area.input(
+                    system_prompt_handler,
+                    inputs=[system_prompt_area],
+                    outputs=[system_prompt],
+                )
+                system_prompt_button = gr.Button("Set")
+            system_prompt_button.click(
+                on_apply_settings_button_click,
+                [system_prompt, model_name, temperature],
+                [chat, message, chatbot, messages],
             )
-            system_prompt_button = gr.Button("Set")
-        system_prompt_button.click(
-            on_apply_settings_button_click,
-            [system_prompt, model_name, temperature],
-            [chat, message, chatbot, messages],
-        )
 
-demo.queue()
-demo.launch()
+    return demo
+
+
+if __name__ == '__main__':
+    demo = main()
+    demo.queue()
+    demo.launch()

bot/web_scrapping/crawler_and_indexer.py

@@ -7,66 +7,69 @@ import pandas as pd
 import requests
 import os
 
-set_api_key(api_key='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')
+set_api_key(api_key='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')
 
 
 def save_content_to_file(url=None, text=None, output_folder=None, file_format=None):
     file_path = os.path.join(output_folder, f"combined_content.{file_format}")
-    if file_format == 'txt':
-        with open(f"{file_path}", "a", encoding="utf-8") as file:
-            for t in text:
-                file.write(f'{t.text}\n')
-        logger.info(f"Content appended to {file_path}")
-    elif file_format == 'pdf':
-        request.urlretrieve(url, file_path)
-        logger.info(f"Content appended to {file_path}")
-    elif file_format == 'csv':
-        df = pd.DataFrame({'Content': [t.text for t in text]})
-        df.to_csv(f"{file_path}", mode='a', index=False, header=False)
-        logger.info(f"Content appended to {file_path}")
-    elif file_format == 'xml':
-        xml_content = ''.join([f'<item>{t.text}</item>' for t in text])
-        with open(f"{file_path}", "a", encoding="utf-8") as file:
-            file.write(xml_content)
+
+    write_functions = {
+        'txt': lambda: write_text(file_path, text),
+        'pdf': lambda: write_pdf(url, file_path),
+        'csv': lambda: write_csv(file_path, text),
+        'xml': lambda: write_xml(file_path, text)
+    }
+
+    write_function = write_functions.get(file_format)
+    if write_function:
+        write_function()
         logger.info(f"Content appended to {file_path}")
     else:
         logger.warning("Invalid file format. Supported formats: txt, pdf, csv, xml")
+
     return file_path
 
 
-def content_crawler_and_index(url, file_format='txt', output_folder='learning_documents'):
-    if url != 'NO_URL':
-        # Send an HTTP GET request to the URL
+def write_text(file_path, text):
+    with open(file_path, "a", encoding="utf-8") as file:
+        for t in text:
+            file.write(f'{t.text}\n')
+
+
+def write_pdf(url, file_path):
+    request.urlretrieve(url, file_path)
+
+
+def write_csv(file_path, text):
+    df = pd.DataFrame({'Content': [t.text for t in text]})
+    df.to_csv(file_path, mode='a', index=False, header=False)
+
+
+def write_xml(file_path, text):
+    xml_content = ''.join([f'<item>{t.text}</item>' for t in text])
+    with open(file_path, "a", encoding="utf-8") as file:
+        file.write(xml_content)
+
+
+def content_crawler_and_index(url, llm, prompt, file_format='txt', output_folder='learning_documents'):
+    if url == 'NO_URL':
+        file_path = output_folder
+    else:
         responses = requests.get(url)
-        # Check if the request was successful
-        if responses.status_code == 200:
-            # Create output folder if it doesn't exist
-            if not os.path.exists(output_folder):
-                os.makedirs(output_folder)
-            # Parse the HTML content using BeautifulSoup
-            soup = BeautifulSoup(responses.text, "html.parser")
-            text = soup.find_all(['h2', 'p', 'i', 'ul'])
-            if text:
-                # Save content based on the specified file format
-                file_path = save_content_to_file(text=text, output_folder=output_folder, file_format=file_format)
-
-                # Create or update the index
-                index = SearchableIndex.embed_index(url, file_path)
-                if os.path.isfile(file_path):
-                    os.remove(file_path)
-                return index
-            else:
-                file_path = save_content_to_file(url=url, output_folder=output_folder, file_format=file_format)
-                index = SearchableIndex.embed_index(url, file_path)
-                if os.path.isfile(file_path):
-                    os.remove(file_path)
-                return index
-
-        else:
+        if responses.status_code != 200:
             logger.warning("Failed to retrieve content from the URL.")
-    else:
-        index = SearchableIndex.embed_index(url=url, path=output_folder)
-        return index
+            return None
+        if not os.path.exists(output_folder):
+            os.makedirs(output_folder)
+        soup = BeautifulSoup(responses.text, "html.parser")
+        text = soup.find_all(['h2', 'p', 'i', 'ul'])
+        file_path = save_content_to_file(text=text, url=url, output_folder=output_folder, file_format=file_format)
+
+    index = SearchableIndex.embed_index(url=url, path=file_path, llm=llm, prompt=prompt)
+    if url != 'NO_URL' and os.path.isfile(file_path):
+        os.remove(file_path)
+
+    return index
 
 
 if __name__ == '__main__':

bot/web_scrapping/default.py

@@ -0,0 +1,6 @@
+from typing import List
+MODELS_NAMES = ["gpt-3.5-turbo"]
+DEFAULT_TEMPERATURE = 0.7
+ChatHistory = List[str]
+default_system_prompt = 'Put your prompt here'
+default_system_format = 'txt'

bot/web_scrapping/searchable_index.py

@@ -16,125 +16,98 @@ import os
 import queue
 
 
+class Query:
+    def __init__(self, question, llm, index):
+        self.question = question
+        self.llm = llm
+        self.index = index
+
+    def query(self):
+        """Query the vectorstore."""
+        llm = self.llm or ChatOpenAI(model_name='gpt-3.5-turbo', temperature=0)
+        chain = RetrievalQA.from_chain_type(
+            llm, retriever=self.index.as_retriever()
+        )
+        return chain.run(self.question)
+
+
 class SearchableIndex:
     def __init__(self, path):
         self.path = path
 
-    def get_text_splits(self):
-        with open(self.path, 'r') as txt:
-            data = txt.read()
-
-        text_split = RecursiveCharacterTextSplitter(chunk_size=1000,
-                                                    chunk_overlap=0,
-                                                    length_function=len)
-        doc_list = text_split.split_text(data)
-        return doc_list
-
-    def get_pdf_splits(self):
-        loader = PyPDFLoader(self.path)
-        pages = loader.load_and_split()
-        text_split = RecursiveCharacterTextSplitter(chunk_size=1000,
-                                                    chunk_overlap=0,
-                                                    length_function=len)
-        doc_list = []
-        for pg in pages:
-            pg_splits = text_split.split_text(pg.page_content)
-            doc_list.extend(pg_splits)
+    @classmethod
+    def get_splits(cls, path, target_col=None, sheet_name=None):
+        extension = os.path.splitext(path)[1].lower()
+        doc_list = None
+        if extension == ".txt":
+            with open(path, 'r') as txt:
+                data = txt.read()
+                text_split = RecursiveCharacterTextSplitter(chunk_size=1000,
+                                                            chunk_overlap=0,
+                                                            length_function=len)
+                doc_list = text_split.split_text(data)
+        elif extension == ".pdf":
+            loader = PyPDFLoader(path)
+            pages = loader.load_and_split()
+            text_split = RecursiveCharacterTextSplitter(chunk_size=1000,
+                                                        chunk_overlap=0,
+                                                        length_function=len)
+            doc_list = []
+            for pg in pages:
+                pg_splits = text_split.split_text(pg.page_content)
+                doc_list.extend(pg_splits)
+        elif extension == ".xml":
+            df = pd.read_excel(io=path, engine='openpyxl', sheet_name=sheet_name)
+            df_loader = DataFrameLoader(df, page_content_column=target_col)
+            doc_list = df_loader.load()
+        elif extension == ".csv":
+            csv_loader = CSVLoader(path)
+            doc_list = csv_loader.load()
+        if doc_list is None:
+            raise ValueError("Unsupported file format")
         return doc_list
 
-    def get_xml_splits(self, target_col, sheet_name):
-        df = pd.read_excel(io=self.path,
-                           engine='openpyxl',
-                           sheet_name=sheet_name)
-
-        df_loader = DataFrameLoader(df,
-                                    page_content_column=target_col)
-
-        excel_docs = df_loader.load()
-
-        return excel_docs
-
-    def get_csv_splits(self):
-        csv_loader = CSVLoader(self.path)
-        csv_docs = csv_loader.load()
-        return csv_docs
-
     @classmethod
     def merge_or_create_index(cls, index_store, faiss_db, embeddings, logger):
         if os.path.exists(index_store):
             local_db = FAISS.load_local(index_store, embeddings)
             local_db.merge_from(faiss_db)
-            logger.info("Merge index completed")
             local_db.save_local(index_store)
-            return local_db
+            logger.info("Merge index completed")
         else:
             faiss_db.save_local(folder_path=index_store)
             logger.info("New store created and loaded...")
             local_db = FAISS.load_local(index_store, embeddings)
-            return local_db
+        return local_db
 
     @classmethod
-    def check_and_load_index(cls, index_files, embeddings, logger, path, result_queue):
+    def check_and_load_index(cls, index_files, embeddings, logger, result_queue):
         if index_files:
             local_db = FAISS.load_local(index_files[0], embeddings)
-            file_to_remove = os.path.join(path, 'combined_content.txt')
-            if os.path.exists(file_to_remove):
-                os.remove(file_to_remove)
         else:
             raise logger.warning("Index store does not exist")
         result_queue.put(local_db)  # Put the result in the queue
 
     @classmethod
-    def embed_index(cls, url, path, target_col=None, sheet_name=None):
+    def embed_index(cls, url, path, llm, prompt, target_col=None, sheet_name=None):
         embeddings = OpenAIEmbeddings()
 
-        def process_docs(queues, extension):
-            nonlocal doc_list
-            instance = cls(path)
-            if extension == ".txt":
-                doc_list = instance.get_text_splits()
-            elif extension == ".pdf":
-                doc_list = instance.get_pdf_splits()
-            elif extension == ".xml":
-                doc_list = instance.get_xml_splits(target_col, sheet_name)
-            elif extension == ".csv":
-                doc_list = instance.get_csv_splits()
-            else:
-                doc_list = None
-            queues.put(doc_list)
-
         if url != 'NO_URL' and path:
-            file_extension = os.path.splitext(path)[1].lower()
-            data_queue = queue.Queue()
-            thread = threading.Thread(target=process_docs, args=(data_queue, file_extension))
-            thread.start()
-            doc_list = data_queue.get()
-            if not doc_list:
-                raise ValueError("Unsupported file format")
-
+            doc_list = cls.get_splits(path, target_col, sheet_name)
             faiss_db = FAISS.from_texts(doc_list, embeddings)
             index_store = os.path.splitext(path)[0] + "_index"
             local_db = cls.merge_or_create_index(index_store, faiss_db, embeddings, logger)
-            return local_db, index_store
+            return Query(prompt, llm, local_db)
         elif url == 'NO_URL' and path:
             index_files = glob.glob(os.path.join(path, '*_index'))
 
             result_queue = queue.Queue()  # Create a queue to store the result
 
             thread = threading.Thread(target=cls.check_and_load_index,
-                                      args=(index_files, embeddings, logger, path, result_queue))
+                                      args=(index_files, embeddings, logger, result_queue))
             thread.start()
             local_db = result_queue.get()  # Retrieve the result from the queue
-            return local_db
-
-    @classmethod
-    def query(cls, question: str, llm, index):
-        """Query the vectorstore."""
-        llm = llm or ChatOpenAI(model_name='gpt-3.5-turbo', temperature=0)
-        chain = RetrievalQA.from_chain_type(
-            llm, retriever=index.as_retriever()
-        )
-        return chain.run(question)
+            return Query(prompt, llm, local_db)
 
 
 if __name__ == '__main__':

learning_documents/combined_content_index/index.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7246a7c708997772e992539fa7cef62d0e33a4a77a03f6483be6a108106a7c1c
+size 100825