mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
feat(guardrails): support logging_only mode for the Akto guardrail (#45461)
* feat(guardrails): support logging_only mode for the Akto guardrail * test(guardrails): wait for the spend row before reading Akto calls and cover mixed modes * test(guardrails): cover Akto logging_only on MCP tool calls * test(guardrails): type the Akto logging_only unit tests and inject the HTTP handler * test(guardrails): cover failing Akto replies, provider failures and a mixed outage burst under logging_only * test(guardrails): cover an unreachable Akto with fail_open under logging_only
This commit is contained in:
yucheng-berri
GitHub
parent
1a42c5abde
commit
51fea825cc
6 changed files with 824 additions and 4 deletions
|
|
@ -243,8 +243,12 @@ def first_value(request_data: Mapping[str, object], key: str) -> object:
|
|||
|
||||
INPUT_HOOKS: Final = MappingProxyType(
|
||||
{
|
||||
"request": frozenset((GuardrailEventHooks.pre_call, GuardrailEventHooks.pre_mcp_call)),
|
||||
"response": frozenset((GuardrailEventHooks.post_call, GuardrailEventHooks.post_mcp_call)),
|
||||
"request": frozenset(
|
||||
(GuardrailEventHooks.pre_call, GuardrailEventHooks.pre_mcp_call, GuardrailEventHooks.logging_only)
|
||||
),
|
||||
"response": frozenset(
|
||||
(GuardrailEventHooks.post_call, GuardrailEventHooks.post_mcp_call, GuardrailEventHooks.logging_only)
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
|
@ -265,6 +269,7 @@ class AktoGuardrail(CustomGuardrail):
|
|||
GuardrailEventHooks.post_call,
|
||||
GuardrailEventHooks.pre_mcp_call,
|
||||
GuardrailEventHooks.post_mcp_call,
|
||||
GuardrailEventHooks.logging_only,
|
||||
]
|
||||
|
||||
def __init__(
|
||||
|
|
|
|||
|
|
@ -186,7 +186,10 @@ _UNSENDABLE: Final[_Classified] = (None, True)
|
|||
|
||||
def request_attachments(request_data: Mapping[str, object]) -> RequestAttachments:
|
||||
# Both, so a decoy "messages" can't hide attachments in a Responses API "input"
|
||||
containers: Final = (_parse(_ITEMS_ADAPTER, request_data.get(key)) or () for key in ("messages", "input"))
|
||||
messages: Final = request_data.get("messages")
|
||||
responses_input: Final = request_data.get("input")
|
||||
sources: Final = (messages,) if responses_input is messages else (messages, responses_input)
|
||||
containers: Final = (_parse(_ITEMS_ADAPTER, source) or () for source in sources)
|
||||
blocks: Final = tuple(chain.from_iterable(_message_blocks(message) for message in chain.from_iterable(containers)))
|
||||
classified: Final = tuple(
|
||||
chain.from_iterable(_block_attachments(block, index) for index, block in enumerate(blocks))
|
||||
|
|
|
|||
133
tests/integration/mcp/test_mcp_akto_logging_only.py
Normal file
133
tests/integration/mcp/test_mcp_akto_logging_only.py
Normal file
|
|
@ -0,0 +1,133 @@
|
|||
"""Akto guardrail in `mode: logging_only` on MCP tool calls through a real proxy with two workers.
|
||||
|
||||
The scripted MCP server and the Akto `/api/http-proxy` service are the only doubles.
|
||||
"""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qs, urlsplit
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.mcp import EntryPoint, McpCaller, McpPeer, echo_tool, register_mcp, scripted_peer, tool_calls
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
||||
AKTO_KEY: Final = "synthetic-akto-mcp-log-key"
|
||||
BLOCK_MARK: Final = "SYNTHETIC-AKTO-BLOCK"
|
||||
RESPONSE_CHECK: Final = {"akto_connector": "litellm", "response_guardrails": "true", "ingest_data": "true"}
|
||||
MCP_SPEND_ROWS: Final = 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND call_type = %s'
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AktoCall:
|
||||
path: str
|
||||
flags: dict[str, str]
|
||||
authorization: str
|
||||
payload: dict[str, object]
|
||||
|
||||
|
||||
def _akto_call(request: Request) -> AktoCall:
|
||||
target: Final = urlsplit(request.target)
|
||||
return AktoCall(
|
||||
path=target.path,
|
||||
flags={name: values[0] for name, values in parse_qs(target.query).items()},
|
||||
authorization=request.headers.get("authorization", ""),
|
||||
payload=json.loads(request.body),
|
||||
)
|
||||
|
||||
|
||||
def _akto_verdict(request: Request) -> Reply:
|
||||
verdict: Final = (
|
||||
{"Allowed": False, "Behaviour": "block", "Reason": "Synthetic Akto policy block"}
|
||||
if BLOCK_MARK.encode() in request.body
|
||||
else {"Allowed": True}
|
||||
)
|
||||
return Reply(body=json.dumps({"data": {"guardrailsResult": verdict}}).encode())
|
||||
|
||||
|
||||
def _config(akto_url: str, root: Path) -> Path:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["guardrails"] = [
|
||||
{
|
||||
"guardrail_name": "akto-mcp-log",
|
||||
"litellm_params": {
|
||||
"guardrail": "akto",
|
||||
"mode": "logging_only",
|
||||
"default_on": True,
|
||||
"akto_base_url": akto_url,
|
||||
"akto_api_key": AKTO_KEY,
|
||||
},
|
||||
}
|
||||
]
|
||||
path: Final = root / "akto-mcp-logging-only.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
proxy: Gateway
|
||||
akto: Wire
|
||||
peer: McpPeer
|
||||
scenario: Scenario
|
||||
alias: str
|
||||
server_id: str
|
||||
|
||||
def key(self) -> str:
|
||||
return self.scenario.key(object_permission={"mcp_servers": [self.server_id]})
|
||||
|
||||
def settled_akto_calls(self, key: str, marker: str) -> tuple[AktoCall, ...]:
|
||||
eventually(
|
||||
lambda: read_rows(MCP_SPEND_ROWS, (sha256(key.encode()).hexdigest(), "call_mcp_tool")),
|
||||
lambda rows: len(rows) >= 1,
|
||||
seconds=70,
|
||||
)
|
||||
eventually(lambda: self.akto.received.qsize(), lambda count: count >= 1, seconds=30)
|
||||
calls: Final = tuple(_akto_call(request) for request in self.akto.drain() if marker.encode() in request.body)
|
||||
return tuple(call for call in calls if call.authorization == AKTO_KEY)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
root: Final = tmp_path_factory.mktemp("akto-mcp-logging-only")
|
||||
alias: Final = "aktolog" + uuid.uuid4().hex[:8]
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
wire_server(_akto_verdict) as akto,
|
||||
scripted_peer(echo_tool("echo")) as peer,
|
||||
owned_proxy_process(gateway, root, {}, config=_config(akto.url, root), workers=2) as owned,
|
||||
owned.gateway.scenario() as scenario,
|
||||
):
|
||||
identity: Final = register_mcp(scenario, peer, alias)
|
||||
peer.drain()
|
||||
akto.drain()
|
||||
yield Rig(owned.gateway, akto, peer, scenario, alias, identity)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entry", ("mcp", "rest"))
|
||||
def test_logging_only_akto_block_verdict_never_blocks_an_mcp_tool_call_and_still_checks_it(
|
||||
rig: Rig, entry: EntryPoint
|
||||
) -> None:
|
||||
marker: Final = "mark-" + uuid.uuid4().hex
|
||||
arguments: Final = {"note": f"{BLOCK_MARK} {marker}"}
|
||||
key: Final = rig.key()
|
||||
|
||||
outcome: Final = McpCaller(rig.proxy, key, entry, rig.alias).call(
|
||||
f"{rig.alias}-echo", arguments, server_id=rig.server_id
|
||||
)
|
||||
|
||||
assert (outcome.error, outcome.text) == (None, json.dumps(arguments, sort_keys=True)), outcome.raw
|
||||
upstream: Final = [call["body"]["params"]["arguments"] for call in tool_calls(rig.peer.drain())]
|
||||
assert upstream == [arguments], upstream
|
||||
calls: Final = rig.settled_akto_calls(key, marker)
|
||||
assert all(call.path == "/api/http-proxy" for call in calls), calls
|
||||
checked: Final = [call for call in calls if call.flags == RESPONSE_CHECK]
|
||||
assert [marker in str(call.payload.get("responsePayload")) for call in checked] == [True], f"{entry}: {calls}"
|
||||
547
tests/integration/observability/test_akto_logging_only.py
Normal file
547
tests/integration/observability/test_akto_logging_only.py
Normal file
|
|
@ -0,0 +1,547 @@
|
|||
"""Akto guardrail in `mode: logging_only`, driven through a real proxy.
|
||||
|
||||
The Akto service is the only guardrail double: an owned wire peer that answers the `/api/http-proxy`
|
||||
verdict protocol. The provider is a second owned peer. The proxy, its guardrail registry, Postgres and
|
||||
Redis run for real with two workers. Every test waits for the spend row, which is written after the
|
||||
logging-only scans finish, before it reads what Akto received.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qs, urlsplit
|
||||
|
||||
import anthropic
|
||||
import openai
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment, object_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
||||
LOG_KEY: Final = "synthetic-akto-log-key"
|
||||
INPUT_KEY: Final = "synthetic-akto-input-key"
|
||||
OUTPUT_KEY: Final = "synthetic-akto-output-key"
|
||||
DOWN_KEY: Final = "synthetic-akto-down-key"
|
||||
DOWN_OPEN_KEY: Final = "synthetic-akto-down-open-key"
|
||||
MIXED_KEY: Final = "synthetic-akto-mixed-key"
|
||||
BLOCK_MARK: Final = "SYNTHETIC-AKTO-BLOCK"
|
||||
BLOCK_REASON: Final = "Synthetic Akto policy block"
|
||||
AKTO_DROP_MARK: Final = "SYNTHETIC-AKTO-DROP"
|
||||
AKTO_ERROR_MARK: Final = "SYNTHETIC-AKTO-500"
|
||||
AKTO_GARBAGE_MARK: Final = "SYNTHETIC-AKTO-GARBAGE"
|
||||
PROVIDER_FAIL_MARK: Final = "SYNTHETIC-PROVIDER-FAIL"
|
||||
REQUEST_CHECK: Final = {"akto_connector": "litellm", "guardrails": "true", "ingest_data": "true"}
|
||||
RESPONSE_CHECK: Final = {"akto_connector": "litellm", "response_guardrails": "true", "ingest_data": "true"}
|
||||
FILE_CHECK: Final = {"akto_connector": "litellm", "file_guardrails": "true"}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AktoCall:
|
||||
path: str
|
||||
flags: dict[str, str]
|
||||
authorization: str
|
||||
payload: dict[str, object]
|
||||
|
||||
def request_text(self) -> str:
|
||||
return str(self.payload.get("requestPayload", ""))
|
||||
|
||||
def response_text(self) -> str:
|
||||
return str(self.payload.get("responsePayload", ""))
|
||||
|
||||
|
||||
def _akto_call(request: Request) -> AktoCall:
|
||||
target: Final = urlsplit(request.target)
|
||||
return AktoCall(
|
||||
path=target.path,
|
||||
flags={name: values[0] for name, values in parse_qs(target.query).items()},
|
||||
authorization=request.headers.get("authorization", ""),
|
||||
payload=json.loads(request.body),
|
||||
)
|
||||
|
||||
|
||||
def _akto_verdict(request: Request) -> Reply:
|
||||
if AKTO_DROP_MARK.encode() in request.body:
|
||||
return Reply(drop_connection=True)
|
||||
if AKTO_ERROR_MARK.encode() in request.body:
|
||||
return Reply(status=500, body=b'{"error": "synthetic akto failure"}')
|
||||
if AKTO_GARBAGE_MARK.encode() in request.body:
|
||||
return Reply(body=b"synthetic akto garbage", content_type="text/plain")
|
||||
verdict: Final = (
|
||||
{"Allowed": False, "Behaviour": "block", "Reason": BLOCK_REASON}
|
||||
if BLOCK_MARK.encode() in request.body
|
||||
else {"Allowed": True}
|
||||
)
|
||||
return Reply(body=json.dumps({"data": {"guardrailsResult": verdict}}).encode())
|
||||
|
||||
|
||||
def _answer(marker: str) -> str:
|
||||
return "synthetic answer " + marker
|
||||
|
||||
|
||||
def _marker_in(body: bytes) -> str:
|
||||
text: Final = body.decode()
|
||||
start: Final = text.find("mark-")
|
||||
assert start >= 0, text
|
||||
return text[start : start + 37]
|
||||
|
||||
|
||||
def _sse(events: tuple[dict[str, object], ...]) -> tuple[bytes, ...]:
|
||||
return tuple(("data: " + json.dumps(event) + "\n\n").encode() for event in events)
|
||||
|
||||
|
||||
def _chat_reply(marker: str, streaming: bool) -> Reply:
|
||||
answer: Final = _answer(marker)
|
||||
if not streaming:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-" + marker,
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": "gpt-5.4-mini",
|
||||
"choices": [
|
||||
{"index": 0, "message": {"role": "assistant", "content": answer}, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
def chunk(delta: dict[str, object], finish: str | None) -> dict[str, object]:
|
||||
return {
|
||||
"id": "chatcmpl-" + marker,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "gpt-5.4-mini",
|
||||
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
|
||||
}
|
||||
|
||||
events: Final = (chunk({"role": "assistant", "content": answer[:9]}, None), chunk({"content": answer[9:]}, "stop"))
|
||||
return Reply(chunks=(*_sse(events), b"data: [DONE]\n\n"), content_type="text/event-stream")
|
||||
|
||||
|
||||
def _messages_reply(marker: str) -> Reply:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "msg_" + marker,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-haiku-5-5",
|
||||
"content": [{"type": "text", "text": _answer(marker)}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _responses_reply(marker: str) -> Reply:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "resp_" + marker,
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-5.4-mini",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msgo_" + marker,
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": _answer(marker), "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _provider(request: Request) -> Reply:
|
||||
if PROVIDER_FAIL_MARK.encode() in request.body:
|
||||
return Reply(status=500, body=b'{"error": {"message": "synthetic provider failure", "type": "server_error"}}')
|
||||
marker: Final = _marker_in(request.body)
|
||||
if request.target.endswith("/v1/messages"):
|
||||
return _messages_reply(marker)
|
||||
if request.target.endswith("/v1/responses"):
|
||||
return _responses_reply(marker)
|
||||
assert request.target.endswith("/v1/chat/completions"), request.target
|
||||
return _chat_reply(marker, bool(json.loads(request.body).get("stream")))
|
||||
|
||||
|
||||
def _guardrail(name: str, key: str, url: str, **params: object) -> dict[str, object]:
|
||||
return {
|
||||
"guardrail_name": name,
|
||||
"litellm_params": {
|
||||
"guardrail": "akto",
|
||||
"mode": "logging_only",
|
||||
"default_on": True,
|
||||
"akto_base_url": url,
|
||||
"akto_api_key": key,
|
||||
**params,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _rig_config(akto_url: str, down_url: str, root: Path) -> Path:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config["litellm_settings"]["cache"] = False
|
||||
config["guardrails"] = [
|
||||
_guardrail("akto-log", LOG_KEY, akto_url),
|
||||
_guardrail("akto-log-input", INPUT_KEY, akto_url, logging_only_scope="input"),
|
||||
_guardrail("akto-log-output", OUTPUT_KEY, akto_url, logging_only_scope="output"),
|
||||
_guardrail("akto-log-down", DOWN_KEY, down_url),
|
||||
_guardrail("akto-log-down-open", DOWN_OPEN_KEY, down_url, unreachable_fallback="fail_open"),
|
||||
_guardrail("akto-mixed", MIXED_KEY, akto_url, mode=["pre_call", "logging_only"], default_on=False),
|
||||
]
|
||||
path: Final = root / "akto-logging-only.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
proxy: Gateway
|
||||
akto: Wire
|
||||
akto_down: Wire
|
||||
provider: Wire
|
||||
chat_model: str
|
||||
claude_model: str
|
||||
responses_model: str
|
||||
|
||||
def base(self) -> str:
|
||||
return str(self.proxy.client.base_url).rstrip("/")
|
||||
|
||||
def auth(self) -> dict[str, str]:
|
||||
return {"Authorization": f"Bearer {self.proxy.key}"}
|
||||
|
||||
def provider_calls(self, marker: str) -> tuple[Request, ...]:
|
||||
return tuple(request for request in self.provider.drain() if marker.encode() in request.body)
|
||||
|
||||
def guardrail_entries(self, response_id: str, name: str) -> tuple[dict[str, object], ...]:
|
||||
rows: Final = spend_rows(response_id)
|
||||
assert len(rows) == 1, rows
|
||||
entries: Final = object_value(rows[0]["metadata"]).get("guardrail_information")
|
||||
assert isinstance(entries, list), rows[0]
|
||||
return tuple(entry for entry in (object_value(item) for item in entries) if entry.get("guardrail_name") == name)
|
||||
|
||||
def settled_akto_calls(self, response_id: str, marker: str, key: str) -> tuple[AktoCall, ...]:
|
||||
self.guardrail_entries(response_id, "akto-log")
|
||||
calls: Final = tuple(_akto_call(request) for request in self.akto.drain() if marker.encode() in request.body)
|
||||
return tuple(call for call in calls if call.authorization == key)
|
||||
|
||||
|
||||
def spend_rows(request_id: str) -> tuple[dict[str, object], ...]:
|
||||
return tuple(
|
||||
eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT request_id, status, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)
|
||||
),
|
||||
lambda values: len(values) >= 1,
|
||||
seconds=70,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
||||
root: Final = tmp_path_factory.mktemp("akto-logging-only")
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
wire_server(_provider) as provider,
|
||||
wire_server(_akto_verdict) as akto,
|
||||
wire_server(lambda _: Reply(drop_connection=True)) as akto_down,
|
||||
owned_proxy_process(gateway, root, {}, config=_rig_config(akto.url, akto_down.url, root), workers=2) as owned,
|
||||
owned.gateway.scenario() as scenario,
|
||||
):
|
||||
chat: Final = scenario.model(
|
||||
model="openai/gpt-5.4-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key"
|
||||
)
|
||||
claude: Final = scenario.model(
|
||||
model="anthropic/claude-haiku-5-5", api_base=provider.url, api_key="synthetic-anthropic-key"
|
||||
)
|
||||
responses: Final = scenario.model(
|
||||
model="openai/gpt-5.4-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key"
|
||||
)
|
||||
yield Rig(owned.gateway, akto, akto_down, provider, chat, claude, responses)
|
||||
|
||||
|
||||
def _marker() -> str:
|
||||
return "mark-" + uuid.uuid4().hex
|
||||
|
||||
|
||||
def _assert_checked_both_ways(calls: tuple[AktoCall, ...], marker: str) -> None:
|
||||
assert [call.flags for call in calls] == [REQUEST_CHECK, RESPONSE_CHECK], calls
|
||||
assert all(call.path == "/api/http-proxy" for call in calls), calls
|
||||
assert marker in calls[0].request_text(), calls[0].payload
|
||||
assert _answer(marker) in calls[1].response_text(), calls[1].payload
|
||||
|
||||
|
||||
def _chat(rig: Rig, content: object, **extra: object) -> dict[str, object]:
|
||||
response: Final = rig.proxy.client.post(
|
||||
"/v1/chat/completions",
|
||||
json={"model": rig.chat_model, "messages": [{"role": "user", "content": content}], **extra},
|
||||
headers=rig.auth(),
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return response.json()
|
||||
|
||||
|
||||
def _content(body: dict[str, object]) -> object:
|
||||
choices: Final = body["choices"]
|
||||
assert isinstance(choices, list), body
|
||||
return object_value(object_value(choices[0])["message"])["content"]
|
||||
|
||||
|
||||
def test_logging_only_checks_a_sync_openai_chat_request_and_response(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0)
|
||||
completion: Final = client.chat.completions.create(
|
||||
model=rig.chat_model, messages=[{"role": "user", "content": "hello " + marker}]
|
||||
)
|
||||
assert completion.choices[0].message.content == _answer(marker)
|
||||
assert len(rig.provider_calls(marker)) == 1
|
||||
|
||||
_assert_checked_both_ways(rig.settled_akto_calls(completion.id, marker, LOG_KEY), marker)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logging_only_checks_an_async_openai_chat_request_and_response(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
client: Final = openai.AsyncOpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0)
|
||||
completion: Final = await client.chat.completions.create(
|
||||
model=rig.chat_model, messages=[{"role": "user", "content": "hello " + marker}]
|
||||
)
|
||||
assert completion.choices[0].message.content == _answer(marker)
|
||||
|
||||
_assert_checked_both_ways(rig.settled_akto_calls(completion.id, marker, LOG_KEY), marker)
|
||||
|
||||
|
||||
def test_logging_only_block_verdict_never_blocks_the_caller(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
body: Final = _chat(rig, f"{BLOCK_MARK} {marker}")
|
||||
assert _content(body) == _answer(marker)
|
||||
assert len(rig.provider_calls(marker)) == 1
|
||||
|
||||
calls: Final = rig.settled_akto_calls(str(body["id"]), marker, LOG_KEY)
|
||||
assert [call.flags for call in calls] == [REQUEST_CHECK], calls
|
||||
entries: Final = rig.guardrail_entries(str(body["id"]), "akto-log")
|
||||
assert [(entry["guardrail_mode"], entry["guardrail_status"]) for entry in entries] == [
|
||||
("logging_only", "guardrail_intervened")
|
||||
], entries
|
||||
|
||||
|
||||
def test_logging_only_scope_input_and_output_each_check_one_direction(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
body: Final = _chat(rig, "scoped " + marker)
|
||||
|
||||
input_calls: Final = rig.settled_akto_calls(str(body["id"]), marker, INPUT_KEY)
|
||||
assert [call.flags for call in input_calls] == [REQUEST_CHECK], input_calls
|
||||
assert marker in input_calls[0].request_text(), input_calls[0].payload
|
||||
entries: Final = rig.guardrail_entries(str(body["id"]), "akto-log-output")
|
||||
assert [entry["guardrail_status"] for entry in entries] == ["success"], entries
|
||||
|
||||
|
||||
def test_logging_only_scope_output_checks_only_the_response(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
body: Final = _chat(rig, "scoped output " + marker)
|
||||
|
||||
output_calls: Final = rig.settled_akto_calls(str(body["id"]), marker, OUTPUT_KEY)
|
||||
assert [call.flags for call in output_calls] == [RESPONSE_CHECK], output_calls
|
||||
assert _answer(marker) in output_calls[0].response_text(), output_calls[0].payload
|
||||
|
||||
|
||||
def test_unreachable_akto_under_logging_only_never_fails_the_caller(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
body: Final = _chat(rig, "outage " + marker)
|
||||
assert _content(body) == _answer(marker)
|
||||
|
||||
entries: Final = rig.guardrail_entries(str(body["id"]), "akto-log-down")
|
||||
assert [entry["guardrail_mode"] for entry in entries] == ["logging_only"], entries
|
||||
dropped: Final = [_akto_call(request) for request in rig.akto_down.drain() if marker.encode() in request.body]
|
||||
assert [call.flags for call in dropped if call.authorization == DOWN_KEY] == [REQUEST_CHECK], dropped
|
||||
|
||||
|
||||
def test_unreachable_akto_with_fail_open_still_attempts_the_response_check(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
body: Final = _chat(rig, "outage open " + marker)
|
||||
assert _content(body) == _answer(marker)
|
||||
|
||||
entries: Final = rig.guardrail_entries(str(body["id"]), "akto-log-down-open")
|
||||
assert [entry["guardrail_mode"] for entry in entries] == ["logging_only", "logging_only"], entries
|
||||
dropped: Final = [_akto_call(request) for request in rig.akto_down.drain() if marker.encode() in request.body]
|
||||
assert [call.flags for call in dropped if call.authorization == DOWN_OPEN_KEY] == [REQUEST_CHECK, RESPONSE_CHECK]
|
||||
|
||||
|
||||
def test_logging_only_streaming_chat_sends_the_assembled_answer(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0)
|
||||
chunks: Final = tuple(
|
||||
client.chat.completions.create(
|
||||
model=rig.chat_model, messages=[{"role": "user", "content": "stream " + marker}], stream=True
|
||||
)
|
||||
)
|
||||
assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == _answer(marker)
|
||||
|
||||
_assert_checked_both_ways(rig.settled_akto_calls(chunks[0].id, marker, LOG_KEY), marker)
|
||||
|
||||
|
||||
def test_logging_only_anthropic_messages_checks_both_directions(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
client: Final = anthropic.Anthropic(base_url=rig.base(), api_key=rig.proxy.key, max_retries=0)
|
||||
message: Final = client.messages.create(
|
||||
model=rig.claude_model, max_tokens=64, messages=[{"role": "user", "content": "claude " + marker}]
|
||||
)
|
||||
assert message.content[0].type == "text" and message.content[0].text == _answer(marker)
|
||||
|
||||
_assert_checked_both_ways(rig.settled_akto_calls(message.id, marker, LOG_KEY), marker)
|
||||
|
||||
|
||||
def test_logging_only_responses_api_checks_both_directions(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0)
|
||||
result: Final = client.responses.create(model=rig.responses_model, input="responses " + marker)
|
||||
assert result.output_text == _answer(marker)
|
||||
|
||||
_assert_checked_both_ways(rig.settled_akto_calls(result.id, marker, LOG_KEY), marker)
|
||||
|
||||
|
||||
def test_logging_only_sends_each_chat_attachment_to_akto_once(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
image_url: Final = f"https://example.com/{marker}.png"
|
||||
body: Final = _chat(
|
||||
rig, [{"type": "text", "text": "describe " + marker}, {"type": "image_url", "image_url": {"url": image_url}}]
|
||||
)
|
||||
|
||||
calls: Final = rig.settled_akto_calls(str(body["id"]), marker, LOG_KEY)
|
||||
file_checks: Final = [call.payload["files"] for call in calls if call.flags == FILE_CHECK]
|
||||
assert file_checks == [[{"filename": f"{marker}.png", "type": "image", "url": image_url}]], calls
|
||||
|
||||
|
||||
def test_logging_only_sends_each_responses_attachment_to_akto_once(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
image_url: Final = f"https://example.com/{marker}.png"
|
||||
client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0)
|
||||
result: Final = client.responses.create(
|
||||
model=rig.responses_model,
|
||||
input=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "input_text", "text": "describe " + marker},
|
||||
{"type": "input_image", "image_url": image_url, "detail": "auto"},
|
||||
],
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
calls: Final = rig.settled_akto_calls(result.id, marker, LOG_KEY)
|
||||
file_checks: Final = [call.payload["files"] for call in calls if call.flags == FILE_CHECK]
|
||||
assert file_checks == [[{"filename": f"{marker}.png", "type": "image", "url": image_url}]], calls
|
||||
|
||||
|
||||
def test_pre_call_plus_logging_only_checks_the_request_inline_and_both_directions_after(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
body: Final = _chat(rig, "mixed " + marker, guardrails=["akto-mixed"])
|
||||
|
||||
flags: Final = [call.flags for call in rig.settled_akto_calls(str(body["id"]), marker, MIXED_KEY)]
|
||||
assert (flags.count(REQUEST_CHECK), flags.count(RESPONSE_CHECK), len(flags)) == (2, 1, 3), flags
|
||||
|
||||
|
||||
def test_dashboard_offers_logging_only_for_akto(rig: Rig) -> None:
|
||||
response: Final = rig.proxy.client.get("/guardrails/ui/add_guardrail_settings", headers=rig.auth())
|
||||
assert response.status_code == 200, response.text
|
||||
settings: Final = response.json()
|
||||
assert "logging_only" in settings["supported_modes_by_provider"]["akto"], settings["supported_modes_by_provider"]
|
||||
assert "akto" not in settings["providers_without_directional_logging_only_scope"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failure", [AKTO_ERROR_MARK, AKTO_GARBAGE_MARK, AKTO_DROP_MARK])
|
||||
def test_failing_akto_under_logging_only_never_fails_the_caller(rig: Rig, failure: str) -> None:
|
||||
marker: Final = _marker()
|
||||
body: Final = _chat(rig, f"{failure} {marker}")
|
||||
assert _content(body) == _answer(marker)
|
||||
|
||||
calls: Final = rig.settled_akto_calls(str(body["id"]), marker, LOG_KEY)
|
||||
assert [call.flags for call in calls] == [REQUEST_CHECK], calls
|
||||
entries: Final = rig.guardrail_entries(str(body["id"]), "akto-log")
|
||||
assert [entry["guardrail_mode"] for entry in entries] == ["logging_only"], entries
|
||||
|
||||
|
||||
def test_provider_failure_under_logging_only_reaches_the_caller_and_checks_nothing(rig: Rig) -> None:
|
||||
marker: Final = _marker()
|
||||
response: Final = rig.proxy.client.post(
|
||||
"/v1/chat/completions",
|
||||
json={"model": rig.chat_model, "messages": [{"role": "user", "content": f"{PROVIDER_FAIL_MARK} {marker}"}]},
|
||||
headers=rig.auth(),
|
||||
)
|
||||
assert response.status_code == 500, response.text
|
||||
assert "synthetic provider failure" in response.text
|
||||
|
||||
rows: Final = spend_rows(response.headers["x-litellm-call-id"])
|
||||
assert [row["status"] for row in rows] == ["failure"], rows
|
||||
assert [request for request in rig.akto.drain() if marker.encode() in request.body] == []
|
||||
|
||||
|
||||
async def _burst_call(rig: Rig, kind: str, content: str) -> tuple[str, str]:
|
||||
if kind == "messages":
|
||||
claude: Final = anthropic.AsyncAnthropic(base_url=rig.base(), api_key=rig.proxy.key, max_retries=0)
|
||||
message: Final = await claude.messages.create(
|
||||
model=rig.claude_model, max_tokens=64, messages=[{"role": "user", "content": content}]
|
||||
)
|
||||
assert message.content[0].type == "text"
|
||||
return message.id, message.content[0].text
|
||||
client: Final = openai.AsyncOpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0)
|
||||
if kind == "responses":
|
||||
result: Final = await client.responses.create(model=rig.responses_model, input=content)
|
||||
return result.id, result.output_text
|
||||
if kind == "stream":
|
||||
stream: Final = await client.chat.completions.create(
|
||||
model=rig.chat_model, messages=[{"role": "user", "content": content}], stream=True
|
||||
)
|
||||
chunks: Final = [chunk async for chunk in stream]
|
||||
return chunks[0].id, "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices)
|
||||
completion: Final = await client.chat.completions.create(
|
||||
model=rig.chat_model, messages=[{"role": "user", "content": content}]
|
||||
)
|
||||
return completion.id, completion.choices[0].message.content or ""
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_burst_with_akto_failing_for_half_the_calls_answers_and_logs_every_call_once(rig: Rig) -> None:
|
||||
kinds: Final = ("chat", "stream", "messages", "responses") * 6
|
||||
markers: Final = tuple(_marker() for _ in kinds)
|
||||
failing: Final = frozenset(markers[::2])
|
||||
results: Final = await asyncio.gather(
|
||||
*(
|
||||
_burst_call(rig, kind, (AKTO_DROP_MARK + " " if marker in failing else "") + "burst " + marker)
|
||||
for kind, marker in zip(kinds, markers, strict=True)
|
||||
)
|
||||
)
|
||||
assert [answer for _, answer in results] == [_answer(marker) for marker in markers]
|
||||
|
||||
for response_id, _ in results:
|
||||
assert [row["status"] for row in spend_rows(response_id)] == ["success"], response_id
|
||||
calls: Final = tuple(
|
||||
call for call in (_akto_call(request) for request in rig.akto.drain()) if call.authorization == LOG_KEY
|
||||
)
|
||||
observed: Final = {
|
||||
marker: [call.flags for call in calls if marker in json.dumps(call.payload)] for marker in markers
|
||||
}
|
||||
expected: Final = {
|
||||
marker: [REQUEST_CHECK] if marker in failing else [REQUEST_CHECK, RESPONSE_CHECK] for marker in markers
|
||||
}
|
||||
assert observed == expected
|
||||
assert rig.proxy.client.get("/health/liveliness").status_code == 200
|
||||
|
|
@ -5,6 +5,8 @@ import os
|
|||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
|
@ -20,7 +22,8 @@ from litellm.proxy.guardrails.guardrail_registry import (
|
|||
guardrail_class_registry,
|
||||
guardrail_initializer_registry,
|
||||
)
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs, ModelResponse
|
||||
|
||||
|
||||
def test_akto_in_guardrail_initializer_registry():
|
||||
|
|
@ -1869,3 +1872,121 @@ async def test_a_modified_verdict_that_changed_no_text_blocks(akto_pre_call, sam
|
|||
inputs=sample_inputs, request_data=sample_request_data, input_type="request"
|
||||
)
|
||||
assert exc_info.value.message == "Content masked by Akto guardrail policy could not be applied"
|
||||
|
||||
|
||||
REQUEST_CHECK: Final = {"akto_connector": "litellm", "guardrails": "true", "ingest_data": "true"}
|
||||
RESPONSE_CHECK: Final = {"akto_connector": "litellm", "response_guardrails": "true", "ingest_data": "true"}
|
||||
|
||||
|
||||
def _logging_only_akto(
|
||||
post: AsyncMock, unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed"
|
||||
) -> AktoGuardrail:
|
||||
handler: Final = MagicMock(spec=AsyncHTTPHandler)
|
||||
handler.post = post
|
||||
return AktoGuardrail(
|
||||
async_handler=handler,
|
||||
akto_base_url="http://localhost:9090",
|
||||
akto_api_key="test-token",
|
||||
guardrail_name="test-logging_only",
|
||||
event_hook="logging_only",
|
||||
unreachable_fallback=unreachable_fallback,
|
||||
)
|
||||
|
||||
|
||||
def _logged_call(text: str = "Hello, how are you?") -> dict[str, object]:
|
||||
return {
|
||||
"model": "gpt-5.5",
|
||||
"messages": [{"role": "user", "content": text}],
|
||||
"litellm_call_id": "call-1",
|
||||
"litellm_params": {"metadata": {"user_api_key_request_route": "/v1/chat/completions"}},
|
||||
"standard_logging_object": {"guardrail_information": []},
|
||||
}
|
||||
|
||||
|
||||
def _logged_response(text: str = "Fine, thanks") -> ModelResponse:
|
||||
return ModelResponse(id="resp-1", choices=[{"message": {"role": "assistant", "content": text}}])
|
||||
|
||||
|
||||
def _recorded_entries(logged_kwargs: dict[str, object]) -> list[dict[str, object]]:
|
||||
standard_logging_object: Final = logged_kwargs["standard_logging_object"]
|
||||
assert isinstance(standard_logging_object, dict), logged_kwargs
|
||||
entries: Final = standard_logging_object["guardrail_information"]
|
||||
assert isinstance(entries, list), standard_logging_object
|
||||
return entries
|
||||
|
||||
|
||||
def test_logging_only_is_a_supported_mode() -> None:
|
||||
assert GuardrailEventHooks.logging_only in AktoGuardrail.get_supported_event_hooks()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("input_type", "flags"), [("request", REQUEST_CHECK), ("response", RESPONSE_CHECK)])
|
||||
async def test_logging_only_handles_both_directions(
|
||||
input_type: Literal["request", "response"], flags: dict[str, str]
|
||||
) -> None:
|
||||
guardrail: Final = _logging_only_akto(AsyncMock(return_value=_mock_allowed_response()))
|
||||
request_data: Final = _with_complete_response({}) if input_type == "response" else {}
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=GenericGuardrailAPIInputs(texts=["hi"]), request_data=request_data, input_type=input_type
|
||||
)
|
||||
|
||||
assert [params for params, _ in _calls(guardrail)] == [flags]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logging_only_checks_and_records_the_logged_request_and_response() -> None:
|
||||
guardrail: Final = _logging_only_akto(AsyncMock(return_value=_mock_allowed_response()))
|
||||
|
||||
await guardrail.async_logging_hook(_logged_call(), _logged_response(), "acompletion")
|
||||
|
||||
sent: Final = _calls(guardrail)
|
||||
assert [params for params, _ in sent] == [REQUEST_CHECK, RESPONSE_CHECK]
|
||||
assert "Hello, how are you?" in sent[0][1]["requestPayload"]
|
||||
assert "Fine, thanks" in sent[1][1]["responsePayload"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logging_only_block_verdict_is_recorded_without_raising() -> None:
|
||||
guardrail: Final = _logging_only_akto(AsyncMock(return_value=_mock_blocked_response("Rejected")))
|
||||
response: Final = _logged_response()
|
||||
|
||||
out_kwargs, out_result = await guardrail.async_logging_hook(_logged_call(), response, "acompletion")
|
||||
|
||||
assert out_result is response
|
||||
assert [params for params, _ in _calls(guardrail)] == [REQUEST_CHECK]
|
||||
[entry] = _recorded_entries(out_kwargs)
|
||||
assert (entry["guardrail_name"], entry["guardrail_mode"], entry["guardrail_status"]) == (
|
||||
"test-logging_only",
|
||||
"logging_only",
|
||||
"guardrail_intervened",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logging_only_ignores_an_unreachable_akto_even_when_fail_closed() -> None:
|
||||
guardrail: Final = _logging_only_akto(AsyncMock(side_effect=httpx.ConnectError("refused")), "fail_closed")
|
||||
response: Final = _logged_response()
|
||||
|
||||
out_kwargs, out_result = await guardrail.async_logging_hook(_logged_call(), response, "acompletion")
|
||||
|
||||
assert out_result is response
|
||||
[entry] = _recorded_entries(out_kwargs)
|
||||
assert (entry["guardrail_mode"], entry["guardrail_response"]) == (
|
||||
"logging_only",
|
||||
"Akto guardrail service unreachable",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_logging_only_sends_each_attachment_once() -> None:
|
||||
guardrail: Final = _logging_only_akto(_file_verdict({"Allowed": True}))
|
||||
image: Final = {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}
|
||||
call: Final = {**_logged_call(), "messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}, image]}]}
|
||||
|
||||
await guardrail.async_logging_hook(call, _logged_response(), "acompletion")
|
||||
|
||||
[file_call] = _file_calls(guardrail)
|
||||
assert json.loads(file_call.kwargs["data"])["files"] == [
|
||||
{"filename": "a.png", "type": "image", "url": "https://example.com/a.png"}
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import base64
|
||||
import json
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -471,3 +472,13 @@ def test_text_that_isnt_valid_utf8_is_still_sent(block):
|
|||
|
||||
[attachment] = request_attachments(request_data).attachments
|
||||
assert base64.b64decode(attachment.content or "") == "a\ud800".encode(errors="surrogatepass")
|
||||
|
||||
|
||||
def test_request_attachments_reads_a_list_shared_by_messages_and_input_once() -> None:
|
||||
shared: Final = [
|
||||
{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}]}
|
||||
]
|
||||
|
||||
assert request_attachments({"messages": shared, "input": shared}).attachments == (
|
||||
Attachment("a.png", "image", url="https://example.com/a.png"),
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue