From 45bf294117ec7a52ff52d2f22ccf5c4bb7f0a5f6 Mon Sep 17 00:00:00 2001 From: Stephen Hu Date: Mon, 4 Aug 2025 13:54:18 +0800 Subject: [PATCH] Refactor: support config strong test (#9198) ### What problem does this PR solve? https://github.com/infiniflow/ragflow/issues/9189#issuecomment-3148920950 ### Type of change - [x] Refactoring Co-authored-by: Kevin Hu --- api/settings.py | 1 + api/utils/api_utils.py | 8 +++++++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/api/settings.py b/api/settings.py index 32fc7bd69..5b3783b45 100644 --- a/api/settings.py +++ b/api/settings.py @@ -70,6 +70,7 @@ REGISTER_ENABLED = 1 # sandbox-executor-manager SANDBOX_ENABLED = 0 SANDBOX_HOST = None +STRONG_TEST_COUNT = int(os.environ.get("STRONG_TEST_COUNT", "32")) BUILTIN_EMBEDDING_MODELS = ["BAAI/bge-large-zh-v1.5@BAAI", "maidalun1020/bce-embedding-base_v1@Youdao"] diff --git a/api/utils/api_utils.py b/api/utils/api_utils.py index 48be1ecae..27ab6201d 100644 --- a/api/utils/api_utils.py +++ b/api/utils/api_utils.py @@ -687,6 +687,12 @@ def timeout(seconds: float | int = None, attempts: int = 2, *, exception: Option async def is_strong_enough(chat_model, embedding_model): + count = settings.STRONG_TEST_COUNT + if not chat_model or not embedding_model: + return + if isinstance(count, int) and count <= 0: + return + @timeout(60, 2) async def _is_strong_enough(): nonlocal chat_model, embedding_model @@ -701,5 +707,5 @@ async def is_strong_enough(chat_model, embedding_model): # Pressure test for GraphRAG task async with trio.open_nursery() as nursery: - for _ in range(32): + for _ in range(count): nursery.start_soon(_is_strong_enough)