JobSmithManipulation
Kevin Hu
commited on
Commit
·
74bda08
1
Parent(s):
016dc40
update document sdk (#2445)
Browse files### What problem does this PR solve?
### Type of change
- [x] New Feature (non-breaking change which adds functionality)
---------
Co-authored-by: Kevin Hu <[email protected]>
- api/apps/sdk/doc.py +194 -15
- sdk/python/ragflow/modules/chunk.py +33 -17
- sdk/python/ragflow/modules/document.py +12 -11
- sdk/python/ragflow/ragflow.py +61 -3
- sdk/python/test/t_document.py +48 -12
api/apps/sdk/doc.py
CHANGED
|
@@ -84,15 +84,28 @@ def upload(dataset_id, tenant_id):
|
|
| 84 |
@token_required
|
| 85 |
def docinfos(tenant_id):
|
| 86 |
req = request.args
|
|
|
|
|
|
|
|
|
|
|
|
|
| 87 |
if "id" in req:
|
| 88 |
doc_id = req["id"]
|
| 89 |
-
e, doc = DocumentService.get_by_id(doc_id)
|
| 90 |
-
return get_json_result(data=doc.to_json())
|
| 91 |
if "name" in req:
|
| 92 |
doc_name = req["name"]
|
| 93 |
doc_id = DocumentService.get_doc_id_by_doc_name(doc_name)
|
| 94 |
-
|
| 95 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 96 |
|
| 97 |
|
| 98 |
@manager.route('/save', methods=['POST'])
|
|
@@ -246,7 +259,7 @@ def rename():
|
|
| 246 |
req["doc_id"], {"name": req["name"]}):
|
| 247 |
return get_data_error_result(
|
| 248 |
retmsg="Database error (Document rename)!")
|
| 249 |
-
|
| 250 |
informs = File2DocumentService.get_by_document_id(req["doc_id"])
|
| 251 |
if informs:
|
| 252 |
e, file = FileService.get_by_id(informs[0].file_id)
|
|
@@ -259,7 +272,7 @@ def rename():
|
|
| 259 |
|
| 260 |
@manager.route("/<document_id>", methods=["GET"])
|
| 261 |
@token_required
|
| 262 |
-
def download_document(dataset_id, document_id):
|
| 263 |
try:
|
| 264 |
# Check whether there is this document
|
| 265 |
exist, document = DocumentService.get_by_id(document_id)
|
|
@@ -313,7 +326,21 @@ def list_docs(dataset_id, tenant_id):
|
|
| 313 |
try:
|
| 314 |
docs, tol = DocumentService.get_by_kb_id(
|
| 315 |
kb_id, page_number, items_per_page, orderby, desc, keywords)
|
| 316 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 317 |
except Exception as e:
|
| 318 |
return server_error_response(e)
|
| 319 |
|
|
@@ -436,6 +463,8 @@ def list_chunk(tenant_id):
|
|
| 436 |
query["available_int"] = int(req["available_int"])
|
| 437 |
sres = retrievaler.search(query, search.index_name(tenant_id), highlight=True)
|
| 438 |
res = {"total": sres.total, "chunks": [], "doc": doc.to_dict()}
|
|
|
|
|
|
|
| 439 |
for id in sres.ids:
|
| 440 |
d = {
|
| 441 |
"chunk_id": id,
|
|
@@ -455,7 +484,21 @@ def list_chunk(tenant_id):
|
|
| 455 |
poss.append([float(d["positions"][i]), float(d["positions"][i + 1]), float(d["positions"][i + 2]),
|
| 456 |
float(d["positions"][i + 3]), float(d["positions"][i + 4])])
|
| 457 |
d["positions"] = poss
|
| 458 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 459 |
return get_json_result(data=res)
|
| 460 |
except Exception as e:
|
| 461 |
if str(e).find("not_found") > 0:
|
|
@@ -471,8 +514,9 @@ def create(tenant_id):
|
|
| 471 |
req = request.json
|
| 472 |
md5 = hashlib.md5()
|
| 473 |
md5.update((req["content_with_weight"] + req["doc_id"]).encode("utf-8"))
|
| 474 |
-
|
| 475 |
-
|
|
|
|
| 476 |
"content_with_weight": req["content_with_weight"]}
|
| 477 |
d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
|
| 478 |
d["important_kwd"] = req.get("important_kwd", [])
|
|
@@ -503,20 +547,33 @@ def create(tenant_id):
|
|
| 503 |
|
| 504 |
DocumentService.increment_chunk_num(
|
| 505 |
doc.id, doc.kb_id, c, 1, 0)
|
| 506 |
-
|
| 507 |
-
#
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 508 |
except Exception as e:
|
| 509 |
return server_error_response(e)
|
| 510 |
|
| 511 |
-
|
| 512 |
@manager.route('/chunk/rm', methods=['POST'])
|
| 513 |
@token_required
|
| 514 |
@validate_request("chunk_ids", "doc_id")
|
| 515 |
-
def rm_chunk():
|
| 516 |
req = request.json
|
| 517 |
try:
|
| 518 |
if not ELASTICSEARCH.deleteByQuery(
|
| 519 |
-
Q("ids", values=req["chunk_ids"]), search.index_name(
|
| 520 |
return get_data_error_result(retmsg="Index updating failure")
|
| 521 |
e, doc = DocumentService.get_by_id(req["doc_id"])
|
| 522 |
if not e:
|
|
@@ -526,4 +583,126 @@ def rm_chunk():
|
|
| 526 |
DocumentService.decrement_chunk_num(doc.id, doc.kb_id, 1, chunk_number, 0)
|
| 527 |
return get_json_result(data=True)
|
| 528 |
except Exception as e:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 529 |
return server_error_response(e)
|
|
|
|
| 84 |
@token_required
|
| 85 |
def docinfos(tenant_id):
|
| 86 |
req = request.args
|
| 87 |
+
if "id" not in req and "name" not in req:
|
| 88 |
+
return get_data_error_result(
|
| 89 |
+
retmsg="Id or name should be provided")
|
| 90 |
+
doc_id=None
|
| 91 |
if "id" in req:
|
| 92 |
doc_id = req["id"]
|
|
|
|
|
|
|
| 93 |
if "name" in req:
|
| 94 |
doc_name = req["name"]
|
| 95 |
doc_id = DocumentService.get_doc_id_by_doc_name(doc_name)
|
| 96 |
+
e, doc = DocumentService.get_by_id(doc_id)
|
| 97 |
+
#rename key's name
|
| 98 |
+
key_mapping = {
|
| 99 |
+
"chunk_num": "chunk_count",
|
| 100 |
+
"kb_id": "knowledgebase_id",
|
| 101 |
+
"token_num": "token_count",
|
| 102 |
+
}
|
| 103 |
+
renamed_doc = {}
|
| 104 |
+
for key, value in doc.to_dict().items():
|
| 105 |
+
new_key = key_mapping.get(key, key)
|
| 106 |
+
renamed_doc[new_key] = value
|
| 107 |
+
|
| 108 |
+
return get_json_result(data=renamed_doc)
|
| 109 |
|
| 110 |
|
| 111 |
@manager.route('/save', methods=['POST'])
|
|
|
|
| 259 |
req["doc_id"], {"name": req["name"]}):
|
| 260 |
return get_data_error_result(
|
| 261 |
retmsg="Database error (Document rename)!")
|
| 262 |
+
|
| 263 |
informs = File2DocumentService.get_by_document_id(req["doc_id"])
|
| 264 |
if informs:
|
| 265 |
e, file = FileService.get_by_id(informs[0].file_id)
|
|
|
|
| 272 |
|
| 273 |
@manager.route("/<document_id>", methods=["GET"])
|
| 274 |
@token_required
|
| 275 |
+
def download_document(dataset_id, document_id,tenant_id):
|
| 276 |
try:
|
| 277 |
# Check whether there is this document
|
| 278 |
exist, document = DocumentService.get_by_id(document_id)
|
|
|
|
| 326 |
try:
|
| 327 |
docs, tol = DocumentService.get_by_kb_id(
|
| 328 |
kb_id, page_number, items_per_page, orderby, desc, keywords)
|
| 329 |
+
|
| 330 |
+
# rename key's name
|
| 331 |
+
renamed_doc_list = []
|
| 332 |
+
for doc in docs:
|
| 333 |
+
key_mapping = {
|
| 334 |
+
"chunk_num": "chunk_count",
|
| 335 |
+
"kb_id": "knowledgebase_id",
|
| 336 |
+
"token_num": "token_count",
|
| 337 |
+
}
|
| 338 |
+
renamed_doc = {}
|
| 339 |
+
for key, value in doc.items():
|
| 340 |
+
new_key = key_mapping.get(key, key)
|
| 341 |
+
renamed_doc[new_key] = value
|
| 342 |
+
renamed_doc_list.append(renamed_doc)
|
| 343 |
+
return get_json_result(data={"total": tol, "docs": renamed_doc_list})
|
| 344 |
except Exception as e:
|
| 345 |
return server_error_response(e)
|
| 346 |
|
|
|
|
| 463 |
query["available_int"] = int(req["available_int"])
|
| 464 |
sres = retrievaler.search(query, search.index_name(tenant_id), highlight=True)
|
| 465 |
res = {"total": sres.total, "chunks": [], "doc": doc.to_dict()}
|
| 466 |
+
|
| 467 |
+
origin_chunks=[]
|
| 468 |
for id in sres.ids:
|
| 469 |
d = {
|
| 470 |
"chunk_id": id,
|
|
|
|
| 484 |
poss.append([float(d["positions"][i]), float(d["positions"][i + 1]), float(d["positions"][i + 2]),
|
| 485 |
float(d["positions"][i + 3]), float(d["positions"][i + 4])])
|
| 486 |
d["positions"] = poss
|
| 487 |
+
|
| 488 |
+
origin_chunks.append(d)
|
| 489 |
+
##rename keys
|
| 490 |
+
for chunk in origin_chunks:
|
| 491 |
+
key_mapping = {
|
| 492 |
+
"chunk_id": "id",
|
| 493 |
+
"content_with_weight": "content",
|
| 494 |
+
"doc_id": "document_id",
|
| 495 |
+
"important_kwd": "important_keywords",
|
| 496 |
+
}
|
| 497 |
+
renamed_chunk = {}
|
| 498 |
+
for key, value in chunk.items():
|
| 499 |
+
new_key = key_mapping.get(key, key)
|
| 500 |
+
renamed_chunk[new_key] = value
|
| 501 |
+
res["chunks"].append(renamed_chunk)
|
| 502 |
return get_json_result(data=res)
|
| 503 |
except Exception as e:
|
| 504 |
if str(e).find("not_found") > 0:
|
|
|
|
| 514 |
req = request.json
|
| 515 |
md5 = hashlib.md5()
|
| 516 |
md5.update((req["content_with_weight"] + req["doc_id"]).encode("utf-8"))
|
| 517 |
+
|
| 518 |
+
chunk_id = md5.hexdigest()
|
| 519 |
+
d = {"id": chunk_id, "content_ltks": rag_tokenizer.tokenize(req["content_with_weight"]),
|
| 520 |
"content_with_weight": req["content_with_weight"]}
|
| 521 |
d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
|
| 522 |
d["important_kwd"] = req.get("important_kwd", [])
|
|
|
|
| 547 |
|
| 548 |
DocumentService.increment_chunk_num(
|
| 549 |
doc.id, doc.kb_id, c, 1, 0)
|
| 550 |
+
d["chunk_id"] = chunk_id
|
| 551 |
+
#rename keys
|
| 552 |
+
key_mapping = {
|
| 553 |
+
"chunk_id": "id",
|
| 554 |
+
"content_with_weight": "content",
|
| 555 |
+
"doc_id": "document_id",
|
| 556 |
+
"important_kwd": "important_keywords",
|
| 557 |
+
"kb_id":"knowledge_base_id",
|
| 558 |
+
}
|
| 559 |
+
renamed_chunk = {}
|
| 560 |
+
for key, value in d.items():
|
| 561 |
+
new_key = key_mapping.get(key, key)
|
| 562 |
+
renamed_chunk[new_key] = value
|
| 563 |
+
|
| 564 |
+
return get_json_result(data={"chunk": renamed_chunk})
|
| 565 |
+
# return get_json_result(data={"chunk_id": chunk_id})
|
| 566 |
except Exception as e:
|
| 567 |
return server_error_response(e)
|
| 568 |
|
|
|
|
| 569 |
@manager.route('/chunk/rm', methods=['POST'])
|
| 570 |
@token_required
|
| 571 |
@validate_request("chunk_ids", "doc_id")
|
| 572 |
+
def rm_chunk(tenant_id):
|
| 573 |
req = request.json
|
| 574 |
try:
|
| 575 |
if not ELASTICSEARCH.deleteByQuery(
|
| 576 |
+
Q("ids", values=req["chunk_ids"]), search.index_name(tenant_id)):
|
| 577 |
return get_data_error_result(retmsg="Index updating failure")
|
| 578 |
e, doc = DocumentService.get_by_id(req["doc_id"])
|
| 579 |
if not e:
|
|
|
|
| 583 |
DocumentService.decrement_chunk_num(doc.id, doc.kb_id, 1, chunk_number, 0)
|
| 584 |
return get_json_result(data=True)
|
| 585 |
except Exception as e:
|
| 586 |
+
return server_error_response(e)
|
| 587 |
+
|
| 588 |
+
@manager.route('/chunk/set', methods=['POST'])
|
| 589 |
+
@token_required
|
| 590 |
+
@validate_request("doc_id", "chunk_id", "content_with_weight",
|
| 591 |
+
"important_kwd")
|
| 592 |
+
def set(tenant_id):
|
| 593 |
+
req = request.json
|
| 594 |
+
d = {
|
| 595 |
+
"id": req["chunk_id"],
|
| 596 |
+
"content_with_weight": req["content_with_weight"]}
|
| 597 |
+
d["content_ltks"] = rag_tokenizer.tokenize(req["content_with_weight"])
|
| 598 |
+
d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"])
|
| 599 |
+
d["important_kwd"] = req["important_kwd"]
|
| 600 |
+
d["important_tks"] = rag_tokenizer.tokenize(" ".join(req["important_kwd"]))
|
| 601 |
+
if "available_int" in req:
|
| 602 |
+
d["available_int"] = req["available_int"]
|
| 603 |
+
|
| 604 |
+
try:
|
| 605 |
+
tenant_id = DocumentService.get_tenant_id(req["doc_id"])
|
| 606 |
+
if not tenant_id:
|
| 607 |
+
return get_data_error_result(retmsg="Tenant not found!")
|
| 608 |
+
|
| 609 |
+
embd_id = DocumentService.get_embd_id(req["doc_id"])
|
| 610 |
+
embd_mdl = TenantLLMService.model_instance(
|
| 611 |
+
tenant_id, LLMType.EMBEDDING.value, embd_id)
|
| 612 |
+
|
| 613 |
+
e, doc = DocumentService.get_by_id(req["doc_id"])
|
| 614 |
+
if not e:
|
| 615 |
+
return get_data_error_result(retmsg="Document not found!")
|
| 616 |
+
|
| 617 |
+
if doc.parser_id == ParserType.QA:
|
| 618 |
+
arr = [
|
| 619 |
+
t for t in re.split(
|
| 620 |
+
r"[\n\t]",
|
| 621 |
+
req["content_with_weight"]) if len(t) > 1]
|
| 622 |
+
if len(arr) != 2:
|
| 623 |
+
return get_data_error_result(
|
| 624 |
+
retmsg="Q&A must be separated by TAB/ENTER key.")
|
| 625 |
+
q, a = rmPrefix(arr[0]), rmPrefix(arr[1])
|
| 626 |
+
d = beAdoc(d, arr[0], arr[1], not any(
|
| 627 |
+
[rag_tokenizer.is_chinese(t) for t in q + a]))
|
| 628 |
+
|
| 629 |
+
v, c = embd_mdl.encode([doc.name, req["content_with_weight"]])
|
| 630 |
+
v = 0.1 * v[0] + 0.9 * v[1] if doc.parser_id != ParserType.QA else v[1]
|
| 631 |
+
d["q_%d_vec" % len(v)] = v.tolist()
|
| 632 |
+
ELASTICSEARCH.upsert([d], search.index_name(tenant_id))
|
| 633 |
+
return get_json_result(data=True)
|
| 634 |
+
except Exception as e:
|
| 635 |
+
return server_error_response(e)
|
| 636 |
+
|
| 637 |
+
@manager.route('/retrieval_test', methods=['POST'])
|
| 638 |
+
@token_required
|
| 639 |
+
@validate_request("kb_id", "question")
|
| 640 |
+
def retrieval_test(tenant_id):
|
| 641 |
+
req = request.json
|
| 642 |
+
page = int(req.get("page", 1))
|
| 643 |
+
size = int(req.get("size", 30))
|
| 644 |
+
question = req["question"]
|
| 645 |
+
kb_id = req["kb_id"]
|
| 646 |
+
if isinstance(kb_id, str): kb_id = [kb_id]
|
| 647 |
+
doc_ids = req.get("doc_ids", [])
|
| 648 |
+
similarity_threshold = float(req.get("similarity_threshold", 0.2))
|
| 649 |
+
vector_similarity_weight = float(req.get("vector_similarity_weight", 0.3))
|
| 650 |
+
top = int(req.get("top_k", 1024))
|
| 651 |
+
|
| 652 |
+
try:
|
| 653 |
+
tenants = UserTenantService.query(user_id=tenant_id)
|
| 654 |
+
for kid in kb_id:
|
| 655 |
+
for tenant in tenants:
|
| 656 |
+
if KnowledgebaseService.query(
|
| 657 |
+
tenant_id=tenant.tenant_id, id=kid):
|
| 658 |
+
break
|
| 659 |
+
else:
|
| 660 |
+
return get_json_result(
|
| 661 |
+
data=False, retmsg=f'Only owner of knowledgebase authorized for this operation.',
|
| 662 |
+
retcode=RetCode.OPERATING_ERROR)
|
| 663 |
+
|
| 664 |
+
e, kb = KnowledgebaseService.get_by_id(kb_id[0])
|
| 665 |
+
if not e:
|
| 666 |
+
return get_data_error_result(retmsg="Knowledgebase not found!")
|
| 667 |
+
|
| 668 |
+
embd_mdl = TenantLLMService.model_instance(
|
| 669 |
+
kb.tenant_id, LLMType.EMBEDDING.value, llm_name=kb.embd_id)
|
| 670 |
+
|
| 671 |
+
rerank_mdl = None
|
| 672 |
+
if req.get("rerank_id"):
|
| 673 |
+
rerank_mdl = TenantLLMService.model_instance(
|
| 674 |
+
kb.tenant_id, LLMType.RERANK.value, llm_name=req["rerank_id"])
|
| 675 |
+
|
| 676 |
+
if req.get("keyword", False):
|
| 677 |
+
chat_mdl = TenantLLMService.model_instance(kb.tenant_id, LLMType.CHAT)
|
| 678 |
+
question += keyword_extraction(chat_mdl, question)
|
| 679 |
+
|
| 680 |
+
retr = retrievaler if kb.parser_id != ParserType.KG else kg_retrievaler
|
| 681 |
+
ranks = retr.retrieval(question, embd_mdl, kb.tenant_id, kb_id, page, size,
|
| 682 |
+
similarity_threshold, vector_similarity_weight, top,
|
| 683 |
+
doc_ids, rerank_mdl=rerank_mdl, highlight=req.get("highlight"))
|
| 684 |
+
for c in ranks["chunks"]:
|
| 685 |
+
if "vector" in c:
|
| 686 |
+
del c["vector"]
|
| 687 |
+
|
| 688 |
+
##rename keys
|
| 689 |
+
renamed_chunks=[]
|
| 690 |
+
for chunk in ranks["chunks"]:
|
| 691 |
+
key_mapping = {
|
| 692 |
+
"chunk_id": "id",
|
| 693 |
+
"content_with_weight": "content",
|
| 694 |
+
"doc_id": "document_id",
|
| 695 |
+
"important_kwd": "important_keywords",
|
| 696 |
+
}
|
| 697 |
+
rename_chunk={}
|
| 698 |
+
for key, value in chunk.items():
|
| 699 |
+
new_key = key_mapping.get(key, key)
|
| 700 |
+
rename_chunk[new_key] = value
|
| 701 |
+
renamed_chunks.append(rename_chunk)
|
| 702 |
+
ranks["chunks"] = renamed_chunks
|
| 703 |
+
return get_json_result(data=ranks)
|
| 704 |
+
except Exception as e:
|
| 705 |
+
if str(e).find("not_found") > 0:
|
| 706 |
+
return get_json_result(data=False, retmsg=f'No chunk found! Check the chunk status please!',
|
| 707 |
+
retcode=RetCode.DATA_ERROR)
|
| 708 |
return server_error_response(e)
|
sdk/python/ragflow/modules/chunk.py
CHANGED
|
@@ -3,32 +3,48 @@ from .base import Base
|
|
| 3 |
|
| 4 |
class Chunk(Base):
|
| 5 |
def __init__(self, rag, res_dict):
|
| 6 |
-
# 初始化类的属性
|
| 7 |
self.id = ""
|
| 8 |
-
self.
|
| 9 |
-
self.
|
| 10 |
-
self.content_sm_ltks = []
|
| 11 |
-
self.important_kwd = []
|
| 12 |
-
self.important_tks = []
|
| 13 |
self.create_time = ""
|
| 14 |
self.create_timestamp_flt = 0.0
|
| 15 |
-
self.
|
| 16 |
-
self.
|
| 17 |
-
self.
|
| 18 |
-
self.q_vec = []
|
| 19 |
self.status = "1"
|
| 20 |
-
for k
|
| 21 |
-
if
|
| 22 |
-
|
| 23 |
-
|
| 24 |
super().__init__(rag, res_dict)
|
|
|
|
| 25 |
def delete(self) -> bool:
|
| 26 |
"""
|
| 27 |
Delete the chunk in the document.
|
| 28 |
"""
|
| 29 |
-
res = self.
|
| 30 |
-
|
| 31 |
res = res.json()
|
| 32 |
if res.get("retmsg") == "success":
|
| 33 |
return True
|
| 34 |
-
raise Exception(res["retmsg"])
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3 |
|
| 4 |
class Chunk(Base):
|
| 5 |
def __init__(self, rag, res_dict):
|
|
|
|
| 6 |
self.id = ""
|
| 7 |
+
self.content = ""
|
| 8 |
+
self.important_keywords = []
|
|
|
|
|
|
|
|
|
|
| 9 |
self.create_time = ""
|
| 10 |
self.create_timestamp_flt = 0.0
|
| 11 |
+
self.knowledgebase_id = None
|
| 12 |
+
self.document_name = ""
|
| 13 |
+
self.document_id = ""
|
|
|
|
| 14 |
self.status = "1"
|
| 15 |
+
for k in list(res_dict.keys()):
|
| 16 |
+
if k not in self.__dict__:
|
| 17 |
+
res_dict.pop(k)
|
|
|
|
| 18 |
super().__init__(rag, res_dict)
|
| 19 |
+
|
| 20 |
def delete(self) -> bool:
|
| 21 |
"""
|
| 22 |
Delete the chunk in the document.
|
| 23 |
"""
|
| 24 |
+
res = self.post('/doc/chunk/rm',
|
| 25 |
+
{"doc_id": self.document_id, 'chunk_ids': [self.id]})
|
| 26 |
res = res.json()
|
| 27 |
if res.get("retmsg") == "success":
|
| 28 |
return True
|
| 29 |
+
raise Exception(res["retmsg"])
|
| 30 |
+
|
| 31 |
+
def save(self) -> bool:
|
| 32 |
+
"""
|
| 33 |
+
Save the document details to the server.
|
| 34 |
+
"""
|
| 35 |
+
res = self.post('/doc/chunk/set',
|
| 36 |
+
{"chunk_id": self.id,
|
| 37 |
+
"kb_id": self.knowledgebase_id,
|
| 38 |
+
"name": self.document_name,
|
| 39 |
+
"content_with_weight": self.content,
|
| 40 |
+
"important_kwd": self.important_keywords,
|
| 41 |
+
"create_time": self.create_time,
|
| 42 |
+
"create_timestamp_flt": self.create_timestamp_flt,
|
| 43 |
+
"doc_id": self.document_id,
|
| 44 |
+
"status": self.status,
|
| 45 |
+
})
|
| 46 |
+
res = res.json()
|
| 47 |
+
if res.get("retmsg") == "success":
|
| 48 |
+
return True
|
| 49 |
+
raise Exception(res["retmsg"])
|
| 50 |
+
|
sdk/python/ragflow/modules/document.py
CHANGED
|
@@ -9,15 +9,15 @@ class Document(Base):
|
|
| 9 |
self.id = ""
|
| 10 |
self.name = ""
|
| 11 |
self.thumbnail = None
|
| 12 |
-
self.
|
| 13 |
self.parser_method = ""
|
| 14 |
self.parser_config = {"pages": [[1, 1000000]]}
|
| 15 |
self.source_type = "local"
|
| 16 |
self.type = ""
|
| 17 |
self.created_by = ""
|
| 18 |
self.size = 0
|
| 19 |
-
self.
|
| 20 |
-
self.
|
| 21 |
self.progress = 0.0
|
| 22 |
self.progress_msg = ""
|
| 23 |
self.process_begin_at = None
|
|
@@ -34,10 +34,10 @@ class Document(Base):
|
|
| 34 |
Save the document details to the server.
|
| 35 |
"""
|
| 36 |
res = self.post('/doc/save',
|
| 37 |
-
{"id": self.id, "name": self.name, "thumbnail": self.thumbnail, "kb_id": self.
|
| 38 |
"parser_id": self.parser_method, "parser_config": self.parser_config.to_json(),
|
| 39 |
"source_type": self.source_type, "type": self.type, "created_by": self.created_by,
|
| 40 |
-
"size": self.size, "token_num": self.
|
| 41 |
"progress": self.progress, "progress_msg": self.progress_msg,
|
| 42 |
"process_begin_at": self.process_begin_at, "process_duation": self.process_duration
|
| 43 |
})
|
|
@@ -177,8 +177,10 @@ class Document(Base):
|
|
| 177 |
if res.status_code == 200:
|
| 178 |
res_data = res.json()
|
| 179 |
if res_data.get("retmsg") == "success":
|
| 180 |
-
chunks
|
| 181 |
-
|
|
|
|
|
|
|
| 182 |
return chunks
|
| 183 |
else:
|
| 184 |
raise Exception(f"Error fetching chunks: {res_data.get('retmsg')}")
|
|
@@ -187,10 +189,9 @@ class Document(Base):
|
|
| 187 |
|
| 188 |
def add_chunk(self, content: str):
|
| 189 |
res = self.post('/doc/chunk/create', {"doc_id": self.id, "content_with_weight":content})
|
| 190 |
-
|
| 191 |
-
# 假设返回的 response 包含 chunk 的信息
|
| 192 |
if res.status_code == 200:
|
| 193 |
-
|
| 194 |
-
|
|
|
|
| 195 |
else:
|
| 196 |
raise Exception(f"Failed to add chunk: {res.status_code} {res.text}")
|
|
|
|
| 9 |
self.id = ""
|
| 10 |
self.name = ""
|
| 11 |
self.thumbnail = None
|
| 12 |
+
self.knowledgebase_id = None
|
| 13 |
self.parser_method = ""
|
| 14 |
self.parser_config = {"pages": [[1, 1000000]]}
|
| 15 |
self.source_type = "local"
|
| 16 |
self.type = ""
|
| 17 |
self.created_by = ""
|
| 18 |
self.size = 0
|
| 19 |
+
self.token_count = 0
|
| 20 |
+
self.chunk_count = 0
|
| 21 |
self.progress = 0.0
|
| 22 |
self.progress_msg = ""
|
| 23 |
self.process_begin_at = None
|
|
|
|
| 34 |
Save the document details to the server.
|
| 35 |
"""
|
| 36 |
res = self.post('/doc/save',
|
| 37 |
+
{"id": self.id, "name": self.name, "thumbnail": self.thumbnail, "kb_id": self.knowledgebase_id,
|
| 38 |
"parser_id": self.parser_method, "parser_config": self.parser_config.to_json(),
|
| 39 |
"source_type": self.source_type, "type": self.type, "created_by": self.created_by,
|
| 40 |
+
"size": self.size, "token_num": self.token_count, "chunk_num": self.chunk_count,
|
| 41 |
"progress": self.progress, "progress_msg": self.progress_msg,
|
| 42 |
"process_begin_at": self.process_begin_at, "process_duation": self.process_duration
|
| 43 |
})
|
|
|
|
| 177 |
if res.status_code == 200:
|
| 178 |
res_data = res.json()
|
| 179 |
if res_data.get("retmsg") == "success":
|
| 180 |
+
chunks=[]
|
| 181 |
+
for chunk_data in res_data["data"].get("chunks", []):
|
| 182 |
+
chunk=Chunk(self.rag,chunk_data)
|
| 183 |
+
chunks.append(chunk)
|
| 184 |
return chunks
|
| 185 |
else:
|
| 186 |
raise Exception(f"Error fetching chunks: {res_data.get('retmsg')}")
|
|
|
|
| 189 |
|
| 190 |
def add_chunk(self, content: str):
|
| 191 |
res = self.post('/doc/chunk/create', {"doc_id": self.id, "content_with_weight":content})
|
|
|
|
|
|
|
| 192 |
if res.status_code == 200:
|
| 193 |
+
res_data = res.json().get("data")
|
| 194 |
+
chunk_data = res_data.get("chunk")
|
| 195 |
+
return Chunk(self.rag,chunk_data)
|
| 196 |
else:
|
| 197 |
raise Exception(f"Failed to add chunk: {res.status_code} {res.text}")
|
sdk/python/ragflow/ragflow.py
CHANGED
|
@@ -20,6 +20,8 @@ import requests
|
|
| 20 |
from .modules.assistant import Assistant
|
| 21 |
from .modules.dataset import DataSet
|
| 22 |
from .modules.document import Document
|
|
|
|
|
|
|
| 23 |
|
| 24 |
class RAGFlow:
|
| 25 |
def __init__(self, user_key, base_url, version='v1'):
|
|
@@ -143,7 +145,7 @@ class RAGFlow:
|
|
| 143 |
return result_list
|
| 144 |
raise Exception(res["retmsg"])
|
| 145 |
|
| 146 |
-
def create_document(self, ds:DataSet, name: str, blob: bytes) -> bool:
|
| 147 |
url = f"/doc/dataset/{ds.id}/documents/upload"
|
| 148 |
files = {
|
| 149 |
'file': (name, blob)
|
|
@@ -164,6 +166,7 @@ class RAGFlow:
|
|
| 164 |
raise Exception(f"Upload failed: {response.json().get('retmsg')}")
|
| 165 |
|
| 166 |
return False
|
|
|
|
| 167 |
def get_document(self, id: str = None, name: str = None) -> Document:
|
| 168 |
res = self.get("/doc/infos", {"id": id, "name": name})
|
| 169 |
res = res.json()
|
|
@@ -204,8 +207,6 @@ class RAGFlow:
|
|
| 204 |
if not doc_ids or not isinstance(doc_ids, list):
|
| 205 |
raise ValueError("doc_ids must be a non-empty list of document IDs")
|
| 206 |
data = {"doc_ids": doc_ids, "run": 2}
|
| 207 |
-
|
| 208 |
-
|
| 209 |
res = self.post(f'/doc/run', data)
|
| 210 |
|
| 211 |
if res.status_code != 200:
|
|
@@ -217,4 +218,61 @@ class RAGFlow:
|
|
| 217 |
print(f"Error occurred during canceling parsing for documents: {str(e)}")
|
| 218 |
raise
|
| 219 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 220 |
|
|
|
|
| 20 |
from .modules.assistant import Assistant
|
| 21 |
from .modules.dataset import DataSet
|
| 22 |
from .modules.document import Document
|
| 23 |
+
from .modules.chunk import Chunk
|
| 24 |
+
|
| 25 |
|
| 26 |
class RAGFlow:
|
| 27 |
def __init__(self, user_key, base_url, version='v1'):
|
|
|
|
| 145 |
return result_list
|
| 146 |
raise Exception(res["retmsg"])
|
| 147 |
|
| 148 |
+
def create_document(self, ds: DataSet, name: str, blob: bytes) -> bool:
|
| 149 |
url = f"/doc/dataset/{ds.id}/documents/upload"
|
| 150 |
files = {
|
| 151 |
'file': (name, blob)
|
|
|
|
| 166 |
raise Exception(f"Upload failed: {response.json().get('retmsg')}")
|
| 167 |
|
| 168 |
return False
|
| 169 |
+
|
| 170 |
def get_document(self, id: str = None, name: str = None) -> Document:
|
| 171 |
res = self.get("/doc/infos", {"id": id, "name": name})
|
| 172 |
res = res.json()
|
|
|
|
| 207 |
if not doc_ids or not isinstance(doc_ids, list):
|
| 208 |
raise ValueError("doc_ids must be a non-empty list of document IDs")
|
| 209 |
data = {"doc_ids": doc_ids, "run": 2}
|
|
|
|
|
|
|
| 210 |
res = self.post(f'/doc/run', data)
|
| 211 |
|
| 212 |
if res.status_code != 200:
|
|
|
|
| 218 |
print(f"Error occurred during canceling parsing for documents: {str(e)}")
|
| 219 |
raise
|
| 220 |
|
| 221 |
+
def retrieval(self,
|
| 222 |
+
question,
|
| 223 |
+
datasets=None,
|
| 224 |
+
documents=None,
|
| 225 |
+
offset=0,
|
| 226 |
+
limit=6,
|
| 227 |
+
similarity_threshold=0.1,
|
| 228 |
+
vector_similarity_weight=0.3,
|
| 229 |
+
top_k=1024):
|
| 230 |
+
"""
|
| 231 |
+
Perform document retrieval based on the given parameters.
|
| 232 |
+
|
| 233 |
+
:param question: The query question.
|
| 234 |
+
:param datasets: A list of datasets (optional, as documents may be provided directly).
|
| 235 |
+
:param documents: A list of documents (if specific documents are provided).
|
| 236 |
+
:param offset: Offset for the retrieval results.
|
| 237 |
+
:param limit: Maximum number of retrieval results.
|
| 238 |
+
:param similarity_threshold: Similarity threshold.
|
| 239 |
+
:param vector_similarity_weight: Weight of vector similarity.
|
| 240 |
+
:param top_k: Number of top most similar documents to consider (for pre-filtering or ranking).
|
| 241 |
+
|
| 242 |
+
Note: This is a hypothetical implementation and may need adjustments based on the actual backend service API.
|
| 243 |
+
"""
|
| 244 |
+
try:
|
| 245 |
+
data = {
|
| 246 |
+
"question": question,
|
| 247 |
+
"datasets": datasets if datasets is not None else [],
|
| 248 |
+
"documents": [doc.id if hasattr(doc, 'id') else doc for doc in
|
| 249 |
+
documents] if documents is not None else [],
|
| 250 |
+
"offset": offset,
|
| 251 |
+
"limit": limit,
|
| 252 |
+
"similarity_threshold": similarity_threshold,
|
| 253 |
+
"vector_similarity_weight": vector_similarity_weight,
|
| 254 |
+
"top_k": top_k,
|
| 255 |
+
"kb_id": datasets,
|
| 256 |
+
}
|
| 257 |
+
|
| 258 |
+
# Send a POST request to the backend service (using requests library as an example, actual implementation may vary)
|
| 259 |
+
res = self.post(f'/doc/retrieval_test', data)
|
| 260 |
+
|
| 261 |
+
# Check the response status code
|
| 262 |
+
if res.status_code == 200:
|
| 263 |
+
res_data = res.json()
|
| 264 |
+
if res_data.get("retmsg") == "success":
|
| 265 |
+
chunks = []
|
| 266 |
+
for chunk_data in res_data["data"].get("chunks", []):
|
| 267 |
+
chunk = Chunk(self, chunk_data)
|
| 268 |
+
chunks.append(chunk)
|
| 269 |
+
return chunks
|
| 270 |
+
else:
|
| 271 |
+
raise Exception(f"Error fetching chunks: {res_data.get('retmsg')}")
|
| 272 |
+
else:
|
| 273 |
+
raise Exception(f"API request failed with status code {res.status_code}")
|
| 274 |
+
|
| 275 |
+
except Exception as e:
|
| 276 |
+
print(f"An error occurred during retrieval: {e}")
|
| 277 |
+
raise
|
| 278 |
|
sdk/python/test/t_document.py
CHANGED
|
@@ -41,6 +41,7 @@ class TestDocument(TestSdk):
|
|
| 41 |
def test_update_document_with_success(self):
|
| 42 |
"""
|
| 43 |
Test updating a document with success.
|
|
|
|
| 44 |
"""
|
| 45 |
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
| 46 |
doc = rag.get_document(name="TestDocument.txt")
|
|
@@ -60,7 +61,7 @@ class TestDocument(TestSdk):
|
|
| 60 |
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
| 61 |
|
| 62 |
# Retrieve a document
|
| 63 |
-
doc = rag.get_document(name="
|
| 64 |
|
| 65 |
# Check if the retrieved document is of type Document
|
| 66 |
if isinstance(doc, Document):
|
|
@@ -147,14 +148,16 @@ class TestDocument(TestSdk):
|
|
| 147 |
ds = rag.create_dataset(name="God4")
|
| 148 |
|
| 149 |
# Define the document name and path
|
| 150 |
-
name3 = '
|
| 151 |
-
path = 'test_data/
|
|
|
|
| 152 |
|
| 153 |
# Create a document in the dataset using the file path
|
| 154 |
rag.create_document(ds, name=name3, blob=open(path, "rb").read())
|
| 155 |
|
| 156 |
# Retrieve the document by name
|
| 157 |
-
doc = rag.get_document(name="
|
|
|
|
| 158 |
|
| 159 |
# Initiate asynchronous parsing
|
| 160 |
doc.async_parse()
|
|
@@ -185,9 +188,9 @@ class TestDocument(TestSdk):
|
|
| 185 |
|
| 186 |
# Prepare a list of file names and paths
|
| 187 |
documents = [
|
| 188 |
-
{'name': '
|
| 189 |
-
{'name': '
|
| 190 |
-
{'name': '
|
| 191 |
]
|
| 192 |
|
| 193 |
# Create documents in bulk
|
|
@@ -248,6 +251,7 @@ class TestDocument(TestSdk):
|
|
| 248 |
print(c)
|
| 249 |
assert c is not None, "Chunk is None"
|
| 250 |
assert "rag" in c['content_with_weight'].lower(), f"Keyword 'rag' not found in chunk content: {c.content}"
|
|
|
|
| 251 |
def test_add_chunk_to_chunk_list(self):
|
| 252 |
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
| 253 |
doc = rag.get_document(name='story.txt')
|
|
@@ -258,12 +262,44 @@ class TestDocument(TestSdk):
|
|
| 258 |
def test_delete_chunk_of_chunk_list(self):
|
| 259 |
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
| 260 |
doc = rag.get_document(name='story.txt')
|
| 261 |
-
|
| 262 |
chunk = doc.add_chunk(content="assss")
|
| 263 |
assert chunk is not None, "Chunk is None"
|
| 264 |
assert isinstance(chunk, Chunk), "Chunk was not added to chunk list"
|
| 265 |
-
|
|
|
|
| 266 |
chunk.delete()
|
| 267 |
-
|
| 268 |
-
|
| 269 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 41 |
def test_update_document_with_success(self):
|
| 42 |
"""
|
| 43 |
Test updating a document with success.
|
| 44 |
+
Update name or parser_method are supported
|
| 45 |
"""
|
| 46 |
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
| 47 |
doc = rag.get_document(name="TestDocument.txt")
|
|
|
|
| 61 |
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
| 62 |
|
| 63 |
# Retrieve a document
|
| 64 |
+
doc = rag.get_document(name="manual.txt")
|
| 65 |
|
| 66 |
# Check if the retrieved document is of type Document
|
| 67 |
if isinstance(doc, Document):
|
|
|
|
| 148 |
ds = rag.create_dataset(name="God4")
|
| 149 |
|
| 150 |
# Define the document name and path
|
| 151 |
+
name3 = 'westworld.pdf'
|
| 152 |
+
path = 'test_data/westworld.pdf'
|
| 153 |
+
|
| 154 |
|
| 155 |
# Create a document in the dataset using the file path
|
| 156 |
rag.create_document(ds, name=name3, blob=open(path, "rb").read())
|
| 157 |
|
| 158 |
# Retrieve the document by name
|
| 159 |
+
doc = rag.get_document(name="westworld.pdf")
|
| 160 |
+
|
| 161 |
|
| 162 |
# Initiate asynchronous parsing
|
| 163 |
doc.async_parse()
|
|
|
|
| 188 |
|
| 189 |
# Prepare a list of file names and paths
|
| 190 |
documents = [
|
| 191 |
+
{'name': 'test1.txt', 'path': 'test_data/test1.txt'},
|
| 192 |
+
{'name': 'test2.txt', 'path': 'test_data/test2.txt'},
|
| 193 |
+
{'name': 'test3.txt', 'path': 'test_data/test3.txt'}
|
| 194 |
]
|
| 195 |
|
| 196 |
# Create documents in bulk
|
|
|
|
| 251 |
print(c)
|
| 252 |
assert c is not None, "Chunk is None"
|
| 253 |
assert "rag" in c['content_with_weight'].lower(), f"Keyword 'rag' not found in chunk content: {c.content}"
|
| 254 |
+
|
| 255 |
def test_add_chunk_to_chunk_list(self):
|
| 256 |
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
| 257 |
doc = rag.get_document(name='story.txt')
|
|
|
|
| 262 |
def test_delete_chunk_of_chunk_list(self):
|
| 263 |
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
| 264 |
doc = rag.get_document(name='story.txt')
|
|
|
|
| 265 |
chunk = doc.add_chunk(content="assss")
|
| 266 |
assert chunk is not None, "Chunk is None"
|
| 267 |
assert isinstance(chunk, Chunk), "Chunk was not added to chunk list"
|
| 268 |
+
doc = rag.get_document(name='story.txt')
|
| 269 |
+
chunk_count_before=doc.chunk_count
|
| 270 |
chunk.delete()
|
| 271 |
+
doc = rag.get_document(name='story.txt')
|
| 272 |
+
assert doc.chunk_count == chunk_count_before-1, "Chunk was not deleted"
|
| 273 |
+
|
| 274 |
+
def test_update_chunk_content(self):
|
| 275 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
| 276 |
+
doc = rag.get_document(name='story.txt')
|
| 277 |
+
chunk = doc.add_chunk(content="assssd")
|
| 278 |
+
assert chunk is not None, "Chunk is None"
|
| 279 |
+
assert isinstance(chunk, Chunk), "Chunk was not added to chunk list"
|
| 280 |
+
chunk.content = "ragflow123"
|
| 281 |
+
res=chunk.save()
|
| 282 |
+
assert res is True, f"Failed to update chunk, error: {res}"
|
| 283 |
+
|
| 284 |
+
def test_retrieval_chunks(self):
|
| 285 |
+
rag = RAGFlow(API_KEY, HOST_ADDRESS)
|
| 286 |
+
ds = rag.create_dataset(name="God8")
|
| 287 |
+
name = 'ragflow_test.txt'
|
| 288 |
+
path = 'test_data/ragflow_test.txt'
|
| 289 |
+
rag.create_document(ds, name=name, blob=open(path, "rb").read())
|
| 290 |
+
doc = rag.get_document(name=name)
|
| 291 |
+
doc.async_parse()
|
| 292 |
+
# Wait for parsing to complete and get progress updates using join
|
| 293 |
+
for progress, msg in doc.join(interval=5, timeout=30):
|
| 294 |
+
print(progress, msg)
|
| 295 |
+
assert 0 <= progress <= 100, f"Invalid progress: {progress}"
|
| 296 |
+
assert msg, "Message should not be empty"
|
| 297 |
+
for c in rag.retrieval(question="What's ragflow?",
|
| 298 |
+
datasets=[ds.id], documents=[doc],
|
| 299 |
+
offset=0, limit=6, similarity_threshold=0.1,
|
| 300 |
+
vector_similarity_weight=0.3,
|
| 301 |
+
top_k=1024
|
| 302 |
+
):
|
| 303 |
+
print(c)
|
| 304 |
+
assert c is not None, "Chunk is None"
|
| 305 |
+
assert "ragflow" in c.content.lower(), f"Keyword 'rag' not found in chunk content: {c.content}"
|