mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
test(e2e): restore langfuse tool_permission guardrail coverage
The guardrail create/delete client helpers come back without dedicated request/response models; the test is unchanged from its previous revision
This commit is contained in:
parent
1b2222783e
commit
9d3672a1da
2 changed files with 115 additions and 1 deletions
|
|
@ -271,6 +271,15 @@ def observation_mentions_tool(obs: LangfuseObservation, tool_name: str) -> bool:
|
|||
return tool_name in blob
|
||||
|
||||
|
||||
def observation_has_guardrail(obs: LangfuseObservation, *, guardrail_name: str) -> bool:
|
||||
blob = json.dumps(obs.metadata, default=str) if obs.metadata is not None else ""
|
||||
if guardrail_name in blob or "guardrail" in blob.lower():
|
||||
return True
|
||||
if obs.name is not None and "guardrail" in obs.name.lower():
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class LoggingClient:
|
||||
gateway: Gateway
|
||||
|
|
@ -390,6 +399,44 @@ class LoggingClient:
|
|||
f"POST /team/{team_id}/callback must return status=success; got {response.status!r}"
|
||||
)
|
||||
|
||||
def create_tool_permission_guardrail(self, name: str, *, allowed_tool: str) -> str:
|
||||
"""Register a tool_permission guardrail that allows one tool and denies the rest."""
|
||||
response = unwrap(
|
||||
self.gateway.transport.post(
|
||||
"/guardrails",
|
||||
headers=self.gateway.transport.master,
|
||||
json=RootModel[JsonValue].model_validate(
|
||||
{
|
||||
"guardrail": {
|
||||
"guardrail_name": name,
|
||||
"litellm_params": {
|
||||
"guardrail": "tool_permission",
|
||||
"mode": "post_call",
|
||||
"default_on": False,
|
||||
"default_action": "deny",
|
||||
"on_disallowed_action": "block",
|
||||
"rules": [{"id": "allow-named-tool", "tool_name": allowed_tool, "decision": "allow"}],
|
||||
},
|
||||
}
|
||||
}
|
||||
),
|
||||
response_type=RootModel[JsonValue],
|
||||
)
|
||||
).root
|
||||
guardrail_id = response.get("guardrail_id") if isinstance(response, dict) else None
|
||||
assert isinstance(guardrail_id, str) and guardrail_id, (
|
||||
f"POST /guardrails must return a guardrail_id; got {response!r}"
|
||||
)
|
||||
return guardrail_id
|
||||
|
||||
def delete_guardrail(self, guardrail_id: str) -> None:
|
||||
_ = self.gateway.transport.delete(
|
||||
f"/guardrails/{guardrail_id}",
|
||||
headers=self.gateway.transport.master,
|
||||
json=NoBody(),
|
||||
response_type=NoBody,
|
||||
)
|
||||
|
||||
def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str:
|
||||
return self.gateway.create_model(model_name, litellm_params)
|
||||
|
||||
|
|
@ -417,6 +464,7 @@ class LoggingClient:
|
|||
stream: bool = False,
|
||||
tools: list[ChatTool] | None = None,
|
||||
tool_choice: str | None = None,
|
||||
guardrails: list[str] | None = None,
|
||||
max_tokens: int = 64,
|
||||
) -> StreamingResponse:
|
||||
body = ChatBody(
|
||||
|
|
@ -426,6 +474,7 @@ class LoggingClient:
|
|||
stream=stream,
|
||||
tools=tools,
|
||||
tool_choice=tool_choice,
|
||||
guardrails=guardrails,
|
||||
)
|
||||
if stream:
|
||||
return self.gateway.chat_stream(key, body)
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ Dynamic credentials by product surface:
|
|||
- user/key: key metadata.logging with callback_name=langfuse_otel
|
||||
- org: organization + team under it + team callback (no org-level callback API)
|
||||
|
||||
Extra success paths assert tool calls land on the trace.
|
||||
Extra success paths assert tool calls and applied guardrails land on the trace.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -34,6 +34,7 @@ from logging_client import (
|
|||
LoggingClient,
|
||||
completion_response_id,
|
||||
costs_agree,
|
||||
observation_has_guardrail,
|
||||
observation_mentions_tool,
|
||||
observation_spend,
|
||||
)
|
||||
|
|
@ -284,6 +285,70 @@ class TestLangfuseTeamLogging:
|
|||
scope="team-tools",
|
||||
)
|
||||
|
||||
@pytest.mark.covers("logging.langfuse.success.logs_spend", exercised_on=["chat_completions"])
|
||||
def test_tool_permission_guardrail_logged(
|
||||
self,
|
||||
client: LoggingClient,
|
||||
resources: ResourceManager,
|
||||
langfuse_creds: LangfuseCreds,
|
||||
) -> None:
|
||||
"""tool_permission post_call guardrail must appear on the Langfuse trace
|
||||
(StandardLogging guardrail_information -> Langfuse guardrail span)."""
|
||||
marker = unique_marker()
|
||||
guardrail_name = f"e2e-lf-tool-perm-{marker}"
|
||||
guardrail_id = client.create_tool_permission_guardrail(
|
||||
guardrail_name, allowed_tool="get_weather"
|
||||
)
|
||||
resources.defer(lambda: client.delete_guardrail(guardrail_id))
|
||||
|
||||
_, key, key_alias = self._team_key(
|
||||
client, resources, langfuse_creds, models=[DRIVER_MODEL]
|
||||
)
|
||||
prompt_marker = unique_marker()
|
||||
outcome = client.chat_raw(
|
||||
key,
|
||||
DRIVER_MODEL,
|
||||
f"Use get_weather for Berlin. marker={prompt_marker}",
|
||||
tools=[WEATHER_TOOL],
|
||||
tool_choice="required",
|
||||
guardrails=[guardrail_name],
|
||||
max_tokens=128,
|
||||
)
|
||||
require_successful_call(outcome)
|
||||
|
||||
observations = client.poll_langfuse_trace_observations(
|
||||
langfuse_creds, key_alias=key_alias, prompt_marker=prompt_marker
|
||||
)
|
||||
assert observations, (
|
||||
f"team+guardrail: no Langfuse observations for key_alias={key_alias!r}"
|
||||
)
|
||||
gen = next(
|
||||
(
|
||||
o
|
||||
for o in observations
|
||||
if prompt_marker in _json_blob(o.input)
|
||||
or key_alias in _json_blob(o.metadata)
|
||||
or o.name in (f"litellm:{key_alias}", "litellm_request")
|
||||
),
|
||||
observations[0],
|
||||
)
|
||||
_assert_logs_spend(
|
||||
client,
|
||||
key=key,
|
||||
outcome=outcome,
|
||||
obs_cost=observation_spend(gen),
|
||||
scope="team-guardrail",
|
||||
)
|
||||
assert any(
|
||||
observation_has_guardrail(o, guardrail_name=guardrail_name)
|
||||
or (o.name is not None and "guardrail" in o.name.lower())
|
||||
for o in observations
|
||||
), (
|
||||
f"Langfuse trace must include applied guardrail {guardrail_name!r}; "
|
||||
f"observation names={[o.name for o in observations]}"
|
||||
)
|
||||
|
||||
|
||||
class TestLangfuseUserKeyLogging:
|
||||
"""User-owned key with metadata.logging (key-level dynamic Langfuse credentials).
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue