mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
application name over application ID
This commit is contained in:
Shalom Jamil
parent
af9436f44a
commit
e9e1fc6bcb
2 changed files with 227 additions and 35 deletions
|
|
@ -89,6 +89,21 @@ class ResolvedRouting(NamedTuple):
|
|||
)
|
||||
|
||||
|
||||
class CheckpointTarget(NamedTuple):
|
||||
"""How a checkpoint call addresses the application.
|
||||
|
||||
When application_name is set the call sends the name and the direction, and the tracker resolves
|
||||
and evaluates in one request; resolution can create the application, and ids from a separate
|
||||
resolve call may name one the tracker's config has not caught up with, which it reports as an
|
||||
uninspected allow rather than an error. application_id is what the session id groups on, and is
|
||||
what gets sent when the deployment pins the application in config and reads no alias.
|
||||
"""
|
||||
|
||||
application_id: str
|
||||
input_type: str
|
||||
application_name: str | None = None
|
||||
|
||||
|
||||
def _coerce_bool(value: bool | str) -> bool:
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
|
|
@ -253,27 +268,35 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
checkpoint_id: str,
|
||||
actor: str,
|
||||
session_id: str,
|
||||
application_id: str,
|
||||
target: CheckpointTarget,
|
||||
) -> Mapping[str, Any]:
|
||||
"""Call the Ovalix Tracker checkpoint API and return the JSON response."""
|
||||
if not application_id or not checkpoint_id:
|
||||
"""Call the Ovalix Tracker checkpoint API and return the JSON response.
|
||||
|
||||
Both routes live on the tracker's /beta litellm router, which accepts the api key this
|
||||
guardrail already sends. The name and id routing forms are mutually exclusive, so exactly one
|
||||
reaches the wire; the name form also sends the direction, which is what the tracker selects
|
||||
the pre or post checkpoint by.
|
||||
"""
|
||||
if not target.application_name and (not target.application_id or not checkpoint_id):
|
||||
raise ValueError("Ovalix: application_id or checkpoint_id not resolved")
|
||||
|
||||
url = (
|
||||
f"{self._tracker_api_base}/tracking/beta/file_checkpoint"
|
||||
if data_type == "FILE"
|
||||
else f"{self._tracker_api_base}/tracking/custom_application/checkpoint"
|
||||
route = "file_checkpoint" if data_type == "FILE" else "checkpoint"
|
||||
routing = (
|
||||
{"application_name": target.application_name, "input_type": target.input_type}
|
||||
if target.application_name
|
||||
else {"application_id": target.application_id, "checkpoint_id": checkpoint_id}
|
||||
)
|
||||
payload = {
|
||||
"application_id": application_id,
|
||||
"checkpoint_id": checkpoint_id,
|
||||
"actor": actor,
|
||||
"session_id": session_id,
|
||||
"data_type": data_type,
|
||||
"data": data,
|
||||
"tool": "LiteLLM",
|
||||
**routing,
|
||||
}
|
||||
response = await self._async_handler.post(url, headers=self._tracker_headers, json=payload)
|
||||
response = await self._async_handler.post(
|
||||
f"{self._tracker_api_base}/tracking/beta/{route}", headers=self._tracker_headers, json=payload
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
|
@ -287,11 +310,11 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
checkpoint_id: str,
|
||||
actor: str,
|
||||
session_id: str,
|
||||
application_id: str,
|
||||
target: CheckpointTarget,
|
||||
escalation_reason: str,
|
||||
) -> str | None:
|
||||
try:
|
||||
resp = await self._call_checkpoint(data_type, data, checkpoint_id, actor, session_id, application_id)
|
||||
resp = await self._call_checkpoint(data_type, data, checkpoint_id, actor, session_id, target)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Ovalix checkpoint call failed: %s", e)
|
||||
raise GuardrailRaisedException(
|
||||
|
|
@ -312,12 +335,12 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
checkpoint_id: str,
|
||||
actor: str,
|
||||
session_id: str,
|
||||
application_id: str,
|
||||
target: CheckpointTarget,
|
||||
escalation_reason: str,
|
||||
) -> str | None:
|
||||
for data_type, data in items:
|
||||
reason = await self._block_reason_for_item(
|
||||
data_type, data, checkpoint_id, actor, session_id, application_id, escalation_reason
|
||||
data_type, data, checkpoint_id, actor, session_id, target, escalation_reason
|
||||
)
|
||||
if reason is not None:
|
||||
return reason
|
||||
|
|
@ -329,12 +352,12 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
checkpoint_id: str,
|
||||
actor: str,
|
||||
session_id: str,
|
||||
application_id: str,
|
||||
target: CheckpointTarget,
|
||||
) -> str | None:
|
||||
for part in sorted(file_parts, key=lambda p: p.message_index, reverse=True):
|
||||
data = await self._file_part_to_data(part)
|
||||
reason = await self._block_reason_for_item(
|
||||
"FILE", data, checkpoint_id, actor, session_id, application_id, _FILE_BLOCK_ESCALATION_REASON
|
||||
"FILE", data, checkpoint_id, actor, session_id, target, _FILE_BLOCK_ESCALATION_REASON
|
||||
)
|
||||
if reason is not None:
|
||||
return reason
|
||||
|
|
@ -354,6 +377,11 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
actor = self._get_actor(request_data)
|
||||
session_id = self._get_session_id_for_application(request_data, routing.application_id)
|
||||
is_response = input_type == "response"
|
||||
target = CheckpointTarget(
|
||||
application_id=routing.application_id,
|
||||
input_type=input_type,
|
||||
application_name=await self._checkpoint_routing_name(request_data),
|
||||
)
|
||||
|
||||
prompt_checkpoint = routing.checkpoint_id_post if is_response else routing.checkpoint_id_pre
|
||||
file_checkpoint = (
|
||||
|
|
@ -379,9 +407,7 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
if is_response
|
||||
else extract_file_parts_from_messages(structured_messages, size_limit=_DEFAULT_FILE_SIZE_LIMIT)
|
||||
)
|
||||
file_block = await self._check_files_for_block(
|
||||
file_parts, file_checkpoint, actor, session_id, routing.application_id
|
||||
)
|
||||
file_block = await self._check_files_for_block(file_parts, file_checkpoint, actor, session_id, target)
|
||||
if file_block is not None:
|
||||
self._block_current_message(file_block)
|
||||
|
||||
|
|
@ -401,7 +427,7 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
prompt_checkpoint,
|
||||
actor,
|
||||
session_id,
|
||||
routing.application_id,
|
||||
target,
|
||||
_TOOL_BLOCK_ESCALATION_REASON,
|
||||
)
|
||||
if tool_block is not None:
|
||||
|
|
@ -414,7 +440,7 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
prompt_checkpoint,
|
||||
actor,
|
||||
session_id,
|
||||
routing.application_id,
|
||||
target,
|
||||
_TOOL_RESULT_BLOCK_ESCALATION_REASON,
|
||||
)
|
||||
if tool_result_block is not None:
|
||||
|
|
@ -428,7 +454,7 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
prompt_checkpoint,
|
||||
actor,
|
||||
session_id,
|
||||
routing.application_id,
|
||||
target,
|
||||
tool_result_text_indices(structured_messages, texts),
|
||||
)
|
||||
if output_texts is None:
|
||||
|
|
@ -451,7 +477,7 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
checkpoint_id: str,
|
||||
actor: str,
|
||||
session_id: str,
|
||||
application_id: str,
|
||||
target: CheckpointTarget,
|
||||
skip_indices: frozenset[int],
|
||||
) -> list[str] | None:
|
||||
output = list(texts)
|
||||
|
|
@ -465,7 +491,7 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
content = texts[original_index]
|
||||
try:
|
||||
resp = await self._call_checkpoint(
|
||||
"TEXT", {"content": content}, checkpoint_id, actor, session_id, application_id
|
||||
"TEXT", {"content": content}, checkpoint_id, actor, session_id, target
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Ovalix checkpoint call failed: %s", e)
|
||||
|
|
@ -580,6 +606,19 @@ class OvalixGuardrail(CustomGuardrail):
|
|||
should_wrap_with_default_message=False,
|
||||
)
|
||||
|
||||
async def _checkpoint_routing_name(self, request_data: Mapping[str, Any]) -> str | None:
|
||||
"""The application name to route checkpoints by, or None to route by resolved ids.
|
||||
|
||||
None when the deployment pins an application in config, or when no name can be read from the
|
||||
api key alias. Only reached after _resolve_routing has already fetched and cached the regex.
|
||||
"""
|
||||
if self._application_id:
|
||||
return None
|
||||
alias = self._get_key_alias(request_data)
|
||||
if not alias:
|
||||
return None
|
||||
return self._extract_application_name(alias, await self._get_app_name_regex())
|
||||
|
||||
async def _resolve_routing(self, request_data: Mapping[str, Any]) -> ResolvedRouting | None:
|
||||
if self._application_id:
|
||||
return ResolvedRouting(
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import base64
|
|||
import gzip
|
||||
import json as json_lib
|
||||
import os
|
||||
from typing import Any, List
|
||||
import re
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -15,6 +15,7 @@ import pytest
|
|||
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm.proxy.guardrails.guardrail_hooks.ovalix.ovalix import (
|
||||
CheckpointTarget,
|
||||
OvalixGuardrail,
|
||||
OvalixGuardrailBlockedException,
|
||||
OvalixGuardrailMissingSecrets,
|
||||
|
|
@ -276,13 +277,13 @@ class TestOvalixGuardrail:
|
|||
checkpoint_id="pre-1",
|
||||
actor="a1b2c3d4",
|
||||
session_id="session-1",
|
||||
application_id="app-1",
|
||||
target=CheckpointTarget("app-1", "request"),
|
||||
)
|
||||
|
||||
assert result == TRACKER_RESPONSE_ALLOW
|
||||
mock_post.assert_called_once()
|
||||
call_args = mock_post.call_args
|
||||
assert call_args.args[0] == ("https://tracker.test/tracking/custom_application/checkpoint")
|
||||
assert call_args.args[0] == ("https://tracker.test/tracking/beta/checkpoint")
|
||||
body = call_args.kwargs["json"]
|
||||
assert body["application_id"] == "app-1"
|
||||
assert body["checkpoint_id"] == "pre-1"
|
||||
|
|
@ -746,15 +747,23 @@ def _routing_body(pre, post, pre_file, post_file, application_id="app-9"):
|
|||
}
|
||||
|
||||
|
||||
def _checkpoint_bodies(mock_post):
|
||||
"""Bodies of the tracker checkpoint calls only, excluding regex/resolve traffic."""
|
||||
def _checkpoint_calls(mock_post):
|
||||
"""(body, url) of the tracker checkpoint calls only, excluding regex/resolve traffic."""
|
||||
return [
|
||||
c.kwargs["json"]
|
||||
(c.kwargs["json"], c.args[0])
|
||||
for c in mock_post.call_args_list
|
||||
if c.args and c.args[0].endswith(("/checkpoint", "/file_checkpoint"))
|
||||
]
|
||||
|
||||
|
||||
def _checkpoint_bodies(mock_post):
|
||||
return [body for body, _ in _checkpoint_calls(mock_post)]
|
||||
|
||||
|
||||
def _route_of(url):
|
||||
return url.rsplit("/", 1)[-1]
|
||||
|
||||
|
||||
def _mock_handler(g, routing=None):
|
||||
get_resp = MagicMock()
|
||||
get_resp.json.return_value = {"regex": _REGEX}
|
||||
|
|
@ -979,7 +988,9 @@ async def test_file_checkpoint_call_routes_to_litellm_file_endpoint():
|
|||
return r
|
||||
|
||||
with patch.object(g._async_handler, "post", new=_post):
|
||||
await g._call_checkpoint("FILE", {"name": "f.txt", "content": "x"}, "file-1", "a", "s", "app-1")
|
||||
await g._call_checkpoint(
|
||||
"FILE", {"name": "f.txt", "content": "x"}, "file-1", "a", "s", CheckpointTarget("app-1", "request")
|
||||
)
|
||||
assert seen["url"] == "https://t/tracking/beta/file_checkpoint"
|
||||
|
||||
|
||||
|
|
@ -1251,7 +1262,7 @@ def test_enable_routing_cache_from_env_string(monkeypatch):
|
|||
async def test_call_checkpoint_requires_application_and_checkpoint():
|
||||
g = _static_guardrail()
|
||||
with pytest.raises(ValueError):
|
||||
await g._call_checkpoint("TEXT", {"content": "x"}, "", "actor", "sess", "app-1")
|
||||
await g._call_checkpoint("TEXT", {"content": "x"}, "", "actor", "sess", CheckpointTarget("app-1", "request"))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1283,7 +1294,7 @@ async def test_discovery_resolved_without_any_checkpoint_raises():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("input_type, inspected", [("request", ["pre-9"]), ("response", [])])
|
||||
@pytest.mark.parametrize("input_type, inspected", [("request", ["request"]), ("response", [])])
|
||||
async def test_one_sided_discovery_inspects_configured_direction_only(input_type, inspected):
|
||||
"""Discovery registers both hooks speculatively, so a direction the app does not inspect must pass through."""
|
||||
g = _discovery_guardrail(enable_cache=False)
|
||||
|
|
@ -1293,7 +1304,7 @@ async def test_one_sided_discovery_inspects_configured_direction_only(input_type
|
|||
inputs=inputs, request_data=_alias_request_data(), input_type=input_type, logging_obj=None
|
||||
)
|
||||
assert result["texts"] == ["hi"]
|
||||
assert [b["checkpoint_id"] for b in _checkpoint_bodies(mock_post)] == inspected
|
||||
assert [b["input_type"] for b in _checkpoint_bodies(mock_post)] == inspected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1312,7 +1323,7 @@ async def test_file_only_checkpoint_inspects_files_and_skips_text():
|
|||
inputs=inputs, request_data=_alias_request_data(), input_type="request", logging_obj=None
|
||||
)
|
||||
assert result["texts"] == ["hi"]
|
||||
assert [(b["data_type"], b["checkpoint_id"]) for b in _checkpoint_bodies(mock_post)] == [("FILE", "pre-file-9")]
|
||||
assert [(b["data_type"], _route_of(c)) for b, c in _checkpoint_calls(mock_post)] == [("FILE", "file_checkpoint")]
|
||||
|
||||
|
||||
def test_initialize_guardrail_wires_new_params(monkeypatch):
|
||||
|
|
@ -1592,3 +1603,145 @@ async def test_cached_404_still_raises_when_failing_closed(monkeypatch):
|
|||
with pytest.raises(GuardrailRaisedException):
|
||||
await g._resolve_routing(_alias_request_data("[Ghost App] prod"))
|
||||
assert mock_post.call_count == 1
|
||||
|
||||
|
||||
def _alias_guardrail():
|
||||
return OvalixGuardrail(
|
||||
tracker_api_base="https://t",
|
||||
tracker_api_key="k",
|
||||
guardrail_name="o",
|
||||
event_hook="pre_call",
|
||||
default_on=True,
|
||||
)
|
||||
|
||||
|
||||
def _capturing_post(response=None):
|
||||
seen = {}
|
||||
|
||||
async def _post(url, headers=None, json=None):
|
||||
seen["url"] = url
|
||||
seen["body"] = json
|
||||
r = MagicMock()
|
||||
r.json.return_value = response if response is not None else _ALLOW
|
||||
r.raise_for_status = MagicMock()
|
||||
return r
|
||||
|
||||
return seen, _post
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"data_type,expected_route",
|
||||
[("TEXT", "checkpoint"), ("TOOL", "checkpoint"), ("FILE", "file_checkpoint")],
|
||||
)
|
||||
async def test_call_checkpoint_by_name_sends_name_and_direction_not_ids(data_type, expected_route):
|
||||
"""Name routing must put application_name and input_type on the wire and no ids at all.
|
||||
|
||||
The two forms are mutually exclusive server-side, so leaking an id alongside the name is rejected
|
||||
"""
|
||||
g = _alias_guardrail()
|
||||
seen, post = _capturing_post()
|
||||
|
||||
with patch.object(g._async_handler, "post", new=post):
|
||||
await g._call_checkpoint(
|
||||
data_type,
|
||||
{"content": "x"},
|
||||
"unused-cp",
|
||||
"actor",
|
||||
"sess",
|
||||
CheckpointTarget("app-9", "response", application_name="Weather App"),
|
||||
)
|
||||
|
||||
assert seen["url"] == f"https://t/tracking/beta/{expected_route}"
|
||||
assert seen["body"]["application_name"] == "Weather App"
|
||||
assert seen["body"]["input_type"] == "response"
|
||||
assert "application_id" not in seen["body"]
|
||||
assert "checkpoint_id" not in seen["body"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_checkpoint_by_ids_sends_ids_and_no_name_or_direction():
|
||||
g = _static_guardrail()
|
||||
seen, post = _capturing_post()
|
||||
|
||||
with patch.object(g._async_handler, "post", new=post):
|
||||
await g._call_checkpoint(
|
||||
"TEXT", {"content": "x"}, "pre-1", "actor", "sess", CheckpointTarget("app-1", "request")
|
||||
)
|
||||
|
||||
assert seen["body"]["application_id"] == "app-1"
|
||||
assert seen["body"]["checkpoint_id"] == "pre-1"
|
||||
assert "application_name" not in seen["body"]
|
||||
assert "input_type" not in seen["body"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_checkpoint_by_name_does_not_require_a_checkpoint_id():
|
||||
"""The tracker chooses the checkpoint under name routing, so an empty id must not be rejected."""
|
||||
g = _alias_guardrail()
|
||||
seen, post = _capturing_post()
|
||||
|
||||
with patch.object(g._async_handler, "post", new=post):
|
||||
await g._call_checkpoint(
|
||||
"TEXT",
|
||||
{"content": "x"},
|
||||
"",
|
||||
"actor",
|
||||
"sess",
|
||||
CheckpointTarget("", "request", application_name="Weather App"),
|
||||
)
|
||||
|
||||
assert seen["body"]["application_name"] == "Weather App"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoint_routing_name_is_none_when_application_is_pinned():
|
||||
"""A deployment that pins application_id reads no alias, so it must keep routing by ids."""
|
||||
g = _static_guardrail()
|
||||
|
||||
assert await g._checkpoint_routing_name({"metadata": {"user_api_key_alias": "[Weather App] k"}}) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoint_routing_name_extracted_from_key_alias():
|
||||
g = _alias_guardrail()
|
||||
with patch.object(g, "_get_app_name_regex", new=AsyncMock(return_value=re.compile(r"^\s*\[([^\]]+)\]"))):
|
||||
name = await g._checkpoint_routing_name({"metadata": {"user_api_key_alias": "[Weather App] free text"}})
|
||||
|
||||
assert name == "Weather App"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoint_routing_name_is_none_without_an_alias():
|
||||
g = _alias_guardrail()
|
||||
|
||||
assert await g._checkpoint_routing_name({"metadata": {}}) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_resolved_by_alias_routes_checkpoints_by_name():
|
||||
"""End to end: an alias-resolved application sends its name per checkpoint, never the resolved id.
|
||||
|
||||
This is what keeps a just-created application inspectable: the ids from resolution may name an
|
||||
application the tracker's process-wide config has not picked up yet
|
||||
"""
|
||||
g = _alias_guardrail()
|
||||
seen, post = _capturing_post()
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
g, "_resolve_routing", new=AsyncMock(return_value=ResolvedRouting("app-9", "pre-9", "post-9", None, None))
|
||||
),
|
||||
patch.object(g, "_get_app_name_regex", new=AsyncMock(return_value=re.compile(r"^\s*\[([^\]]+)\]"))),
|
||||
patch.object(g._async_handler, "post", new=post),
|
||||
):
|
||||
await g.apply_guardrail(
|
||||
GenericGuardrailAPIInputs(texts=["hello"]),
|
||||
{"metadata": {"user_api_key_alias": "[Weather App] k", "user_api_key_user_email": "u@e.com"}},
|
||||
"request",
|
||||
)
|
||||
|
||||
assert seen["url"] == "https://t/tracking/beta/checkpoint"
|
||||
assert seen["body"]["application_name"] == "Weather App"
|
||||
assert seen["body"]["input_type"] == "request"
|
||||
assert "application_id" not in seen["body"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue