forked from phoenix/litellm-mirror
feat(proxy_server.py): enable llm api based prompt injection checks
run user calls through an llm api to check for prompt injection attacks. This happens in parallel to th e actual llm call using `async_moderation_hook`
This commit is contained in:
parent
f24d3ffdb6
commit
d91f9a9f50
11 changed files with 271 additions and 24 deletions
|
@ -96,6 +96,9 @@ class _ENTERPRISE_GoogleTextModeration(CustomLogger):
|
|||
async def async_moderation_hook(
|
||||
self,
|
||||
data: dict,
|
||||
call_type: (
|
||||
Literal["completion"] | Literal["embeddings"] | Literal["image_generation"]
|
||||
),
|
||||
):
|
||||
"""
|
||||
- Calls Google's Text Moderation API
|
||||
|
|
|
@ -99,6 +99,9 @@ class _ENTERPRISE_LlamaGuard(CustomLogger):
|
|||
async def async_moderation_hook(
|
||||
self,
|
||||
data: dict,
|
||||
call_type: (
|
||||
Literal["completion"] | Literal["embeddings"] | Literal["image_generation"]
|
||||
),
|
||||
):
|
||||
"""
|
||||
- Calls the Llama Guard Endpoint
|
||||
|
|
|
@ -22,6 +22,7 @@ from litellm.utils import (
|
|||
)
|
||||
from datetime import datetime
|
||||
import aiohttp, asyncio
|
||||
from litellm.utils import get_formatted_prompt
|
||||
|
||||
litellm.set_verbose = True
|
||||
|
||||
|
@ -94,6 +95,9 @@ class _ENTERPRISE_LLMGuard(CustomLogger):
|
|||
async def async_moderation_hook(
|
||||
self,
|
||||
data: dict,
|
||||
call_type: (
|
||||
Literal["completion"] | Literal["embeddings"] | Literal["image_generation"]
|
||||
),
|
||||
):
|
||||
"""
|
||||
- Calls the LLM Guard Endpoint
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue