mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
feat: integrate Repelloai Argus guardrail (#30673)
* feat(guardrails): add RepelloAI Argus guardrail integration (#1) * feat(guardrails): add RepelloAI Argus guardrail integration Add a new guardrail hook backed by RepelloAI Argus, with dashboard-managed asset policies enforced via an asset_id and X-API-Key auth. * fix(guardrails): harden RepelloAI Argus guardrail - scan streaming responses on output (was bypassing the guardrail) - log blocked verdicts as guardrail_intervened instead of success - treat auth/config errors (401/403/404/422) as misconfiguration that always blocks, not a fail-open-able unreachable error - default unreachable_fallback to fail_closed and read it directly; block on unknown/malformed verdicts so an API change can't silently disable enforcement - type unreachable_fallback as a Literal, drop the duplicate config model, expose unreachable_fallback in the config schema, and stop leaking the raw provider response / exception strings to the client * fix(guardrails): address RepelloAI Argus review feedback - support ARGUS_API_KEY (with REPELLOAI_API_KEY fallback) - make asset_id required in the config model - normalize unreachable_fallback so only fail_open opens; block on 400 misconfig - correct the shared unreachable_fallback field description * docs(guardrails): add RepelloAI Argus docs page and dashboard listing - add docs page covering config, env vars, modes, verdicts, failure semantics - list RepelloAI Argus in the Guardrail Garden with provider/logo mappings - add a regression test for the provider logo and display-name resolution * fix(guardrails): keep RepelloAI asset_id optional in config model A required asset_id leaked onto the shared LitellmParams (which inherits RepelloAIGuardrailConfigModel), breaking validation for every other guardrail. Keep it optional like sibling models; the guardrail __init__ still raises when asset_id is missing, which is the real enforcement. * Add comment for last user turn scanning * feat(guardrails): harden repelloai scanning * feat(guardrails): expand repelloai scanning to include tool definitions Add extraction of tool definitions and tool call arguments to the RepelloAI guardrail scanning. Improves detection coverage by including function schemas and parameters in the prompt sent to the guardrail service. Also captures detailed error responses in logs and adds guardrail header to streaming responses. * refactor(guardrails): fix and harden repelloai schema text extraction - Fix duplicate text in _iter_schema_text: previously all dict values were re-queued onto the stack even after scalar/list keys were already extracted explicitly, causing names/descriptions to appear twice in the scanned prompt - Extract schema key frozensets to module-level constants so they are not reconstructed on every call - Change _iter_schema_text from @classmethod to @staticmethod (cls unused) - Narrow _call_analyze stage param from str to Literal["prompt", "response"] - Add HttpxResponse type annotation to _raise_for_config_error - Add LLMResponseTypes annotation to async_post_call_success_hook response param * fix(guardrails): resolve pyright type errors in repelloai guardrail - Narrow async_handler.post return from Response|None to Response with explicit None guard before calling raise_for_status/json - Fix list comprehension returning str|None by switching to explicit loop with isinstance guard so pyright tracks the narrowing - Cast model_dump() result to Dict since hasattr does not narrow object type in pyright * fix(guardrails/repello): include Responses API instructions field in prompt scan The /v1/responses top-level `instructions` field was not included in _extract_prompt_text, allowing a caller to bypass guardrail policy checks by putting blocked content in `instructions` while keeping `input` benign. * feat: add api_key to config model and read prompt from data dict * fix(guardrails/repello): plug input_text and tool-call response bypass gaps Responses API input content parts with type 'input_text' were silently dropped by build_inspection_messages (which only handles type='text'), allowing callers to send blocked content via that path without triggering the pre-call scan. Fix: add _extract_input_text_parts to RepelloAIGuardrail and call it when walking the Responses API input messages. Post-call scanning skipped responses whose choices contained only tool_calls or function_call (message.content=None), letting models put blocked output in function arguments undetected. Fix: _extract_chat_completion_text now calls _extract_tool_call_args_from_message on each choice message. Also replace typing.Dict/List with builtin dict/list to clear TID251 strict ruff violations introduced by this file. * fix(guardrails/repello): scan Responses API function_call output arguments Output items with type 'function_call' in a /v1/responses response were skipped by _extract_responses_api_text; only 'message' items were walked. A model could return blocked content in function_call.arguments undetected. Now extract arguments from function_call output items before scanning. * refactor(guardrails/repello): clean up typing and remove lint-any workarounds - Replace Optional[X]/Union[X,Y] with X|None/X|Y union syntax throughout - Use dict[str, object] instead of bare dict in all signatures - Remove **kwargs from __init__; declare guardrail_name, event_hook, default_on explicitly - Replace getattr(litellm_params, ...) with direct attribute access now that LitellmParams inherits RepelloAIGuardrailConfigModel - Add _event_hook_from_mode() to convert str|list[str]|Mode to typed GuardrailEventHooks - Use TypeAdapter.validate_json() instead of response.json() + manual dict construction - Add _is_object_dict/_is_object_list TypeGuard helpers to narrow object types without Any - Remove cast() workarounds and typed intermediate variables that existed only for the now-removed lint-any CI check - Drop _AddLiteLLMCallback Protocol; budget has sufficient slack for the one reportUnknownMemberType - Fix GuardrailConfigModel missing type arg: GuardrailConfigModel[BaseModel] * fix(guardrails/repello): suppress LIT007 on TypeGuard helpers and add streaming scan-skip warning - Add guard-ok suppressions to _is_object_dict and _is_object_list to satisfy the LIT007 hard-zero budget gate - Emit verbose_proxy_logger.warning when the streaming hook finds no inspectable text after assembly, matching observability of pre/post hooks * refactor: modifications for lint check
This commit is contained in:
Lavish Bansal
GitHub
parent
50f34e0b15
commit
cd4bd92c0a
10 changed files with 1913 additions and 1 deletions
|
|
@ -0,0 +1,51 @@
|
|||
from typing import TYPE_CHECKING, Union
|
||||
|
||||
from litellm.types.guardrails import (
|
||||
GuardrailEventHooks,
|
||||
Mode,
|
||||
SupportedGuardrailIntegrations,
|
||||
)
|
||||
|
||||
from .repelloai import RepelloAIGuardrail
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
|
||||
def _event_hook_from_mode(
|
||||
mode: str | list[str] | Mode,
|
||||
) -> Union[GuardrailEventHooks, list[GuardrailEventHooks], Mode]:
|
||||
if isinstance(mode, Mode):
|
||||
return mode
|
||||
if isinstance(mode, list):
|
||||
return [GuardrailEventHooks(item) for item in mode]
|
||||
return GuardrailEventHooks(mode)
|
||||
|
||||
|
||||
def initialize_guardrail(
|
||||
litellm_params: "LitellmParams", guardrail: "Guardrail"
|
||||
) -> RepelloAIGuardrail:
|
||||
import litellm
|
||||
|
||||
_repelloai_callback = RepelloAIGuardrail(
|
||||
guardrail_name=guardrail["guardrail_name"],
|
||||
api_key=litellm_params.api_key,
|
||||
api_base=litellm_params.api_base,
|
||||
asset_id=litellm_params.asset_id,
|
||||
unreachable_fallback=litellm_params.unreachable_fallback,
|
||||
event_hook=_event_hook_from_mode(litellm_params.mode),
|
||||
default_on=litellm_params.default_on or False,
|
||||
)
|
||||
litellm.logging_callback_manager.add_litellm_callback(_repelloai_callback)
|
||||
|
||||
return _repelloai_callback
|
||||
|
||||
|
||||
guardrail_initializer_registry = {
|
||||
SupportedGuardrailIntegrations.REPELLOAI.value: initialize_guardrail,
|
||||
}
|
||||
|
||||
|
||||
guardrail_class_registry = {
|
||||
SupportedGuardrailIntegrations.REPELLOAI.value: RepelloAIGuardrail,
|
||||
}
|
||||
613
litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py
Normal file
613
litellm/proxy/guardrails/guardrail_hooks/repelloai/repelloai.py
Normal file
|
|
@ -0,0 +1,613 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from typing import AsyncGenerator, Literal
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import TypeGuard
|
||||
|
||||
from fastapi import HTTPException
|
||||
from httpx import HTTPError, Response as HttpxResponse
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client, # pyright: ignore[reportUnknownVariableType]
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType]
|
||||
)
|
||||
from litellm.proxy.guardrails._content_utils import build_inspection_messages
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.guardrails import GuardrailEventHooks, Mode
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.repelloai import (
|
||||
RepelloAIAnalyzeResponse,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
CallTypesLiteral,
|
||||
GuardrailStatus,
|
||||
LLMResponseTypes,
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
)
|
||||
|
||||
DEFAULT_REPELLOAI_API_BASE = "https://argusapi.repello.ai/sdk/v1"
|
||||
DEFAULT_REPELLOAI_TIMEOUT = 30.0
|
||||
BLOCKED_VERDICT = "blocked"
|
||||
FLAGGED_VERDICT = "flagged"
|
||||
PASSED_VERDICT = "passed"
|
||||
|
||||
# Argus returns these for a permanently broken guardrail (bad key, unknown
|
||||
# asset_id, malformed payload), not a transient outage. They must always
|
||||
# block, never honour fail_open.
|
||||
CONFIG_ERROR_STATUS_CODES = frozenset({400, 401, 403, 404, 422})
|
||||
_SCHEMA_SCALAR_KEYS = frozenset(("name", "description", "title", "const", "default"))
|
||||
_SCHEMA_LIST_KEYS = frozenset(("enum", "examples"))
|
||||
_SCHEMA_EXTRACTED_KEYS = _SCHEMA_SCALAR_KEYS | _SCHEMA_LIST_KEYS
|
||||
|
||||
|
||||
class RepelloAIGuardrailMissingSecrets(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def _is_object_dict(value: object) -> TypeGuard[dict[str, object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip
|
||||
return isinstance(value, dict)
|
||||
|
||||
|
||||
def _is_object_list(value: object) -> TypeGuard[list[object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip
|
||||
return isinstance(value, list)
|
||||
|
||||
|
||||
class RepelloAIGuardrail(CustomGuardrail):
|
||||
@staticmethod
|
||||
def _get_field(obj: object, key: str) -> object:
|
||||
if _is_object_dict(obj):
|
||||
return obj.get(key)
|
||||
return getattr(obj, key, None)
|
||||
|
||||
@classmethod
|
||||
def _extract_tool_call_args_from_message(cls, message: object) -> list[str]:
|
||||
args: list[str] = []
|
||||
|
||||
tool_calls = cls._get_field(message, "tool_calls")
|
||||
if _is_object_list(tool_calls):
|
||||
for tool_call in tool_calls:
|
||||
function = cls._get_field(tool_call, "function")
|
||||
arguments = cls._get_field(function, "arguments")
|
||||
if isinstance(arguments, str) and arguments.strip():
|
||||
args.append(arguments)
|
||||
|
||||
function_call = cls._get_field(message, "function_call")
|
||||
arguments = cls._get_field(function_call, "arguments")
|
||||
if isinstance(arguments, str) and arguments.strip():
|
||||
args.append(arguments)
|
||||
|
||||
return args
|
||||
|
||||
@staticmethod
|
||||
def _iter_schema_text(node: object) -> list[str]:
|
||||
texts: list[str] = []
|
||||
stack: list[object] = [node]
|
||||
|
||||
while stack:
|
||||
current = stack.pop()
|
||||
if _is_object_dict(current):
|
||||
for key in _SCHEMA_SCALAR_KEYS:
|
||||
value = current.get(key)
|
||||
if isinstance(value, str) and value:
|
||||
texts.append(value)
|
||||
for key in _SCHEMA_LIST_KEYS:
|
||||
items = current.get(key)
|
||||
if _is_object_list(items):
|
||||
for item in items:
|
||||
if isinstance(item, str) and item:
|
||||
texts.append(item)
|
||||
remaining: list[object] = [
|
||||
v for k, v in current.items() if k not in _SCHEMA_EXTRACTED_KEYS
|
||||
]
|
||||
stack.extend(reversed(remaining))
|
||||
elif _is_object_list(current):
|
||||
stack.extend(reversed(current))
|
||||
|
||||
return texts
|
||||
|
||||
@classmethod
|
||||
def _extract_tool_definition_text(cls, data: dict[str, object]) -> list[str]:
|
||||
texts: list[str] = []
|
||||
|
||||
tools = data.get("tools")
|
||||
for tool in tools if _is_object_list(tools) else []:
|
||||
if not _is_object_dict(tool):
|
||||
continue
|
||||
function = tool.get("function")
|
||||
if _is_object_dict(function):
|
||||
texts.extend(cls._iter_schema_text(function))
|
||||
|
||||
functions = data.get("functions")
|
||||
for function in functions if _is_object_list(functions) else []:
|
||||
if _is_object_dict(function):
|
||||
texts.extend(cls._iter_schema_text(function))
|
||||
|
||||
return texts
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
asset_id: str | None = None,
|
||||
unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed",
|
||||
guardrail_name: str | None = None,
|
||||
event_hook: (
|
||||
GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None
|
||||
) = None,
|
||||
default_on: bool = False,
|
||||
):
|
||||
self.repelloai_api_key = (
|
||||
api_key
|
||||
or get_secret_str("ARGUS_API_KEY")
|
||||
or get_secret_str("REPELLOAI_API_KEY")
|
||||
or ""
|
||||
)
|
||||
if not self.repelloai_api_key:
|
||||
raise RepelloAIGuardrailMissingSecrets(
|
||||
"Couldn't get Repello API key. Set `ARGUS_API_KEY` in the environment "
|
||||
"or pass `api_key` to the guardrail in the config file."
|
||||
)
|
||||
|
||||
self.asset_id = asset_id
|
||||
if not self.asset_id:
|
||||
raise ValueError(
|
||||
"Repello guardrail requires an `asset_id`. Create an asset in the Repello "
|
||||
"dashboard and set `asset_id` on the guardrail in the config file."
|
||||
)
|
||||
|
||||
self.api_base = (
|
||||
api_base
|
||||
or get_secret_str("REPELLOAI_API_BASE")
|
||||
or DEFAULT_REPELLOAI_API_BASE
|
||||
)
|
||||
self.unreachable_fallback: Literal["fail_closed", "fail_open"] = (
|
||||
"fail_open" if unreachable_fallback == "fail_open" else "fail_closed"
|
||||
)
|
||||
self.async_handler = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.GuardrailCallback,
|
||||
params={"timeout": DEFAULT_REPELLOAI_TIMEOUT},
|
||||
)
|
||||
super().__init__( # pyright: ignore[reportUnknownMemberType]
|
||||
guardrail_name=guardrail_name,
|
||||
event_hook=event_hook,
|
||||
default_on=default_on,
|
||||
)
|
||||
|
||||
async def _call_analyze(
|
||||
self,
|
||||
text: str,
|
||||
stage: Literal["prompt", "response"],
|
||||
request_data: dict[str, object],
|
||||
event_type: GuardrailEventHooks,
|
||||
) -> RepelloAIAnalyzeResponse | None:
|
||||
endpoint = f"{self.api_base}/analyze/{stage}"
|
||||
request: dict[str, object] = {
|
||||
"asset_id": self.asset_id or "",
|
||||
"scan_data": {stage: text},
|
||||
}
|
||||
|
||||
status: GuardrailStatus = "success"
|
||||
guardrail_json_response: str | dict[str, object] | list[dict[str, object]] = ""
|
||||
start_time: datetime = datetime.now()
|
||||
repelloai_response: RepelloAIAnalyzeResponse | None = None
|
||||
try:
|
||||
verbose_proxy_logger.debug("RepelloAI Argus request: %s", request)
|
||||
raw_response: HttpxResponse | None = (
|
||||
await self.async_handler.post( # pyright: ignore[reportUnknownMemberType]
|
||||
url=endpoint,
|
||||
headers={"X-API-Key": self.repelloai_api_key},
|
||||
json=request,
|
||||
)
|
||||
)
|
||||
if raw_response is None:
|
||||
raise ValueError("RepelloAI Argus returned no response")
|
||||
response: HttpxResponse = raw_response
|
||||
self._raise_for_config_error(response)
|
||||
response.raise_for_status()
|
||||
try:
|
||||
repelloai_response = TypeAdapter(
|
||||
RepelloAIAnalyzeResponse
|
||||
).validate_json(response.text)
|
||||
except ValidationError as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={
|
||||
"error": "RepelloAI Argus guardrail returned invalid JSON",
|
||||
"status_code": response.status_code,
|
||||
},
|
||||
) from e
|
||||
verbose_proxy_logger.debug(
|
||||
"RepelloAI Argus response: %s", repelloai_response
|
||||
)
|
||||
if self._verdict_blocks(repelloai_response):
|
||||
status = "guardrail_intervened"
|
||||
return repelloai_response
|
||||
except HTTPException as e:
|
||||
status = "guardrail_failed_to_respond"
|
||||
guardrail_json_response = str(e.detail) if not isinstance(e.detail, (dict, list)) else e.detail # type: ignore[assignment]
|
||||
raise
|
||||
except HTTPError as e:
|
||||
status = "guardrail_failed_to_respond"
|
||||
guardrail_json_response = str(e)
|
||||
return self._handle_unreachable(e)
|
||||
except Exception as e:
|
||||
status = "guardrail_failed_to_respond"
|
||||
guardrail_json_response = str(e)
|
||||
raise HTTPException(
|
||||
status_code=500, detail={"error": "RepelloAI Argus guardrail failed"}
|
||||
) from e
|
||||
finally:
|
||||
end_time = datetime.now()
|
||||
if repelloai_response is not None:
|
||||
guardrail_json_response = dict(repelloai_response)
|
||||
self.add_standard_logging_guardrail_information_to_request_data( # pyright: ignore[reportUnknownMemberType]
|
||||
guardrail_json_response=guardrail_json_response,
|
||||
guardrail_status=status,
|
||||
request_data=request_data,
|
||||
start_time=start_time.timestamp(),
|
||||
end_time=end_time.timestamp(),
|
||||
duration=(end_time - start_time).total_seconds(),
|
||||
masked_entity_count={},
|
||||
event_type=event_type,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _raise_for_config_error(response: HttpxResponse) -> None:
|
||||
if response.status_code in CONFIG_ERROR_STATUS_CODES:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={
|
||||
"error": "RepelloAI Argus guardrail is misconfigured",
|
||||
"status_code": response.status_code,
|
||||
},
|
||||
)
|
||||
|
||||
def _verdict_blocks(
|
||||
self, repelloai_response: RepelloAIAnalyzeResponse | None
|
||||
) -> bool:
|
||||
if repelloai_response is None:
|
||||
return False
|
||||
verdict = repelloai_response.get("verdict")
|
||||
if verdict == BLOCKED_VERDICT:
|
||||
return True
|
||||
if verdict in (PASSED_VERDICT, FLAGGED_VERDICT):
|
||||
return False
|
||||
verbose_proxy_logger.warning(
|
||||
"RepelloAI Argus returned an unrecognized verdict (%s) - blocking.",
|
||||
verdict,
|
||||
)
|
||||
return True
|
||||
|
||||
def _handle_unreachable(self, error: Exception) -> RepelloAIAnalyzeResponse | None:
|
||||
verbose_proxy_logger.warning("RepelloAI Argus unreachable: %s", str(error))
|
||||
if self.unreachable_fallback == "fail_closed":
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={"error": "RepelloAI Argus guardrail unreachable"},
|
||||
)
|
||||
return None
|
||||
|
||||
def _raise_if_blocked(
|
||||
self, repelloai_response: RepelloAIAnalyzeResponse | None
|
||||
) -> None:
|
||||
if repelloai_response is None:
|
||||
return
|
||||
if self._verdict_blocks(repelloai_response):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=self._format_blocked_detail(repelloai_response),
|
||||
)
|
||||
self._log_flagged_verdict(repelloai_response)
|
||||
|
||||
@classmethod
|
||||
def _format_blocked_detail(
|
||||
cls, repelloai_response: RepelloAIAnalyzeResponse
|
||||
) -> str:
|
||||
policies = repelloai_response.get("policies_violated")
|
||||
if not isinstance(policies, list) or not policies:
|
||||
return "Blocked by RepelloAI Argus guardrail."
|
||||
|
||||
formatted_policies: list[str] = []
|
||||
for policy in policies:
|
||||
policy_name = policy.get("policy_name") or "unknown_policy"
|
||||
details: list[str] = []
|
||||
action_taken = policy.get("action_taken")
|
||||
if action_taken:
|
||||
details.append(f"action: {action_taken}")
|
||||
policy_details = policy.get("details")
|
||||
if isinstance(policy_details, dict):
|
||||
score = policy_details.get("score")
|
||||
if score is not None:
|
||||
details.append(f"score: {score}")
|
||||
suffix = f" ({', '.join(details)})" if details else ""
|
||||
formatted_policies.append(f"{policy_name}{suffix}")
|
||||
|
||||
if not formatted_policies:
|
||||
return "Blocked by RepelloAI Argus guardrail."
|
||||
return f"Blocked by RepelloAI Argus guardrail. Policies violated: {'; '.join(formatted_policies)}."
|
||||
|
||||
@staticmethod
|
||||
def _log_flagged_verdict(repelloai_response: RepelloAIAnalyzeResponse) -> None:
|
||||
if repelloai_response.get("verdict") == FLAGGED_VERDICT:
|
||||
verbose_proxy_logger.warning(
|
||||
"RepelloAI Argus flagged content (allowed): %s",
|
||||
repelloai_response.get("policies_violated"),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _extract_prompt_message_text(data: dict[str, object]) -> list[str]:
|
||||
messages = build_inspection_messages(data)
|
||||
return [
|
||||
content
|
||||
for message in messages
|
||||
if isinstance(content := message.get("content"), str) and content
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _extract_input_text_parts(content: object) -> list[str]:
|
||||
if not _is_object_list(content):
|
||||
return []
|
||||
return [
|
||||
text
|
||||
for part in content
|
||||
if _is_object_dict(part) and part.get("type") == "input_text"
|
||||
if isinstance(text := part.get("text"), str) and text
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _extract_prompt_field_text(data: dict[str, object]) -> list[str]:
|
||||
prompt = data.get("prompt")
|
||||
if isinstance(prompt, str) and prompt:
|
||||
return [prompt]
|
||||
if _is_object_list(prompt):
|
||||
return [item for item in prompt if isinstance(item, str) and item]
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def _extract_prompt_text(cls, data: dict[str, object]) -> str | None:
|
||||
texts = cls._extract_prompt_message_text(data)
|
||||
texts.extend(cls._extract_prompt_field_text(data))
|
||||
|
||||
instructions = data.get("instructions")
|
||||
if isinstance(instructions, str) and instructions:
|
||||
texts.append(instructions)
|
||||
|
||||
raw_messages = data.get("messages")
|
||||
if _is_object_list(raw_messages):
|
||||
for message in raw_messages:
|
||||
texts.extend(cls._extract_tool_call_args_from_message(message))
|
||||
|
||||
raw_input = data.get("input")
|
||||
if _is_object_list(raw_input):
|
||||
for item in raw_input:
|
||||
if _is_object_dict(item):
|
||||
if "role" not in item:
|
||||
continue
|
||||
texts.extend(cls._extract_tool_call_args_from_message(item))
|
||||
texts.extend(cls._extract_input_text_parts(item.get("content")))
|
||||
|
||||
texts.extend(cls._extract_tool_definition_text(data))
|
||||
return "\n".join(text for text in texts if text) if texts else None
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: litellm.DualCache,
|
||||
data: dict[str, object],
|
||||
call_type: CallTypesLiteral,
|
||||
) -> Exception | str | dict[str, object] | None:
|
||||
verbose_proxy_logger.debug("RepelloAI Argus: pre_call_hook")
|
||||
|
||||
event_type = GuardrailEventHooks.pre_call
|
||||
if (
|
||||
self.should_run_guardrail( # pyright: ignore[reportUnknownMemberType]
|
||||
data=data, event_type=event_type
|
||||
)
|
||||
is not True
|
||||
):
|
||||
return data
|
||||
|
||||
text = self._extract_prompt_text(data)
|
||||
if not text:
|
||||
verbose_proxy_logger.warning(
|
||||
"RepelloAI Argus: no inspectable prompt text in data - skipping."
|
||||
)
|
||||
return data
|
||||
|
||||
repelloai_response = await self._call_analyze(
|
||||
text=text,
|
||||
stage="prompt",
|
||||
request_data=data,
|
||||
event_type=event_type,
|
||||
)
|
||||
self._raise_if_blocked(repelloai_response)
|
||||
|
||||
add_guardrail_to_applied_guardrails_header(
|
||||
request_data=data, guardrail_name=self.guardrail_name
|
||||
)
|
||||
return data
|
||||
|
||||
async def async_post_call_success_hook(
|
||||
self,
|
||||
data: dict[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response: LLMResponseTypes,
|
||||
):
|
||||
verbose_proxy_logger.debug("RepelloAI Argus: post_call_success_hook")
|
||||
|
||||
event_type = GuardrailEventHooks.post_call
|
||||
if (
|
||||
self.should_run_guardrail( # pyright: ignore[reportUnknownMemberType]
|
||||
data=data, event_type=event_type
|
||||
)
|
||||
is not True
|
||||
):
|
||||
return response
|
||||
|
||||
text = self._extract_response_text(response)
|
||||
if not text:
|
||||
verbose_proxy_logger.warning(
|
||||
"RepelloAI Argus: no inspectable response text - skipping."
|
||||
)
|
||||
return response
|
||||
|
||||
repelloai_response = await self._call_analyze(
|
||||
text=text,
|
||||
stage="response",
|
||||
request_data=data,
|
||||
event_type=event_type,
|
||||
)
|
||||
self._raise_if_blocked(repelloai_response)
|
||||
|
||||
add_guardrail_to_applied_guardrails_header(
|
||||
request_data=data, guardrail_name=self.guardrail_name
|
||||
)
|
||||
return response
|
||||
|
||||
async def async_post_call_streaming_iterator_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response: AsyncGenerator[ModelResponseStream, None],
|
||||
request_data: dict[str, object],
|
||||
) -> AsyncGenerator[ModelResponseStream, None]:
|
||||
from litellm import main as litellm_main
|
||||
|
||||
event_type = GuardrailEventHooks.post_call
|
||||
if (
|
||||
self.should_run_guardrail( # pyright: ignore[reportUnknownMemberType]
|
||||
data=request_data, event_type=event_type
|
||||
)
|
||||
is not True
|
||||
):
|
||||
async for chunk in response:
|
||||
yield chunk
|
||||
return
|
||||
|
||||
chunks: list[ModelResponseStream] = []
|
||||
async for chunk in response:
|
||||
chunks.append(chunk)
|
||||
|
||||
assembled = litellm_main.stream_chunk_builder( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
|
||||
chunks=chunks
|
||||
)
|
||||
text = (
|
||||
self._extract_response_text(assembled)
|
||||
if isinstance(assembled, ModelResponse)
|
||||
else None
|
||||
)
|
||||
if text:
|
||||
repelloai_response = await self._call_analyze(
|
||||
text=text,
|
||||
stage="response",
|
||||
request_data=request_data,
|
||||
event_type=event_type,
|
||||
)
|
||||
if repelloai_response is not None:
|
||||
self._log_flagged_verdict(repelloai_response)
|
||||
if self._verdict_blocks(repelloai_response):
|
||||
from litellm.proxy.proxy_server import StreamingCallbackError
|
||||
|
||||
raise StreamingCallbackError("Blocked by RepelloAI Argus guardrail")
|
||||
add_guardrail_to_applied_guardrails_header(
|
||||
request_data=request_data, guardrail_name=self.guardrail_name
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.warning(
|
||||
"RepelloAI Argus: no inspectable text in streamed response; skipping scan. "
|
||||
"guardrail=%s assembled_type=%s",
|
||||
self.guardrail_name,
|
||||
type(assembled).__name__,
|
||||
)
|
||||
|
||||
for chunk in chunks:
|
||||
yield chunk
|
||||
|
||||
@staticmethod
|
||||
def _extract_response_text(response: object) -> str | None:
|
||||
if _is_object_dict(response):
|
||||
response_dict = response
|
||||
elif isinstance(response, ModelResponse):
|
||||
response_dict = (
|
||||
response.model_dump() # pyright: ignore[reportUnknownMemberType]
|
||||
)
|
||||
else:
|
||||
output_text = getattr(response, "output_text", None)
|
||||
if isinstance(output_text, str) and output_text:
|
||||
return output_text
|
||||
response_dict = {}
|
||||
|
||||
text = RepelloAIGuardrail._extract_chat_completion_text(response_dict)
|
||||
if text:
|
||||
return text
|
||||
return RepelloAIGuardrail._extract_responses_api_text(response_dict)
|
||||
|
||||
@classmethod
|
||||
def _extract_chat_completion_text(
|
||||
cls, response_dict: dict[str, object]
|
||||
) -> str | None:
|
||||
choices = response_dict.get("choices")
|
||||
if not _is_object_list(choices):
|
||||
return None
|
||||
parts: list[str] = []
|
||||
for choice in choices:
|
||||
if not _is_object_dict(choice):
|
||||
continue
|
||||
message = choice.get("message")
|
||||
if _is_object_dict(message):
|
||||
content = message.get("content")
|
||||
if isinstance(content, str) and content:
|
||||
parts.append(content)
|
||||
parts.extend(cls._extract_tool_call_args_from_message(message))
|
||||
text = choice.get("text")
|
||||
if isinstance(text, str) and text:
|
||||
parts.append(text)
|
||||
return "\n".join(parts) if parts else None
|
||||
|
||||
@staticmethod
|
||||
def _extract_responses_api_text(response_dict: dict[str, object]) -> str | None:
|
||||
output = response_dict.get("output")
|
||||
if not _is_object_list(output):
|
||||
return None
|
||||
texts: list[str] = []
|
||||
for output_item in output:
|
||||
if not _is_object_dict(output_item):
|
||||
continue
|
||||
item_type = output_item.get("type")
|
||||
if item_type == "function_call":
|
||||
arguments = output_item.get("arguments")
|
||||
if isinstance(arguments, str) and arguments:
|
||||
texts.append(arguments)
|
||||
continue
|
||||
if item_type != "message":
|
||||
continue
|
||||
content = output_item.get("content")
|
||||
if not _is_object_list(content):
|
||||
continue
|
||||
for content_item in content:
|
||||
if not _is_object_dict(content_item):
|
||||
continue
|
||||
if content_item.get("type") not in ("output_text", "text"):
|
||||
continue
|
||||
text = content_item.get("text")
|
||||
if isinstance(text, str) and text:
|
||||
texts.append(text)
|
||||
return "".join(texts) if texts else None
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> type[GuardrailConfigModel[BaseModel]] | None:
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.repelloai import (
|
||||
RepelloAIGuardrailConfigModel,
|
||||
)
|
||||
|
||||
return RepelloAIGuardrailConfigModel
|
||||
|
|
@ -44,6 +44,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import (
|
|||
from litellm.types.proxy.guardrails.guardrail_hooks.qohash import (
|
||||
QostodianNexusConfigModel,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.repelloai import (
|
||||
RepelloAIGuardrailConfigModel,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import (
|
||||
VigilGuardGuardrailConfigModel,
|
||||
)
|
||||
|
|
@ -115,6 +118,7 @@ class SupportedGuardrailIntegrations(Enum):
|
|||
QOSTODIAN_NEXUS = "qostodian_nexus"
|
||||
RUBRIK = "rubrik"
|
||||
VIGIL_GUARD = "vigil_guard"
|
||||
REPELLOAI = "repelloai"
|
||||
|
||||
|
||||
class Role(Enum):
|
||||
|
|
@ -758,7 +762,7 @@ class BaseLitellmParams(
|
|||
default="fail_closed",
|
||||
description=(
|
||||
"Behavior when a guardrail endpoint is unreachable due to network errors. "
|
||||
"NOTE: This is currently only implemented by guardrail='generic_guardrail_api'. "
|
||||
"Implemented by guardrail='generic_guardrail_api', 'akto', 'vigil_guard', and 'repelloai'. "
|
||||
"'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed."
|
||||
),
|
||||
)
|
||||
|
|
@ -856,6 +860,7 @@ class LitellmParams(
|
|||
PresidioConfigModel,
|
||||
BedrockGuardrailConfigModel,
|
||||
LakeraV2GuardrailConfigModel,
|
||||
RepelloAIGuardrailConfigModel,
|
||||
LassoGuardrailConfigModel,
|
||||
PillarGuardrailConfigModel,
|
||||
GraySwanGuardrailConfigModel,
|
||||
|
|
|
|||
65
litellm/types/proxy/guardrails/guardrail_hooks/repelloai.py
Normal file
65
litellm/types/proxy/guardrails/guardrail_hooks/repelloai.py
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
from typing import List, Literal, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from .base import GuardrailConfigModel
|
||||
|
||||
|
||||
class RepelloAIGuardrailConfigModel(GuardrailConfigModel[BaseModel]):
|
||||
"""Config model for the RepelloAI Argus guardrail."""
|
||||
|
||||
api_key: Optional[str] = Field(
|
||||
default=None,
|
||||
description="API key for the RepelloAI Argus service. Falls back to ARGUS_API_KEY or REPELLOAI_API_KEY.",
|
||||
)
|
||||
api_base: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Base URL for the RepelloAI Argus API. Defaults to https://argusapi.repello.ai/sdk/v1",
|
||||
)
|
||||
asset_id: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Repello asset ID whose dashboard policies are enforced. Required; the guardrail raises at init if it is missing.",
|
||||
)
|
||||
unreachable_fallback: Literal["fail_closed", "fail_open"] = Field(
|
||||
default="fail_closed",
|
||||
description="What to do when the RepelloAI Argus API is unreachable. 'fail_closed' = block (default), 'fail_open' = allow.",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "RepelloAI Argus"
|
||||
|
||||
|
||||
class RepelloAIScanData(TypedDict, total=False):
|
||||
"""The text payload sent to the RepelloAI Argus analyze endpoints.
|
||||
Only one of 'prompt' or 'response' is set per request.
|
||||
"""
|
||||
|
||||
prompt: Optional[str]
|
||||
response: Optional[str]
|
||||
|
||||
|
||||
class RepelloAIAnalyzeRequest(TypedDict, total=False):
|
||||
"""Request body for POST {api_base}/analyze/{prompt|response}."""
|
||||
|
||||
asset_id: str
|
||||
scan_data: RepelloAIScanData
|
||||
|
||||
|
||||
class RepelloAIViolatedPolicy(TypedDict, total=False):
|
||||
policy_name: Optional[str]
|
||||
policy_id: Optional[str]
|
||||
action_taken: Optional[str]
|
||||
scope: Optional[str]
|
||||
details: Optional[dict[str, object]]
|
||||
masked_result: Optional[str]
|
||||
|
||||
|
||||
class RepelloAIAnalyzeResponse(TypedDict, total=False):
|
||||
"""Response body returned by the RepelloAI Argus analyze endpoints."""
|
||||
|
||||
verdict: Optional[str] # "blocked" | "flagged" | "passed"
|
||||
request_id: Optional[str]
|
||||
policies_violated: Optional[List[RepelloAIViolatedPolicy]]
|
||||
policies_applied: Optional[List[dict[str, object]]]
|
||||
File diff suppressed because it is too large
Load diff
BIN
ui/litellm-dashboard/public/assets/logos/repelloai.png
Normal file
BIN
ui/litellm-dashboard/public/assets/logos/repelloai.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 14 KiB |
|
|
@ -294,4 +294,10 @@ export const GUARDRAIL_PRESETS: Record<string, GuardrailPreset> = {
|
|||
mode: "pre_call",
|
||||
defaultOn: false,
|
||||
},
|
||||
repelloai: {
|
||||
provider: "Repelloai",
|
||||
guardrailNameSuggestion: "RepelloAI Argus",
|
||||
mode: "pre_call",
|
||||
defaultOn: false,
|
||||
},
|
||||
};
|
||||
|
|
|
|||
|
|
@ -432,6 +432,16 @@ export const PARTNER_GUARDRAIL_CARDS: GuardrailCardInfo[] = [
|
|||
tags: ["Security", "Policy", "Grounding", "RAG"],
|
||||
providerKey: "Xecguard",
|
||||
},
|
||||
{
|
||||
id: "repelloai",
|
||||
name: "RepelloAI Argus",
|
||||
description:
|
||||
"RepelloAI Argus scans prompts and responses against policies configured per asset in the Repello dashboard.",
|
||||
category: "partner",
|
||||
logo: `${ASSET_PREFIX}repelloai.png`,
|
||||
tags: ["Security", "Policy", "Prompt Injection"],
|
||||
providerKey: "Repelloai",
|
||||
},
|
||||
];
|
||||
|
||||
export const ALL_CARDS = [...LITELLM_CONTENT_FILTER_CARDS, ...PARTNER_GUARDRAIL_CARDS];
|
||||
|
|
|
|||
|
|
@ -194,6 +194,20 @@ describe("guardrail_info_helpers", () => {
|
|||
expect(result.displayName).toBe("Noma Security");
|
||||
expect(result.logo).toContain("noma_security.png");
|
||||
});
|
||||
|
||||
it("should resolve RepelloAI Argus logo and display name", () => {
|
||||
populateGuardrailProviders({
|
||||
repelloai: { ui_friendly_name: "RepelloAI Argus" },
|
||||
});
|
||||
populateGuardrailProviderMap({
|
||||
repelloai: { ui_friendly_name: "RepelloAI Argus" },
|
||||
});
|
||||
|
||||
const result = getGuardrailLogoAndName("repelloai");
|
||||
|
||||
expect(result.displayName).toBe("RepelloAI Argus");
|
||||
expect(result.logo).toContain("repelloai.png");
|
||||
});
|
||||
});
|
||||
|
||||
describe("skipSystemMessageToChoice / choiceToSkipSystemForCreate", () => {
|
||||
|
|
|
|||
|
|
@ -53,6 +53,7 @@ export const guardrail_provider_map: Record<string, string> = {
|
|||
LlmAsAJudge: "llm_as_a_judge",
|
||||
Xecguard: "xecguard",
|
||||
QostodianNexus: "qostodian_nexus",
|
||||
Repelloai: "repelloai",
|
||||
};
|
||||
|
||||
// Function to populate provider map from API response - updates the original map
|
||||
|
|
@ -142,6 +143,7 @@ export const guardrailLogoMap: Record<string, string> = {
|
|||
"LiteLLM LLM as a Judge": `${asset_logos_folder}litellm_logo.jpg`,
|
||||
Akto: `${asset_logos_folder}akto.svg`,
|
||||
"Qostodian Nexus": `${asset_logos_folder}qohash.jpg`,
|
||||
"RepelloAI Argus": `${asset_logos_folder}repelloai.png`,
|
||||
};
|
||||
|
||||
export const getGuardrailLogoAndName = (guardrailValue: string): { logo: string; displayName: string } => {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue