fix(proxy): derive prompt injection heuristics thread count from CPU count with env override

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng 2026-09-17 05:30:41 +00:00
parent 719d7a1983
commit d6f6f64c0f
2 changed files with 18 additions and 1 deletions

View file

@ -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"
)

View file

@ -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")