fix(proxy): run prompt injection heuristics on a dedicated bounded executor

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/constants.py

@@ -602,6 +602,7 @@ 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
 DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE: Final = os.getenv(
     "DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield"
 )

litellm/proxy/hooks/prompt_injection_detection.py

@@ -8,6 +8,7 @@
 
 
 import asyncio
+from concurrent.futures import ThreadPoolExecutor
 from difflib import SequenceMatcher
 from typing import Final, Literal
 
@@ -16,7 +17,10 @@ from fastapi import HTTPException
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.caching import DualCache
-from litellm.constants import DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD
+from litellm.constants import (
+    DEFAULT_PROMPT_INJECTION_SIMILARITY_THRESHOLD,
+    PROMPT_INJECTION_HEURISTICS_MAX_THREADS,
+)
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.prompt_templates.factory import (
     prompt_injection_detection_default_pt,
@@ -25,6 +29,10 @@ from litellm.proxy._types import LiteLLMPromptInjectionParams, UserAPIKeyAuth
 from litellm.router import Router
 from litellm.utils import get_formatted_prompt
 
+HEURISTICS_EXECUTOR: Final = ThreadPoolExecutor(
+    max_workers=PROMPT_INJECTION_HEURISTICS_MAX_THREADS, thread_name_prefix="prompt-injection-heuristics"
+)
+
 
 class _OPTIONAL_PromptInjectionDetection(CustomLogger):
     enforces_request_content: bool = True
@@ -107,6 +115,11 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger):
                         combinations.append(phrase.lower())
         return combinations
 
+    async def check_user_input_similarity_off_loop(self, user_input: str) -> bool:
+        return await asyncio.get_running_loop().run_in_executor(
+            HEURISTICS_EXECUTOR, self.check_user_input_similarity, user_input
+        )
+
     def check_user_input_similarity(
         self,
         user_input: str,
@@ -168,7 +181,7 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger):
             if self.prompt_injection_params is not None:
                 # 1. check if heuristics check turned on
                 if self.prompt_injection_params.heuristics_check is True:
-                    is_prompt_attack = await asyncio.to_thread(self.check_user_input_similarity, formatted_prompt)
+                    is_prompt_attack = await self.check_user_input_similarity_off_loop(formatted_prompt)
                     if is_prompt_attack is True:
                         raise HTTPException(
                             status_code=400,
@@ -178,7 +191,7 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger):
                 if self.prompt_injection_params.vector_db_check is True:
                     pass
             else:
-                is_prompt_attack = await asyncio.to_thread(self.check_user_input_similarity, formatted_prompt)
+                is_prompt_attack = await self.check_user_input_similarity_off_loop(formatted_prompt)
 
             if is_prompt_attack is True:
                 raise HTTPException(

tests/test_litellm/proxy/hooks/test_prompt_injection_detection.py

@@ -1,5 +1,6 @@
 import asyncio
 import time
+from concurrent.futures import ThreadPoolExecutor
 
 import pytest
 from fastapi import HTTPException
@@ -13,6 +14,8 @@ from litellm.proxy.hooks.prompt_injection_detection import (
 from litellm.proxy.utils import ProxyLogging
 from litellm.router import Router
 
+LONG_SAFE_PROMPT = "Summarize the quarterly revenue report for the finance team. " * 3
+
 
 def _moderation_detector(verdict: str) -> _OPTIONAL_PromptInjectionDetection:
     detector = _OPTIONAL_PromptInjectionDetection(
@@ -93,8 +96,7 @@ async def test_heuristics_check_keeps_event_loop_responsive():
     detector = _OPTIONAL_PromptInjectionDetection(
         prompt_injection_params=LiteLLMPromptInjectionParams(heuristics_check=True)
     )
-    long_safe_prompt = "Summarize the quarterly revenue report for the finance team. " * 3
-    data = {"model": "test-model", "messages": [{"role": "user", "content": long_safe_prompt}]}
+    data = {"model": "test-model", "messages": [{"role": "user", "content": LONG_SAFE_PROMPT}]}
     ticks_during_scan: list[float] = []
     scan_done = asyncio.Event()
 
@@ -120,6 +122,36 @@ async def test_heuristics_check_keeps_event_loop_responsive():
     assert len(ticks_before_finish) >= int((finished - started) / 0.05)
 
 
+@pytest.mark.asyncio
+async def test_heuristics_check_does_not_occupy_default_executor():
+    detector = _OPTIONAL_PromptInjectionDetection(
+        prompt_injection_params=LiteLLMPromptInjectionParams(heuristics_check=True)
+    )
+    data = {"model": "test-model", "messages": [{"role": "user", "content": LONG_SAFE_PROMPT}]}
+    loop = asyncio.get_running_loop()
+    single_worker_default_executor = ThreadPoolExecutor(max_workers=1)
+    loop.set_default_executor(single_worker_default_executor)
+
+    scan = asyncio.create_task(
+        detector.async_pre_call_hook(
+            user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
+            cache=DualCache(),
+            data=data,
+            call_type="acompletion",
+        )
+    )
+    await asyncio.sleep(0.05)
+    started = time.perf_counter()
+    await loop.run_in_executor(None, time.sleep, 0)
+    unrelated_work_wait = time.perf_counter() - started
+    result = await scan
+    scan_wall = time.perf_counter() - started
+    single_worker_default_executor.shutdown(wait=False)
+
+    assert result == data
+    assert unrelated_work_wait < scan_wall / 4
+
+
 @pytest.mark.asyncio
 async def test_moderation_hook_rejects_unsafe_llm_verdict():
     detector = _moderation_detector(verdict="UNSAFE")