systex_phase2/file_loader/add_vectordb.py

from dotenv import load_dotenv
load_dotenv("../.env")

from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.document_loaders import Docx2txtLoader
from langchain.document_loaders import CSVLoader

import os
import glob

from langchain_community.vectorstores import SupabaseVectorStore
from langchain_openai import OpenAIEmbeddings
from supabase.client import Client, create_client

load_dotenv()
supabase_url = os.environ.get("SUPABASE_URL")
supabase_key = os.environ.get("SUPABASE_KEY")
document_table = "documents2"
supabase: Client = create_client(supabase_url, supabase_key)
def get_data_list(data_list=None, path=None, extension=None, update=False):
    files = data_list or glob.glob(os.path.join(path, f"*.{extension}"))
    if update:    
        doc = files.copy()
    else:
        existed_data = check_existed_data(supabase)
        doc = []
        for file_path in files:
            filename = os.path.basename(file_path)
            if filename not in existed_data:
                doc.append(file_path)

    return doc


def read_and_split_files(data_list=None, path=None, extension=None, update=False):

    def read_csv(path):
        extension = "csv"
        # path = r"./Phase2/"
        files = glob.glob(os.path.join(path, f"*.{extension}"))

        if not files:
            return None

        documents = []
        for file_path in files:
            print(file_path)

            loader = CSVLoader(file_path, encoding="utf-8")
            doc = loader.load()

            documents.extend(doc)
        
        return documents
    
    def load_and_split(file_list):
        chunks = []
        for file in file_list:
            if file.endswith(".txt") or file.endswith(".md"):
                loader = TextLoader(file, encoding='utf-8')
                docs = loader.load()
            elif file.endswith(".pdf"):
                loader = PyPDFLoader(file)
                docs = loader.load()
            elif file.endswith(".docx"):
                loader = Docx2txtLoader(file)
                docs = loader.load()
            else:
                print(f"Unsupported file extension: {file}")
                continue



            # Split
            rules = ['低碳產品獎勵辦法.docx', '公私場所固定污染源空氣污染物排放量申報管理辦法.docx', '氣候變遷因應法.docx', '氣候變遷因應法施行細則.docx', 
                     '淘汰老舊機車換購電動機車溫室氣體減量獎勵辦法訂定總說明及逐條說明.docx', '溫室氣體抵換專案管理辦法.docx', '溫室氣體排放源符合效能標準獎勵辦法.docx', 
                     '溫室氣體排放量增量抵換管理辦法.docx', '溫室氣體排放量增量抵換管理辦法訂定總說明及逐條說明.docx', '溫室氣體排放量盤查登錄及查驗管理辦法修正條文.docx', 
                     '溫室氣體排放量盤查登錄管理辦法(溫室氣體排放量盤查登錄及查驗管理辦法修正條文前身).docx', '溫室氣體自願減量專案管理辦法.docx', 
                     '溫室氣體自願減量專案管理辦法中華民國112年10月12日訂定總說明及逐條說明.docx', '溫室氣體認證機構及查驗機構管理辦法.docx', 
                     '溫室氣體階段管制目標及管制方式作業準則.docx', '碳足跡產品類別規則訂定、引用及修訂指引.docx', 
                     '老舊汽車汰舊換新溫室氣體減量獎勵辦法中華民國112年1月11日訂定總說明及逐條說明.docx']
            print(os.path.basename(file))
            if file.endswith(".docx") and os.path.basename(file) in rules:
                separators = ['\u25cb\s*第.*?條', '\u25cf\s*第.*?條']
                text_splitter = RecursiveCharacterTextSplitter(is_separator_regex=True, separators=separators, chunk_size=500, chunk_overlap=0)
                splits = text_splitter.split_documents(docs)
            elif os.path.basename(file) in ["new_information.docx"]:
                print(file)
                separators = ['###']
                text_splitter = RecursiveCharacterTextSplitter(is_separator_regex=True, separators=separators, chunk_size=500, chunk_overlap=0)
                splits = text_splitter.split_documents(docs)

            elif file.endswith(".md"):
                headers_to_split_on = [
                    ("#", "Header 1"),
                    ("##", "Header 2"),
                    ("###", "Header 3"),
                ]

                markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
                splits = markdown_splitter.split_text(docs[0].page_content)
                for split in splits:
                    split
                
            else:
                text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=500, chunk_overlap=0)
                splits = text_splitter.split_documents(docs)

            chunks.extend(splits)
            # doc = read_csv(path)
            # chunks.extend(doc)

        return chunks
    


    # specific data type
    doc = get_data_list(data_list=data_list, path=path, extension=extension, update=update)

    # web url

    # csv
    # doc = read_csv(path)

    # Index
    docs = load_and_split(doc)

    return docs

def create_ids(docs):
    # Create a dictionary to count occurrences of each page in each document
    page_counter = {}

    # List to store the resulting IDs
    document_ids = []

    # Generate IDs
    for doc in [docs[i].metadata for i in range(len(docs))]:
        if "source" in doc.keys():
            source = doc['source']
            file_name = os.path.basename(source).split('.')[0]
        else:
            source = "supplement"
            file_name = "supplement"

        if "page" in doc.keys():
            page = doc['page']
            key = f"{source}_{page}"
        else:
            key = f"{source}"

        if key not in page_counter:
            page_counter[key] = 1
        else:
            page_counter[key] += 1
        
        if "page" in doc.keys():
            doc_id = f"{file_name} | page {page} | chunk {page_counter[key]}"
        else:
            doc_id = f"{file_name} | chunk {page_counter[key]}"

        
        document_ids.append(doc_id)

    return document_ids

def get_document(data_list=None, path=None, extension=None, update=False):
    docs = read_and_split_files(data_list=data_list, path=path, extension=extension, update=update)
    document_ids = create_ids(docs)

    for doc in docs:
        doc.metadata['source'] = os.path.basename(doc.metadata['source']) if 'source' in doc.metadata else "supplement.md"
        # print(doc.metadata)

    # document_metadatas = [{'source': doc.metadata['source'], 'page': doc.metadata['page'], 'chunk': int(id.split("chunk ")[-1])} for doc, id in zip(docs, document_ids)]
    document_metadatas = []

    for doc, id in zip(docs, document_ids):
        chunk_number = int(id.split("chunk ")[-1])
        doc.metadata['chunk'] = chunk_number
        doc.metadata['extension'] = os.path.basename(doc.metadata['source']).split(".")[-1]
        document_metadatas.append(doc.metadata)

    documents = [docs.metadata['source'].split(".")[0] + docs.page_content for docs in docs]

    return document_ids, documents, document_metadatas

def check_existed_data(supabase):
    response = supabase.table(document_table).select("id, metadata").execute()
    existed_data = list(set([data['metadata']['source'] for data in response.data]))
    # existed_data = [(data['id'], data['metadata']['source']) for data in response.data]
    return existed_data


class GetVectorStore(SupabaseVectorStore):
    def __init__(self, embeddings, supabase, table_name):
        super().__init__(embedding=embeddings, client=supabase, table_name=table_name, query_name="match_documents")

    def insert(self, documents, document_metadatas):
        self.add_texts(
            texts=documents,
            metadatas=document_metadatas,
        )

    def delete(self, file_list):
        for file in file_list:
            file_name = os.path.basename(file)
            self._client.table(self.table_name).delete().eq('metadata->>source', file_name).execute()

    def update(self, documents, document_metadatas, update_existing_data=False):
        if not document_metadatas:  # no new data
            return

        if update_existing_data:
            file_list = list(set(metadata['source'] for metadata in document_metadatas))
            self.delete(file_list)

        self.insert(documents, document_metadatas)

if __name__ == "__main__":

    load_dotenv("../.env")
    supabase_url = os.environ.get("SUPABASE_URL")
    supabase_key = os.environ.get("SUPABASE_KEY")
    document_table = "documents"
    supabase: Client = create_client(supabase_url, supabase_key)

    embeddings = OpenAIEmbeddings()

    ###################################################################################
    # get vector store
    vector_store = GetVectorStore(embeddings, supabase, document_table)

    ###################################################################################
    # update data (old + new / all new / all old)
    path = "/home/mia/systex/Documents"
    extension = "pdf"
    # file = None

    # file_list = ["溫室氣體排放量盤查作業指引113年版.pdf"]
    # file = [os.path.join(path, file) for file in file_list]
    # file_list = glob.glob(os.path.join(path, "*"))
    file_list =glob.glob(os.path.join(path, f"*.{extension}"))
    # print(file_list)
    
    # update = False
    # document_ids, documents, document_metadatas = get_document(data_list=file_list, path=path, extension=extension, update=update)
    # vector_store.update(documents, document_metadatas, update_existing_data=update)

    ###################################################################################
    # insert new data (all new)
    # vector_store.insert(documents, document_metadatas)

    ###################################################################################
    # delete data
    # file_list = ["溫室氣體排放量盤查作業指引113年版.pdf"]
    # file_list = glob.glob(os.path.join(path, f"*.docx"))
    # file_list = [os.path.basename(file_path) for file_path in file_list]
    # print(file_list)
    # vector_store.delete(file_list)

    ###################################################################################
    # get retriver
    # retriever = vector_store.as_retriever(search_kwargs={"k": 6})