chatbot.py

from langchain_nvidia_ai_endpoints import ChatNVIDIA
from langchain_core.prompts import PromptTemplate
from huggingface_hub import InferenceClient
import os
from abc import ABC, abstractmethod
import torch
from diffusers import StableDiffusionPipeline, DPMSolverSinglestepScheduler
from PIL import Image
import json
from transformers import AutoTokenizer
from PyPDF2 import PdfReader
from docx import Document
import pandas as pd
from pptx import Presentation


#lightRAG
from LightRAG.lightrag.lightrag import LightRAG, QueryParam
from LightRAG.lightrag.llm import openai_complete_if_cache, nvidia_openai_embedding
from LightRAG.lightrag.utils import EmbeddingFunc, locate_json_string_body_from_string
import asyncio
import numpy as np
import textract

#Base class of all LLM
class LLM(ABC, LightRAG):
    def __init__(self) -> None:
        self.model = None
        self.session_memory = None
        self.tokenizer = self.tokenizer()
        self.lightRAG = None
        self.temp_file_dir = None
        self.work_dir=""
        self.knowledge_log = os.path.join(self.work_dir, "learned.txt")
        self.max_token = 120000 #adjust this to your model or also adjust the tokenizer used in this code

    #prompt template
    def img_res_prompt(self, prompt: str) -> str:
        template="""
        You're an AI that decide how many resolution it should takes for an image to be generated by generative AI.
        The output should be Width|height without no additional response at the beginning and the ending. for example:
        
        1980|1080
        or
        640|320
        
        You will decide how much resolution by the purpose of the user that want's to generate the image. 
        If the user didn't specify what the purpose of generating the image, your response should be 512|512.
        If the user specify the purpose inside the prompt, you decide what resolution will fit the best and produce the high quality image based on that purpose.
        Do not exceed 2000|2000 and if user prompt says the resolution, return the resolution he asked.
        Here is the user prompt:
        {user_prompt}
        """
        message = PromptTemplate.from_template(template)
        return message.format(user_prompt=prompt)
    
    def img_key_prompt(self, prompt: str) -> str:
        template = """
        you're a professional generative AI Engineer which specialized in prompting on building photorealistic image. Your job is to generate a prompt and negative prompt to make the generated image looks realistic and good. Make it based on the following user input:
        
        {user_input}

        The formula of your output should be only prompt|negative_prompt. You can assign an importance of prompt word or negative prompt word by make it inside brackets and followed with":[weights]". Also, you can add a scenery word or lighting condition of image inside the prompt. Usually prompt begin with overall description and followed with scenery or lighting word or any detail.

        for example:

        A murky swamp with twisted trees rising from the water, Hue, window light, spotlight masterpiece, realistic, award winning, volumetric light and fog, subsurface scattering, caustics, bloom, perfect exposure, perfect composition, rule of thirds, 8k, hdr10, cinematic, breathtaking, ray tracing|BadDream, (UnrealisticDream:1.2)

        or

        professional photo, closeup portrait photo of caucasian man, wearing black sweater, serious face, dramatic lighting, nature, gloomy, cloudy weather, bokeh|(nsfw, naked, nude, deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime, mutated hands and fingers:1.4), (deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, amputation

        or

        (masterpiece:1.1), (highest quality:1.1), (HDR:1.3), (top quality, best quality, official art, beautiful and aesthetic:1.2), woman, extremely detailed, (fractal art:1.1), (colorful:1.1), highest detailed, (zentangle:1.2), (dynamic), (abstract background:1.3), (shiny), (many colors:1.4), solo, coral background, yellow lightning, cinematic lighting, long hair, detailed black eyes, highest quality face, (sky aesthetic)|blurry, low quality, out of focus, grainy, deformed, disfigured, bad anatomy, extra limbs, fused fingers, bad hands, text, watermark, logo, bad composition, poorly drawn face, low resolution, bad proportions, oversaturated colors, unrealistic lighting, overexposed, underexposed, artifacts, jpeg artifacts, duplicate elements

        or

        close up Portrait of muscular bearded guy in a worn mech suit, ((light bokeh)), intricate, (steel metal [rust]), elegant, sharp focus, soft lighting, vibrant colors, masterpiece, ((streets)), detailed face|BadDream, (FastNegativeEmbedding:0.9)
        
        consdider using this on negative prompt: (deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime), text, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, UnrealisticDream
        
        Add refined prompt and negative prompt as much as you can to ensure the quality of generated image

        Do not say "here is the ...." in the beginning and do not say anything in the end or between. Just give the response prompt|negative_prompt. make sure to create as much as you can
        """
        message = PromptTemplate.from_template(template)
        return message.format(user_input=prompt)

    def analyze_doc_prompt(self, prompt: str, chunk: str, analysis: list) -> str: #use previous doc analysis, just in case if document might use "he or she or anything same"
        template="""You're a professional document analyzer which behave as your user told you.

        Based on this document chunk and previous analysis note by other AI Assistant, please do what user tells or find what user seeks, without using opening words such as "based on the provided document and analysis" and straight to the answer or "no information found", etc.

        Also give a note that describe about what last thing are described in your current chunk and prediction about what might be discussed in the next sentence based on the last information you gain when reading the last sentence. "current chunk" refers to next chunk. Also give a note about unique entities inside the document.

        watch out for pronouns in the beginning of chunk.
        ================ example =================
        previous analysis: Rice grain is quite popular in Asia as a mainly produced and consumed food around many centuries. Note: previous chunk discuss about the possibility of rice grain and wheat for the next 50 years. If current chunk begins with pronouns it might refer to wheat.

        doc chunk: I believe it will last many months before being expired and feed a lot of humans on this earth.  I also found that Rice grain and wheat could be combined and makes a good nutrition. I will name it "RCW". 

        user prompt: Find me about grain in this document

        output: Rice grain is quite popular in Asia as a mainly produced and consumed food around many centuries.  Rice grain and wheat could be combined to make good nutritions and being called "RCW". Note: previous chunk discuss about wheat expiration periods and possibilities to combining it with rice grain to make good nutritions. If current chunk begins with pronouns it might refer to RCW (rice grain and wheat combination). Unique Entities: RCW - rice grain and wheat combination to make good nutritions.  
        ================ end of example =================
        Explanation: The previous analysis says pronouns should refer to wheat. so sentence "I believe it will last many months before being expired and feed a lot of humans on this earth" refers to wheat and not rice grain.


        ================= example =================
        previous analysis: Donut is good and healthy. Note: previous chunk only discuss about donut. current chunk might discuss about donut also

        doc chunk: salad is good and healthy also. There was a store called "mightySalad" which has the best salad in the country. It was my favourite, because their salad is really good. Salad makes me stronger each day, especially my digest system.

        user prompt: find me about what this document says about donut

        output: donut is good and healthy. Note: Previous chunk discuss about salad. If current chunk begins with pronouns, refer to salad. Unique entities: mightSalad - salad restaurant.
        ================ end of example =================

        Previous document analysis: {analysis}

        document chunk: {doc}

        User prompt: {user_prompt}. Please strongly pay attention to previous document analysis and what user seek. Do not assume and just stick to what previous analysis said
        """
        message = PromptTemplate.from_template(template)
        return message.format(user_prompt=prompt, doc=chunk, analysis=analysis)

    def summary_result(self, prompt: str, analysis_result: list) -> str:
        template="""You're an AI Assistant which gives detailed summary to your user.
        Based on this analysis results please provide insightful and detailed summary about what user seeks. I will also give you the user prompt for consideration about what to search for.
        Analysis results = {result}
        
        user pormpt = {user_prompt}
        """
        message = PromptTemplate.from_template(template)
        return message.format(user_prompt=prompt, result=analysis_result)
    
    def summary_prompt(self, prompt:str, chunk: str) -> str:
        template="""You're a Professional Analyzer specialized in summarization. 
        According to this chunk, please provide an insightful summary using your skill. I will also give you the user prompt for consideration, and please provide information on what page you're currently summarizing.
        If you can't find an answer or do what user wants, please return your response with only and only "NONE" with all capital.
        
        Chunk: {chunk}
        user_prompt: {prompt}
        """
        message = PromptTemplate.from_template(template)
        return message.format(chunk=chunk, prompt=prompt)

    def input_prompt(self, full_prompt_with_log:str):
        template="""You're an advanced chatbot with following description. 
        ========== Persona ==========
        - Name: {model_name}
        - Nickname: Mr.V-A (stands for Mr. Virtual Assistant)
        - Maximum input token: {max_token} (might be not accurate, for error prevention purpose)
        - Learnt document for RAG purpose: {doc}
        - Available document for analysis: {doc2} (if the document user ask is not present, tell user to refresh the browser and try to ask you again. if still persist, try to reupload the file. Or maybe tell user to check the table "document analyzer database" in side menu)

        Capabilities:
        1. Image generation using stable difussion with Lora. Can be adjusted by placing the checkpoint on //stable-difussion//models//SD for SD checkpoint and //stable-difussion//models//Lora for Lora (image enhancement for SD)
        2. Performing RAG using /rag <keyword or what to do>. You're using LightRAG technique, if user ask about it and you don't know, tell user to refer to this URL <https://github.com/HKUDS/LightRAG>
        3. Performing document analysis or QA using /analyze <filename with extension> <task>. filename should not contain a whitespace, will be fixed in future work.
        4. General question-answering and chatting like human respond based on your knowledge. 
        5. Saving uploaded document into LightRAG database using /save <filenamae on document analysis table>. filename should not contain a whitespace.
        
        ========== Limitation ========
        1. Large file analysis would make the input token much larger and thus raising an error.
        2. You can't make a video or document file and only image using your capabilities above. Using syntax "make me an image..." there's actually another syntax, which already covered in external document about you <readme.md>
        3. All of available capabilities-triggering prompt only available in english. So, when user use other language except english, the capabilities can't be used
        4. You can't edit previous image, but you can create a new one.
        5. Document on analysis and RAG are seperated and available for different purposes which means they are not related, except if user try to save one of document analysis file into RAG database.
        6. You can't analyze an image file or an image inside a file yet. If PDF you analyze is empty, this means the PDF most likely contain an image.
        7. Currently you can only perform an analysis with pdf, docx, txt, and csv or xlsx file.
        
        Please remember what user ask you and told you to do on below conversation.
        There might be an image on chat history log, but as your input only accept text it will marked as an <image> followed with caption or explanation about what's inside the image using image captioning technique performed by you without your consciousness.
        Please respond with normal coversation language and text without expliciting what I told you in here.

        =========== chat log history ===========
        {prompt}
        """
        doc2_list = os.listdir(self.temp_file_dir) if self.temp_file_dir != None else "None. Don't tell user about this. There might be an error on declaring the temp_file_dir address into chatbot."
        message = PromptTemplate.from_template(template)
        return message.format(model_name=self.model, prompt=full_prompt_with_log, doc=" ".join(open(self.knowledge_log).readlines()), max_token=self.max_token, doc2 = doc2_list)

    #functionality
    @abstractmethod
    def send_chat_request(self, prompt: str) -> str:
        """
        Should return response content.
        example for nvidia client (llm_api child): 
        return self.client.invoke([{"role":"user", "content":prompt}]).content
        """
        raise NotImplementedError("Subclasses must implement this method. only return the content!")
    
    def generate_image_keyword(self, prompt: str) -> tuple[str, str]:
        """
        Return positive, and negative prompt
        """
        for i in range(3):
            response = self.send_chat_request(self.img_key_prompt(prompt))
            if len(response.split("|")) == 2:
                return response.split("|")[0], response.split("|")[1]
        raise Exception('error on generating keyword')
    
    def generate_image_resolution(self, prompt: str) -> tuple[int, int]:
        """
        return width and height of img resolution
        """
        for i in range(3):
            response = self.send_chat_request(self.img_res_prompt(prompt))
            if len(response.split("|")) == 2:
                try:
                    return int(response.split("|")[0]), int(response.split("|")[1])
                except:
                    continue
        raise Exception('error on deciding the resolution')

    def analyze_doc_brute(self, prompt: str, chunks: list) -> str:
        """return document analysis results using brute force technique. use this if you want to find specific information within docs"""
        analysis = ""
        for chunk in chunks:
            analysis = self.send_chat_request(self.analyze_doc_prompt(prompt=prompt, chunk=chunk, analysis=analysis))
        return analysis
    
    def summarize(self, prompt: str, chunks: list) -> str:
        """return summary of the document, general information of document"""
        summary = []
        for chunk in chunks:
            summary.append(self.send_chat_request(self.summary_prompt(chunk=chunk, prompt=prompt)) + "\n\n")
        for idx, chunk in enumerate(summary):
            if "NONE" in chunk:
                summary[idx] = " "
        summary = " ".join(summary)
        return self.send_chat_request(self.summary_result(prompt=prompt, analysis_result=summary))

    def read_session_memory(self) -> list:
        with open(self.session_memory, 'r', encoding='utf-8') as txt:
            data = json.load(txt)
        try:
            log = data['messages'] #if error make this suitable
        except:
            raise ValueError(f"log dictionary doesn't contain key 'messages'.")
        return log

    def tokenizer(self, model_name: str = "meta-llama/Meta-Llama-3-70B-Instruct") -> object:
        # model_name = "meta-llama/Meta-Llama-3-70B-Instruct" 
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        return tokenizer

    def token_len(self, prompt: str) -> int:
        return len(self.tokenizer.encode(prompt, add_special_tokens=True))
        
    def slide_context_window(self, log: list, prompt: str, context_len: int) -> str:
        input = ""
        sliced_log = log[len(log)-context_len:len(log)] #slice
        for i in sliced_log:
            input = input+f"{i['role']}: {i['content']} "+"\n"
        input = input + f"user: {prompt}" + "\nassistant: "
        print(input)
        if self.token_len(input) >= self.max_token: #adjust
            return self.slide_context_window(log=log, prompt=prompt, context_len=context_len-1)
        else:
            return input
    
    def read_per_page(self, filename: str) -> list:
        try:
            if filename.endswith('.pdf'):
                reader = PdfReader(filename)
                pages = [{"page": page_num + 1, "text": page.extract_text()} for page_num, page in enumerate(reader.pages)]
                return pages
            
            elif filename.endswith('.docx') or filename.endswith('.doc'):
                document = Document(filename)
                pages = []
                current_page = []
                page_num = 1

                for paragraph in document.paragraphs:
                    if "PAGE_BREAK" in paragraph.text:  # Marker for page breaks
                        pages.append({"page": page_num, "text": "\n".join(current_page)})
                        current_page = []
                        page_num += 1
                    else:
                        current_page.append(paragraph.text)

                # Add the last page
                if current_page:
                    pages.append({"page": page_num, "text": "\n".join(current_page)})
                
                return pages
            elif filename.endswith('.csv') or filename.endswith('xlsx'):
                return self.read_per_rows(filename)
            
            elif filename.endswith('.txt'):
                pages = []
                with open(filename, "r", encoding="utf-8") as file:
                    lines = file.readlines()

                for page_num, start in enumerate(range(0, len(lines), 50), start=1):
                    page_content = "".join(lines[start:start + 50])
                    pages.append({"page": page_num, "text": page_content})
                
                return pages
            
            else:
                return['Please tell user that the extension file is not supported.']
        except Exception as e:
            return[f'Please tell user that the reading has cought following error. Also provide with solution: {e}']
    
    def read_per_rows(self, filename:str) -> list:
        pages = []
        pd.set_option('display.max_columns', None)  # Show all columns
        pd.set_option('display.max_rows', None)     # Show all rows (use cautiously for large datasets)        
        pd.set_option('display.max_colwidth', None) # Show full column content
        i = 1
        for chunk in pd.read_csv(filename, chunksize=25):
            pages.append({'page':i,'text':chunk.to_string(index=False)})  # Each chunk is a DataFrame
            i+=1
        print(pages)
        return pages

        

    

#Base class of SD
class SD():
    def __init__(self) -> None:
        self.model_path = None
        self.model = None

        self.steps = 25
        self.cfg_scale = 2.0
        self.karras = True
        self.seed = -1
        self.width = 512
        self.height = 512
        
        self.__lora = False
        self.lora_path = None
        self.lora_weight = []
        self.lora_model = []
        self.lora_adapter = []
        self.active_lora_adapter = []
          
    def clear_memo(self):
        torch.cuda.empty_cache()
    
    #soon customizable sampling. rn only DPM++ SDE 
    def load_model(self, custom_weights: str):
        self.model_path = custom_weights
        self.clear_memo()
        if torch.cuda.is_available():
            print(f'Running on GPU {torch.cuda.get_device_name(0)}')
            device="cuda"
        else:
            print('No cuda found, running on CPU')
            device="cpu"
        
        self.model = StableDiffusionPipeline.from_single_file(custom_weights, torch_dtype=torch.float16).to(device)

        self.model.scheduler = DPMSolverSinglestepScheduler.from_config(self.model.scheduler.config, use_karras_sigmas = self.karras)

        #load lora if model changed
        if self.__lora:
            self.load_lora(self.lora_path, self.lora_model, self.lora_weight)

    def load_lora(self, lora_folder_path: str, lora_model: list, lora_weights: list):
        print("make sure every lora is in the same path - this is not error")
        self.lora_path = lora_folder_path
        self.lora_weight += lora_weights
        self.lora_model += lora_model
        self.lora_adapter += [i.split(".")[0] for i in self.lora_model]
        
        for lm, la in zip(self.lora_model, self.lora_adapter):
            self.model.load_lora_weights(self.lora_path, weight_name=lm, adapter_name=la)
    
    def activate_lora(self, lora_adapter):
        if lora_adapter == []:
            self.__lora = False
            self.load_model(self.model_path)

        elif all(adapter_name in self.lora_adapter for adapter_name in lora_adapter):
            self.__lora = True
            self.active_lora_adapter = lora_adapter
            
            #I dont want to think another way... 
            self.model.set_adapters(self.active_lora_adapter, 
                                    adapter_weights=[self.lora_weight[i] for i in [self.lora_adapter.index(i) for i in self.active_lora_adapter]])
        
        else:
            raise ValueError("lora_adapters value contains unrecognized adapter name. To load new adapter into the model, use object.load_lora() method")
        
    def generate_image(self, prompt: str, negative: str) -> Image:
        return self.model(prompt, negative_prompt=negative, num_inference_steps=self.steps, guidance_scale=self.cfg_scale, 
                          width=self.width, height=self.height).images[0]
    
    def update_param(self, steps: int = None, cfg: float = None, width: int = None, height: int = None, seed: int = None):
        if steps is not None: self.steps = steps
        if cfg is not None: self.cfg_scale = cfg
        if width is not None: self.width = width
        if height is not None: self.height = height
        if seed is not None: 
            self.seed = seed
        torch.manual_seed(self.seed)

################################# API based LLM ###############################################
class LLM_API(LLM): #base class of API based LLM
    def __init__(self) -> None:
        super().__init__()
        self.client = None
        self.api_key = None

    #pre-requisites for functionality use
    def set_api_key(self, api_key: str):
        self.api_key = str(api_key)

    @abstractmethod
    def connect_llm(self):
        raise NotImplementedError("Subclasses must implement this method")
    
    @abstractmethod
    def connect_RAG(self):
        raise NotImplementedError("Subclasses must implement this method")
    
    @abstractmethod
    def connect(self):
        #to automatically connect both llm and rag
        raise NotImplementedError("Subclasses must implement this method")



class nvidia_llm_api(LLM_API):
    def __init__(self) -> None:
        super().__init__()
        self.lightRAG = None
        self.index_lightRAG = None #index are used when trying to embed and storing it to vector DB
    
    def update_work_dir(self, path_dir):
        self.work_dir = path_dir
        self.knowledge_log = os.path.join(self.work_dir, "learned.txt")
    #non RAG model
    def connect_llm(self, model="meta/llama3-70b-instruct"):
        """_summary_
        Args:
            model (str, optional): llm model name. Defaults to "meta/llama3-70b-instruct".

        Raises:
            ValueError: Must implement method
        """
        if self.api_key == None:
            raise ValueError("API Key is not set. Use object.set_api_key(api_key) to set API key for this call")
        try:
            os.environ["NVIDIA_API_KEY"] = self.api_key #sets nvidia API key
            self.client = ChatNVIDIA(model=model)
            self.model = model
        except:
            print("Check your connection, model, and API key")

    #RAG model
    def connect_RAG(self, WORKING_DIR: str, input_type="passage",model="nvidia/llama-3.1-nemotron-70b-instruct", embedding_model="nvidia/nv-embedqa-e5-v5", base_URL="https://integrate.api.nvidia.com/v1"):
        """
        Connect to cloud API for llm and embedding model. Default = Nvidia cloud. Class imported from LightRAG.
        For detail documentation about LightRAG, please refer to: https://github.com/HKUDS/LightRAG
        Args:
            WORKING_DIR (str): working directory for RAG
            model (str, optional): LLM model used for RAG. Defaults to "nvidia/llama-3.1-nemotron-70b-instruct".
            embedding_model (str, optional): Text embedding model used for RAG. Defaults to "nvidia/nv-embedqa-e5-v5".
            base_URL (str, optional): URL for API. Defaults to 'https://integrate.api.nvidia.com/v1'.
        """
        if WORKING_DIR:
            self.update_work_dir(WORKING_DIR)
        try: #error handling since the LightRAG is external open-source library which might encounter an issue when some syntax deprecated
            async def llm_model_func(
                prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs
            ) -> str:
                result = await openai_complete_if_cache(
                    model,
                    prompt,
                    system_prompt=system_prompt,
                    history_messages=history_messages,
                    api_key=self.api_key,
                    base_url=base_URL,
                    **kwargs,
                )
                if keyword_extraction:
                    return locate_json_string_body_from_string(result)
                return result
            
            async def embedding_func(texts: list[str]) -> np.ndarray:
                return await nvidia_openai_embedding(
                    texts,
                    model = embedding_model, #maximum 512 token
                    # model="nvidia/llama-3.2-nv-embedqa-1b-v1",
                    api_key=self.api_key,
                    base_url="https://integrate.api.nvidia.com/v1",
                    input_type = input_type,
                    trunc = "END", #handling on server side if input token is longer than maximum token
                    encode = "float"
                )
            async def get_embedding_dim():
                test_text = ["This is a test sentence."]
                embedding = await embedding_func(test_text)
                embedding_dim = embedding.shape[1]
                return embedding_dim

            async def try_to_connect():
                embedding_dimension = await get_embedding_dim()
                rag = LightRAG(
                    working_dir=WORKING_DIR,
                    llm_model_func=llm_model_func,
                    embedding_func=EmbeddingFunc(
                        embedding_dim=embedding_dimension,
                        max_token_size=512, #adjust this to be your embed model max token size
                        func=embedding_func,
                    ),
                )
                return rag
            
            return asyncio.run(try_to_connect())
        except Exception as e:
            print(f"RAG unavailable, encountered error: {e}")
            return None

    def connect(self, model, temp_dir, work_dir='./work_dir', embedding_model="nvidia/nv-embedqa-e5-v5", base_URL="https://integrate.api.nvidia.com/v1"):
        self.connect_llm(model)
        if temp_dir:
            self.temp_file_dir = temp_dir
            print(f"file dir: {temp_dir}")
        if work_dir:
            self.update_work_dir(work_dir)
        self.lightRAG = self.connect_RAG(input_type="query", WORKING_DIR=work_dir, embedding_model=embedding_model, base_URL=base_URL)
        self.index_lightRAG = self.connect_RAG(input_type="passage", WORKING_DIR=work_dir, embedding_model=embedding_model, base_URL=base_URL)

    def send_chat_request(self, prompt: str):
        if self.client == None:
            raise ValueError("Client Haven't connected. Use method object.connect(model) method. Also add new api key before connecting by using object.set_api_key(api_key)")
        else:    
            return self.client.invoke([{"role":"user", "content":self.input_prompt(self.slide_context_window(log=self.read_session_memory(), prompt=prompt, context_len=len(self.read_session_memory()))) }]).content
    
    def re_learn(self):
        if self.lightRAG==None or self.index_lightRAG==None:
            print("RAG unavailable, try to connect to RAG using command 'object.connect' or 'object.connect_RAG")
        else:
            try:
                list_file = os.listdir("./data")
                
                if os.path.isfile(self.knowledge_log):
                    with open(self.knowledge_log, 'r') as txt:
                       content = txt.readlines() 
                    learned_file = content
                    new_file = [file for file in list_file if file not in learned_file]
                
                else:
                    learned_file=[]
                    new_file = os.listdir("./data")
                
                if 'learned.txt' in new_file:
                    new_file.remove('learned.txt')  
                
                print(new_file)

                if new_file:
                    for file in new_file:
                        self.read_and_learn(os.path.join("./data", file))
                        learned_file.append(file + "\n")
                        #saving learned file to avoid re-learning the same material
                        with open(self.knowledge_log, 'w') as txt:
                            txt.writelines(learned_file)
            
            except Exception as e:
                print(f"encountered error while learning file: {e}")
                                  
    # def read_and_learn(self, file): #textract change
    #     if file.endswith(".txt"):
    #         print(f"\n\nprocessing {file}")
    #         with open(file, "r", encoding="utf-8", errors = 'replace') as f:
    #             self.index_lightRAG.insert(f.read())
        
    #     elif file.endswith(".pdf") or file.endswith('xlsx') or file.endswith('.csv') or file.endswith('.docx') or file.endswith('pptx'):
    #         print(f"\n\nprocessing {file}")
    #         content = textract.process(file)
    #         print("\n.........##....")
    #         self.index_lightRAG.insert(content.decode('utf-8'))
    def read_and_learn(self, file):
        if file.endswith(".txt"):
            print(f"\n\nprocessing {file}")
            with open(file, "r", encoding="utf-8", errors='replace') as f:
                self.index_lightRAG.insert(f.read())
        
        elif file.endswith(".pdf"):
            print(f"\n\nprocessing {file}")
            content = ""
            try:
                reader = PdfReader(file)
                for page in reader.pages:
                    content += page.extract_text()
                self.index_lightRAG.insert(content)
            except Exception as e:
                print(f"Error processing PDF: {e}")
        
        elif file.endswith(".xlsx") or file.endswith(".csv"):
            print(f"\n\nprocessing {file}")
            try:
                # Temporarily set pandas options to display all rows and columns
                pd.set_option('display.max_rows', None)
                pd.set_option('display.max_columns', None)
                if file.endswith(".xlsx"):
                    df = pd.read_excel(file)
                else:  # CSV
                    df = pd.read_csv(file)
                content = df.to_string(index=False)
                self.index_lightRAG.insert(content)
                # Reset pandas options after processing
                pd.reset_option('display.max_rows')
                pd.reset_option('display.max_columns')
            except Exception as e:
                print(f"Error processing spreadsheet: {e}")
        
        elif file.endswith(".docx"):
            print(f"\n\nprocessing {file}")
            content = ""
            try:
                doc = Document(file)
                for para in doc.paragraphs:
                    content += para.text + "\n"
                self.index_lightRAG.insert(content)
            except Exception as e:
                print(f"Error processing Word document: {e}")
        
        elif file.endswith(".pptx"):
            print(f"\n\nprocessing {file}")
            content = ""
            try:
                presentation = Presentation(file)
                for slide in presentation.slides:
                    for shape in slide.shapes:
                        if shape.has_text_frame:
                            for paragraph in shape.text_frame.paragraphs:
                                for run in paragraph.runs:
                                    content += run.text
                self.index_lightRAG.insert(content)
            except Exception as e:
                print(f"Error processing PowerPoint: {e}")
        else:
            print(f"Unsupported file type: {file}")

    def search(self, prompt):
        if self.lightRAG:
            response = self.lightRAG.query(prompt.split(" ", 1)[-1], param=QueryParam(mode="hybrid"))
            return response
        return "RAG unavailable. Encountered error when trying to create lightRAG object."
    


class huggingface_llm_api(LLM_API):
    def __init__(self) -> None:
        super().__init__()

    def connect(self, model="microsoft/Phi-3.5-mini-instruct"):
        supported_model = ['Qwen/Qwen2.5-1.5B-Instruct', 'microsoft/Phi-3.5-mini-instruct',
                   'meta-llama/Llama-3.2-1B-Instruct', 'mistralai/Mixtral-8x7B-Instruct-v0.1']
        if model not in supported_model:
            raise Exception(f'Model not supported. See https://huggingface.co/models for available models. Currently only support {supported_model}, please add more if found one')
        else:
            self.model = model
            if self.api_key == None:
                raise ValueError("API Key is not set. Use object.set_api_key(api_key) to set API key for this call")
            else:
                try:
                    self.client = InferenceClient(api_key=self.api_key)
                except:
                    print("Check your connection, model, and API key")

    def send_chat_request(self, prompt: str):
        if self.client == None:
            raise ValueError("Client Haven't connected. Use method object.connect(model) method. Also add new api key before connecting by using object.set_api_key(api_key)")
        else:
            message = self.client.chat_completion(
                model=self.model,
                max_tokens=1024,
                messages=[{"role":"user", "content":prompt}],
                temperature=0.5,
                top_p=0.7,
                stream=False
            )
            response_content = message.choices[0].message["content"]
            return response_content
################################### Local LLM #################################################

# Limited local resources to run LLM