knowledge_base.py

import os
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"

import hashlib
from dotenv import load_dotenv
load_dotenv()

import datetime
from langchain_chroma import Chroma
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

import config_data as config


def check_md5(md5_str:str):
    """检查是否被处理过"""
    if not os.path.exists(config.md5_path):
        open(config.md5_path,'w',encoding="utf-8").close()
        return False
    else:
        for line in open(config.md5_path,'r',encoding="utf-8").readlines():
            line = line.strip() #空格回车处理
            if line == md5_str:
                return True
        return False


        pass
def save_md5(md5_str:str):
    """记录md5"""
    with open(config.md5_path,'a',encoding="utf-8") as f:
        f.write(md5_str+'\n')

def get_string_md5(input_str:str,encoding="utf-8"):
    str_bytes=input_str.encode(encoding=encoding)
    md5_obj=hashlib.md5()
    md5_obj.update(str_bytes)
    md5_hex=md5_obj.hexdigest()
    return md5_hex #16进制字符串


class KnowledgeBaseService(object):
    def __init__(self):
        os.makedirs(config.persist_directory,exist_ok=True)

        self.embeddings = DashScopeEmbeddings(
            model="text-embedding-v4"
        )


        self.chroma=Chroma(
            collection_name=config.collection_name, #数据库表明
            embedding_function=self.embeddings,
            persist_directory=config.persist_directory
        )        #向量存储实例Chroma向量库对象
        self.spliter=RecursiveCharacterTextSplitter(
            chunk_size=config.chunk_size,
            chunk_overlap=config.chunk_overlap, #连续文本字符重叠数量
            separators=config.separators,
            length_function=len

        )       #文本分割器对象


    def upload_by_str(self,data,filename):
        """传入的字符串向量化传入数据库"""


        md5_hex=get_string_md5(data)
        if check_md5(md5_hex):
            return "[跳过]内容已经被处理过"

        if len(data)>config.max_split_char_number:
            knowledge_chunks:list[str]=self.spliter.split_text(data)
        else:
            knowledge_chunks: list[str] = [data]

        metadata={
            "source": filename,
            "create_time": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "operator":"白曜石"
        }

        self.chroma.add_texts(
            knowledge_chunks,
            metadatas=[metadata for _ in knowledge_chunks]

        )

        save_md5(md5_hex)
        return "[成功]内容成功上传向量库"

if __name__ == '__main__':

    service=KnowledgeBaseService()
    print(service.upload_by_str("bys", "testfile"))