diff --git a/litellm/constants.py b/litellm/constants.py index 12749a0fce0..02413ee97ab 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -603,7 +603,9 @@ LOGGING_EXECUTOR_MAX_THREADS: Final = get_env_int("LOGGING_EXECUTOR_MAX_THREADS" LOGGING_EXECUTOR_MAX_PENDING_TASKS: Final = get_env_int("LOGGING_EXECUTOR_MAX_PENDING_TASKS", 10_000) LOGGING_EXECUTOR_DROPPED_TASK_LOG_INTERVAL_SECONDS: Final = 30.0 AWS_SIGNING_MAX_THREADS: Final = 16 -PROMPT_INJECTION_HEURISTICS_MAX_THREADS: Final = 4 +PROMPT_INJECTION_HEURISTICS_MAX_THREADS: Final = get_env_int( + "PROMPT_INJECTION_HEURISTICS_MAX_THREADS", os.cpu_count() or 1 +) DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv( "DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield" ) diff --git a/tests/test_litellm/proxy/hooks/test_prompt_injection_detection.py b/tests/test_litellm/proxy/hooks/test_prompt_injection_detection.py index f6016971357..b189ee740fe 100644 --- a/tests/test_litellm/proxy/hooks/test_prompt_injection_detection.py +++ b/tests/test_litellm/proxy/hooks/test_prompt_injection_detection.py @@ -1,4 +1,6 @@ import asyncio +import importlib +import os import time from concurrent.futures import ThreadPoolExecutor @@ -152,6 +154,19 @@ async def test_heuristics_check_does_not_occupy_default_executor(): assert unrelated_work_wait < scan_wall / 4 +@pytest.mark.parametrize( + ("configured", "expected"), + [("3", 3), ("not-an-int", os.cpu_count() or 1)], +) +def test_heuristics_thread_count_config_is_honoured(monkeypatch: pytest.MonkeyPatch, configured: str, expected: int): + monkeypatch.setenv("PROMPT_INJECTION_HEURISTICS_MAX_THREADS", configured) + try: + assert importlib.reload(litellm.constants).PROMPT_INJECTION_HEURISTICS_MAX_THREADS == expected + finally: + monkeypatch.delenv("PROMPT_INJECTION_HEURISTICS_MAX_THREADS") + importlib.reload(litellm.constants) + + @pytest.mark.asyncio async def test_moderation_hook_rejects_unsafe_llm_verdict(): detector = _moderation_detector(verdict="UNSAFE")