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:
mubashir1osmani 2026-07-13 18:55:01 -07:00
parent 1b2222783e
commit 9d3672a1da
2 changed files with 115 additions and 1 deletions

View file

@ -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)

View file

@ -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).