Merge pull request #41541 from BerriAI/litellm_prompt_injection_async_llm_check

fix(proxy): run prompt injection heuristics off the event loop

litellm/constants.py

@@ -605,6 +605,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 = max(1, get_env_int("PROMPT_INJECTION_HEURISTICS_MAX_THREADS", 1))
 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

@@ -7,6 +7,8 @@
 ## Reject a call if it contains a prompt injection attack.
 
 
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
 from difflib import SequenceMatcher
 from typing import Final, Literal
 
@@ -15,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,
@@ -24,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
@@ -106,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,
@@ -167,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 = self.check_user_input_similarity(user_input=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,
@@ -177,7 +191,7 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger):
                 if self.prompt_injection_params.vector_db_check is True:
                     pass
             else:
-                is_prompt_attack = self.check_user_input_similarity(user_input=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,12 +1,21 @@
+import asyncio
+import importlib
+import time
+from collections.abc import AsyncIterator
+from concurrent.futures import ThreadPoolExecutor
+
 import pytest
 from fastapi import HTTPException
 
+import litellm
 from litellm.caching.caching import DualCache
-from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy._types import LiteLLMPromptInjectionParams, UserAPIKeyAuth
 from litellm.proxy.hooks.prompt_injection_detection import (
     _OPTIONAL_PromptInjectionDetection,
 )
 
+LONG_SAFE_PROMPT = "Summarize the quarterly revenue report for the finance team. " * 3
+
 
 @pytest.mark.asyncio
 async def test_acompletion_call_type_rejects_prompt_injection():
@@ -57,3 +66,76 @@ async def test_acompletion_call_type_allows_safe_prompt():
     )
 
     assert result == data
+
+
+@pytest.mark.asyncio
+async def test_heuristics_check_keeps_event_loop_responsive():
+    detector = _OPTIONAL_PromptInjectionDetection(
+        prompt_injection_params=LiteLLMPromptInjectionParams(heuristics_check=True)
+    )
+    data = {"model": "test-model", "messages": [{"role": "user", "content": LONG_SAFE_PROMPT}]}
+
+    async def ticks_until_done(task: asyncio.Task[dict]) -> AsyncIterator[float]:
+        while not task.done():
+            await asyncio.sleep(0.01)
+            yield time.perf_counter()
+
+    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",
+        )
+    )
+    started = time.perf_counter()
+    ticks_during_scan = tuple([tick async for tick in ticks_until_done(scan)])
+    finished = time.perf_counter()
+    result = await scan
+
+    assert result == data
+    assert len(ticks_during_scan) >= 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.parametrize(
+    ("configured", "expected"),
+    [("3", 3), ("not-an-int", 1), ("0", 1), ("-2", 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)
+