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:
Krrish Dholakia 2024-03-20 22:43:42 -07:00
parent f24d3ffdb6
commit d91f9a9f50
11 changed files with 271 additions and 24 deletions

View file

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

View file

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

View file

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