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)