mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
c6ca22ebc5
* fix(guardrails): make hide-secrets honor skip_tool_message_in_guardrail With the flag on, per guardrail or in litellm_settings, hide-secrets leaves role: tool messages and Anthropic tool_result blocks as sent, as the unified guardrails do. Tool call arguments are still masked. * test(guardrails): type the skip-tool walker test and assert the whole request
483 lines
19 KiB
Python
483 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import random
|
|
import string
|
|
import uuid
|
|
from collections.abc import Callable, Iterator, Mapping
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import httpx
|
|
import pytest
|
|
import yaml
|
|
from integration._support.budget import case_timeout_seconds
|
|
from integration._support.client import Gateway
|
|
from integration._support.database import scratch_database
|
|
from integration._support.process import OwnedProxy, owned_proxy_process
|
|
from integration._support.timing import assert_not_delayed_by, assert_within_factor_of_control, timed
|
|
from integration._support.wire import Reply, Request, Wire, wire_server
|
|
from pydantic import JsonValue
|
|
|
|
SECRET: Final = "AKIAIOSFODNN7EXAMPLE"
|
|
TEXT: Final = f"AWS_ACCESS_KEY_ID={SECRET}"
|
|
MASTER_KEY: Final = "sk-integration-hide-secrets-7426-synthetic-master"
|
|
|
|
|
|
def _provider(request: Request) -> Reply:
|
|
if request.method != "POST":
|
|
return Reply()
|
|
body: Final = json.loads(request.body)
|
|
identity: Final = uuid.uuid4().hex
|
|
if request.target.endswith("/messages"):
|
|
response: Final = {
|
|
"id": f"msg_{identity}",
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"model": body["model"],
|
|
"content": [{"type": "text", "text": "accepted"}],
|
|
"stop_reason": "end_turn",
|
|
"stop_sequence": None,
|
|
"usage": {"input_tokens": 10, "output_tokens": 1},
|
|
}
|
|
if body.get("stream"):
|
|
events: Final = (
|
|
("message_start", {"type": "message_start", "message": {**response, "content": []}}),
|
|
(
|
|
"content_block_start",
|
|
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
|
),
|
|
(
|
|
"content_block_delta",
|
|
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "accepted"}},
|
|
),
|
|
("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
|
(
|
|
"message_delta",
|
|
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}},
|
|
),
|
|
("message_stop", {"type": "message_stop"}),
|
|
)
|
|
return Reply(
|
|
content_type="text/event-stream",
|
|
chunks=tuple(f"event: {name}\ndata: {json.dumps(value)}\n\n".encode() for name, value in events),
|
|
)
|
|
return Reply(body=json.dumps(response).encode())
|
|
return Reply(
|
|
body=json.dumps(
|
|
{
|
|
"id": f"chatcmpl_{identity}",
|
|
"object": "chat.completion",
|
|
"created": 1,
|
|
"model": body["model"],
|
|
"choices": [
|
|
{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "accepted"}}
|
|
],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
|
|
}
|
|
).encode()
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Rig:
|
|
proxy: Gateway
|
|
wire: Wire
|
|
database_url: str
|
|
owned: OwnedProxy
|
|
environment: Mapping[str, str]
|
|
|
|
|
|
@contextmanager
|
|
def tool_rig(
|
|
directory: Path,
|
|
respond: Callable[[Request], Reply] = _provider,
|
|
*,
|
|
configure: Callable[[dict[str, JsonValue]], None] | None = None,
|
|
environment: Mapping[str, str] | None = None,
|
|
workers: int = 2,
|
|
) -> Iterator[Rig]:
|
|
assert os.environ.get("LITELLM_LICENSE"), "LITELLM_LICENSE must be set: hide-secrets is an enterprise feature"
|
|
product_root: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3])
|
|
with (
|
|
wire_server(respond) as wire,
|
|
scratch_database("hide_secrets") as database,
|
|
):
|
|
config: Final = directory / "config.yaml"
|
|
configuration: Final = {
|
|
"model_list": [
|
|
{
|
|
"model_name": "claude",
|
|
"litellm_params": {
|
|
"model": "anthropic/claude-sonnet-5",
|
|
"api_key": "synthetic-provider-key",
|
|
"api_base": wire.url,
|
|
"num_retries": 0,
|
|
},
|
|
},
|
|
{
|
|
"model_name": "gpt",
|
|
"litellm_params": {
|
|
"model": "openai/gpt-5",
|
|
"api_key": "synthetic-provider-key",
|
|
"api_base": wire.url,
|
|
"num_retries": 0,
|
|
},
|
|
},
|
|
],
|
|
"general_settings": {
|
|
"master_key": MASTER_KEY,
|
|
"store_prompts_in_spend_logs": True,
|
|
"disable_spend_logs": False,
|
|
"proxy_batch_write_at": 1,
|
|
"proxy_batch_polling_interval": 1,
|
|
},
|
|
"router_settings": {"num_retries": 0},
|
|
"litellm_settings": {
|
|
"cache": True,
|
|
"cache_params": {
|
|
"type": "redis",
|
|
"host": "os.environ/REDIS_HOST",
|
|
"port": "os.environ/REDIS_PORT",
|
|
},
|
|
},
|
|
"guardrails": [
|
|
{
|
|
"guardrail_name": "hide-secrets",
|
|
"litellm_params": {"guardrail": "hide-secrets", "mode": "pre_call", "default_on": True},
|
|
}
|
|
],
|
|
}
|
|
if configure is not None:
|
|
configure(configuration)
|
|
config.write_text(yaml.safe_dump(configuration))
|
|
child_environment: Final = {
|
|
"DATABASE_URL": database,
|
|
"LITELLM_DISABLE_NO_REDIS_WARNING": "true",
|
|
"ANTHROPIC_API_BASE": wire.url,
|
|
"ANTHROPIC_API_KEY": "synthetic-provider-key",
|
|
"LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS": "True",
|
|
"PYTHONPATH": os.pathsep.join(
|
|
(str(product_root / "enterprise"), str(product_root), os.environ.get("PYTHONPATH", ""))
|
|
),
|
|
**(environment or {}),
|
|
}
|
|
with httpx.Client() as client:
|
|
seed: Final = Gateway(client, MASTER_KEY, wire.url)
|
|
with owned_proxy_process(
|
|
seed,
|
|
directory,
|
|
child_environment,
|
|
config=config,
|
|
workers=workers,
|
|
client_timeout=case_timeout_seconds(),
|
|
) as owned:
|
|
yield Rig(owned.gateway, wire, database, owned, child_environment)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
|
with tool_rig(tmp_path_factory.mktemp("hide-secrets-tools")) as value:
|
|
yield value
|
|
|
|
|
|
def _input() -> dict[str, JsonValue]:
|
|
return {
|
|
"env": TEXT,
|
|
"path": 'C:\\tmp\\"config"',
|
|
"count": 3,
|
|
"enabled": False,
|
|
"optional": None,
|
|
"nested": ["safe", {"value": TEXT}],
|
|
}
|
|
|
|
|
|
def _messages(shape: str) -> dict[str, JsonValue]:
|
|
user: Final[dict[str, JsonValue]] = {"role": "user", "content": "read .env"}
|
|
tool: Final[dict[str, JsonValue]] = {
|
|
"type": "tool_use",
|
|
"id": "tool_1",
|
|
"name": "read_file",
|
|
"input": {"path": ".env"},
|
|
}
|
|
prefix: Final[list[JsonValue]] = [user, {"role": "assistant", "content": [tool]}]
|
|
if shape == "plain":
|
|
return {"model": "claude", "max_tokens": 32, "messages": [{"role": "user", "content": TEXT}]}
|
|
if shape.startswith("result"):
|
|
content: Final[JsonValue] = (
|
|
TEXT
|
|
if shape == "result-string"
|
|
else [
|
|
{"type": "text", "text": TEXT},
|
|
{"type": "text", "text": "keep me"},
|
|
]
|
|
)
|
|
return {
|
|
"model": "claude",
|
|
"max_tokens": 32,
|
|
"messages": [
|
|
*prefix,
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "tool_result",
|
|
"tool_use_id": "tool_1",
|
|
"content": content,
|
|
"cache_control": {"type": "ephemeral"},
|
|
},
|
|
{"type": "text", "text": "summarize"},
|
|
],
|
|
},
|
|
],
|
|
}
|
|
if shape.startswith("input"):
|
|
arguments: Final = (
|
|
_input()
|
|
if shape == "input-dict"
|
|
else {"command": "export API_TOKEN=0123456789abcdef0123456789abcdef01234567", "enabled": False}
|
|
)
|
|
return {
|
|
"model": "claude",
|
|
"max_tokens": 32,
|
|
"messages": [
|
|
user,
|
|
{"role": "assistant", "content": [{**tool, "input": arguments}]},
|
|
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "tool_1", "content": "done"}]},
|
|
],
|
|
}
|
|
return {
|
|
"model": "claude",
|
|
"max_tokens": 32,
|
|
"messages": [user],
|
|
"system": TEXT
|
|
if shape == "system-string"
|
|
else [{"type": "text", "text": TEXT, "cache_control": {"type": "ephemeral"}}],
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"shape",
|
|
("plain", "result-string", "result-nested", "input-dict", "input-assignment", "system-string", "system-blocks"),
|
|
)
|
|
@pytest.mark.parametrize("stream", (False, True))
|
|
def test_messages_secret_never_reaches_provider(rig: Rig, shape: str, stream: bool) -> None:
|
|
rig.wire.drain()
|
|
request: Final = {**_messages(shape), "stream": stream, "cache": {"no-cache": True}}
|
|
response: Final = rig.proxy.request("POST", "/v1/messages", request)
|
|
assert response.status_code == 200, response.text
|
|
assert "accepted" in response.text
|
|
sent: Final = rig.wire.drain()
|
|
assert len(sent) == 1
|
|
assert SECRET.encode() not in sent[0].body
|
|
assert b"0123456789abcdef0123456789abcdef01234567" not in sent[0].body
|
|
assert b"REDACTED" in sent[0].body
|
|
observed: Final = json.loads(sent[0].body)
|
|
if shape == "input-dict":
|
|
value: Final = observed["messages"][1]["content"][0]["input"]
|
|
assert value.keys() == _input().keys()
|
|
assert value["path"] == _input()["path"]
|
|
assert value["count"] == 3
|
|
assert value["enabled"] is False
|
|
assert value["optional"] is None
|
|
|
|
|
|
def test_chat_tool_arguments_remain_valid_json(rig: Rig) -> None:
|
|
rig.wire.drain()
|
|
response: Final = rig.proxy.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": "gpt",
|
|
"cache": {"no-cache": True},
|
|
"messages": [
|
|
{"role": "user", "content": "read .env"},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_1",
|
|
"type": "function",
|
|
"function": {"name": "read_file", "arguments": json.dumps(_input())},
|
|
}
|
|
],
|
|
},
|
|
{"role": "tool", "tool_call_id": "call_1", "content": "done"},
|
|
],
|
|
},
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
assert response.json()["choices"][0]["message"]["content"] == "accepted"
|
|
sent: Final = rig.wire.drain()
|
|
assert len(sent) == 1
|
|
assert SECRET.encode() not in sent[0].body
|
|
arguments: Final = json.loads(json.loads(sent[0].body)["messages"][1]["tool_calls"][0]["function"]["arguments"])
|
|
assert arguments.keys() == _input().keys()
|
|
assert arguments["count"] == 3
|
|
assert arguments["path"] == _input()["path"]
|
|
assert arguments["enabled"] is False
|
|
assert arguments["optional"] is None
|
|
|
|
|
|
def test_deeply_nested_tool_arguments_do_not_stall_the_worker(rig: Rig) -> None:
|
|
def call(depth: int) -> float:
|
|
rig.wire.drain()
|
|
arguments: Final = '{"k":' * depth + json.dumps(TEXT) + "}" * depth
|
|
sent: Final = timed(
|
|
lambda: rig.proxy.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": "gpt",
|
|
"cache": {"no-cache": True},
|
|
"messages": [
|
|
{"role": "user", "content": "read .env"},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_1",
|
|
"type": "function",
|
|
"function": {"name": "read", "arguments": arguments},
|
|
}
|
|
],
|
|
},
|
|
{"role": "tool", "tool_call_id": "call_1", "content": "done"},
|
|
],
|
|
},
|
|
)
|
|
)
|
|
assert sent.value.status_code == 200, sent.value.text
|
|
upstream: Final = rig.wire.drain()
|
|
assert len(upstream) == 1
|
|
assert SECRET.encode() not in upstream[0].body
|
|
return sent.seconds
|
|
|
|
control: Final = tuple(call(1) for _ in range(3))
|
|
assert_not_delayed_by((call(300),), control=control, injected_delay=5.0)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def single_worker(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
|
|
with tool_rig(tmp_path_factory.mktemp("hide-secrets-one-worker"), workers=1) as value:
|
|
yield value
|
|
|
|
|
|
def _tool_history(turns: int, seed: int, shape: str = "anthropic") -> tuple[str, dict[str, JsonValue]]:
|
|
generator: Final = random.Random(seed)
|
|
alphabet: Final = string.ascii_letters + string.digits
|
|
messages: Final[list[JsonValue]] = [{"role": "user", "content": "read the project"}]
|
|
for turn in range(turns):
|
|
result = "\n".join("".join(generator.choices(alphabet, k=60)) for _ in range(64))
|
|
if shape == "anthropic":
|
|
call: JsonValue = {
|
|
"type": "tool_use",
|
|
"id": f"tool_{turn}",
|
|
"name": "read",
|
|
"input": {"path": f"{turn}.py"},
|
|
}
|
|
messages.append({"role": "assistant", "content": [call]})
|
|
messages.append(
|
|
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": f"tool_{turn}", "content": result}]}
|
|
)
|
|
else:
|
|
function: JsonValue = {"name": "read", "arguments": json.dumps({"path": f"{turn}.py"})}
|
|
messages.append(
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [{"id": f"call_{turn}", "type": "function", "function": function}],
|
|
}
|
|
)
|
|
messages.append({"role": "tool", "tool_call_id": f"call_{turn}", "content": result})
|
|
if shape == "anthropic":
|
|
return "/v1/messages", {"model": "claude", "max_tokens": 32, "messages": messages, "cache": {"no-cache": True}}
|
|
return "/v1/chat/completions", {"model": "gpt", "messages": messages, "cache": {"no-cache": True}}
|
|
|
|
|
|
@pytest.mark.parametrize("shape", ("anthropic", "chat"))
|
|
def test_resent_tool_history_is_scanned_once(single_worker: Rig, shape: str) -> None:
|
|
path, body = _tool_history(150, seed=2 if shape == "anthropic" else 4, shape=shape)
|
|
turns: Final = body["messages"]
|
|
assert isinstance(turns, list)
|
|
first: Final = timed(lambda: single_worker.proxy.request("POST", path, body))
|
|
following: Final = {
|
|
**body,
|
|
"messages": [*turns, {"role": "assistant", "content": "ok"}, {"role": "user", "content": "next"}],
|
|
}
|
|
second: Final = timed(lambda: single_worker.proxy.request("POST", path, following))
|
|
assert first.value.status_code == 200 and second.value.status_code == 200, (first.value.text, second.value.text)
|
|
assert_within_factor_of_control(second.seconds, control=first.seconds, factor=0.5)
|
|
|
|
|
|
@pytest.mark.parametrize("placement", ("guardrail", "global"))
|
|
def test_skip_tool_message_in_guardrail_forwards_tool_results(tmp_path: Path, placement: str) -> None:
|
|
def configure(configuration: dict[str, JsonValue]) -> None:
|
|
guardrails: Final = configuration["guardrails"]
|
|
settings: Final = configuration["litellm_settings"]
|
|
assert isinstance(guardrails, list) and isinstance(guardrails[0], dict) and isinstance(settings, dict)
|
|
params: Final = guardrails[0]["litellm_params"]
|
|
assert isinstance(params, dict)
|
|
(params if placement == "guardrail" else settings)["skip_tool_message_in_guardrail"] = True
|
|
|
|
with tool_rig(tmp_path, configure=configure) as rig:
|
|
rig.wire.drain()
|
|
replies: Final = (
|
|
rig.proxy.request(
|
|
"POST",
|
|
"/v1/messages",
|
|
{
|
|
"model": "claude",
|
|
"max_tokens": 32,
|
|
"cache": {"no-cache": True},
|
|
"messages": [
|
|
{"role": "user", "content": TEXT},
|
|
{
|
|
"role": "assistant",
|
|
"content": [
|
|
{"type": "tool_use", "id": "tool_1", "name": "read_file", "input": {"env": TEXT}}
|
|
],
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "tool_use_id": "tool_1", "content": TEXT}],
|
|
},
|
|
],
|
|
},
|
|
),
|
|
rig.proxy.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": "gpt",
|
|
"cache": {"no-cache": True},
|
|
"messages": [
|
|
{"role": "user", "content": TEXT},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_1",
|
|
"type": "function",
|
|
"function": {"name": "read_file", "arguments": json.dumps({"env": TEXT})},
|
|
}
|
|
],
|
|
},
|
|
{"role": "tool", "tool_call_id": "call_1", "content": TEXT},
|
|
],
|
|
},
|
|
),
|
|
)
|
|
assert all(reply.status_code == 200 for reply in replies), [reply.text for reply in replies]
|
|
sent: Final = rig.wire.drain()
|
|
assert len(sent) == 2
|
|
for request in sent:
|
|
messages = json.loads(request.body)["messages"]
|
|
assert SECRET not in json.dumps(messages[:2])
|
|
assert SECRET in json.dumps(messages[2])
|