""" 使用向量库和全文检索两个来构建知识库 """ from pathlib import Path from tika import parser from langchain.text_splitter import RecursiveCharacterTextSplitter from datetime import datetime from full_index_search import create_full_index,create_full_index2,delete_fs_index_by_base_path from k_database import Doc,BaiKe,BaiKe_Catalog import json import json5 import time #from llm_embedding2 import embeddings #from milvus import default_server #from langchain_milvus import Milvus from init import gcfg from ff import smb_scan_directory,sftp_scan_directory,sftp_conn,sftp_close,ftp_scan_directory def extract_text(content): #raw = parser.from_file(file_path) try: raw = parser.from_buffer(content) except: raw ={'metadata':{"Content-Type":""},'content':''} return raw["metadata"],raw['content'] def extract_text_path(file_path): try: raw = parser.from_file(file_path) except: raw ={'metadata':{"Content-Type":""},'content':''} return raw["metadata"],raw['content'] #遍历目录下的所有文件和子目录,返回文件path和meta信息 def list_files(directory,pre="*"): for path in Path(directory).rglob(pre): if path.is_file(): yield path,path.stat() # 创建一个 RecursiveCharacterTextSplitter 实例 text_splitter = RecursiveCharacterTextSplitter( chunk_size=gcfg["full_index"]["chunk_size"], # 设置每个文本块的目标大小 chunk_overlap=gcfg["full_index"]["chunk_overlap"] # 设置相邻文本块之间的重叠字符数 ) # vector_db=Milvus( # embeddings, # connection_args={"host":"127.0.0.1","port": default_server.listen_port}, # collection_name='test', # auto_id=True, # ) #忽略的文件类型 ignore_file_type=gcfg["full_index"]["ignore_ftype"].split(",") #忽略的文件大小 max_file_size = int(gcfg["full_index"]["ignore_fsize"])*1024*1024 #文件的扩展名 def get_file_ext(fname): return fname[fname.rfind("."):] import os import pwd def build_linux_base(base,cfg): count=0 error=0 for path,stat in list_files(cfg["path"]): count +=1 size = stat.st_size print(str(path)) if get_file_ext(path.name) not in ignore_file_type and size < max_file_size and path.name[0:2] not in [".~","~$"]: try: doc = Doc.get_or_none( (Doc.abs_path==str(path)) & (Doc.base==cfg["name"]) ) stat_info = os.stat(str(path)) uid = stat_info.st_uid username = pwd.getpwuid(uid).pw_name if not doc or doc.f_up_at.timestamp() < stat.st_mtime: meta,text = extract_text_path(str(path)) if text==None: text="" build_one_doc(base,str(path),path.name,meta,text,size,stat.st_mtime,username) except Exception as e : print(f"{path} 构建出错:{e}") error +=1 if error >0: raise Exception(f"本次扫描有{error}个文档错误") return count import smbclient #构建smb的知识库 def build_smb_base(base,cfg): count=0 smbclient.register_session(cfg["address"],cfg["user"],cfg["password"]) for path,fname,mtime,size in smb_scan_directory(cfg["address"],cfg["user"],cfg["password"],cfg["path"]): count +=1 print(count,path,size) if get_file_ext(fname) not in ignore_file_type and size < max_file_size and fname[0:2] not in [".~","~$"]: doc = Doc.get_or_none( (Doc.abs_path==str(path)) & (Doc.base==cfg["name"]) ) if not doc or doc.f_up_at.timestamp() < mtime: try: with smbclient.open_file(path, mode="rb") as fd: content = fd.read() except: #出错重连，继续 smbclient.register_session(cfg["address"],cfg["user"],cfg["password"]) continue meta,text = extract_text(content) if text==None: text="" build_one_doc(base,path,fname,meta,text,size,mtime) return count #构建sftp的知识库 def build_sftp_base(base,cfg): count=0 ssh,sftp = sftp_conn(cfg["address"],cfg["user"],cfg["password"]) for path,fname,mtime,size,ownername in sftp_scan_directory(cfg["address"],cfg["user"],cfg["password"],cfg["path"]): count +=1 print(count,path,size,ownername) if get_file_ext(fname) not in ignore_file_type and size < max_file_size and fname[0:2] not in [".~","~$"]: doc = Doc.get_or_none( (Doc.abs_path==str(path)) & (Doc.base==cfg["name"]) ) if not doc or doc.f_up_at.timestamp() < mtime: try: with sftp.file(path, 'r') as remote_file: content = remote_file.read() except: #出错重连，继续 sftp_close(ssh,sftp) ssh,sftp = sftp_conn(cfg["address"],cfg["user"],cfg["password"]) continue meta,text = extract_text(content) if text==None: text="" build_one_doc(base,path,fname,meta,text,size,mtime,ownername) sftp_close(ssh,sftp) return count from ftplib import FTP from io import BytesIO def get_binary_file_content(ftp, remote_path): """ 获取二进制文件的内容。 :param ftp: 已连接并登录的 FTP 对象 :param remote_path: 远程文件路径 :return: 文件内容 (bytes) """ binary_data = BytesIO() ftp.retrbinary(f'RETR {remote_path}', binary_data.write) return binary_data.getvalue() #构建ftp的知识库 def build_ftp_base(base,cfg): count=0 for entry in ftp_scan_directory(cfg["address"],cfg["user"],cfg["password"],cfg["path"]): path = entry["path"] fname = entry["name"] size = int(entry["size"]) mtime = entry["mtime"] count +=1 print(count,path,size,max_file_size) if get_file_ext(fname) not in ignore_file_type and size < max_file_size and fname[0:2] not in [".~","~$"]: doc = Doc.get_or_none( (Doc.abs_path==str(path)) & (Doc.base==cfg["name"]) ) if not doc or doc.f_up_at.timestamp() < mtime: # 连接到FTP服务器 with FTP(cfg["address"]) as ftp: ftp.login(user=cfg["user"], passwd=cfg["password"]) content = get_binary_file_content(ftp,path) meta,text = extract_text(content) if text==None: text="" build_one_doc(base,path,fname,meta,text,size,mtime) return count from dingpan import ding_scan_directory,ding_get_file_content #构建钉盘上文件目录的知识库 def build_ding_pan(base,cfg): count=0 for entry in ding_scan_directory(cfg["address"],cfg["password"],cfg["path"]): path = entry["path"] fname = entry["fname"] file_id = entry["id"] mtime = entry["mtime"] spaceId = entry["spaceId"] unionid = entry["unionid"] access_token = entry["access_token"] count +=1 print(count,path,file_id) if get_file_ext(fname) not in ignore_file_type and fname[0:2] not in [".~","~$"]: doc = Doc.get_or_none( (Doc.abs_path==str(path)) & (Doc.base==cfg["name"]) ) if not doc or doc.f_up_at < mtime: content = ding_get_file_content(spaceId,file_id,unionid,access_token) if content and len(content) < max_file_size: meta,text = extract_text(content) if text: build_one_doc(base,path,fname,{"id":file_id,"spaceId":entry["spaceId"]},text,len(content),mtime.timestamp()) return count #构建单个文档 def build_one_doc(base,path,file_name,meta,text,st_size,st_mtime,ownername="admin"): #向量数据库的处理 # #分割文档，包含文件名称的meta信息 #metas=[{"name":file_name}]*len(split_texts) #documents = text_splitter.create_documents(split_texts,metas) # try: # vector_db.add_documents(documents) # except: # pass # query = Doc.insert(base=base,abs_path=path, # f_name=file_name, # f_size=st_size, # f_up_at = datetime.fromtimestamp(st_mtime), # ctype= meta["Content-Type"], # catalog="内部", # meta=json.dumps(meta)[:1000] # ).on_conflict( # conflict_target=[Doc.abs_path], # update={Doc.f_up_at:datetime.fromtimestamp(st_mtime)} # ) #query.execute() try: Doc.create(base=base,abs_path=path, f_name=file_name, f_size=st_size, f_up_at = datetime.fromtimestamp(st_mtime), ctype= meta["Content-Type"], catalog="内部", meta=json.dumps(meta)[:1000], author = ownername, ) except: # 更新，同时删除原有全文索引 Doc.update({Doc.f_up_at:datetime.fromtimestamp(st_mtime),Doc.base:base, Doc.meta:json.dumps(meta)[:1000],Doc.author:ownername, }).where(Doc.abs_path==path).execute() delete_fs_index_by_base_path(base,path) #内容分片 split_texts = text_splitter.split_text(text) print(path,len(text),len(split_texts)) split_texts = list(map(lambda x: f"文件: {path}\n"+x,split_texts)) create_full_index(split_texts,file_name,path,st_mtime,base) #print(file,len(text),len(documents)) #百科分类的统计 def baike_catalog_total(): from peewee import fn query = (BaiKe. select( BaiKe.catalog,fn.COUNT(BaiKe.id).alias('count') ) .group_by(BaiKe.catalog) ) for record in query: BaiKe_Catalog.update(doc_count= record.count).where(BaiKe_Catalog.catalog == record.catalog).execute() # 清理过期session def clear_session_id(): import glob import os session_ids = glob.glob(f"/tmp/k3gpt/session_id_*", recursive=True) for file in session_ids: try: with open(file,"r") as f: s_timestamp = f.read() if int(time.time())-int(s_timestamp) > 3600*12: os.remove(file) except Exception as e: print (e) if __name__=="__main__": print("开始知识库构建工作") while True: try: with open(f"{gcfg['fs']['path']}/conf/source.conf") as conf: source = json5.load(conf) except: source={"测试":{"name":"测试","type":"FTP","address":"192.168.0.1:21","user":"admin","password":"1","path":"home","scan":1,"count":1024}} bases=list(source.values()) for base in bases: try: now = int(time.time()) #扫描周期 scant = int(float(source[base["name"]]["scan"])*3600) #扫描时间的判断,不到时间 if "scan_date" in base and now < source[base["name"]]["scan_date"] + scant: print("扫描时间未到",base["name"],datetime.now()) continue print("扫描开始",base["name"],datetime.now()) count=0 if base["type"]=="SFTP": count = build_sftp_base(base["name"],base) elif base["type"]=="本地": count = build_linux_base(base["name"],base) elif base["type"]=="Windows共享": count= build_smb_base(base["name"],base) elif base["type"]=="FTP": count= build_ftp_base(base["name"],base) elif base["type"]=="钉盘": count= build_ding_pan(base["name"],base) source[base["name"]]["count"] = count source[base["name"]]["scan_date"] = int(time.time()) print("扫描结束",base["name"],datetime.now()) except Exception as e : print(e) #更新 #end for #写回 #最新的配置 with open(f"{gcfg['fs']['path']}/conf/source.conf") as conf: source_new = json5.load(conf) #更新文件数和扫描时间 for k,v in source.items(): if "count" in v : source_new[k]["count"] = v["count"] if "scan_date" in v : source_new[k]["scan_date"] = v["scan_date"] with open(f"{gcfg['fs']['path']}/conf/source.conf","w+") as conf: json5.dump(source_new,conf,ensure_ascii=False) try: baike_catalog_total() except Exception as e: print("百科分类统计",e) clear_session_id() #休息30秒 time.sleep(30)