phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 11:14:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(prompt_injection_detection.py): fix type check

Browse source
This commit is contained in:
Krrish Dholakia 2024-03-21 08:56:13 -07:00
parent 42d62cf99b
commit 0521e8a1d9
3 changed files with 56 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

12
litellm/proxy/hooks/prompt_injection_detection.py
View file

@ -82,7 +82,7 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger):
"PromptInjectionDetection: Model List not set. Required for Prompt Injection detection."
)
verbose_proxy_logger.debug(
self.print_verbose(
f"model_names: {self.llm_router.model_names}; self.prompt_injection_params.llm_api_name: {self.prompt_injection_params.llm_api_name}"
)
if (
@ -201,7 +201,7 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger):
data: dict,
call_type: Literal["completion", "embeddings", "image_generation"],
):
verbose_proxy_logger.debug(
self.print_verbose(
f"IN ASYNC MODERATION HOOK - self.prompt_injection_params = {self.prompt_injection_params}"
)
@ -235,10 +235,12 @@ class _OPTIONAL_PromptInjectionDetection(CustomLogger):
],
)
verbose_proxy_logger.debug(f"Received LLM Moderation response: {response}")
self.print_verbose(f"Received LLM Moderation response: {response}")
self.print_verbose(
f"llm_api_fail_call_string: {self.prompt_injection_params.llm_api_fail_call_string}"
)
if isinstance(response, litellm.ModelResponse) and isinstance(
response.choices, litellm.Choices
response.choices[0], litellm.Choices
):
if self.prompt_injection_params.llm_api_fail_call_string in response.choices[0].message.content: # type: ignore
is_prompt_attack = True