Skip to content

Commit 00f08a9

Browse files
asiroliuasiroliu
committed
Fix: Refactor parser config handling and add GraphRAG defaults
- Update `get_parser_config` to merge provided configs with defaults - Add GraphRAG configuration defaults for all chunk methods - Make raptor and graphrag fields non-nullable in ParserConfig schema - Update related test cases to reflect config changes - Ensure backward compatibility while adding new GraphRAG support - #8396
1 parent 131fc10 commit 00f08a9

File tree

7 files changed

+194
-30
lines changed

7 files changed

+194
-30
lines changed

‎api/utils/api_utils.py

Lines changed: 39 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -351,28 +351,47 @@ def generate_confirmation_token(tenant_id):
351351

352352

353353
def get_parser_config(chunk_method, parser_config):
354-
if parser_config:
355-
return parser_config
356354
if not chunk_method:
357355
chunk_method = "naive"
356+
357+
# Define default configurations for each chunk method
358358
key_mapping = {
359-
"naive": {"chunk_token_num": 512, "delimiter": r"\n", "html4excel": False, "layout_recognize": "DeepDOC", "raptor": {"use_raptor": False}},
360-
"qa": {"raptor": {"use_raptor": False}},
359+
"naive": {"chunk_token_num": 512, "delimiter": r"\n", "html4excel": False, "layout_recognize": "DeepDOC", "raptor": {"use_raptor": False}, "graphrag": {"use_graphrag": False}},
360+
"qa": {"raptor": {"use_raptor": False}, "graphrag": {"use_graphrag": False}},
361361
"tag": None,
362362
"resume": None,
363-
"manual": {"raptor": {"use_raptor": False}},
363+
"manual": {"raptor": {"use_raptor": False}, "graphrag": {"use_graphrag": False}},
364364
"table": None,
365-
"paper": {"raptor": {"use_raptor": False}},
366-
"book": {"raptor": {"use_raptor": False}},
367-
"laws": {"raptor": {"use_raptor": False}},
368-
"presentation": {"raptor": {"use_raptor": False}},
365+
"paper": {"raptor": {"use_raptor": False}, "graphrag": {"use_graphrag": False}},
366+
"book": {"raptor": {"use_raptor": False}, "graphrag": {"use_graphrag": False}},
367+
"laws": {"raptor": {"use_raptor": False}, "graphrag": {"use_graphrag": False}},
368+
"presentation": {"raptor": {"use_raptor": False}, "graphrag": {"use_graphrag": False}},
369369
"one": None,
370-
"knowledge_graph": {"chunk_token_num": 8192, "delimiter": r"\n", "entity_types": ["organization", "person", "location", "event", "time"]},
370+
"knowledge_graph": {
371+
"chunk_token_num": 8192,
372+
"delimiter": r"\n",
373+
"entity_types": ["organization", "person", "location", "event", "time"],
374+
"raptor": {"use_raptor": False},
375+
"graphrag": {"use_graphrag": False},
376+
},
371377
"email": None,
372378
"picture": None,
373379
}
374-
parser_config = key_mapping[chunk_method]
375-
return parser_config
380+
381+
default_config = key_mapping[chunk_method]
382+
383+
# If no parser_config provided, return default
384+
if not parser_config:
385+
return default_config
386+
387+
# If parser_config is provided, merge with defaults to ensure required fields exist
388+
if default_config is None:
389+
return parser_config
390+
391+
# Ensure raptor and graphrag fields have default values if not provided
392+
merged_config = deep_merge(default_config, parser_config)
393+
394+
return merged_config
376395

377396

378397
def get_data_openai(
@@ -602,17 +621,14 @@ def get_mcp_tools(mcp_servers: list, timeout: float | int = 10) -> tuple[dict, s
602621

603622
TimeoutException = Union[Type[BaseException], BaseException]
604623
OnTimeoutCallback = Union[Callable[..., Any], Coroutine[Any, Any, Any]]
605-
def timeout(
606-
seconds: float |int = None,
607-
attempts: int = 2,
608-
*,
609-
exception: Optional[TimeoutException] = None,
610-
on_timeout: Optional[OnTimeoutCallback] = None
611-
):
624+
625+
626+
def timeout(seconds: float | int = None, attempts: int = 2, *, exception: Optional[TimeoutException] = None, on_timeout: Optional[OnTimeoutCallback] = None):
612627
def decorator(func):
613628
@wraps(func)
614629
def wrapper(*args, **kwargs):
615630
result_queue = queue.Queue(maxsize=1)
631+
616632
def target():
617633
try:
618634
result = func(*args, **kwargs)
@@ -644,7 +660,7 @@ async def async_wrapper(*args, **kwargs) -> Any:
644660
with trio.fail_after(seconds):
645661
return await func(*args, **kwargs)
646662
except trio.TooSlowError:
647-
if a < attempts -1:
663+
if a < attempts - 1:
648664
continue
649665
if on_timeout is not None:
650666
if callable(on_timeout):
@@ -668,11 +684,11 @@ async def async_wrapper(*args, **kwargs) -> Any:
668684
if asyncio.iscoroutinefunction(func):
669685
return async_wrapper
670686
return wrapper
687+
671688
return decorator
672689

673690

674691
async def is_strong_enough(chat_model, embedding_model):
675-
676692
@timeout(30, 2)
677693
async def _is_strong_enough():
678694
nonlocal chat_model, embedding_model
@@ -681,11 +697,11 @@ async def _is_strong_enough():
681697
_ = await trio.to_thread.run_sync(lambda: embedding_model.encode(["Are you strong enough!?"]))
682698
if chat_model:
683699
with trio.fail_after(30):
684-
res = await trio.to_thread.run_sync(lambda: chat_model.chat("Nothing special.", [{"role":"user", "content": "Are you strong enough!?"}], {}))
700+
res = await trio.to_thread.run_sync(lambda: chat_model.chat("Nothing special.", [{"role": "user", "content": "Are you strong enough!?"}], {}))
685701
if res.find("**ERROR**") >= 0:
686702
raise Exception(res)
687703

688704
# Pressure test for GraphRAG task
689705
async with trio.open_nursery() as nursery:
690706
for _ in range(32):
691-
nursery.start_soon(_is_strong_enough)
707+
nursery.start_soon(_is_strong_enough)

‎api/utils/validation_utils.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -365,10 +365,10 @@ class ParserConfig(Base):
365365
auto_questions: int = Field(default=0, ge=0, le=10)
366366
chunk_token_num: int = Field(default=512, ge=1, le=2048)
367367
delimiter: str = Field(default=r"\n", min_length=1)
368-
graphrag: GraphragConfig | None = None
368+
graphrag: GraphragConfig = Field(default_factory=lambda: GraphragConfig(use_graphrag=False))
369369
html4excel: bool = False
370370
layout_recognize: str = "DeepDOC"
371-
raptor: RaptorConfig | None = None
371+
raptor: RaptorConfig = Field(default_factory=lambda: RaptorConfig(use_raptor=False))
372372
tag_kb_ids: list[str] = Field(default_factory=list)
373373
topn_tags: int = Field(default=1, ge=1, le=10)
374374
filename_embd_weight: float | None = Field(default=0.1, ge=0.0, le=1.0)

‎test/testcases/test_http_api/test_dataset_mangement/test_create_dataset.py

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -644,6 +644,7 @@ def test_parser_config_empty(self, HttpApiAuth):
644644
"html4excel": False,
645645
"layout_recognize": "DeepDOC",
646646
"raptor": {"use_raptor": False},
647+
"graphrag": {"use_graphrag": False},
647648
}, res
648649

649650
@pytest.mark.p2
@@ -657,6 +658,7 @@ def test_parser_config_unset(self, HttpApiAuth):
657658
"html4excel": False,
658659
"layout_recognize": "DeepDOC",
659660
"raptor": {"use_raptor": False},
661+
"graphrag": {"use_graphrag": False},
660662
}, res
661663

662664
@pytest.mark.p3
@@ -670,6 +672,7 @@ def test_parser_config_none(self, HttpApiAuth):
670672
"html4excel": False,
671673
"layout_recognize": "DeepDOC",
672674
"raptor": {"use_raptor": False},
675+
"graphrag": {"use_graphrag": False},
673676
}, res
674677

675678
@pytest.mark.p2
@@ -695,3 +698,64 @@ def test_unsupported_field(self, HttpApiAuth, payload):
695698
res = create_dataset(HttpApiAuth, payload)
696699
assert res["code"] == 101, res
697700
assert "Extra inputs are not permitted" in res["message"], res
701+
702+
703+
@pytest.mark.usefixtures("clear_datasets")
704+
class TestParserConfigBugFix:
705+
@pytest.mark.p1
706+
def test_parser_config_missing_raptor_and_graphrag(self, HttpApiAuth):
707+
payload = {"name": "test_parser_config_missing_fields", "parser_config": {"chunk_token_num": 1024}}
708+
res = create_dataset(HttpApiAuth, payload)
709+
assert res["code"] == 0, res
710+
711+
parser_config = res["data"]["parser_config"]
712+
assert "raptor" in parser_config, "raptor field should be present"
713+
assert "graphrag" in parser_config, "graphrag field should be present"
714+
assert parser_config["raptor"]["use_raptor"] is False, "raptor.use_raptor should default to False"
715+
assert parser_config["graphrag"]["use_graphrag"] is False, "graphrag.use_graphrag should default to False"
716+
assert parser_config["chunk_token_num"] == 1024, "User-provided chunk_token_num should be preserved"
717+
718+
@pytest.mark.p1
719+
def test_parser_config_with_only_raptor(self, HttpApiAuth):
720+
payload = {"name": "test_parser_config_only_raptor", "parser_config": {"chunk_token_num": 1024, "raptor": {"use_raptor": True}}}
721+
res = create_dataset(HttpApiAuth, payload)
722+
assert res["code"] == 0, res
723+
724+
parser_config = res["data"]["parser_config"]
725+
assert parser_config["raptor"]["use_raptor"] is True, "User-provided raptor.use_raptor should be preserved"
726+
assert "graphrag" in parser_config, "graphrag field should be present"
727+
assert parser_config["graphrag"]["use_graphrag"] is False, "graphrag.use_graphrag should default to False"
728+
729+
@pytest.mark.p1
730+
def test_parser_config_with_only_graphrag(self, HttpApiAuth):
731+
payload = {"name": "test_parser_config_only_graphrag", "parser_config": {"chunk_token_num": 1024, "graphrag": {"use_graphrag": True}}}
732+
res = create_dataset(HttpApiAuth, payload)
733+
assert res["code"] == 0, res
734+
735+
parser_config = res["data"]["parser_config"]
736+
assert "raptor" in parser_config, "raptor field should be present"
737+
assert parser_config["raptor"]["use_raptor"] is False, "raptor.use_raptor should default to False"
738+
assert parser_config["graphrag"]["use_graphrag"] is True, "User-provided graphrag.use_graphrag should be preserved"
739+
740+
@pytest.mark.p1
741+
def test_parser_config_with_both_fields(self, HttpApiAuth):
742+
payload = {"name": "test_parser_config_both_fields", "parser_config": {"chunk_token_num": 1024, "raptor": {"use_raptor": True}, "graphrag": {"use_graphrag": True}}}
743+
res = create_dataset(HttpApiAuth, payload)
744+
assert res["code"] == 0, res
745+
746+
parser_config = res["data"]["parser_config"]
747+
assert parser_config["raptor"]["use_raptor"] is True, "User-provided raptor.use_raptor should be preserved"
748+
assert parser_config["graphrag"]["use_graphrag"] is True, "User-provided graphrag.use_graphrag should be preserved"
749+
750+
@pytest.mark.p2
751+
@pytest.mark.parametrize("chunk_method", ["qa", "manual", "paper", "book", "laws", "presentation"])
752+
def test_parser_config_different_chunk_methods(self, HttpApiAuth, chunk_method):
753+
payload = {"name": f"test_parser_config_{chunk_method}", "chunk_method": chunk_method, "parser_config": {"chunk_token_num": 512}}
754+
res = create_dataset(HttpApiAuth, payload)
755+
assert res["code"] == 0, res
756+
757+
parser_config = res["data"]["parser_config"]
758+
assert "raptor" in parser_config, f"raptor field should be present for {chunk_method}"
759+
assert "graphrag" in parser_config, f"graphrag field should be present for {chunk_method}"
760+
assert parser_config["raptor"]["use_raptor"] is False, f"raptor.use_raptor should default to False for {chunk_method}"
761+
assert parser_config["graphrag"]["use_graphrag"] is False, f"graphrag.use_graphrag should default to False for {chunk_method}"

‎test/testcases/test_http_api/test_dataset_mangement/test_update_dataset.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -755,6 +755,7 @@ def test_parser_config_empty(self, HttpApiAuth, add_dataset_func):
755755
"html4excel": False,
756756
"layout_recognize": "DeepDOC",
757757
"raptor": {"use_raptor": False},
758+
"graphrag": {"use_graphrag": False},
758759
}, res
759760

760761
@pytest.mark.p3
@@ -772,6 +773,7 @@ def test_parser_config_none(self, HttpApiAuth, add_dataset_func):
772773
"html4excel": False,
773774
"layout_recognize": "DeepDOC",
774775
"raptor": {"use_raptor": False},
776+
"graphrag": {"use_graphrag": False},
775777
}, res
776778

777779
@pytest.mark.p3
@@ -783,7 +785,7 @@ def test_parser_config_empty_with_chunk_method_change(self, HttpApiAuth, add_dat
783785

784786
res = list_datasets(HttpApiAuth)
785787
assert res["code"] == 0, res
786-
assert res["data"][0]["parser_config"] == {"raptor": {"use_raptor": False}}, res
788+
assert res["data"][0]["parser_config"] == {"raptor": {"use_raptor": False}, "graphrag": {"use_graphrag": False}}, res
787789

788790
@pytest.mark.p3
789791
def test_parser_config_unset_with_chunk_method_change(self, HttpApiAuth, add_dataset_func):
@@ -794,7 +796,7 @@ def test_parser_config_unset_with_chunk_method_change(self, HttpApiAuth, add_dat
794796

795797
res = list_datasets(HttpApiAuth)
796798
assert res["code"] == 0, res
797-
assert res["data"][0]["parser_config"] == {"raptor": {"use_raptor": False}}, res
799+
assert res["data"][0]["parser_config"] == {"raptor": {"use_raptor": False}, "graphrag": {"use_graphrag": False}}, res
798800

799801
@pytest.mark.p3
800802
def test_parser_config_none_with_chunk_method_change(self, HttpApiAuth, add_dataset_func):
@@ -805,7 +807,7 @@ def test_parser_config_none_with_chunk_method_change(self, HttpApiAuth, add_data
805807

806808
res = list_datasets(HttpApiAuth, {"id": dataset_id})
807809
assert res["code"] == 0, res
808-
assert res["data"][0]["parser_config"] == {"raptor": {"use_raptor": False}}, res
810+
assert res["data"][0]["parser_config"] == {"raptor": {"use_raptor": False}, "graphrag": {"use_graphrag": False}}, res
809811

810812
@pytest.mark.p2
811813
@pytest.mark.parametrize(

‎test/testcases/test_http_api/test_file_management_within_dataset/test_update_document.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -540,6 +540,7 @@ def test_parser_config(
540540
"html4excel": False,
541541
"layout_recognize": "DeepDOC",
542542
"raptor": {"use_raptor": False},
543+
"graphrag": {"use_graphrag": False},
543544
}
544545
else:
545546
for k, v in parser_config.items():

‎test/testcases/test_sdk_api/test_dataset_mangement/test_create_dataset.py

Lines changed: 80 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -588,11 +588,26 @@ def test_parser_config_empty(self, client):
588588
excepted_value = DataSet.ParserConfig(
589589
client,
590590
{
591+
"auto_keywords": 0,
592+
"auto_questions": 0,
591593
"chunk_token_num": 512,
592-
"delimiter": r"\n",
594+
"delimiter": "\\n",
595+
"filename_embd_weight": 0.1,
596+
"graphrag": {"community": False, "entity_types": ["organization", "person", "geo", "event", "category"], "method": "light", "resolution": False, "use_graphrag": False},
593597
"html4excel": False,
594598
"layout_recognize": "DeepDOC",
595-
"raptor": {"use_raptor": False},
599+
"pages": None,
600+
"raptor": {
601+
"max_cluster": 64,
602+
"max_token": 256,
603+
"prompt": "Please summarize the following paragraphs. Be careful with the numbers, do not make things up. Paragraphs as following:\n {cluster_content}\nThe above is the content you need to summarize.",
604+
"random_seed": 0,
605+
"threshold": 0.1,
606+
"use_raptor": False,
607+
},
608+
"tag_kb_ids": [],
609+
"task_page_size": None,
610+
"topn_tags": 1,
596611
},
597612
)
598613
parser_config_o = DataSet.ParserConfig(client, {})
@@ -610,6 +625,7 @@ def test_parser_config_unset(self, client):
610625
"html4excel": False,
611626
"layout_recognize": "DeepDOC",
612627
"raptor": {"use_raptor": False},
628+
"graphrag": {"use_graphrag": False},
613629
},
614630
)
615631
payload = {"name": "parser_config_unset"}
@@ -626,6 +642,7 @@ def test_parser_config_none(self, client):
626642
"html4excel": False,
627643
"layout_recognize": "DeepDOC",
628644
"raptor": {"use_raptor": False},
645+
"graphrag": {"use_graphrag": False},
629646
},
630647
)
631648
payload = {"name": "parser_config_empty", "parser_config": None}
@@ -655,3 +672,64 @@ def test_unsupported_field(self, client, payload):
655672
with pytest.raises(Exception) as excinfo:
656673
client.create_dataset(**payload)
657674
assert "got an unexpected keyword argument" in str(excinfo.value), str(excinfo.value)
675+
676+
677+
@pytest.mark.usefixtures("clear_datasets")
678+
class TestParserConfigBugFix:
679+
@pytest.mark.p1
680+
def test_parser_config_missing_raptor_and_graphrag(self, client):
681+
parser_config = DataSet.ParserConfig(client, {"chunk_token_num": 1024})
682+
payload = {"name": "test_parser_config_missing_fields_sdk", "parser_config": parser_config}
683+
dataset = client.create_dataset(**payload)
684+
685+
config = dataset.parser_config
686+
assert hasattr(config, "raptor"), "raptor field should be present"
687+
assert hasattr(config, "graphrag"), "graphrag field should be present"
688+
assert config.raptor.use_raptor is False, "raptor.use_raptor should default to False"
689+
assert config.graphrag.use_graphrag is False, "graphrag.use_graphrag should default to False"
690+
assert config.chunk_token_num == 1024, "User-provided chunk_token_num should be preserved"
691+
692+
@pytest.mark.p1
693+
def test_parser_config_with_only_raptor(self, client):
694+
parser_config = DataSet.ParserConfig(client, {"chunk_token_num": 1024, "raptor": {"use_raptor": True}})
695+
payload = {"name": "test_parser_config_only_raptor_sdk", "parser_config": parser_config}
696+
dataset = client.create_dataset(**payload)
697+
698+
config = dataset.parser_config
699+
assert config.raptor.use_raptor is True, "User-provided raptor.use_raptor should be preserved"
700+
assert hasattr(config, "graphrag"), "graphrag field should be present"
701+
assert config.graphrag.use_graphrag is False, "graphrag.use_graphrag should default to False"
702+
703+
@pytest.mark.p1
704+
def test_parser_config_with_only_graphrag(self, client):
705+
parser_config = DataSet.ParserConfig(client, {"chunk_token_num": 1024, "graphrag": {"use_graphrag": True}})
706+
payload = {"name": "test_parser_config_only_graphrag_sdk", "parser_config": parser_config}
707+
dataset = client.create_dataset(**payload)
708+
709+
config = dataset.parser_config
710+
assert hasattr(config, "raptor"), "raptor field should be present"
711+
assert config.raptor.use_raptor is False, "raptor.use_raptor should default to False"
712+
assert config.graphrag.use_graphrag is True, "User-provided graphrag.use_graphrag should be preserved"
713+
714+
@pytest.mark.p1
715+
def test_parser_config_with_both_fields(self, client):
716+
parser_config = DataSet.ParserConfig(client, {"chunk_token_num": 1024, "raptor": {"use_raptor": True}, "graphrag": {"use_graphrag": True}})
717+
payload = {"name": "test_parser_config_both_fields_sdk", "parser_config": parser_config}
718+
dataset = client.create_dataset(**payload)
719+
720+
config = dataset.parser_config
721+
assert config.raptor.use_raptor is True, "User-provided raptor.use_raptor should be preserved"
722+
assert config.graphrag.use_graphrag is True, "User-provided graphrag.use_graphrag should be preserved"
723+
724+
@pytest.mark.p2
725+
@pytest.mark.parametrize("chunk_method", ["qa", "manual", "paper", "book", "laws", "presentation"])
726+
def test_parser_config_different_chunk_methods(self, client, chunk_method):
727+
parser_config = DataSet.ParserConfig(client, {"chunk_token_num": 512})
728+
payload = {"name": f"test_parser_config_{chunk_method}_sdk", "chunk_method": chunk_method, "parser_config": parser_config}
729+
dataset = client.create_dataset(**payload)
730+
731+
config = dataset.parser_config
732+
assert hasattr(config, "raptor"), f"raptor field should be present for {chunk_method}"
733+
assert hasattr(config, "graphrag"), f"graphrag field should be present for {chunk_method}"
734+
assert config.raptor.use_raptor is False, f"raptor.use_raptor should default to False for {chunk_method}"
735+
assert config.graphrag.use_graphrag is False, f"graphrag.use_graphrag should default to False for {chunk_method}"

0 commit comments

Comments
 (0)