mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-23 00:41:40 +00:00
refactor: slim upstream adapter — import from conduct-litellm-guard PyPI
The full adapter (response parser, session-ID chain, fail-mode logic,
HTTP client) lives in the conduct-litellm-guard package on PyPI. The
upstream tree hosts a thin re-export + the LiteLLM registration wiring.
Matches the Aporia / Lakera pattern — vendor SDK on PyPI, upstream
integration is a tiny adapter.
Benefits:
- Passes ruff-strict-budget and type-discipline-budget without new
violations.
- Users get the same install experience as any other guardrail vendor:
pip install conduct-litellm-guard
- Vendor keeps ownership of the parser + fail-mode semantics; upstream
keeps a stable interface.
Tests slimmed to smoke coverage (imports work, class is a
CustomGuardrail, enum + registries wired, missing-package error path).
Full behavioural coverage stays in the PyPI package.
Local runs of both scripts/ruff_strict_gate.py and
scripts/type_discipline_gate.py against upstream/litellm_internal_staging:
both pass.
This commit is contained in:
Sudhi Seshachala
Conduct AI
parent
bb74e4e562
commit
73f98dd1df
3 changed files with 76 additions and 390 deletions
|
|
@ -9,11 +9,17 @@ if TYPE_CHECKING:
|
|||
|
||||
|
||||
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"):
|
||||
"""Initialize the Conduct guardrail from LiteLLM's config block.
|
||||
|
||||
Maps LiteLLM's idiomatic ``api_base`` / ``api_key`` to Conduct's
|
||||
``api_url`` / ``agent_token`` constructor arguments. All other
|
||||
settings pass through unchanged.
|
||||
"""
|
||||
import litellm
|
||||
|
||||
_conduct_callback: Final = ConductGuardrail(
|
||||
api_base=getattr(litellm_params, "api_base", None),
|
||||
api_key=getattr(litellm_params, "api_key", None),
|
||||
api_url=getattr(litellm_params, "api_base", None),
|
||||
agent_token=getattr(litellm_params, "api_key", None),
|
||||
workspace_id=getattr(litellm_params, "workspace_id", None),
|
||||
fail_mode=getattr(litellm_params, "fail_mode", "fail_closed"),
|
||||
tool_name=getattr(litellm_params, "tool_name", "llm_call"),
|
||||
|
|
|
|||
|
|
@ -1,299 +1,26 @@
|
|||
# +-------------------------------------------------------------+
|
||||
#
|
||||
# Use Conduct Guard for your LLM calls
|
||||
#
|
||||
# Runtime policy enforcement — block / warn / audit / approval
|
||||
# Signed configuration + hash-chained audit + 20+ compliance packs
|
||||
# Docs: https://conductai.ai/guard
|
||||
#
|
||||
# +-------------------------------------------------------------+
|
||||
"""Conduct Guard as a LiteLLM guardrail.
|
||||
|
||||
Thin re-export. The adapter, response-envelope parser, session-ID chain,
|
||||
and fail-mode logic all live in the `conduct-litellm-guard` PyPI package,
|
||||
which is where issues, versioning, and standalone-user support live.
|
||||
|
||||
Install: `pip install conduct-litellm-guard`
|
||||
Source: https://github.com/sseshachala/conductai/tree/main/packages/conduct-litellm-guard
|
||||
Docs: https://conductai.ai/guard
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Literal
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
GUARDRAIL_NAME: Final = "conduct"
|
||||
try:
|
||||
from conduct_litellm_guard import ConductGuard as ConductGuardrail
|
||||
from conduct_litellm_guard.guardrail import (
|
||||
ConductGuardBlocked as ConductGuardrailBlocked,
|
||||
)
|
||||
from conduct_litellm_guard.guardrail import GuardDecision
|
||||
except ImportError as _e:
|
||||
raise ImportError(
|
||||
"conduct-litellm-guard is required for the Conduct guardrail. Install with: pip install conduct-litellm-guard"
|
||||
) from _e
|
||||
|
||||
|
||||
Verdict = Literal["allow", "advisory", "warning", "block", "approval", "unknown"]
|
||||
FailMode = Literal["fail_open", "fail_closed"]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GuardDecision:
|
||||
"""Structured view of what ``guard_check`` returned. The raw text is
|
||||
kept so audit / logging surfaces can quote it verbatim."""
|
||||
|
||||
verdict: Verdict
|
||||
raw: str
|
||||
rule_id: str | None = None
|
||||
message: str | None = None
|
||||
|
||||
@classmethod
|
||||
def parse(cls, text: str) -> GuardDecision:
|
||||
"""Map the ``guard_check`` string envelope to a verdict.
|
||||
|
||||
Response contract from Conduct:
|
||||
* ``"ok"`` or empty → allow silently
|
||||
* ``"advisory: ..."`` → allow but log
|
||||
* ``"WARNING — ..."`` → allow but surface
|
||||
* ``"BLOCKED — ..."`` → hard block
|
||||
* ``"PENDING approval — ..."`` → HITL — treat as block
|
||||
"""
|
||||
stripped = (text or "").strip()
|
||||
if not stripped or stripped.lower().startswith("ok"):
|
||||
return cls(verdict="allow", raw=stripped)
|
||||
if stripped.startswith("BLOCKED"):
|
||||
return cls(
|
||||
verdict="block",
|
||||
raw=stripped,
|
||||
rule_id=_extract_rule_id(stripped),
|
||||
message=_strip_prefix(stripped, "BLOCKED"),
|
||||
)
|
||||
if stripped.startswith("PENDING approval"):
|
||||
return cls(
|
||||
verdict="approval",
|
||||
raw=stripped,
|
||||
rule_id=_extract_rule_id(stripped),
|
||||
message=_strip_prefix(stripped, "PENDING approval"),
|
||||
)
|
||||
if stripped.startswith("WARNING"):
|
||||
return cls(
|
||||
verdict="warning",
|
||||
raw=stripped,
|
||||
rule_id=_extract_rule_id(stripped),
|
||||
message=_strip_prefix(stripped, "WARNING"),
|
||||
)
|
||||
if stripped.startswith("advisory"):
|
||||
return cls(
|
||||
verdict="advisory",
|
||||
raw=stripped,
|
||||
rule_id=_extract_rule_id(stripped),
|
||||
message=_strip_prefix(stripped, "advisory"),
|
||||
)
|
||||
return cls(verdict="unknown", raw=stripped)
|
||||
|
||||
|
||||
def _strip_prefix(text: str, prefix: str) -> str | None:
|
||||
remainder = text[len(prefix) :].strip()
|
||||
return remainder.lstrip(":—- ").strip() or None
|
||||
|
||||
|
||||
def _extract_rule_id(text: str) -> str | None:
|
||||
marker = "[rule:"
|
||||
idx = text.find(marker)
|
||||
if idx < 0:
|
||||
return None
|
||||
tail = text[idx + len(marker) :]
|
||||
end = tail.find("]")
|
||||
return tail[:end].strip() if end >= 0 else None
|
||||
|
||||
|
||||
class ConductGuardrailBlocked(Exception):
|
||||
"""Raised inside the pre-call hook to abort a LiteLLM request. LiteLLM
|
||||
surfaces the message to the caller."""
|
||||
|
||||
def __init__(self, decision: GuardDecision) -> None:
|
||||
self.decision = decision
|
||||
super().__init__(decision.message or decision.raw or "Blocked by Conduct Guard")
|
||||
|
||||
|
||||
class ConductGuardrail(CustomGuardrail):
|
||||
"""Conduct Guard as a LiteLLM ``CustomGuardrail``.
|
||||
|
||||
Reads config from LiteLLM's guardrail block. Every pre-call hook
|
||||
invocation calls Conduct's ``guard_check`` MCP tool using the
|
||||
supplied agent token. On block, raises so the LiteLLM proxy returns
|
||||
an error to the caller instead of forwarding to the model."""
|
||||
|
||||
@classmethod
|
||||
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
|
||||
return [
|
||||
GuardrailEventHooks.pre_call,
|
||||
GuardrailEventHooks.during_call,
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
api_base: str | None = None,
|
||||
api_key: str | None = None,
|
||||
workspace_id: str | None = None,
|
||||
fail_mode: FailMode = "fail_closed",
|
||||
tool_name: str = "llm_call",
|
||||
timeout: float = 8.0,
|
||||
**kwargs: object,
|
||||
) -> None:
|
||||
kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
|
||||
super().__init__(**kwargs)
|
||||
|
||||
self._api_url = (api_base or os.environ.get("CONDUCT_API_URL", "https://api.conductai.ai")).rstrip("/")
|
||||
token = api_key or os.environ.get("CONDUCT_AGENT_TOKEN")
|
||||
if not token:
|
||||
raise ValueError(
|
||||
"ConductGuardrail: agent token is required. Set CONDUCT_AGENT_TOKEN "
|
||||
"in the environment or pass api_key in the guardrail config."
|
||||
)
|
||||
self._agent_token = token
|
||||
self._workspace_id = workspace_id or os.environ.get("CONDUCT_WORKSPACE_ID")
|
||||
self._fail_mode: FailMode = fail_mode
|
||||
self._tool_name = tool_name
|
||||
self._timeout = timeout
|
||||
self._async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
|
||||
|
||||
# ── LiteLLM contract ───────────────────────────────────────────────
|
||||
|
||||
@log_guardrail_information
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: object,
|
||||
data: dict,
|
||||
call_type: Literal[
|
||||
"completion",
|
||||
"text_completion",
|
||||
"embeddings",
|
||||
"image_generation",
|
||||
"moderation",
|
||||
"audio_transcription",
|
||||
],
|
||||
) -> dict | None:
|
||||
decision = await self._check(data=data, call_type=call_type)
|
||||
|
||||
if decision.verdict in ("block", "approval"):
|
||||
raise ConductGuardrailBlocked(decision)
|
||||
|
||||
data.setdefault("metadata", {}).setdefault("conduct_guard", {}).update(
|
||||
{"verdict": decision.verdict, "rule_id": decision.rule_id}
|
||||
)
|
||||
return data
|
||||
|
||||
# ── Guard check ─────────────────────────────────────────────────
|
||||
|
||||
async def _check(self, *, data: dict, call_type: str) -> GuardDecision:
|
||||
tool_input = _build_tool_input(data, call_type)
|
||||
session_id = _extract_session_id(data)
|
||||
prompt = _extract_prompt_text(data)
|
||||
|
||||
arguments: dict = {
|
||||
"tool_name": self._tool_name,
|
||||
"tool_input": tool_input,
|
||||
}
|
||||
if prompt is not None:
|
||||
arguments["prompt"] = prompt
|
||||
|
||||
payload = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": str(uuid.uuid4()),
|
||||
"method": "tools/call",
|
||||
"params": {"name": "guard_check", "arguments": arguments},
|
||||
}
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self._agent_token}",
|
||||
"Content-Type": "application/json",
|
||||
"User-Agent": "litellm-conduct-guardrail/1.0",
|
||||
"X-Claude-Surface": "litellm",
|
||||
}
|
||||
if self._workspace_id:
|
||||
headers["X-Workspace-Id"] = self._workspace_id
|
||||
if session_id:
|
||||
headers["X-Conduct-Session-Id"] = session_id
|
||||
|
||||
try:
|
||||
response = await self._async_handler.post(
|
||||
f"{self._api_url}/guard/mcp",
|
||||
json=payload,
|
||||
headers=headers,
|
||||
timeout=self._timeout,
|
||||
)
|
||||
response.raise_for_status()
|
||||
body = response.json()
|
||||
|
||||
if "error" in body:
|
||||
err = body["error"]
|
||||
verbose_proxy_logger.warning("conduct_guard: eval error %s — applying %s", err, self._fail_mode)
|
||||
if self._fail_mode == "fail_closed":
|
||||
return GuardDecision(
|
||||
verdict="block",
|
||||
raw=str(err),
|
||||
message="Conduct Guard policy-eval error (fail_closed).",
|
||||
)
|
||||
return GuardDecision(verdict="allow", raw="fail_open")
|
||||
|
||||
result = body.get("result") or {}
|
||||
for item in result.get("content", []) or []:
|
||||
if item.get("type") == "text":
|
||||
return GuardDecision.parse(item.get("text", ""))
|
||||
return GuardDecision(verdict="allow", raw="")
|
||||
except Exception as e: # noqa: BLE001 — transport failure fallback path is intentionally broad
|
||||
verbose_proxy_logger.warning("conduct_guard: transport error %s — applying %s", e, self._fail_mode)
|
||||
if self._fail_mode == "fail_closed":
|
||||
return GuardDecision(
|
||||
verdict="block",
|
||||
raw=str(e),
|
||||
message="Conduct Guard is unreachable (fail_closed).",
|
||||
)
|
||||
return GuardDecision(verdict="allow", raw="fail_open")
|
||||
|
||||
|
||||
# ── Helpers ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _extract_session_id(data: dict) -> str | None:
|
||||
metadata = data.get("litellm_metadata") or data.get("metadata") or {}
|
||||
for key in ("trace_id", "X-Conduct-Session-Id", "conduct_session_id"):
|
||||
val = metadata.get(key)
|
||||
if val:
|
||||
return str(val)
|
||||
|
||||
user = data.get("user") or metadata.get("user") or ""
|
||||
first_msg = ""
|
||||
for m in data.get("messages") or []:
|
||||
if isinstance(m, dict) and m.get("role") == "user":
|
||||
first_msg = str(m.get("content", ""))[:512]
|
||||
break
|
||||
if not user and not first_msg:
|
||||
return None
|
||||
digest = hashlib.sha256((user + "|" + first_msg).encode("utf-8")).hexdigest()
|
||||
return f"litellm-{digest[:16]}"
|
||||
|
||||
|
||||
def _extract_prompt_text(data: dict) -> str | None:
|
||||
for m in reversed(data.get("messages") or []):
|
||||
if isinstance(m, dict) and m.get("role") == "user":
|
||||
content = m.get("content")
|
||||
if isinstance(content, str):
|
||||
return content[:4000]
|
||||
if isinstance(content, list):
|
||||
parts = [p.get("text", "") for p in content if isinstance(p, dict) and p.get("type") == "text"]
|
||||
return " ".join(parts)[:4000] or None
|
||||
return None
|
||||
|
||||
|
||||
def _build_tool_input(data: dict, call_type: str) -> dict:
|
||||
messages = data.get("messages") or []
|
||||
return {
|
||||
"model": data.get("model"),
|
||||
"call_type": call_type,
|
||||
"message_count": len(messages),
|
||||
"temperature": data.get("temperature"),
|
||||
"max_tokens": data.get("max_tokens"),
|
||||
"stream": bool(data.get("stream")),
|
||||
"content": _extract_prompt_text(data) or "",
|
||||
}
|
||||
__all__ = ["ConductGuardrail", "ConductGuardrailBlocked", "GuardDecision"]
|
||||
|
|
|
|||
|
|
@ -1,114 +1,67 @@
|
|||
"""Unit tests for the Conduct guardrail.
|
||||
"""Smoke tests for the Conduct guardrail integration.
|
||||
|
||||
Mocked transport — no real network. Verifies the response-envelope
|
||||
parser, pre-call hook behavior (allow / block / approval), fail-mode
|
||||
handling, and session-ID resolution chain.
|
||||
The adapter itself is tested in the ``conduct-litellm-guard`` PyPI
|
||||
package. Here we only verify:
|
||||
* the LiteLLM-tree module imports cleanly when the standalone package
|
||||
is installed
|
||||
* the enum + registry entries are wired
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock
|
||||
import importlib
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.guardrails.guardrail_hooks.conduct.conduct import (
|
||||
ConductGuardrail,
|
||||
ConductGuardrailBlocked,
|
||||
GuardDecision,
|
||||
)
|
||||
|
||||
def test_import_module() -> None:
|
||||
"""The wrapper module imports without side effects."""
|
||||
module = importlib.import_module("litellm.proxy.guardrails.guardrail_hooks.conduct")
|
||||
assert module.ConductGuardrail is not None
|
||||
|
||||
|
||||
# ── Decision parsing ────────────────────────────────────────────────────
|
||||
def test_class_is_custom_guardrail_subclass() -> None:
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.proxy.guardrails.guardrail_hooks.conduct import ConductGuardrail
|
||||
|
||||
assert issubclass(ConductGuardrail, CustomGuardrail)
|
||||
|
||||
|
||||
class TestGuardDecisionParse:
|
||||
@pytest.mark.parametrize("raw", ["ok", "OK", "", " ok "])
|
||||
def test_ok_variants_are_allow(self, raw: str) -> None:
|
||||
assert GuardDecision.parse(raw).verdict == "allow"
|
||||
def test_enum_value_registered() -> None:
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
|
||||
def test_blocked_extracts_rule_id(self) -> None:
|
||||
d = GuardDecision.parse("BLOCKED — command touches /etc/passwd [rule: no-etc-passwd]")
|
||||
assert d.verdict == "block"
|
||||
assert d.rule_id == "no-etc-passwd"
|
||||
|
||||
def test_pending_approval_treated_as_block(self) -> None:
|
||||
d = GuardDecision.parse("PENDING approval — HITL required [rule: prod-deploy-gate]")
|
||||
assert d.verdict == "approval"
|
||||
assert d.rule_id == "prod-deploy-gate"
|
||||
|
||||
def test_warning_is_warning(self) -> None:
|
||||
d = GuardDecision.parse("WARNING — high-risk model [rule: model-tier]")
|
||||
assert d.verdict == "warning"
|
||||
|
||||
def test_advisory_is_advisory(self) -> None:
|
||||
d = GuardDecision.parse("advisory: policy eval error: boom")
|
||||
assert d.verdict == "advisory"
|
||||
|
||||
def test_unknown_prefix_marked_unknown(self) -> None:
|
||||
assert GuardDecision.parse("wat").verdict == "unknown"
|
||||
assert SupportedGuardrailIntegrations.CONDUCT.value == "conduct"
|
||||
|
||||
|
||||
# ── Pre-call hook ──────────────────────────────────────────────────────
|
||||
def test_registries_populated() -> None:
|
||||
from litellm.proxy.guardrails.guardrail_hooks.conduct import (
|
||||
guardrail_class_registry,
|
||||
guardrail_initializer_registry,
|
||||
)
|
||||
|
||||
assert "conduct" in guardrail_class_registry
|
||||
assert "conduct" in guardrail_initializer_registry
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _agent_token_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("CONDUCT_AGENT_TOKEN", "cond_agt_test_placeholder")
|
||||
def test_missing_standalone_package_raises_helpful_error(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""When ``conduct-litellm-guard`` is not installed, the import fails
|
||||
with a message pointing users at the ``pip install`` command."""
|
||||
# Ensure the module is re-imported without the standalone package.
|
||||
for name in list(sys.modules):
|
||||
if name.startswith(("conduct_litellm_guard", "litellm.proxy.guardrails.guardrail_hooks.conduct")):
|
||||
monkeypatch.delitem(sys.modules, name, raising=False)
|
||||
|
||||
real_import = __builtins__["__import__"] if isinstance(__builtins__, dict) else __builtins__.__import__
|
||||
|
||||
def _guard() -> ConductGuardrail:
|
||||
return ConductGuardrail()
|
||||
def _fake_import(name: str, *args: object, **kwargs: object) -> object:
|
||||
if name.startswith("conduct_litellm_guard"):
|
||||
raise ImportError("simulated missing package")
|
||||
return real_import(name, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr("builtins.__import__", _fake_import)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestPreCallHook:
|
||||
async def test_allow_returns_data_with_metadata_tag(self) -> None:
|
||||
g = _guard()
|
||||
g._check = AsyncMock(return_value=GuardDecision(verdict="allow", raw="ok"))
|
||||
data = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}
|
||||
result = await g.async_pre_call_hook(None, None, data, "completion")
|
||||
assert result is data
|
||||
assert result["metadata"]["conduct_guard"]["verdict"] == "allow"
|
||||
|
||||
async def test_block_raises(self) -> None:
|
||||
g = _guard()
|
||||
g._check = AsyncMock(
|
||||
return_value=GuardDecision(
|
||||
verdict="block",
|
||||
raw="BLOCKED — no secrets [rule: no-prod-secrets]",
|
||||
rule_id="no-prod-secrets",
|
||||
)
|
||||
)
|
||||
with pytest.raises(ConductGuardrailBlocked) as exc:
|
||||
await g.async_pre_call_hook(None, None, {}, "completion")
|
||||
assert exc.value.decision.rule_id == "no-prod-secrets"
|
||||
|
||||
async def test_pending_approval_also_raises(self) -> None:
|
||||
g = _guard()
|
||||
g._check = AsyncMock(return_value=GuardDecision(verdict="approval", raw="PENDING approval — review"))
|
||||
with pytest.raises(ConductGuardrailBlocked):
|
||||
await g.async_pre_call_hook(None, None, {}, "completion")
|
||||
|
||||
|
||||
class TestConfig:
|
||||
def test_missing_token_raises_at_construction(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.delenv("CONDUCT_AGENT_TOKEN", raising=False)
|
||||
with pytest.raises(ValueError, match="agent token"):
|
||||
ConductGuardrail()
|
||||
|
||||
def test_config_api_key_wins_over_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("CONDUCT_AGENT_TOKEN", "env-token")
|
||||
g = ConductGuardrail(api_key="config-token")
|
||||
assert g._agent_token == "config-token"
|
||||
|
||||
def test_config_api_base_wins_over_default(self) -> None:
|
||||
g = ConductGuardrail(api_base="https://conduct.example.com/")
|
||||
assert g._api_url == "https://conduct.example.com"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
raise SystemExit(subprocess.call([sys.executable, "-m", "pytest", __file__, "-v"]))
|
||||
with pytest.raises(ImportError, match="pip install conduct-litellm-guard"):
|
||||
importlib.import_module("litellm.proxy.guardrails.guardrail_hooks.conduct.conduct")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue