pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

go through smoke test of all API (#12)

Browse Source 
* add field progress msg into docinfo; add file processing procedure

* go through upload, create kb, add doc to kb

* smoke test for all API

* smoke test for all API
This commit is contained in:
KevinHuSh
2023-12-22 17:57:27 +08:00
committed by GitHub
parent 72b7b5fae5
commit 1eb186a25f
27 changed files with 921 additions and 281 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
python/conf/sys.cnf
View File

@ -1,7 +1,10 @@
[infiniflow]
es=127.0.0.1:9200
es=http://127.0.0.1:9200
pgdb_usr=root
pgdb_pwd=infiniflow_docgpt
pgdb_host=127.0.0.1
pgdb_port=5455
minio_host=127.0.0.1:9000
minio_usr=infiniflow
minio_pwd=infiniflow_docgpt

14
python/nlp/huchunk.py
View File

@ -2,6 +2,7 @@ import re
import os
import copy
import base64
import magic
from dataclasses import dataclass
from typing import List
import numpy as np
@ -373,6 +374,7 @@ class PptChunker(HuChunker):
from pptx import Presentation
ppt = Presentation(fnm)
flds = self.Fields()
flds.text_chunks = []
for slide in ppt.slides:
for shape in slide.shapes:
if hasattr(shape, "text"):
@ -391,11 +393,21 @@ class TextChunker(HuChunker):
def __init__(self):
super().__init__()
@staticmethod
def is_binary_file(file_path):
mime = magic.Magic(mime=True)
file_type = mime.from_file(file_path)
if 'text' in file_type:
return False
else:
return True
def __call__(self, fnm):
flds = self.Fields()
if self.is_binary_file(fnm):return flds
with open(fnm, "r") as f:
txt = f.read()
flds.text_chunks = self.naive_text_chunk(txt)
flds.text_chunks = [(c, None) for c in self.naive_text_chunk(txt)]
flds.table_chunks = []
return flds

187
python/svr/parse_user_docs.py
View File

@ -1,10 +1,15 @@
import json, re, sys, os, hashlib, copy, glob, util, time, random
from util.es_conn import HuEs, Postgres
import json, os, sys, hashlib, copy, time, random, re, logging, torch
from os.path import dirname, realpath
sys.path.append(dirname(realpath(__file__)) + "/../")
from util.es_conn import HuEs
from util.db_conn import Postgres
from util.minio_conn import HuMinio
from util import rmSpace, findMaxDt
from FlagEmbedding import FlagModel
from nlp import huchunk, huqie
import base64, hashlib
from io import BytesIO
import pandas as pd
from elasticsearch_dsl import Q
from parser import (
PdfParser,
@ -22,73 +27,115 @@ from nlp.huchunk import (
ES = HuEs("infiniflow")
BATCH_SIZE = 64
PG = Postgres("infiniflow", "docgpt")
MINIO = HuMinio("infiniflow")
PDF = PdfChunker(PdfParser())
DOC = DocxChunker(DocxParser())
EXC = ExcelChunker(ExcelParser())
PPT = PptChunker()
UPLOAD_LOCATION = os.environ.get("UPLOAD_LOCATION", "./")
logging.warning(f"The files are stored in {UPLOAD_LOCATION}, please check it!")
def chuck_doc(name):
name = os.path.split(name)[-1].lower().split(".")[-1]
if name.find("pdf") >= 0: return PDF(name)
if name.find("doc") >= 0: return DOC(name)
if name.find("xlsx") >= 0: return EXC(name)
if name.find("ppt") >= 0: return PDF(name)
if name.find("pdf") >= 0: return PPT(name)
suff = os.path.split(name)[-1].lower().split(".")[-1]
if suff.find("pdf") >= 0: return PDF(name)
if suff.find("doc") >= 0: return DOC(name)
if re.match(r"(xlsx|xlsm|xltx|xltm)", suff): return EXC(name)
if suff.find("ppt") >= 0: return PPT(name)
if re.match(r"(txt|csv)", name): return TextChunker(name)
return TextChunker()(name)
def collect(comm, mod, tm):
sql = f"""
select
id as kb2doc_id,
kb_id,
did,
updated_at,
is_deleted
from kb2_doc
where
updated_at >= '{tm}'
and kb_progress = 0
and MOD(did, {comm}) = {mod}
order by updated_at asc
limit 1000
"""
kb2doc = PG.select(sql)
if len(kb2doc) == 0:return pd.DataFrame()
sql = """
select
did,
uid,
doc_name,
location,
updated_at
from docinfo
size
from doc_info
where
updated_at >= '{tm}'
and kb_progress = 0
and type = 'doc'
and MOD(uid, {comm}) = {mod}
order by updated_at asc
limit 1000
"""
df = PG.select(sql)
df = df.fillna("")
mtm = str(df["updated_at"].max())[:19]
print("TOTAL:", len(df), "To: ", mtm)
return df, mtm
did in (%s)
"""%",".join([str(i) for i in kb2doc["did"].unique()])
docs = PG.select(sql)
docs = docs.fillna("")
docs = docs.join(kb2doc.set_index("did"), on="did", how="left")
mtm = str(docs["updated_at"].max())[:19]
print("TOTAL:", len(docs), "To: ", mtm)
return docs
def set_progress(did, prog, msg):
def set_progress(kb2doc_id, prog, msg="Processing..."):
sql = f"""
update docinfo set kb_progress={prog}, kb_progress_msg='{msg}' where did={did}
update kb2_doc set kb_progress={prog}, kb_progress_msg='{msg}'
where
id={kb2doc_id}
"""
PG.update(sql)
def build(row):
if row["size"] > 256000000:
set_progress(row["did"], -1, "File size exceeds( <= 256Mb )")
set_progress(row["kb2doc_id"], -1, "File size exceeds( <= 256Mb )")
return []
res = ES.search(Q("term", doc_id=row["did"]))
if ES.getTotal(res) > 0:
ES.updateScriptByQuery(Q("term", doc_id=row["did"]),
scripts="""
if(!ctx._source.kb_id.contains('%s'))
ctx._source.kb_id.add('%s');
"""%(str(row["kb_id"]), str(row["kb_id"])),
idxnm = index_name(row["uid"])
)
set_progress(row["kb2doc_id"], 1, "Done")
return []
random.seed(time.time())
set_progress(row["kb2doc_id"], random.randint(0, 20)/100., "Finished preparing! Start to slice file!")
try:
obj = chuck_doc(os.path.join(UPLOAD_LOCATION, row["location"]))
except Exception as e:
if re.search("(No such file|not found)", str(e)):
set_progress(row["kb2doc_id"], -1, "Can not find file <%s>"%row["doc_name"])
else:
set_progress(row["kb2doc_id"], -1, f"Internal system error: %s"%str(e).replace("'", ""))
return []
print(row["doc_name"], obj)
if not obj.text_chunks and not obj.table_chunks:
set_progress(row["kb2doc_id"], 1, "Nothing added! Mostly, file type unsupported yet.")
return []
set_progress(row["kb2doc_id"], random.randint(20, 60)/100., "Finished slicing files. Start to embedding the content.")
doc = {
"doc_id": row["did"],
"kb_id": [str(row["kb_id"])],
"title_tks": huqie.qie(os.path.split(row["location"])[-1]),
"updated_at": row["updated_at"]
"updated_at": str(row["updated_at"]).replace("T", " ")[:19]
}
random.seed(time.time())
set_progress(row["did"], random.randint(0, 20)/100., "Finished preparing! Start to slice file!")
obj = chuck_doc(row["location"])
if not obj:
set_progress(row["did"], -1, "Unsuported file type.")
return []
set_progress(row["did"], random.randint(20, 60)/100.)
output_buffer = BytesIO()
docs = []
md5 = hashlib.md5()
@ -97,12 +144,11 @@ def build(row):
md5.update((txt + str(d["doc_id"])).encode("utf-8"))
d["_id"] = md5.hexdigest()
d["content_ltks"] = huqie.qie(txt)
d["docnm_kwd"] = rmSpace(d["docnm_tks"])
if not img:
docs.append(d)
continue
img.save(output_buffer, format='JPEG')
d["img_bin"] = base64.b64encode(output_buffer.getvalue())
d["img_bin"] = str(output_buffer.getvalue())
docs.append(d)
for arr, img in obj.table_chunks:
@ -115,9 +161,11 @@ def build(row):
docs.append(d)
continue
img.save(output_buffer, format='JPEG')
d["img_bin"] = base64.b64encode(output_buffer.getvalue())
MINIO.put("{}-{}".format(row["uid"], row["kb_id"]), d["_id"],
output_buffer.getvalue())
d["img_id"] = "{}-{}".format(row["uid"], row["kb_id"])
docs.append(d)
set_progress(row["did"], random.randint(60, 70)/100., "Finished slicing. Start to embedding the content.")
set_progress(row["kb2doc_id"], random.randint(60, 70)/100., "Continue embedding the content.")
return docs
@ -127,7 +175,7 @@ def index_name(uid):return f"docgpt_{uid}"
def init_kb(row):
idxnm = index_name(row["uid"])
if ES.indexExist(idxnm): return
return ES.createIdx(idxnm, json.load(open("res/mapping.json", "r")))
return ES.createIdx(idxnm, json.load(open("conf/mapping.json", "r")))
model = None
@ -138,27 +186,59 @@ def embedding(docs):
vects = 0.1 * tts + 0.9 * cnts
assert len(vects) == len(docs)
for i,d in enumerate(docs):d["q_vec"] = vects[i].tolist()
for d in docs:
set_progress(d["doc_id"], random.randint(70, 95)/100.,
"Finished embedding! Start to build index!")
def rm_doc_from_kb(df):
if len(df) == 0:return
for _,r in df.iterrows():
ES.updateScriptByQuery(Q("term", doc_id=r["did"]),
scripts="""
if(ctx._source.kb_id.contains('%s'))
ctx._source.kb_id.remove(
ctx._source.kb_id.indexOf('%s')
);
"""%(str(r["kb_id"]),str(r["kb_id"])),
idxnm = index_name(r["uid"])
)
if len(df) == 0:return
sql = """
delete from kb2_doc where id in (%s)
"""%",".join([str(i) for i in df["kb2doc_id"]])
PG.update(sql)
def main(comm, mod):
global model
from FlagEmbedding import FlagModel
model = FlagModel('/opt/home/kevinhu/data/bge-large-zh-v1.5/',
query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章：",
use_fp16=torch.cuda.is_available())
tm_fnm = f"res/{comm}-{mod}.tm"
tmf = open(tm_fnm, "a+")
tm = findMaxDt(tm_fnm)
rows, tm = collect(comm, mod, tm)
for r in rows:
if r["is_deleted"]:
ES.deleteByQuery(Q("term", dock_id=r["did"]), index_name(r["uid"]))
continue
rows = collect(comm, mod, tm)
if len(rows) == 0:return
rm_doc_from_kb(rows.loc[rows.is_deleted == True])
rows = rows.loc[rows.is_deleted == False].reset_index(drop=True)
if len(rows) == 0:return
tmf = open(tm_fnm, "a+")
for _, r in rows.iterrows():
cks = build(r)
if not cks:
tmf.write(str(r["updated_at"]) + "\n")
continue
## TODO: exception handler
## set_progress(r["did"], -1, "ERROR: ")
embedding(cks)
if cks: init_kb(r)
ES.bulk(cks, index_name(r["uid"]))
set_progress(r["kb2doc_id"], random.randint(70, 95)/100.,
"Finished embedding! Start to build index!")
init_kb(r)
es_r = ES.bulk(cks, index_name(r["uid"]))
if es_r:
set_progress(r["kb2doc_id"], -1, "Index failure!")
print(es_r)
else: set_progress(r["kb2doc_id"], 1., "Done!")
tmf.write(str(r["updated_at"]) + "\n")
tmf.close()
@ -166,6 +246,5 @@ def main(comm, mod):
if __name__ == "__main__":
from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
main(comm, rank)
main(comm.Get_size(), comm.Get_rank())

4
python/util/config.py
View File

@ -14,9 +14,9 @@ class Config:
self.env = env
if env == "spark":CF.read("./cv.cnf")
def get(self, key):
def get(self, key, default=None):
global CF
return CF.get(self.env, key)
return CF[self.env].get(key, default)
def init(env):
return Config(env)

4
python/util/db_conn.py
View File

@ -49,7 +49,11 @@ class Postgres(object):
cur = self.conn.cursor()
cur.execute(sql)
updated_rows = cur.rowcount
<<<<<<< HEAD
self.conn.commit()
=======
conn.commit()
>>>>>>> upstream/main
cur.close()
return updated_rows
except Exception as e:

36
python/util/es_conn.py
View File

@ -5,10 +5,10 @@ import time
import copy
import elasticsearch
from elasticsearch import Elasticsearch
from elasticsearch_dsl import UpdateByQuery, Search, Index
from elasticsearch_dsl import UpdateByQuery, Search, Index, Q
from util import config
print("Elasticsearch version: ", elasticsearch.__version__)
logging.info("Elasticsearch version: ", elasticsearch.__version__)
def instance(env):
@ -20,7 +20,7 @@ def instance(env):
timeout=600
)
print("ES: ", ES_DRESS, ES.info())
logging.info("ES: ", ES_DRESS, ES.info())
return ES
@ -31,7 +31,7 @@ class HuEs:
self.info = {}
self.config = config.init(env)
self.conn()
self.idxnm = self.config.get("idx_nm","")
self.idxnm = self.config.get("idx_nm", "")
if not self.es.ping():
raise Exception("Can't connect to ES cluster")
@ -46,6 +46,7 @@ class HuEs:
break
except Exception as e:
logging.error("Fail to connect to es: " + str(e))
time.sleep(1)
def version(self):
v = self.info.get("version", {"number": "5.6"})
@ -121,7 +122,6 @@ class HuEs:
acts.append(
{"update": {"_id": id, "_index": ids[id]["_index"]}, "retry_on_conflict": 100})
acts.append({"doc": d, "doc_as_upsert": "true"})
logging.info("bulk upsert: %s" % id)
res = []
for _ in range(100):
@ -148,7 +148,6 @@ class HuEs:
return res
except Exception as e:
logging.warn("Fail to bulk: " + str(e))
print(e)
if re.search(r"(Timeout|time out)", str(e), re.IGNORECASE):
time.sleep(3)
continue
@ -229,7 +228,7 @@ class HuEs:
return False
def search(self, q, idxnm=None, src=False, timeout="2s"):
print(json.dumps(q, ensure_ascii=False))
if not isinstance(q, dict): q = Search().query(q).to_dict()
for i in range(3):
try:
res = self.es.search(index=(self.idxnm if not idxnm else idxnm),
@ -271,9 +270,31 @@ class HuEs:
str(e) + "【Q】：" + str(q.to_dict()))
if str(e).find("Timeout") > 0 or str(e).find("Conflict") > 0:
continue
self.conn()
return False
def updateScriptByQuery(self, q, scripts, idxnm=None):
ubq = UpdateByQuery(index=self.idxnm if not idxnm else idxnm).using(self.es).query(q)
ubq = ubq.script(source=scripts)
ubq = ubq.params(refresh=True)
ubq = ubq.params(slices=5)
ubq = ubq.params(conflicts="proceed")
for i in range(3):
try:
r = ubq.execute()
return True
except Exception as e:
logging.error("ES updateByQuery exception: " +
str(e) + "【Q】：" + str(q.to_dict()))
if str(e).find("Timeout") > 0 or str(e).find("Conflict") > 0:
continue
self.conn()
return False
def deleteByQuery(self, query, idxnm=""):
for i in range(3):
try:
@ -307,7 +328,6 @@ class HuEs:
routing=routing, refresh=False) # , doc_type="_doc")
return True
except Exception as e:
print(e)
logging.error("ES update exception: " + str(e) + " id：" + str(id) + ", version:" + str(self.version()) +
json.dumps(script, ensure_ascii=False))
if str(e).find("Timeout") > 0:

73
python/util/minio_conn.py Normal file
View File

@ -0,0 +1,73 @@
import logging
import time
from util import config
from minio import Minio
from io import BytesIO
class HuMinio(object):
def __init__(self, env):
self.config = config.init(env)
self.conn = None
self.__open__()
def __open__(self):
try:
if self.conn:self.__close__()
except Exception as e:
pass
try:
self.conn = Minio(self.config.get("minio_host"),
access_key=self.config.get("minio_usr"),
secret_key=self.config.get("minio_pwd"),
secure=False
)
except Exception as e:
logging.error("Fail to connect %s "%self.config.get("minio_host") + str(e))
def __close__(self):
del self.conn
self.conn = None
def put(self, bucket, fnm, binary):
for _ in range(10):
try:
if not self.conn.bucket_exists(bucket):
self.conn.make_bucket(bucket)
r = self.conn.put_object(bucket, fnm,
BytesIO(binary),
len(binary)
)
return r
except Exception as e:
logging.error(f"Fail put {bucket}/{fnm}: "+str(e))
self.__open__()
time.sleep(1)
def get(self, bucket, fnm):
for _ in range(10):
try:
r = self.conn.get_object(bucket, fnm)
return r.read()
except Exception as e:
logging.error(f"Fail get {bucket}/{fnm}: "+str(e))
self.__open__()
time.sleep(1)
return
if __name__ == "__main__":
conn = HuMinio("infiniflow")
fnm = "/opt/home/kevinhu/docgpt/upload/13/11-408.jpg"
from PIL import Image
img = Image.open(fnm)
buff = BytesIO()
img.save(buff, format='JPEG')
print(conn.put("test", "11-408.jpg", buff.getvalue()))
bts = conn.get("test", "11-408.jpg")
img = Image.open(BytesIO(bts))
img.save("test.jpg")