scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-04 07:27:51 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 299

Merge pull request #49962 from BerriAI/litellm_rc108_move_e2e_to_integration

Browse source 
test(e2e): repair rc/1.108.0 release qualification
This commit is contained in:
yuneng-jiang 2026-11-03 15:35:41 -08:00 • committed by GitHub
commit 49a2e97f84
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 1104 additions and 1865 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
tests/e2e/CONTRIBUTING.md
View file

@ -90,18 +90,7 @@ The suites run against a live proxy, so bring one up first by running the litell
They also need a proxy whose bundled UI contains the change under test, so run the proxy from your branch (an editable install serves the UI your checkout builds)
Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and point `PHOENIX_COLLECTOR_HTTP_ENDPOINT` at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy. The presidio guardrail tests need a running Presidio analyzer and anonymizer the proxy can reach, addressed by `PRESIDIO_ANALYZER_API_BASE` / `PRESIDIO_ANONYMIZER_API_BASE`
The Presidio spend-log audit test also requires prompt storage on the proxy, because its assertions inspect the detected entities and their scores. Add the following to the proxy config before starting it:
```yaml
general_settings:
store_prompts_in_spend_logs: true
```
Alternatively, start the proxy with `STORE_PROMPTS_IN_SPEND_LOGS=true litellm --config <your-e2e-config>.yml --port 4000`. Setting the variable only on the pytest process does not configure the proxy. With prompt storage disabled, the proxy correctly redacts `guardrail_response`, so that configuration cannot exercise this test's entity-detail assertions. Enable this only on a test stack using synthetic prompts
A couple of logging destinations are configured on the proxy rather than by the test. The Weave tests scope their callback to the key they create, but litellm builds the `weave_otel` logger from `WANDB_API_KEY` and `WANDB_PROJECT_ID` before it applies the per-key vars, so the proxy needs both in its own environment or the key-scoped callback never initializes and nothing ships
Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and start the proxy with `LITELLM_OTEL_V2=true`, the `otel` callback, and `OTEL_EXPORTER_OTLP_ENDPOINT` pointing at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy.
### Secret manager lanes

128
tests/e2e/a2a/a2a_client.py
View file

@ -18,7 +18,7 @@ from dataclasses import dataclass
from pydantic import BaseModel, ConfigDict, Field
from e2e_config import settle_propagation
from e2e_http import NoBody, Result, Success, get_external, is_ok
from e2e_http import NoBody, Result, Success, is_ok
from e2e_metadata import STEP_FRAMES, step
from proxy_client import ProxyClient
@ -36,11 +36,6 @@ class A2ASkill(BaseModel):
examples: list[str] | None = None
class A2AProvider(BaseModel):
organization: str
url: str
class AgentCardParams(BaseModel):
"""The upstream agent card an admin registers. `protocolVersion` is the field the
proxy validates against SUPPORTED_A2A_PROTOCOL_VERSIONS on registration."""
@ -57,29 +52,6 @@ class AgentCardParams(BaseModel):
preferred_transport: str | None = Field(default=None, serialization_alias="preferredTransport")
class UpstreamAgentCard(BaseModel):
"""A real published agent card parsed from a public /.well-known endpoint. Keys on
the A2A wire aliases so `model_validate_json` reads the served JSON and
`model_dump(by_alias=True)` re-emits it unchanged for verbatim registration; it is
only ever fetched-and-validated, never hand-constructed, so aliasing on the wire
names does not affect any call site."""
model_config = ConfigDict(populate_by_name=True)
protocol_version: str = Field(alias="protocolVersion")
name: str
description: str
version: str
url: str
provider: A2AProvider | None = None
documentation_url: str | None = Field(default=None, alias="documentationUrl")
capabilities: A2ACapabilities = A2ACapabilities()
skills: list[A2ASkill]
default_input_modes: list[str] = Field(default=["text"], alias="defaultInputModes")
default_output_modes: list[str] = Field(default=["text"], alias="defaultOutputModes")
preferred_transport: str | None = Field(default=None, alias="preferredTransport")
class A2ABridgeParams(BaseModel):
"""litellm_params that route the agent through the completion bridge: an A2A
message/send is transformed into a litellm.acompletion against this provider."""
@ -93,7 +65,7 @@ class A2ABridgeParams(BaseModel):
class AgentRegisterBody(BaseModel):
agent_name: str
agent_card_params: AgentCardParams | UpstreamAgentCard
agent_card_params: AgentCardParams
litellm_params: A2ABridgeParams | None = None
@ -136,32 +108,9 @@ class A2ATextPart(BaseModel):
text: str
class A2ASearchPropertiesParams(BaseModel):
"""The strict param schema of the published property agent's `search_properties`
skill (unknown keys are rejected upstream), so a natural-language query like
"properties for sale in SF under $2M" is expressed as typed fields."""
un_locode: str | None = None
service_type: str | None = None
property_type: str | None = None
bedrooms_min: int | None = None
asking_price_max: float | None = None
limit: int | None = None
class A2ASkillInvocation(BaseModel):
skill: str
params: A2ASearchPropertiesParams
class A2ADataPart(BaseModel):
kind: str = "data"
data: A2ASkillInvocation
class A2AOutboundMessage(BaseModel):
role: str = "user"
parts: list[A2ATextPart | A2ADataPart]
parts: list[A2ATextPart]
message_id: str = Field(serialization_alias="messageId")
@ -189,57 +138,10 @@ class A2AResponseMessage(BaseModel):
parts: list[A2AResponsePart] = []
class A2ATaskStatus(BaseModel):
state: str | None = None
message: A2AResponseMessage | None = None
class A2AListingLocation(BaseModel):
"""Only the location fields a test reads back off a returned listing."""
un_locode: str | None = None
class A2AListing(BaseModel):
"""A single property card from the agent's `search_results` artifact; only the
identity/location fields a test asserts on are modelled."""
raia_id: str
property_type: str | None = None
service_type: str | None = None
location: A2AListingLocation = A2AListingLocation()
class A2ASearchResults(BaseModel):
"""The DataPart payload the property agent's `search_properties` skill returns:
the run count plus the listing cards themselves. Proof the tool actually ran and
matched, not just that the task completed with some text."""
total: int
count: int
listings: list[A2AListing] = []
class A2AArtifactPart(BaseModel):
kind: str | None = None
data: A2ASearchResults | None = None
class A2AArtifact(BaseModel):
model_config = ConfigDict(populate_by_name=True)
artifact_id: str | None = Field(default=None, alias="artifactId")
name: str | None = None
parts: list[A2AArtifactPart] = []
class A2AResult(BaseModel):
"""A message/send result. In 0.3 the message fields sit directly on the result
(`kind`/`role`/`parts`); in 1.0 they are nested under `message`; a real agent that
runs a task replies with a `task` whose agent text lives on `status.message` and
whose tool output lives on `artifacts`. `text` reads the agent's reply from
whichever shape the served version produced; `search_results` reads the tool's
structured output when the agent ran a skill."""
(`kind`/`role`/`parts`); in 1.0 they are nested under `message`. `text` reads the
agent's reply from whichever shape the served version produced."""
model_config = ConfigDict(populate_by_name=True)
@ -248,8 +150,6 @@ class A2AResult(BaseModel):
message_id: str | None = Field(default=None, alias="messageId")
parts: list[A2AResponsePart] = []
message: A2AResponseMessage | None = None
status: A2ATaskStatus | None = None
artifacts: list[A2AArtifact] = []
@property
def text(self) -> str:
@ -257,8 +157,6 @@ class A2AResult(BaseModel):
parts = self.message.parts
elif self.parts:
parts = self.parts
elif self.status is not None and self.status.message is not None:
parts = self.status.message.parts
else:
parts = []
return "".join(part.text or "" for part in parts)
@ -267,14 +165,6 @@ class A2AResult(BaseModel):
def is_nested_v1_shape(self) -> bool:
return self.message is not None
@property
def search_results(self) -> A2ASearchResults | None:
for artifact in self.artifacts:
for part in artifact.parts:
if part.data is not None:
return part.data
return None
class A2AError(BaseModel):
code: int
@ -380,11 +270,3 @@ class A2AClient:
def build_a2a_client(proxy: ProxyClient) -> A2AClient:
return A2AClient(proxy=proxy)
@step("Fetch a published A2A agent card from its /.well-known endpoint")
def fetch_agent_card(url: str, *, timeout: float = 20.0) -> Result[UpstreamAgentCard]:
"""Fetch a live A2A agent card from its /.well-known endpoint and parse it into the
registration model, so a test can register a real published card verbatim rather
than a hand-rolled one."""
return get_external(url, response_type=UpstreamAgentCard, timeout=timeout)

49
tests/e2e/a2a/test_a2a_agent_e2e.py
View file

@ -17,18 +17,14 @@ import pytest
from a2a_client import (
A2ABridgeParams,
A2AClient,
A2ADataPart,
A2AJsonRpcRequest,
A2AMessageSendParams,
A2AOutboundMessage,
A2ASearchPropertiesParams,
A2ASkill,
A2ASkillInvocation,
A2ATextPart,
AgentCardParams,
AgentRegisterBody,
AgentResponse,
fetch_agent_card,
)
from e2e_config import unique_marker
from e2e_http import Result, UnknownApiError, unwrap
@ -50,9 +46,6 @@ BRIDGE = A2ABridgeParams(
model=BRIDGE_MODEL,
)
MOVEHOME_AGENT_CARD_URL = "https://movehome.org/.well-known/agent.json"
MOVEHOME_ORIGIN = "https://movehome.org"
pytestmark = pytest.mark.e2e
@ -135,48 +128,6 @@ class TestA2AAgentLifecycle:
assert result is not None
assert result.text != ""
@pytest.mark.covers("other.a2a.message_send.real_world_agent_replies")
@meta(
Subject(
domain=Domain.AGENTS_API,
route=Route.A2A,
)
)
def test_real_world_agent_replies_to_property_query(self, client: A2AClient, resources: ResourceManager, scoped_key: str) -> None:
upstream = unwrap(fetch_agent_card(MOVEHOME_AGENT_CARD_URL)).model_copy(update={"url": MOVEHOME_ORIGIN})
assert upstream.protocol_version == "0.3.0"
marker = unique_marker()
body = AgentRegisterBody(agent_name=f"e2e-a2a-real-{marker}", agent_card_params=upstream)
agent = unwrap(client.register_agent(body))
resources.defer(lambda: client.delete_agent(agent.agent_id))
assert agent.agent_card_params.protocol_version == "0.3"
location = "GBLON"
request = A2AJsonRpcRequest(
id=f"e2e-{unique_marker()}",
params=A2AMessageSendParams(
message=A2AOutboundMessage(
parts=[
A2ADataPart(
data=A2ASkillInvocation(
skill="search_properties",
params=A2ASearchPropertiesParams(un_locode=location, service_type="long_term", limit=3),
)
)
],
message_id=unique_marker(),
)
),
)
response = unwrap(client.send_message(agent.agent_id, scoped_key, request))
assert response.error is None
assert response.result is not None
results = response.result.search_results
assert results is not None, "agent returned no search_results artifact; skill did not run"
assert results.total > 0
assert results.listings, "search_properties matched nothing; agent returned no property cards"
assert all(listing.raia_id for listing in results.listings)
assert all(listing.location.un_locode == location for listing in results.listings)
@pytest.mark.covers("other.a2a.discovery.proxy_fronted_card")
@meta(
Subject(

5
tests/e2e/coverage_registry/guardrail.yaml
View file

@ -1,11 +1,6 @@
# Guardrail enforcement (behavior features). Grounded in litellm/proxy/guardrails/guardrail_hooks/.
# Rolls up into the "Logging & Guardrails" dashboard module together with logging.*
- {id: guardrail.presidio.pre_call.masks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "PII masking pre-call; data-leak blast radius"}
- {id: guardrail.presidio.post_call.masks, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Mask PII in model output"}
- {id: guardrail.presidio.post_call.masks_generated_output, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, chat_completions_stream, anthropic_messages_stream], source: "guardrail_hooks/presidio.py", rationale: "Mask model-generated credit-card output with the UI default scope"}
- {id: guardrail.presidio.post_call.spend_log_stores_masked_output, module: guardrail, tier: P0, hook_point: post_call, assertions: [masks], exercised_on: [chat_completions, chat_completions_stream, messages, anthropic_messages_stream, responses], source: "guardrail_hooks/presidio.py", fail_before_fix: proven, rationale: "When an output guardrail masks the response, the spend log stores the masked text the caller received rather than the raw model output, on every endpoint and both stream modes (LIT-8325)"}
- {id: guardrail.presidio.logging_only.masks, module: guardrail, tier: P0, hook_point: logging_only, assertions: [masks], exercised_on: [chat_completions, messages], source: "guardrail_hooks/presidio.py", rationale: "Redact in logs without blocking"}
- {id: guardrail.presidio.pre_call.logs_masked_entities, module: guardrail, tier: P0, hook_point: pre_call, assertions: [logs_masked_entities], exercised_on: [chat_completions], source: "guardrail_hooks/presidio.py", rationale: "A masking run must record itself on the spend log: the dashboard's guardrail panel renders the masked-entity counts and per-entity scores straight off metadata.guardrail_information, so a run that masks but records nothing leaves an operator unable to audit it"}
- {id: guardrail.bedrock.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions, messages, responses], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "AWS content guardrail blocks harmful input"}
- {id: guardrail.litellm_content_filter.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Local content-filter default-on blocks banned keyword pre-call"}
- {id: guardrail.litellm_content_filter.pre_call.blocks_video, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [videos], source: "test_key_guardrail_video_e2e.py", fail_before_fix: proven, rationale: "A content-filter guardrail attached to a key (metadata.guardrails) blocks a banned prompt on POST /v1/videos before the provider is called; before the fix the route's call type was unknown to the unified guardrail hook and the prompt went to the provider unscanned (LIT-6685)"}

5
tests/e2e/coverage_registry/llm_conversational.yaml
View file

@ -133,11 +133,6 @@
- {id: llm.responses.ollama.basic.stream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/responses streams text deltas, usage and a terminal event"}
- {id: llm.responses.ollama.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/responses returns one addressable get_weather call"}
- {id: llm.responses.ollama.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: responses, route: ollama, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_ollama_e2e.py", rationale: "Ollama /api/generate (prompt-based tools) over /v1/responses tool result round trip reaches the model"}
- {id: llm.chat_completions.sail.service_tier.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: sail, capability: service_tier, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_sail_e2e.py", rationale: "service_tier balanced and auto map to Sail completion windows and bill the matching price columns; Sail serves flex only to background responses and Batch"}
- {id: llm.chat_completions.sail.service_tier.nonstream.rejects_unknown_tier, module: llm, tier: P1, subject_endpoint: chat_completions, route: sail, capability: service_tier, streaming: nonstream, assertions: [rejects_unknown_tier], source: "llm_translation/test_sail_e2e.py", rationale: "A service_tier Sail has no completion window for is a 400 without drop_params"}
- {id: llm.chat_completions.sail.service_tier.nonstream.drops_unknown_tier_and_bills_asap, module: llm, tier: P1, subject_endpoint: chat_completions, route: sail, capability: service_tier, streaming: nonstream, assertions: [drops_unknown_tier_and_bills_asap], source: "llm_translation/test_sail_e2e.py", rationale: "An unknown service_tier under drop_params is dropped and billed at asap in both the cost header and spend log"}
- {id: llm.responses.sail.service_tier.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: responses, route: sail, capability: service_tier, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_sail_e2e.py", rationale: "A caller metadata.completion_window of balanced on /v1/responses bills Sail balanced rates"}
- {id: llm.messages.sail.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: sail, capability: basic, streaming: nonstream, assertions: [works], source: "llm_translation/test_sail_e2e.py", rationale: "Sail over /v1/messages"}
- {id: llm.chat_completions.anthropic.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: chat_completions, route: anthropic, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic over /chat/completions: cost header and spend row agree"}
- {id: llm.chat_completions.anthropic.multi_turn.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: anthropic, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic tool result round trip over /chat/completions"}
- {id: llm.messages.anthropic.multi_turn.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: anthropic, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_conversational_matrix_e2e.py", rationale: "Anthropic tool result round trip over /v1/messages"}

1
tests/e2e/coverage_registry/other.yaml
View file

@ -55,7 +55,6 @@
- {id: other.a2a.register.persists, module: other, tier: P1, area: a2a, assertions: [persists], source: "agent_endpoints/endpoints.py:325-443", rationale: "POST /v1/agents registers an agent card; GET /v1/agents/{id} reads it back"}
- {id: other.a2a.register.unsupported_version_rejected, module: other, tier: P1, area: a2a, assertions: [unsupported_version_rejected], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a protocolVersion outside SUPPORTED_A2A_PROTOCOL_VERSIONS is refused with 400"}
- {id: other.a2a.register.semver_version_accepted, module: other, tier: P1, area: a2a, assertions: [semver_version_accepted], source: "agent_endpoints/endpoints.py _validate_protocol_version", rationale: "A card pinning a patch-level semver like 0.3.0 (what the Google A2A SDK emits) registers, stores and serves the canonical 0.3 rather than 400ing; regression guard for the v1.92 report"}
- {id: other.a2a.message_send.real_world_agent_replies, module: other, tier: P1, area: a2a, assertions: [real_world_agent_replies], source: "agent_endpoints/a2a_endpoints.py asend_message", rationale: "A real published a2a agent fetched live from a public /.well-known endpoint (pinning the full semver 0.3.0 the a2a-sdk emits) registers, serves the canonical 0.3, and a message/send skill invocation proxies to the live upstream and returns the agent's reply"}
- {id: other.a2a.register.malformed_version_rejected, module: other, tier: P1, area: a2a, assertions: [malformed_version_rejected], source: "a2a/agent_card.py normalize_protocol_version", rationale: "A malformed protocolVersion like 0.3.garbage fails full-string semver validation and is refused with 400 instead of truncating to a supported family"}
- {id: other.a2a.discovery.proxy_fronted_card, module: other, tier: P1, area: a2a, assertions: [proxy_fronted_card], source: "agent_endpoints/a2a_endpoints.py get_agent_card", rationale: "/.well-known/agent-card.json serves the proxy url + supportedInterfaces and the LiteLLM virtual-key bearer scheme, not the upstream"}
- {id: other.a2a.message_send.bridge_invokes, module: other, tier: P1, area: a2a, assertions: [bridge_invokes], source: "a2a_protocol/litellm_completion_bridge/handler.py", rationale: "A2A message/send routes through the completion bridge to a real provider and logs an asend_message spend row"}

19
tests/e2e/guardrails/guardrails_client.py
View file

@ -41,8 +41,6 @@ from pydantic import BaseModel, Field
GUARDRAIL_BACKEND: Final = "gemini/gemini-3.8-flash"
GuardrailMode = Literal["pre_call", "post_call", "during_call", "logging_only"]
PiiEntity = Literal["EMAIL_ADDRESS", "PHONE_NUMBER", "PERSON", "CREDIT_CARD", "US_SSN"]
PiiAction = Literal["MASK", "BLOCK"]
BlockedWordAction = Literal["BLOCK", "MASK"]
@ -80,22 +78,6 @@ class BlockCodeExecutionParamsBody(GuardrailParamsBase):
guardrail: Literal["block_code_execution"] = "block_code_execution"
class PresidioParamsBody(GuardrailParamsBase):
"""Presidio PII guardrail params. `presidio_filter_scope="input"` keeps the
registration to a single callback on the configured mode; the default
("both") also registers a second post_call output-masking callback, which a
pre_call- or logging_only-scoped test must not drag in. `output_parse_pii`
stays unset/False: True would unmask the response back to the caller."""
guardrail: Literal["presidio"] = "presidio"
presidio_analyzer_api_base: str
presidio_anonymizer_api_base: str
presidio_filter_scope: Literal["input", "output", "both"] | None = None
presidio_language: str | None = None
output_parse_pii: bool | None = None
pii_entities_config: dict[PiiEntity, PiiAction] | None = None
class ToolPermissionRuleBody(BaseModel):
"""One tool_permission rule: a decision for the tool named by `tool_name`."""
@ -130,7 +112,6 @@ GuardrailParamsBody = (
| BedrockGuardrailParamsBody
| OpenAIModerationParamsBody
| BlockCodeExecutionParamsBody
| PresidioParamsBody
| ToolPermissionParamsBody
| CustomCodeParamsBody
)

857
tests/e2e/guardrails/test_presidio_masking_e2e.py
View file

@ -1,857 +0,0 @@
"""Live e2e: the Presidio PII guardrail masks, per its configured hook point.
pre_call: the guardrail calls the Presidio analyzer/anonymizer on the request
messages BEFORE the model runs, so the model only ever sees placeholders like
<EMAIL_ADDRESS>. A prompt asking the model to repeat a fake email + phone back
must come back with the placeholders echoed and the raw PII absent, on
/chat/completions and on /v1/messages (Anthropic format).
post_call: the mirror hook. The request reaches the model unmasked and the
MODEL OUTPUT is what gets anonymized, so the caller never receives raw PII the
model repeated back. The two hooks are told apart behaviorally rather than by
configuration: the post_call prompt asks for a value derived from the raw email
(its local part, which is not itself an entity Presidio masks) alongside the
address itself, so the answer proves the model saw the raw address while the
address in the same response comes back as <EMAIL_ADDRESS>.
The analyzer/anonymizer endpoints come from PRESIDIO_ANALYZER_API_BASE /
PRESIDIO_ANONYMIZER_API_BASE; missing env is a hard failure, never a skip.
Each guardrail registers with an explicit presidio_filter_scope so only the
configured hook's callback exists (the default "both" registers input masking
AND a post_call output masker), and is deleted on teardown.
The spend-log audit test requires general_settings.store_prompts_in_spend_logs:
true in the proxy config, or STORE_PROMPTS_IN_SPEND_LOGS=true in the proxy
process environment before startup. Setting it only on pytest has no effect.
Without this opt-in, redacting guardrail_response is expected proxy behavior;
this suite deliberately requires the detected-entity details to remain visible.
"""
from __future__ import annotations
import os
import re
import time
from collections.abc import Callable, Iterator
from typing import Final, Literal
import pytest
from pydantic import BaseModel, JsonValue, TypeAdapter
from e2e_config import unique_marker
from e2e_http import Result, StreamingResponse, Success
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
from guardrails_client import GuardrailMode, GuardrailsClient, PiiAction, PiiEntity, PresidioParamsBody
from lifecycle import ResourceManager
from models import (
AnthropicMessagesResponse,
ChatResponse,
GuardrailEntityMatch,
GuardrailRunRecord,
SpendLogRow,
)
pytestmark = pytest.mark.e2e
MODEL = "gemini-2.5-flash"
# A guardrail created via POST /guardrails reaches the worker that served the
# create immediately, but every other worker only picks it up on its next
# periodic DB sync (~30s), so the first requests can be served unguarded.
GUARDRAIL_PROPAGATION_DEADLINE_SECONDS = 40.0
GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS = 5.0
# Presidio's anonymizer replaces a detected entity with its unnumbered type
# placeholder, e.g. <EMAIL_ADDRESS>. The pre_call assertions match on the bare
# token because the model is echoing the masked prompt and may not preserve the
# angle brackets; the logged payload keeps the placeholder verbatim.
MASKED_EMAIL_TOKEN = "EMAIL_ADDRESS"
MASKED_PHONE_TOKEN = "PHONE_NUMBER"
MASKED_CREDIT_CARD_TOKEN = "CREDIT_CARD"
# Fictional NANP 555 number; a standard format Presidio's phone recognizer detects.
FAKE_PHONE = "+1 415-555-0134"
FAKE_VISA_TEST_CARD = "4111 1111 1111 1111"
_CARD_DIGIT_RUN: Final = re.compile(r"(?:\d[ -]?){13,19}")
_CONTENT_KEYS: Final = frozenset({"content", "text"})
def _presidio_bases() -> tuple[str, str]:
analyzer = os.environ.get("PRESIDIO_ANALYZER_API_BASE", "").strip()
anonymizer = os.environ.get("PRESIDIO_ANONYMIZER_API_BASE", "").strip()
if not analyzer or not anonymizer:
pytest.fail(
"Presidio e2e requires PRESIDIO_ANALYZER_API_BASE and PRESIDIO_ANONYMIZER_API_BASE "
"(the running Presidio analyzer/anonymizer services); missing env is a hard failure, not a skip"
)
return analyzer, anonymizer
def _register_presidio(
client: GuardrailsClient,
resources: ResourceManager,
*,
name: str,
mode: GuardrailMode | list[GuardrailMode] = "pre_call",
filter_scope: Literal["input", "output", "both"] | None = "input",
entities: dict[PiiEntity, PiiAction] | None = None,
) -> None:
analyzer, anonymizer = _presidio_bases()
guardrail_id = client.register(
name,
PresidioParamsBody(
mode=mode,
default_on=False,
presidio_analyzer_api_base=analyzer,
presidio_anonymizer_api_base=anonymizer,
presidio_filter_scope=filter_scope,
pii_entities_config=entities,
),
)
resources.defer(lambda: client.delete_guardrail(guardrail_id))
def _fake_email() -> str:
return f"jane.doe.{unique_marker()}@example.com"
def _pii_prompt(marker: str, email: str) -> str:
return (
f"{marker} Repeat this sentence back to me exactly, word for word: "
f"My email address is {email} and my phone number is {FAKE_PHONE}."
)
def _first_content(response: ChatResponse) -> str:
if not response.choices:
return ""
message = response.choices[0].message
return (message.content if message else None) or ""
class _StreamDelta(BaseModel):
content: str | None = None
class _StreamChoice(BaseModel):
delta: _StreamDelta
class _StreamChunk(BaseModel):
choices: tuple[_StreamChoice, ...] = ()
class _AnthropicStreamDelta(BaseModel):
type: str | None = None
text: str | None = None
class _AnthropicStreamEvent(BaseModel):
type: str
delta: _AnthropicStreamDelta | None = None
def _credit_card_prompt(marker: str) -> str:
return (
f"{marker} Reply with only the well known Visa sandbox test card number that starts with 4111, "
"the 16 digits grouped in fours separated by spaces, and nothing else."
)
def _passes_luhn(digits: str) -> bool:
checksum = sum(
digit if position % 2 == 0 else (digit * 2 - 9 if digit * 2 > 9 else digit * 2)
for position, digit in enumerate(int(char) for char in reversed(digits))
)
return checksum % 10 == 0
def _contains_card_number(text: str) -> bool:
"""Presidio's CREDIT_CARD recognizer only reports Luhn-valid digit runs, so a
Luhn-invalid number the model hallucinates is not something masking can catch."""
return any(
13 <= len(digits) <= 19 and _passes_luhn(digits)
for digits in (re.sub(r"[ -]", "", match.group()) for match in _CARD_DIGIT_RUN.finditer(text))
)
def _stream_content(result: StreamingResponse) -> str:
return "".join(
choice.delta.content
for event in result.stream_events
if event != "[DONE]"
for choice in _StreamChunk.model_validate_json(event).choices[:1]
if choice.delta.content
)
def _anthropic_stream_content(result: StreamingResponse) -> str:
return "".join(
event.delta.text
for payload in result.stream_events
for event in [_AnthropicStreamEvent.model_validate_json(payload)]
if event.type == "content_block_delta"
and event.delta is not None
and event.delta.type == "text_delta"
and event.delta.text
)
def _messages_text(response: AnthropicMessagesResponse) -> str:
"""The text of a /v1/messages answer, whichever shape the proxy produced
(Anthropic-native content blocks or OpenAI-normalized choices)."""
parts: list[str] = []
for block in response.content or []:
if block.text:
parts.append(block.text)
for choice in response.choices or []:
if choice.message and choice.message.content:
parts.append(choice.message.content)
return "\n".join(parts)
def _assert_eventually_masked[R: BaseModel](
fetch: Callable[[], Result[R]], extract: Callable[[R], str], *, email: str
) -> None:
"""Retry the call until the response comes back masked, to the propagation
deadline. An unmasked early response is in-flight guardrail propagation, not
a failure, and neither is a transient non-Success (a replica that has not
reloaded the guardrail answers 404, the live model can rate-limit) - only a
response that still carries the raw PII at the deadline is."""
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
last: str = "<no successful response yet>"
while True:
result = fetch()
match result:
case Success(data=data):
content = extract(data)
last = content
masked = MASKED_EMAIL_TOKEN in content and MASKED_PHONE_TOKEN in content and email not in content
if masked:
assert FAKE_PHONE not in content, (
f"the raw phone number must be masked before the model sees it, but the "
f"response echoed it: {content[:300]!r}"
)
return
case _:
last = f"<non-Success result: {result}>"
if time.monotonic() >= deadline:
pytest.fail(
f"presidio pre_call guardrail never masked the PII within "
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
)
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
class TestPresidioPreCallMasking:
@pytest.mark.covers(
"guardrail.presidio.pre_call.masks",
exercised_on=["chat_completions"],
)
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_pre_call_masks_pii_on_chat_completions(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
name = f"e2e-presidio-pre-chat-{unique_marker()}"
_register_presidio(client, resources, name=name)
email = _fake_email()
prompt = _pii_prompt(unique_marker(), email)
_assert_eventually_masked(
lambda: client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128),
_first_content,
email=email,
)
@pytest.mark.covers(
"guardrail.presidio.pre_call.masks",
exercised_on=["messages"],
)
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.MESSAGES,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_pre_call_masks_pii_on_messages(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
name = f"e2e-presidio-pre-msg-{unique_marker()}"
_register_presidio(client, resources, name=name)
email = _fake_email()
prompt = _pii_prompt(unique_marker(), email)
_assert_eventually_masked(
lambda: client.messages(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128),
_messages_text,
email=email,
)
#: Room for the model's reasoning tokens plus the three-line answer; a lower cap
#: truncates the response before the address it is supposed to mask.
_POST_CALL_MAX_TOKENS = 512
#: The post_call scenario masks these two entities and nothing else. Left
#: unscoped, Presidio's broader recognizers claim the local part too (a random
#: marker reads as an NRP), which would erase the very token that tells output
#: masking apart from input masking.
_POST_CALL_ENTITIES: dict[PiiEntity, PiiAction] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"}
def _post_call_prompt(marker: str, local_part: str) -> str:
"""Ask for the local part and the full address in one answer. Presidio masks
an EMAIL_ADDRESS entity and a bare local part is not one, so the two land
differently in the same response and pin the hook point behaviorally."""
return (
f"{marker} My email address is {local_part}@example.com and my phone number is {FAKE_PHONE}. "
"Reply with exactly three lines and nothing else. "
"Line 1: the part of the email address before the @ sign. "
"Line 2: the full email address. "
"Line 3: the phone number."
)
class TestPresidioPostCallMasking:
@pytest.mark.covers(
"guardrail.presidio.post_call.masks",
exercised_on=["chat_completions"],
)
@meta(
Subject(
domain=Domain.GUARDRAILS,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_post_call_masks_pii_in_model_output(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
"""A guardrail scoped to the output must anonymize the PII the model
repeats back, so a caller (or a downstream log of the response) never
receives it, while the request itself reaches the model untouched.
Both facts are asserted from one response: the local part comes back raw,
which is only possible if the model saw the real address, and the address
itself comes back as <EMAIL_ADDRESS> in the same answer.
"""
name = f"e2e-presidio-post-chat-{unique_marker()}"
_register_presidio(
client,
resources,
name=name,
mode="post_call",
filter_scope="output",
entities=_POST_CALL_ENTITIES,
)
local_part = f"e2euser{unique_marker()}"
email = f"{local_part}@example.com"
prompt = _post_call_prompt(unique_marker(), local_part)
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
last = "<no successful response yet>"
while True:
result = client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=_POST_CALL_MAX_TOKENS)
match result:
case Success(data=data):
last = _first_content(data)
if MASKED_EMAIL_TOKEN in last and email not in last:
assert local_part in last, (
"the model must have seen the RAW address (it is asked for the local "
"part, which Presidio does not mask); the local part is missing, so "
f"this response cannot tell post_call masking from pre_call: {last[:300]!r}"
)
assert MASKED_PHONE_TOKEN in last and FAKE_PHONE not in last, (
f"the phone number in the model's answer must be masked too, got: {last[:300]!r}"
)
return
case _:
last = f"<non-Success result: {result}>"
if time.monotonic() >= deadline:
pytest.fail(
f"presidio post_call guardrail never masked the model's output within "
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
)
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
def _assert_eventually_masks_generated_card(fetch: Callable[[], str | None]) -> None:
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
last: str = "<no successful response yet>"
while True:
content = fetch()
if content is not None:
last = content
if _contains_card_number(content):
pytest.fail(
"the post_call output masking let a card number through: "
f"{content[:300]!r}"
)
if MASKED_CREDIT_CARD_TOKEN in content:
return
if time.monotonic() >= deadline:
pytest.fail(
"presidio post_call output masking never masked the generated card within "
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
)
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
class TestPresidioCreditCardOutputMasking:
"""Proves the UI-default Presidio scope masks model-generated card output."""
@pytest.mark.covers(
"guardrail.presidio.post_call.masks_generated_output",
exercised_on=["chat_completions"],
)
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_ui_default_scope_masks_a_card_number_the_model_generates_on_chat_completions(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
name = f"e2e-presidio-card-chat-{unique_marker()}"
_register_presidio(
client,
resources,
name=name,
mode=["pre_call", "post_call"],
filter_scope=None,
entities={"CREDIT_CARD": "MASK"},
)
prompt: Final = _credit_card_prompt(unique_marker())
def fetch() -> str | None:
result: Final = client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=512)
match result:
case Success(data=data):
return _first_content(data)
case _:
return None
_assert_eventually_masks_generated_card(fetch)
@pytest.mark.covers(
"guardrail.presidio.post_call.masks_generated_output",
exercised_on=["chat_completions_stream"],
)
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.STREAM,
)
)
def test_ui_default_scope_masks_a_card_number_the_model_generates_on_streaming_chat_completions(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
name = f"e2e-presidio-card-stream-{unique_marker()}"
_register_presidio(
client,
resources,
name=name,
mode=["pre_call", "post_call"],
filter_scope=None,
entities={"CREDIT_CARD": "MASK"},
)
prompt: Final = _credit_card_prompt(unique_marker())
def fetch() -> str | None:
result: Final = client.chat_stream_raw(
scoped_key,
MODEL,
prompt,
guardrails=[name],
max_tokens=512,
)
if not result.ok or result.stream_error:
return None
return _stream_content(result)
_assert_eventually_masks_generated_card(fetch)
@pytest.mark.covers(
"guardrail.presidio.post_call.masks_generated_output",
exercised_on=["anthropic_messages_stream"],
)
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.MESSAGES,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.STREAM,
)
)
def test_ui_default_scope_masks_a_card_number_the_model_generates_on_streaming_anthropic_messages(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
name = f"e2e-presidio-card-messages-stream-{unique_marker()}"
_register_presidio(
client,
resources,
name=name,
mode=["pre_call", "post_call"],
filter_scope=None,
entities={"CREDIT_CARD": "MASK"},
)
prompt: Final = _credit_card_prompt(unique_marker())
def fetch() -> str | None:
result: Final = client.messages_stream_raw(
scoped_key,
MODEL,
prompt,
guardrails=[name],
max_tokens=512,
)
if not result.ok or result.stream_error:
return None
return _anthropic_stream_content(result)
_assert_eventually_masks_generated_card(fetch)
def _wire_text(outcome: StreamingResponse) -> str:
return "\n".join(outcome.stream_events) if outcome.is_streaming else outcome.body
def _poll_until_generated_card_masked(fetch: Callable[[], StreamingResponse]) -> StreamingResponse:
"""The raw HTTP outcome once the output masker is in effect on the serving
worker: whichever wire shape the endpoint speaks, a masked body carries the
CREDIT_CARD placeholder and no Luhn-valid card run. A raw card is a worker
that has not loaded the guardrail yet, so it is polled through like any
other unmasked answer."""
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
last: str = "<no successful response yet>"
while True:
outcome = fetch()
if outcome.ok and not outcome.stream_error:
wire = _wire_text(outcome)
last = wire
if MASKED_CREDIT_CARD_TOKEN in wire and not _contains_card_number(wire):
return outcome
if time.monotonic() >= deadline:
pytest.fail(
"presidio post_call output masking never masked the generated card within "
f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}"
)
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
def _spend_log_response_text(client: GuardrailsClient, key: str, call_id: str) -> str:
rows = client.proxy.poll_logs_for_key(
key,
predicate=lambda logged: any(row.litellm_call_id == call_id for row in logged),
)
row = next((row for row in rows if row.litellm_call_id == call_id), None)
assert row is not None, f"no spend log row ever appeared for x-litellm-call-id {call_id}"
return "\n".join(_stored_content(row.response))
def _stored_content(node: JsonValue, key: str | None = None) -> Iterator[str]:
match node:
case str() if key in _CONTENT_KEYS:
yield node
case dict():
for child_key, child in node.items():
yield from _stored_content(child, child_key)
case list():
for item in node:
yield from _stored_content(item, key)
case _:
return
class TestPresidioSpendLogStoresMaskedOutput:
"""The spend log stores the response the caller received, on every endpoint
and both stream modes, when an output-only post_call Presidio guardrail masks
a card number the model generated."""
_CELL: Final = "guardrail.presidio.post_call.spend_log_stores_masked_output"
def _assert_spend_log_is_masked(
self,
client: GuardrailsClient,
resources: ResourceManager,
key: str,
*,
name: str,
fetch: Callable[[str, str], StreamingResponse],
) -> None:
_register_presidio(
client,
resources,
name=name,
mode="post_call",
filter_scope="output",
entities={"CREDIT_CARD": "MASK"},
)
prompt: Final = _credit_card_prompt(unique_marker())
outcome = _poll_until_generated_card_masked(lambda: fetch(prompt, name))
assert outcome.call_id, f"the served response must carry x-litellm-call-id: {dict(outcome.headers)}"
logged = _spend_log_response_text(client, key, outcome.call_id)
assert not _contains_card_number(logged), (
"the caller got the masked response but the spend log stored the raw model output: "
f"{logged[:400]!r}"
)
assert MASKED_CREDIT_CARD_TOKEN in logged, (
f"the spend log response carries neither the card nor the placeholder: {logged[:400]!r}"
)
@pytest.mark.covers(_CELL, exercised_on=["chat_completions"])
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_spend_log_stores_masked_output_on_chat_completions(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
self._assert_spend_log_is_masked(
client,
resources,
scoped_key,
name=f"e2e-presidio-log-card-chat-{unique_marker()}",
fetch=lambda prompt, guardrail: client.chat_raw(
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
),
)
@pytest.mark.covers(_CELL, exercised_on=["chat_completions_stream"])
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.STREAM,
)
)
def test_spend_log_stores_masked_output_on_streaming_chat_completions(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
self._assert_spend_log_is_masked(
client,
resources,
scoped_key,
name=f"e2e-presidio-log-card-chat-stream-{unique_marker()}",
fetch=lambda prompt, guardrail: client.chat_stream_raw(
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
),
)
@pytest.mark.covers(_CELL, exercised_on=["messages"])
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.MESSAGES,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_spend_log_stores_masked_output_on_anthropic_messages(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
self._assert_spend_log_is_masked(
client,
resources,
scoped_key,
name=f"e2e-presidio-log-card-messages-{unique_marker()}",
fetch=lambda prompt, guardrail: client.messages_raw(
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
),
)
@pytest.mark.covers(_CELL, exercised_on=["anthropic_messages_stream"])
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.MESSAGES,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.STREAM,
)
)
def test_spend_log_stores_masked_output_on_streaming_anthropic_messages(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
self._assert_spend_log_is_masked(
client,
resources,
scoped_key,
name=f"e2e-presidio-log-card-messages-stream-{unique_marker()}",
fetch=lambda prompt, guardrail: client.messages_stream_raw(
scoped_key, MODEL, prompt, guardrails=[guardrail], max_tokens=512
),
)
@pytest.mark.covers(_CELL, exercised_on=["responses"])
@meta(
Subject(
domain=Domain.GUARDRAILS,
route=Route.RESPONSES,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_spend_log_stores_masked_output_on_responses(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
self._assert_spend_log_is_masked(
client,
resources,
scoped_key,
name=f"e2e-presidio-log-card-responses-{unique_marker()}",
fetch=lambda prompt, guardrail: client.responses(scoped_key, MODEL, prompt, guardrails=[guardrail]),
)
_LOGGED_ENTITIES: dict[PiiEntity, PiiAction] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"}
_ENTITY_LIST_ADAPTER: Final = TypeAdapter(list[GuardrailEntityMatch])
def _applied_guardrails(outcome: StreamingResponse) -> str:
return outcome.headers.get("x-litellm-applied-guardrails", "")
def _guardrail_records(row: SpendLogRow) -> tuple[GuardrailRunRecord, ...]:
metadata = row.metadata
if metadata is None:
return ()
return tuple(metadata.guardrail_information or ())
def _poll_until_guardrail_applied(
client: GuardrailsClient, key: str, guardrail_name: str, prompt: str
) -> StreamingResponse:
deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS
last = client.chat_raw(key, MODEL, prompt, guardrails=[guardrail_name], max_tokens=128)
while time.monotonic() < deadline:
if last.ok and guardrail_name in _applied_guardrails(last):
return last
time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS)
last = client.chat_raw(key, MODEL, prompt, guardrails=[guardrail_name], max_tokens=128)
return last
class TestPresidioSpendLogRecord:
@pytest.mark.covers(
"guardrail.presidio.pre_call.logs_masked_entities",
exercised_on=["chat_completions"],
)
@meta(
Subject(
domain=Domain.GUARDRAILS,
providers=(Provider.GEMINI,),
models=(MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_masking_run_is_recorded_on_the_spend_log(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
name = f"e2e-presidio-log-{unique_marker()}"
_register_presidio(client, resources, name=name, entities=_LOGGED_ENTITIES)
email = _fake_email()
prompt = _pii_prompt(unique_marker(), email)
outcome = _poll_until_guardrail_applied(client, scoped_key, name, prompt)
assert outcome.ok, f"the guarded call must be served, got {outcome.status_code}: {outcome.body[:400]}"
assert name in _applied_guardrails(outcome), (
"the response must carry x-litellm-applied-guardrails naming the guardrail; without it "
f"the 200 only proves the guardrail never attached. Got {_applied_guardrails(outcome)!r}"
)
request_id = ChatResponse.model_validate_json(outcome.body).id
assert request_id, f"the served response must carry an id to look the spend log up by: {outcome.body[:400]}"
rows = client.proxy.poll_logs_for_request_id(
request_id,
predicate=lambda logged: bool(_guardrail_records(logged[0])),
)
assert rows, f"no spend log row ever appeared for request {request_id}"
records = _guardrail_records(rows[0])
assert records, (
f"the spend log for {request_id} carries no guardrail_information, so the dashboard's "
"guardrail panel would render nothing for a request the guardrail demonstrably ran on"
)
record = next(
(entry for entry in records if entry.guardrail_name == name and entry.guardrail_mode == "pre_call"),
None,
)
assert record is not None, (
"guardrail_information carries no pre_call record for this guardrail, only "
f"{[(entry.guardrail_name, entry.guardrail_mode) for entry in records]}"
)
assert record.guardrail_status == "success", (
f"the recorded status must be success for a run that masked and served, got {record.guardrail_status!r}"
)
assert record.guardrail_provider == "presidio", (
f"the record must attribute the run to presidio so the dashboard picks the right "
f"renderer, got {record.guardrail_provider!r}"
)
counts = record.masked_entity_count or {}
assert _LOGGED_ENTITIES.keys() <= counts.keys(), (
f"every entity the guardrail was configured to mask must appear in masked_entity_count, got {counts}"
)
assert all(counts[entity] >= 1 for entity in _LOGGED_ENTITIES), (
f"each masked entity must be counted at least once, got {counts}"
)
assert not isinstance(record.guardrail_response, str), (
"This audit test requires general_settings.store_prompts_in_spend_logs: true in the proxy config "
"or STORE_PROMPTS_IN_SPEND_LOGS=true in the proxy process environment before startup "
"(not just the pytest environment). Default prompt redaction is valid proxy behavior, "
f"but prevents entity-detail assertions; got guardrail_response={record.guardrail_response!r}"
)
entities = _ENTITY_LIST_ADAPTER.validate_python(record.guardrail_response)
assert entities, (
"guardrail_response must carry the detected entities; the dashboard's Detected Entities "
"list and its per-entity scores are rendered from exactly this array"
)
assert {entity.entity_type for entity in entities} >= _LOGGED_ENTITIES.keys(), (
f"the detected entities must cover what was masked, got "
f"{sorted(entity.entity_type for entity in entities)}"
)
for entity in entities:
assert 0.0 < entity.score <= 1.0, f"{entity.entity_type} carries an out-of-range score: {entity.score}"
assert 0 <= entity.start < entity.end, (
f"{entity.entity_type} carries a degenerate span: {entity.start}-{entity.end}"
)

256
tests/e2e/llm_translation/test_sail_e2e.py
View file

@ -1,256 +0,0 @@
"""Live e2e: Sail through the gateway, where LiteLLM turns ``service_tier`` into Sail's
``metadata.completion_window`` and bills the price columns of the window it sent.
The deployment carries its own base, balanced and flex rates, each distinct, so a bill at
the wrong tier cannot pass. They are registered on the deployment instead of read from the
proxy's cost map, because a stack that loads the published map has no ``sail/`` rows until
this provider ships. Requires SAIL_API_KEY on the proxy; no skip gate.
"""
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final, Literal
import pytest
from e2e_config import SLOW_PROVIDER_TIMEOUT_SECONDS, unique_marker
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
from lifecycle import ResourceManager
from models import LiteLLMParamsBody, SpendLogRow
from openai import OpenAI
from proxy_client import ProxyClient
from sdk_clients import NO_PROXY_CACHE, SdkClients, response_header
pytestmark = pytest.mark.e2e
BACKEND: Final = "sail/zai-org/GLM-5.3"
PricedTier = Literal["base", "balanced", "flex"]
PRICED_TIERS: Final[tuple[PricedTier, ...]] = ("base", "balanced", "flex")
PROMPT: Final = "Reply with one word."
MAX_TOKENS: Final = 512
@dataclass(frozen=True, slots=True)
class _Rates:
input: float
output: float
cache_read: float
RATES: Final[Mapping[PricedTier, _Rates]] = {
"base": _Rates(input=3e-06, output=9e-06, cache_read=1e-06),
"balanced": _Rates(input=2e-06, output=6e-06, cache_read=7e-07),
"flex": _Rates(input=1e-06, output=3e-06, cache_read=4e-07),
}
@dataclass(frozen=True, slots=True)
class _Tokens:
prompt: int
cached: int
completion: int
def _approx_equal(actual: float, expected: float) -> bool:
return abs(actual - expected) <= max(1e-12, abs(expected) * 1e-2)
def _cost(rates: _Rates, tokens: _Tokens) -> float:
return (
(tokens.prompt - tokens.cached) * rates.input
+ tokens.cached * rates.cache_read
+ tokens.completion * rates.output
)
def _register(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, str]:
model: Final = f"e2e-sail-{unique_marker()}"
model_id: Final = proxy.create_model(
model,
LiteLLMParamsBody(
model=BACKEND,
api_key="os.environ/SAIL_API_KEY",
input_cost_per_token=RATES["base"].input,
output_cost_per_token=RATES["base"].output,
cache_read_input_token_cost=RATES["base"].cache_read,
input_cost_per_token_balanced=RATES["balanced"].input,
output_cost_per_token_balanced=RATES["balanced"].output,
cache_read_input_token_cost_balanced=RATES["balanced"].cache_read,
input_cost_per_token_flex=RATES["flex"].input,
output_cost_per_token_flex=RATES["flex"].output,
cache_read_input_token_cost_flex=RATES["flex"].cache_read,
),
)
resources.defer(lambda: proxy.delete_model(model_id))
return model, resources.key()
def _openai(sdk: SdkClients, key: str) -> OpenAI:
return sdk.openai(key).with_options(timeout=SLOW_PROVIDER_TIMEOUT_SECONDS)
def _assert_billed_at(tier: PricedTier, tokens: _Tokens, header_cost: str | None) -> float:
assert tokens.prompt > 0 and tokens.completion > 0, f"Sail reported no usage, so no cost is real: {tokens}"
assert header_cost is not None, "x-litellm-response-cost header missing"
costs: Final = {priced: _cost(rates, tokens) for priced, rates in RATES.items()}
assert not any(_approx_equal(costs[other], costs[tier]) for other in PRICED_TIERS if other != tier), (
f"{BACKEND} tier rates too close together to tell {tier} apart at {tokens}: {costs}"
)
assert _approx_equal(float(header_cost), costs[tier]), (
f"header cost {header_cost} is not the {tier} price at {tokens}: expected {costs[tier]}, all tiers {costs}"
)
return float(header_cost)
def _assert_spend_row_matches(proxy: ProxyClient, key: str, header_cost: float) -> None:
def priced(rows: list[SpendLogRow]) -> bool:
return any((row.spend or 0) > 0 for row in rows)
rows: Final = [row for row in proxy.poll_logs_for_key(key, predicate=priced) if (row.spend or 0) > 0]
assert rows, f"no priced spend row landed for key {key}"
assert rows[0].custom_llm_provider == "sail", f"spend row misattributed: {rows[0]}"
assert rows[0].spend is not None and _approx_equal(rows[0].spend, header_cost), (
f"logged spend {rows[0].spend} disagrees with the x-litellm-response-cost header {header_cost}"
)
class TestSailChatCompletions:
@pytest.mark.covers("llm.chat_completions.sail.service_tier.nonstream.cost_logged")
@meta(
Subject(
domain=Domain.SPEND_BUDGETS,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.SAIL,),
models=(BACKEND,),
mode=Mode.NONSTREAM,
)
)
@pytest.mark.parametrize(
("service_tier", "billed_tier"), [("balanced", "balanced"), ("auto", "base")]
)
def test_service_tier_bills_the_matching_completion_window(
self,
proxy: ProxyClient,
resources: ResourceManager,
sdk: SdkClients,
service_tier: str,
billed_tier: PricedTier,
) -> None:
model, key = _register(proxy, resources)
raw: Final = _openai(sdk, key).chat.completions.with_raw_response.create(
model=model,
messages=[{"role": "user", "content": f"{PROMPT} {unique_marker()}"}],
max_completion_tokens=MAX_TOKENS,
extra_body={**NO_PROXY_CACHE, "service_tier": service_tier},
)
usage: Final = raw.parse().usage
assert usage is not None, "chat response carries no usage"
details: Final = usage.prompt_tokens_details
tokens: Final = _Tokens(
prompt=usage.prompt_tokens,
cached=(details.cached_tokens or 0) if details else 0,
completion=usage.completion_tokens,
)
header_cost: Final = _assert_billed_at(
billed_tier, tokens, response_header(raw.headers, "x-litellm-response-cost")
)
_assert_spend_row_matches(proxy, key, header_cost)
@pytest.mark.covers("llm.chat_completions.sail.service_tier.nonstream.drops_unknown_tier_and_bills_asap")
@meta(
Subject(
domain=Domain.SPEND_BUDGETS,
route=Route.CHAT_COMPLETIONS,
providers=(Provider.SAIL,),
models=(BACKEND,),
mode=Mode.NONSTREAM,
)
)
@pytest.mark.parametrize("service_tier", ["bogus", 5])
def test_unknown_service_tier_is_dropped_and_billed_asap(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients, service_tier: str | int
) -> None:
model, key = _register(proxy, resources)
raw: Final = _openai(sdk, key).chat.completions.with_raw_response.create(
model=model,
messages=[{"role": "user", "content": f"{PROMPT} {unique_marker()}"}],
max_completion_tokens=MAX_TOKENS,
extra_body={**NO_PROXY_CACHE, "service_tier": service_tier, "drop_params": True},
)
usage: Final = raw.parse().usage
assert usage is not None, "chat response carries no usage"
details: Final = usage.prompt_tokens_details
tokens: Final = _Tokens(
prompt=usage.prompt_tokens,
cached=(details.cached_tokens or 0) if details else 0,
completion=usage.completion_tokens,
)
header_cost: Final = _assert_billed_at("base", tokens, response_header(raw.headers, "x-litellm-response-cost"))
_assert_spend_row_matches(proxy, key, header_cost)
class TestSailResponses:
@pytest.mark.covers("llm.responses.sail.service_tier.nonstream.cost_logged")
@meta(
Subject(
domain=Domain.SPEND_BUDGETS,
route=Route.RESPONSES,
providers=(Provider.SAIL,),
models=(BACKEND,),
mode=Mode.NONSTREAM,
)
)
def test_caller_completion_window_bills_its_rates(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
) -> None:
model, key = _register(proxy, resources)
raw: Final = _openai(sdk, key).responses.with_raw_response.create(
model=model,
input=f"{PROMPT} {unique_marker()}",
max_output_tokens=MAX_TOKENS,
metadata={"completion_window": "balanced"},
extra_body=NO_PROXY_CACHE,
)
usage: Final = raw.parse().usage
assert usage is not None, "responses answer carries no usage"
tokens: Final = _Tokens(
prompt=usage.input_tokens,
cached=usage.input_tokens_details.cached_tokens,
completion=usage.output_tokens,
)
header_cost: Final = _assert_billed_at(
"balanced", tokens, response_header(raw.headers, "x-litellm-response-cost")
)
_assert_spend_row_matches(proxy, key, header_cost)
class TestSailMessages:
@pytest.mark.covers("llm.messages.sail.basic.nonstream.works")
@meta(
Subject(
domain=Domain.LLM_TRANSLATION,
route=Route.MESSAGES,
providers=(Provider.SAIL,),
models=(BACKEND,),
mode=Mode.NONSTREAM,
)
)
def test_plain_call_returns_a_message(
self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients
) -> None:
model, key = _register(proxy, resources)
message: Final = sdk.anthropic(key).messages.create(
model=model,
max_tokens=MAX_TOKENS,
messages=[{"role": "user", "content": PROMPT}],
extra_body=NO_PROXY_CACHE,
)
assert message.role == "assistant" and message.content, f"/v1/messages returned no content: {message}"
assert message.usage.output_tokens > 0, f"/v1/messages reported no output usage: {message.usage}"

39
tests/e2e/logging/logging_client.py
View file

@ -233,45 +233,6 @@ class LangfuseCreds:
)
@dataclass(frozen=True, slots=True)
class WeaveCreds:
"""Weights & Biases Weave credentials for a key-scoped ``weave_otel`` callback.
The proxy still needs WANDB_API_KEY / WANDB_PROJECT_ID in its own environment:
the weave_otel logger is constructed from those before the per-key vars are
applied, so a key-scoped callback on a proxy without them never initializes.
The per-key vars are what direct THIS key's spans at this project.
"""
api_key: str
project_id: str
def key_logging_metadata(self) -> KeyMetadata:
return KeyMetadata(
logging=[
KeyLoggingCallback(
callback_name="weave_otel",
callback_type="success_and_failure",
callback_vars=KeyLoggingCallbackVars(
wandb_api_key=self.api_key,
weave_project_id=self.project_id,
),
)
]
)
def load_weave_creds() -> WeaveCreds:
api_key = os.getenv("WANDB_API_KEY")
project_id = (os.getenv("WEAVE_PROJECT_ID") or os.getenv("WANDB_PROJECT_ID") or "").strip()
if not (api_key and project_id):
pytest.fail(
"Weave e2e requires WANDB_API_KEY and WEAVE_PROJECT_ID (or WANDB_PROJECT_ID, "
"format <entity>/<project>); missing credentials is a hard failure, not a skip"
)
return WeaveCreds(api_key=api_key, project_id=project_id)
def load_langfuse_creds() -> LangfuseCreds:
public_key = os.getenv("LANGFUSE_PUBLIC_KEY")
secret_key = os.getenv("LANGFUSE_SECRET_KEY")

212
tests/e2e/logging/test_weave_log_e2e.py
View file

@ -1,212 +0,0 @@
"""Live e2e: key-scoped Weave (Weights & Biases) delivery, success and failure.
Covers the two `logging.niche_integrations.*.logs_spend` cells with a real member
of that cohort. A key carrying a `weave_otel` callback in its logging metadata
must deliver its calls to the real Weave project, and each call must arrive
exactly once, carrying the same cost the response header reported:
- success: one `litellm_request` call, OTEL status OK, `llm.response.cost` equal
to `x-litellm-response-cost`, and non-zero tokens
- failure: a provider-rejected call arrives too, as one call with OTEL status
ERROR naming the provider exception, and with no cost - a failed call that
silently never reaches the destination is an invisible outage, and a billed
one is worse
Both halves assert the recorded state (the key's callback registration answers
success and the destination holds the call) and the enforced behavior (the
delivered payload's status and cost). Delivery is read back through Weave's own
query API; nothing is mocked.
"""
from __future__ import annotations
import time
from typing import Final
import pytest
from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker
from e2e_http import StreamingResponse
from e2e_metadata import Domain, Mode, Provider, Subject, meta
from lifecycle import ResourceManager
from logging_client import (
INVALID_UPSTREAM_API_KEY,
LoggingClient,
WeaveCreds,
costs_agree,
first_ok,
load_weave_creds,
)
from models import LiteLLMParamsBody
from weave_reader import WeaveCall, WeaveReader, build_weave_reader
pytestmark = pytest.mark.e2e
UNREACHABLE_ANTHROPIC_BACKEND: Final = "anthropic/claude-haiku-4-5"
@pytest.fixture(scope="session")
def weave_creds() -> WeaveCreds:
return load_weave_creds()
@pytest.fixture(scope="session")
def weave_reader() -> WeaveReader:
return build_weave_reader()
#: How far before the request the Weave read-back window opens, to absorb clock
#: skew between this host and Weave. Without it a host running slightly fast
#: would filter out its own call.
_WINDOW_SKEW_SECONDS = 120.0
def _window_start() -> float:
return time.time() - _WINDOW_SKEW_SECONDS
def _exactly_one(calls: tuple[WeaveCall, ...], *, marker: str, what: str) -> WeaveCall:
assert calls, f"no Weave call for the {what} (marker {marker}) reached the project within the deadline"
assert len(calls) == 1, (
f"expected exactly ONE Weave call for the {what} (marker {marker}), got {len(calls)}: "
f"{[call.id for call in calls]} - more than one call for one request is the "
"duplicate-delivery bug"
)
return calls[0]
WEAVE_STAGE_RED_REASON = (
"stage red: product gap, key-scoped weave_otel spans are not delivered when the OTEL v2 callback is active"
)
class TestWeaveLogDelivery:
@pytest.mark.skip(reason=WEAVE_STAGE_RED_REASON)
@pytest.mark.covers("logging.niche_integrations.success.logs_spend", exercised_on=["chat_completions"])
@meta(
Subject(
domain=Domain.OBSERVABILITY,
providers=(Provider.ANTHROPIC,),
models=(CHEAP_ANTHROPIC_MODEL,),
mode=Mode.NONSTREAM,
)
)
def test_chat_completions_delivers_one_call_with_spend(
self,
client: LoggingClient,
weave_creds: WeaveCreds,
weave_reader: WeaveReader,
resources: ResourceManager,
) -> None:
alias = f"weave-key-{unique_marker()}"
key = client.key_with_alias(
alias,
models=[CHEAP_ANTHROPIC_MODEL],
metadata=weave_creds.key_logging_metadata(),
)
resources.defer(lambda: client.delete_key(key))
marker = unique_marker()
since = _window_start()
outcome = first_ok(
client,
lambda: client.chat_raw(key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=64),
)
assert outcome.response_cost is not None and outcome.response_cost > 0, (
f"the response must report x-litellm-response-cost, got {outcome.response_cost!r}"
)
call = _exactly_one(
weave_reader.poll_calls_matching(marker, since=since), marker=marker, what="successful call"
)
assert call.status_code == "OK", f"a successful call must land at OK span status, got {call.status_code!r}"
cost = call.response_cost
assert cost is not None and costs_agree(outcome.response_cost, cost), (
f"the Weave call's llm.response.cost {cost!r} must agree with the header cost "
f"{outcome.response_cost} - a delivered span with the wrong cost is a silent "
"billing-attribution bug"
)
assert call.total_tokens is not None and call.total_tokens > 0, (
f"the delivered call must carry token usage, got {call.total_tokens!r}"
)
@pytest.mark.skip(reason=WEAVE_STAGE_RED_REASON)
@pytest.mark.covers("logging.niche_integrations.failure.logs_spend", exercised_on=["chat_completions"])
@meta(
Subject(
domain=Domain.OBSERVABILITY,
providers=(Provider.ANTHROPIC,),
models=(UNREACHABLE_ANTHROPIC_BACKEND,),
mode=Mode.NONSTREAM,
)
)
def test_failed_chat_completions_delivers_one_error_call(
self,
client: LoggingClient,
weave_creds: WeaveCreds,
weave_reader: WeaveReader,
resources: ResourceManager,
) -> None:
"""A deployment with an invalid upstream key passes proxy auth and fails
at the provider, so exactly one provider failure exists for it. Proxy-side
401s during key propagation never reach the provider and ship no payload,
which is what the retry loop below relies on."""
model_name = f"weave-err-{unique_marker()}"
model_id = client.create_model(
model_name,
LiteLLMParamsBody(model=UNREACHABLE_ANTHROPIC_BACKEND, api_key=INVALID_UPSTREAM_API_KEY),
)
resources.defer(lambda: client.delete_model(model_id))
key = client.key_with_alias(
f"weave-err-key-{unique_marker()}",
models=[model_name],
metadata=weave_creds.key_logging_metadata(),
)
resources.defer(lambda: client.delete_key(key))
marker = unique_marker()
since = _window_start()
outcome = _provoke_provider_failure(client, key, model_name, marker)
call = _exactly_one(weave_reader.poll_calls_matching(marker, since=since), marker=marker, what="failed call")
assert call.status_code == "ERROR", (
f"a failed call must land at ERROR span status, got {call.status_code!r} - "
"Weave's own summary.weave.status reads success either way, which is exactly "
"why the span status is what this asserts on"
)
error = call.error
assert error is not None and error.message is not None and "AnthropicException" in error.message, (
f"the delivered call must carry the provider error, got {error!r}"
)
assert not call.response_cost, f"a failed call must not be billed, got llm.response.cost={call.response_cost!r}"
assert outcome.status_code == 401, (
f"an upstream auth failure must map to 401, got {outcome.status_code}: {outcome.body[:200]}"
)
def _provoke_provider_failure(client: LoggingClient, key: str, model_name: str, marker: str) -> StreamingResponse:
"""Send until the provider (not the proxy) is the one rejecting the call.
A network failure between the test and the proxy is NOT retried: the request
may have been served, and a retry would double-log the failure payload and
falsely trip the exactly-one assertion.
"""
deadline = time.monotonic() + client.proxy.poll_timeout
while True:
outcome = client.chat_raw(key, model_name, f"trigger an upstream auth failure {marker}", max_tokens=16)
assert not outcome.ok, "the call must fail; the deployment's upstream key is invalid"
assert outcome.status_code != -1, (
"network failure between the test and the proxy while provoking the provider failure; "
"retrying now could double-log the failure payload and falsely trip the exactly-one "
f"assertion - fix the rig connectivity first: {outcome.body[:200]}"
)
if "AnthropicException" in outcome.body or time.monotonic() >= deadline:
break
time.sleep(client.proxy.poll_interval)
assert "AnthropicException" in outcome.body, (
"never saw the upstream provider failure before the deadline; the key may still be "
f"propagating - last outcome {outcome.status_code}: {outcome.body[:200]}"
)
return outcome

285
tests/e2e/logging/weave_reader.py
View file

@ -1,285 +0,0 @@
"""Read-back for the Weave (Weights & Biases) logging tests against the real
Weave project.
The proxy ships OTEL spans to https://trace.wandb.ai/otel/v1/traces with the
``weave_otel`` callback, and the tests read the ingested calls back through
Weave's own query API (``POST /calls/stream_query``), which answers JSON Lines:
one JSON object per call, so the body is parsed line by line rather than as one
document.
The project is shared with other traffic, so the read never relies on the target
being among the newest N calls: the query is scoped server-side to the
``litellm_request`` op and to calls that started after the test's own request,
and pages with ``offset`` until the window is exhausted.
Weave's own ``summary.weave.status`` is a rollup that reads "success" even for a
span the exporter marked failed, so status comes from the OTEL span itself
(``attributes.otel_span.status.code``), and the shipped cost from
``attributes.otel_span.attributes.llm.response.cost`` - the StandardLogging
``response_cost``, which is what makes this a spend assertion rather than a
delivery ping.
Missing configuration is a hard failure, never a skip.
"""
from __future__ import annotations
import base64
import json
import os
import time
from dataclasses import dataclass, field
from itertools import count, takewhile
from typing import Final
import pytest
from pydantic import BaseModel, ConfigDict, Field
from e2e_config import POLL_INTERVAL, POLL_TIMEOUT
from e2e_http import URL, AuthHeaders, send
from e2e_metadata import step
_WEAVE_TRACE_API: Final = "https://trace.wandb.ai"
#: The op every litellm LLM call lands under. The proxy also exports a root
#: server span ("Received Proxy Server Request") and management spans; only the
#: LLM call carries the usage and cost this suite asserts on.
LITELLM_REQUEST_OP: Final = "litellm_request"
#: How long to keep re-reading after the first matching call before trusting the
#: exactly-one assertion. The OTEL batch exporter flushes on its own schedule, so
#: a duplicate export can surface well after the first one, and a duplicate IS
#: the bug being guarded against.
WEAVE_SETTLE_SECONDS: Final = 45.0
#: Rows per page. The query is already scoped to this run's time window, so this
#: only bounds one round trip, not what the read can see.
_PAGE_SIZE: Final = 500
class _WeaveSortBy(BaseModel):
field: str
direction: str
class _WeaveOpFilter(BaseModel):
op_names: list[str]
class _WeaveGetField(BaseModel):
get_field: str = Field(serialization_alias="$getField")
class _WeaveLiteral(BaseModel):
literal: float = Field(serialization_alias="$literal")
class _WeaveGreaterThan(BaseModel):
gt: tuple[_WeaveGetField, _WeaveLiteral] = Field(serialization_alias="$gt")
class _WeaveQuery(BaseModel):
expr: _WeaveGreaterThan = Field(serialization_alias="$expr")
class _WeaveQueryBody(BaseModel):
project_id: str
filter: _WeaveOpFilter
query: _WeaveQuery
limit: int = _PAGE_SIZE
offset: int = 0
sort_by: list[_WeaveSortBy] = [_WeaveSortBy(field="started_at", direction="asc")]
class _OtelStatus(BaseModel):
model_config = ConfigDict(extra="ignore")
code: str | None = None
message: str | None = None
class _OtelError(BaseModel):
model_config = ConfigDict(extra="ignore")
code: str | None = None
type: str | None = None
message: str | None = None
class _LlmResponse(BaseModel):
model_config = ConfigDict(extra="ignore")
cost: float | None = None
class _LlmAttributes(BaseModel):
model_config = ConfigDict(extra="ignore")
response: _LlmResponse | None = None
class _OtelSpanAttributes(BaseModel):
model_config = ConfigDict(extra="ignore")
llm: _LlmAttributes | None = None
error: _OtelError | None = None
class _OtelSpan(BaseModel):
model_config = ConfigDict(extra="ignore")
name: str | None = None
status: _OtelStatus | None = None
attributes: _OtelSpanAttributes | None = None
class _CallAttributes(BaseModel):
model_config = ConfigDict(extra="ignore")
otel_span: _OtelSpan | None = None
class _Usage(BaseModel):
model_config = ConfigDict(extra="ignore")
total_tokens: int | None = None
class _WeaveSummary(BaseModel):
model_config = ConfigDict(extra="ignore")
usage: dict[str, _Usage] = {}
class WeaveCall(BaseModel):
"""One ingested Weave call, reduced to what the scenarios assert on."""
model_config = ConfigDict(extra="ignore")
id: str
op_name: str
started_at: str | None = None
inputs: dict[str, object] = {}
attributes: _CallAttributes | None = None
summary: _WeaveSummary | None = Field(default=None)
@property
def op(self) -> str:
"""The bare op name out of ``weave:///<entity>/<project>/op/<op>:<digest>``."""
return self.op_name.split("/op/")[-1].split(":")[0]
@property
def status_code(self) -> str | None:
"""The OTEL span status, not Weave's own rollup (which reads "success"
even for a span the exporter marked ERROR)."""
span = self.attributes.otel_span if self.attributes else None
return span.status.code if span and span.status else None
@property
def error(self) -> _OtelError | None:
span = self.attributes.otel_span if self.attributes else None
return span.attributes.error if span and span.attributes else None
@property
def response_cost(self) -> float | None:
span = self.attributes.otel_span if self.attributes else None
llm = span.attributes.llm if span and span.attributes else None
return llm.response.cost if llm and llm.response else None
@property
def total_tokens(self) -> int | None:
"""Weave keys usage by model, so the total is summed across whatever
models the call reported."""
if not self.summary or not self.summary.usage:
return None
totals = [usage.total_tokens for usage in self.summary.usage.values() if usage.total_tokens is not None]
return sum(totals) if totals else None
def mentions(self, needle: str) -> bool:
return needle in json.dumps(self.inputs, default=str)
@dataclass(frozen=True, slots=True)
class WeaveReader:
project_id: str
api_key: str = field(repr=False)
@property
def _headers(self) -> AuthHeaders:
"""Weave authenticates with HTTP Basic as the fixed user ``api``."""
token = base64.b64encode(f"api:{self.api_key}".encode()).decode()
return AuthHeaders(authorization=f"Basic {token}")
def _query_body(self, *, since: float, offset: int, op: str) -> _WeaveQueryBody:
return _WeaveQueryBody(
project_id=self.project_id,
filter=_WeaveOpFilter(op_names=[f"weave:///{self.project_id}/op/{op}:*"]),
query=_WeaveQuery(
expr=_WeaveGreaterThan(gt=(_WeaveGetField(get_field="started_at"), _WeaveLiteral(literal=since)))
),
offset=offset,
)
def _page(self, *, since: float, offset: int, op: str) -> tuple[WeaveCall, ...]:
outcome = send(
URL(f"{_WEAVE_TRACE_API}/calls/stream_query"),
headers=self._headers,
json=self._query_body(since=since, offset=offset, op=op),
)
if not outcome.ok:
pytest.fail(
f"Weave calls query for project {self.project_id!r} failed "
f"({outcome.status_code}): {outcome.body[:300]}"
)
return tuple(WeaveCall.model_validate_json(line) for line in outcome.body.splitlines() if line.strip())
@step("Read the Weave {op} calls carrying the marker {marker}")
def calls_matching(self, marker: str, *, since: float, op: str = LITELLM_REQUEST_OP) -> tuple[WeaveCall, ...]:
"""Every call under ``op`` started after ``since`` whose inputs carry
``marker``, paging until the window is exhausted.
More than one is the duplicate-delivery bug, so this never collapses to a
single call.
"""
pages = tuple(
takewhile(
bool,
(self._page(since=since, offset=offset, op=op) for offset in count(0, _PAGE_SIZE)),
)
)
return tuple(call for page in pages for call in page if call.mentions(marker))
@step("Wait for Weave to ingest a {op} call carrying the marker {marker}, then watch for duplicates")
def poll_calls_matching(self, marker: str, *, since: float, op: str = LITELLM_REQUEST_OP) -> tuple[WeaveCall, ...]:
"""Poll until the call is readable, then keep re-reading for
WEAVE_SETTLE_SECONDS so a duplicate exported by a later batch flush
cannot hide from the exactly-one assertion. A duplicate ends the settle
early, because more waiting cannot clear it."""
deadline = time.monotonic() + POLL_TIMEOUT
while time.monotonic() < deadline:
calls = self.calls_matching(marker, since=since, op=op)
if calls:
return self._settled(marker, since=since, op=op, first=calls)
time.sleep(POLL_INTERVAL)
return ()
def _settled(self, marker: str, *, since: float, op: str, first: tuple[WeaveCall, ...]) -> tuple[WeaveCall, ...]:
"""A transiently empty re-read never downgrades what was already seen."""
settle_deadline = time.monotonic() + WEAVE_SETTLE_SECONDS
latest = first # rebind-ok: one settle window, re-read per poll interval
while time.monotonic() < settle_deadline and len(latest) <= 1:
time.sleep(POLL_INTERVAL)
latest = self.calls_matching(marker, since=since, op=op) or latest
return latest
def build_weave_reader() -> WeaveReader:
project_id = (os.environ.get("WEAVE_PROJECT_ID") or os.environ.get("WANDB_PROJECT_ID") or "").strip()
api_key = os.environ.get("WANDB_API_KEY", "").strip()
if not project_id or not api_key:
pytest.fail(
"Weave e2e requires WANDB_API_KEY and WEAVE_PROJECT_ID (or WANDB_PROJECT_ID, "
"format <entity>/<project>): the test reads the proxy's weave_otel delivery "
"back from the real Weave project; missing credentials is a hard failure, not a skip"
)
return WeaveReader(project_id=project_id, api_key=api_key)

2
tests/e2e/models.py
View file

@ -45,8 +45,6 @@ class KeyLoggingCallbackVars(BaseModel):
langfuse_public_key: str | None = Field(default=None, repr=False)
langfuse_secret_key: str | None = Field(default=None, repr=False)
langfuse_host: str | None = Field(default=None, repr=False)
wandb_api_key: str | None = Field(default=None, repr=False)
weave_project_id: str | None = None
class KeyLoggingCallback(BaseModel):

53
tests/integration/compatibility/test_a2a_wire_versions.py
View file

@ -755,6 +755,59 @@ def test_a2a_message_send_aliases_and_agent_name_reach_the_same_agent_with_the_s
), posts
def test_a2a_data_part_skill_call_and_data_artifact_reach_each_side_unchanged(gateway: Gateway) -> None:
marker: Final = "a2adata" + uuid.uuid4().hex
invocation: Final[dict[str, JsonValue]] = {"skill": "search_items", "params": {"region": "emea", "limit": 3}}
found: Final[dict[str, JsonValue]] = {
"total": 2,
"count": 2,
"listings": [{"id": marker + "-1", "region": "emea"}, {"id": marker + "-2", "region": "emea"}],
}
task: Final[dict[str, JsonValue]] = {
"kind": "task",
"id": marker + "-task",
"contextId": marker + "-context",
"status": {"state": "completed"},
"artifacts": [
{"artifactId": marker + "-artifact", "name": "search_results", "parts": [{"kind": "data", "data": found}]}
],
}
def upstream(request: Request) -> Reply:
if request.method == "GET":
return Reply(body=json.dumps(_peer_card(wire.url, marker)).encode())
body: Final = json.loads(request.body)
return Reply(body=json.dumps({"jsonrpc": "2.0", "id": body["id"], "result": task}).encode())
with wire_server(upstream) as wire, gateway.scenario() as scenario:
identity: Final = _register_agent(
gateway, scenario, {"agent_name": marker, "agent_card_params": _peer_card(wire.url, marker)}
)
response: Final = gateway.client.post(
f"/a2a/{identity}",
headers={"Authorization": f"Bearer {gateway.key}"},
json={
"jsonrpc": "2.0",
"id": marker,
"method": "message/send",
"params": {
"message": {
"kind": "message",
"role": "user",
"messageId": marker + "-in",
"parts": [{"kind": "data", "data": invocation}],
}
},
},
)
assert response.status_code == 200, response.text
assert response.json() == {"jsonrpc": "2.0", "id": marker, "result": task}, response.text
posts: Final = tuple(json.loads(item.body) for item in wire.drain() if item.method == "POST")
assert len(posts) == 1, posts
assert posts[0]["method"] == "message/send", posts
assert posts[0]["params"]["message"]["parts"] == [{"kind": "data", "data": invocation}], posts
def test_a2a_task_methods_pass_through_verbatim_and_push_callbacks_are_validated_before_forwarding(
gateway: Gateway,
) -> None:

593
tests/integration/observability/test_presidio_masking_endpoints.py Normal file
View file

@ -0,0 +1,593 @@
import json
import re
import uuid
from collections.abc import Callable, Iterator
from contextlib import ExitStack
from dataclasses import dataclass
from itertools import accumulate, chain
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway, eventually, gateway_from_environment, list_value, 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
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter
CARD: Final = "4111 1111 1111 1111"
PHONE: Final = "+1 415-555-0134"
MASKED_CARD: Final = "<CREDIT_CARD>"
MASKED_EMAIL: Final = "<EMAIL_ADDRESS>"
MASKED_PHONE: Final = "<PHONE_NUMBER>"
REPEAT_REQUEST: Final = "Repeat this sentence back to me exactly, word for word: "
THREE_LINE_REQUEST: Final = (
"Reply with exactly three lines and nothing else. "
"Line 1: the part of the email address before the @ sign. "
"Line 2: the full email address. "
"Line 3: the phone number."
)
CARD_REQUEST: Final = "Reply with only the well known Visa sandbox test card number that starts with 4111."
EMAIL_PATTERN: Final = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9-]+(?:\.[A-Za-z0-9-]+)*\.[A-Za-z]{2,}")
PHONE_PATTERN: Final = re.compile(r"(?:\+1 )?\d{3}-\d{3}-\d{4}")
CARD_PATTERN: Final = re.compile(r"\b(?:4\d{3}|5[0-5]\d{2}|6\d{3}|1\d{3}|3\d{3})[- ]?\d{3,4}[- ]?\d{3,4}[- ]?\d{3,5}\b")
RECOGNIZERS: Final = {
"EMAIL_ADDRESS": (EMAIL_PATTERN, 1.0),
"PHONE_NUMBER": (PHONE_PATTERN, 0.4),
"CREDIT_CARD": (CARD_PATTERN, 1.0),
}
JSON: Final = TypeAdapter(JsonValue)
class _Frozen(BaseModel):
model_config = ConfigDict(frozen=True)
class _AnalyzeRequest(_Frozen):
text: str = ""
language: str = ""
entities: tuple[str, ...] | None = None
class _Span(_Frozen):
entity_type: str
start: int
end: int
score: float
class _AnonymizeRequest(_Frozen):
text: str
analyzer_results: tuple[_Span, ...]
class _Part(_Frozen):
text: str
class _Message(_Frozen):
content: str | tuple[_Part, ...]
class _ModelRequest(_Frozen):
messages: tuple[_Message, ...] = ()
input: str | tuple[_Message, ...] | None = None
stream: bool = False
@property
def prompt(self) -> str:
turns: Final = self.messages if self.input is None else self.input
if isinstance(turns, str):
return turns
content: Final = turns[-1].content
return content if isinstance(content, str) else "".join(part.text for part in content)
class _StreamDelta(_Frozen):
content: str | None = None
class _StreamChoice(_Frozen):
delta: _StreamDelta
class _ChatChunk(_Frozen):
choices: tuple[_StreamChoice, ...] = ()
class _AnthropicDelta(_Frozen):
type: str | None = None
text: str | None = None
class _AnthropicEvent(_Frozen):
type: str
delta: _AnthropicDelta | None = None
class _LoggedEntity(_Frozen):
entity_type: str
start: int
end: int
score: float
def _passes_luhn(digits: str) -> bool:
checksum: Final = sum(
digit if position % 2 == 0 else (digit * 2 - 9 if digit * 2 > 9 else digit * 2)
for position, digit in enumerate(int(char) for char in reversed(digits))
)
return checksum % 10 == 0
def _detections(entity: str, text: str) -> tuple[_Span, ...]:
pattern, score = RECOGNIZERS[entity]
return tuple(
_Span(entity_type=entity, start=match.start(), end=match.end(), score=score)
for match in pattern.finditer(text)
if entity != "CREDIT_CARD" or _passes_luhn(re.sub(r"\D", "", match.group()))
)
def _analyze(request: Request) -> Reply:
assert request.target == "/analyze", request.target
body: Final = _AnalyzeRequest.model_validate_json(request.body)
if not body.text or not body.language:
return Reply(status=400, body=json.dumps({"error": "No text provided"}).encode())
requested: Final = body.entities or tuple(RECOGNIZERS)
findings: Final = chain.from_iterable(
_detections(entity, body.text) for entity in requested if entity in RECOGNIZERS
)
return Reply(body=json.dumps([span.model_dump() for span in sorted(findings, key=lambda s: s.start)]).encode())
def _anonymize(request: Request) -> Reply:
assert request.target == "/anonymize", request.target
body: Final = _AnonymizeRequest.model_validate_json(request.body)
spans: Final = sorted(body.analyzer_results, key=lambda span: span.start)
masks: Final = [f"<{span.entity_type}>" for span in spans]
shifts: Final = tuple(
accumulate((len(mask) - (span.end - span.start) for mask, span in zip(masks, spans)), initial=0)
)
items: Final = [
{
"start": span.start + shift,
"end": span.start + shift + len(mask),
"entity_type": span.entity_type,
"text": mask,
"operator": "replace",
}
for span, mask, shift in zip(spans, masks, shifts)
]
pieces: Final = [
body.text[(spans[index - 1].end if index else 0) : span.start] + masks[index]
for index, span in enumerate(spans)
]
tail: Final = body.text[spans[-1].end :] if spans else body.text
return Reply(body=json.dumps({"text": "".join(pieces) + tail, "items": items}).encode())
def _model_answer(prompt: str) -> str:
if REPEAT_REQUEST in prompt:
return prompt.split(REPEAT_REQUEST, 1)[1]
if THREE_LINE_REQUEST in prompt:
email: Final = EMAIL_PATTERN.search(prompt)
phone: Final = PHONE_PATTERN.search(prompt)
assert email is not None and phone is not None, prompt
return "\n".join((email.group().split("@")[0], email.group(), phone.group()))
assert CARD_REQUEST in prompt, prompt
return CARD
def _chat_completion(answer: str) -> dict[str, JsonValue]:
return {
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": answer}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 11, "completion_tokens": 5, "total_tokens": 16},
}
def _chat_frames(answer: str) -> tuple[bytes, ...]:
identity: Final = f"chatcmpl-{uuid.uuid4().hex}"
def frame(delta: dict[str, JsonValue], finish: str | None = None) -> bytes:
payload: Final = {
"id": identity,
"object": "chat.completion.chunk",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
}
return b"data: " + json.dumps(payload).encode() + b"\n\n"
return (
frame({"role": "assistant", "content": ""}),
*(frame({"content": word}) for word in re.findall(r"\S+\s*", answer)),
frame({}, "stop"),
b"data: [DONE]\n\n",
)
def _response_object(answer: str, identity: str, item: str, *, completed: bool = True) -> dict[str, JsonValue]:
return {
"id": identity,
"object": "response",
"created_at": 1,
"status": "completed" if completed else "in_progress",
"model": "gpt-4o-mini",
"output": [_message_item(answer, item)] if completed else [],
"usage": {
"input_tokens": 11,
"output_tokens": 5,
"total_tokens": 16,
"input_tokens_details": {"cached_tokens": 0},
"output_tokens_details": {"reasoning_tokens": 0},
}
if completed
else None,
}
def _message_item(answer: str, item: str, *, completed: bool = True) -> dict[str, JsonValue]:
return {
"type": "message",
"id": item,
"status": "completed" if completed else "in_progress",
"role": "assistant",
"content": [_output_text(answer)] if completed else [],
}
def _output_text(text: str) -> dict[str, JsonValue]:
return {"type": "output_text", "text": text, "annotations": []}
def _response_frames(answer: str) -> tuple[bytes, ...]:
identity: Final = f"resp_{uuid.uuid4().hex}"
item: Final = f"msg_{uuid.uuid4().hex}"
location: Final[dict[str, JsonValue]] = {"item_id": item, "output_index": 0, "content_index": 0}
events: Final[tuple[dict[str, JsonValue], ...]] = (
{"type": "response.created", "response": _response_object("", identity, item, completed=False)},
{
"type": "response.output_item.added",
"output_index": 0,
"item": _message_item("", item, completed=False),
},
{"type": "response.content_part.added", **location, "part": _output_text("")},
*({"type": "response.output_text.delta", **location, "delta": word} for word in re.findall(r"\S+\s*", answer)),
{"type": "response.output_text.done", **location, "text": answer},
{"type": "response.content_part.done", **location, "part": _output_text(answer)},
{"type": "response.output_item.done", "output_index": 0, "item": _message_item(answer, item)},
{"type": "response.completed", "response": _response_object(answer, identity, item)},
)
return tuple(
f"event: {event['type']}\ndata: {json.dumps({**event, 'sequence_number': number})}\n\n".encode()
for number, event in enumerate(events)
)
def _model(request: Request) -> Reply:
body: Final = _ModelRequest.model_validate_json(request.body)
answer: Final = _model_answer(body.prompt)
match request.target:
case "/v1/chat/completions" if body.stream:
return Reply(content_type="text/event-stream", chunks=_chat_frames(answer))
case "/v1/chat/completions":
return Reply(body=json.dumps(_chat_completion(answer)).encode())
case "/v1/responses" if body.stream:
return Reply(content_type="text/event-stream", chunks=_response_frames(answer))
case "/v1/responses":
return Reply(
body=json.dumps(
_response_object(answer, f"resp_{uuid.uuid4().hex}", f"msg_{uuid.uuid4().hex}")
).encode()
)
case _:
raise AssertionError(request.target)
def _presidio_guardrail(
name: str,
analyzer: Wire,
anonymizer: Wire,
*,
mode: str | list[str],
scope: str | None,
entities: dict[str, JsonValue] | None,
) -> dict[str, JsonValue]:
return {
"guardrail_name": name,
"litellm_params": {
"guardrail": "presidio",
"mode": mode,
"default_on": False,
"presidio_analyzer_api_base": analyzer.url,
"presidio_anonymizer_api_base": anonymizer.url,
**({} if scope is None else {"presidio_filter_scope": scope}),
**({} if entities is None else {"pii_entities_config": entities}),
},
}
@dataclass(frozen=True, slots=True)
class Guardrails:
input_scope: str
logged: str
output_scope: str
default_scope: str
card_output_scope: str
@dataclass(frozen=True, slots=True)
class Rig:
gateway: Gateway
model: Wire
guardrails: Guardrails
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
directory: Final = tmp_path_factory.mktemp("presidio-masking-endpoints")
suffix: Final = uuid.uuid4().hex[:8]
names: Final = Guardrails(
input_scope=f"presidio-input-{suffix}",
logged=f"presidio-logged-{suffix}",
output_scope=f"presidio-output-{suffix}",
default_scope=f"presidio-default-{suffix}",
card_output_scope=f"presidio-card-output-{suffix}",
)
contact_entities: Final[dict[str, JsonValue]] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"}
card_entities: Final[dict[str, JsonValue]] = {"CREDIT_CARD": "MASK"}
with ExitStack() as stack:
model: Final = stack.enter_context(wire_server(_model))
analyzer: Final = stack.enter_context(wire_server(_analyze))
anonymizer: Final = stack.enter_context(wire_server(_anonymize))
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["general_settings"]["store_prompts_in_spend_logs"] = True
config["guardrails"] = [
_presidio_guardrail(names.input_scope, analyzer, anonymizer, mode="pre_call", scope="input", entities=None),
_presidio_guardrail(
names.logged, analyzer, anonymizer, mode="pre_call", scope="input", entities=contact_entities
),
_presidio_guardrail(
names.output_scope, analyzer, anonymizer, mode="post_call", scope="output", entities=contact_entities
),
_presidio_guardrail(
names.default_scope,
analyzer,
anonymizer,
mode=["pre_call", "post_call"],
scope=None,
entities=card_entities,
),
_presidio_guardrail(
names.card_output_scope, analyzer, anonymizer, mode="post_call", scope="output", entities=card_entities
),
]
path: Final = directory / "presidio-masking.yaml"
path.write_text(yaml.safe_dump(config))
shared: Final = stack.enter_context(gateway_from_environment())
owned: Final = stack.enter_context(owned_proxy_process(shared, directory, {}, config=path))
yield Rig(owned.gateway, model, names)
@pytest.fixture
def model(rig: Rig) -> Iterator[str]:
with rig.gateway.scenario() as scenario:
yield scenario.model(api_base=rig.model.url + "/v1")
def _content_strings(node: JsonValue, key: str | None = None) -> Iterator[str]:
match node:
case str() if key in {"content", "text"}:
yield node
case dict():
for child_key, child in node.items():
yield from _content_strings(child, child_key)
case list():
for item in node:
yield from _content_strings(item, key)
case _:
return
def _json_text(raw: str) -> str:
return "\n".join(_content_strings(JSON.validate_json(raw)))
def _data_events(raw: str) -> tuple[str, ...]:
return tuple(
line.removeprefix("data: ") for line in raw.splitlines() if line.startswith("data: ") and line != "data: [DONE]"
)
def _chat_stream_text(raw: str) -> str:
return "".join(
choice.delta.content
for event in _data_events(raw)
for choice in _ChatChunk.model_validate_json(event).choices[:1]
if choice.delta.content
)
def _messages_stream_text(raw: str) -> str:
return "".join(
event.delta.text
for payload in _data_events(raw)
for event in [_AnthropicEvent.model_validate_json(payload)]
if event.type == "content_block_delta"
and event.delta is not None
and event.delta.type == "text_delta"
and event.delta.text
)
@dataclass(frozen=True, slots=True)
class Surface:
path: str
stream: bool
takes_input: bool
read: Callable[[str], str]
CHAT: Final = Surface("/v1/chat/completions", False, False, _json_text)
CHAT_STREAM: Final = Surface("/v1/chat/completions", True, False, _chat_stream_text)
MESSAGES: Final = Surface("/v1/messages", False, False, _json_text)
MESSAGES_STREAM: Final = Surface("/v1/messages", True, False, _messages_stream_text)
RESPONSES: Final = Surface("/v1/responses", False, True, _json_text)
@dataclass(frozen=True, slots=True)
class Answer:
text: str
wire: str
applied_guardrails: str
sent_to_model: tuple[str, ...]
def _request_body(surface: Surface, model: str, guardrail: str, prompt: str) -> dict[str, JsonValue]:
common: Final[dict[str, JsonValue]] = {"model": model, "guardrails": [guardrail], "stream": surface.stream}
if surface.takes_input:
return {**common, "input": prompt}
return {**common, "max_tokens": 512, "messages": [{"role": "user", "content": prompt}]}
def _ask(rig: Rig, model: str, surface: Surface, guardrail: str, prompt: str) -> Answer:
rig.model.drain()
with rig.gateway.client.stream(
"POST",
surface.path,
json=_request_body(surface, model, guardrail, prompt),
headers={"Authorization": f"Bearer {rig.gateway.key}"},
) as response:
wire: Final = response.read().decode()
status: Final = response.status_code
applied: Final = response.headers.get("x-litellm-applied-guardrails", "")
assert status == 200, wire
return Answer(
surface.read(wire),
wire,
applied,
tuple(_ModelRequest.model_validate_json(request.body).prompt for request in rig.model.drain()),
)
def _spend_row(rig: Rig, model: str) -> dict[str, JsonValue]:
rows: Final = eventually(
lambda: read_rows('SELECT response, metadata FROM "LiteLLM_SpendLogs" WHERE model_group = %s', (model,)),
lambda found: len(found) == 1,
seconds=70,
)
return rows[0]
def _marker() -> str:
return uuid.uuid4().hex[:8]
def _repeat_prompt(marker: str, email: str) -> str:
return f"{marker} {REPEAT_REQUEST}My email address is {email} and my phone number is {PHONE}."
def _three_line_prompt(marker: str, local_part: str) -> str:
return f"{marker} My email address is {local_part}@example.com and my phone number is {PHONE}. {THREE_LINE_REQUEST}"
def _card_prompt(marker: str) -> str:
return f"{marker} {CARD_REQUEST}"
@pytest.mark.parametrize("surface", (CHAT, MESSAGES), ids=("chat", "messages"))
def test_pre_call_masks_contact_details_before_the_model_sees_them(rig: Rig, model: str, surface: Surface) -> None:
marker: Final = _marker()
email: Final = f"jane.doe.{marker}@example.com"
answer: Final = _ask(rig, model, surface, rig.guardrails.input_scope, _repeat_prompt(marker, email))
masked_sentence: Final = f"My email address is {MASKED_EMAIL} and my phone number is {MASKED_PHONE}."
assert answer.sent_to_model == (f"{marker} {REPEAT_REQUEST}{masked_sentence}",), answer.sent_to_model
assert answer.text == masked_sentence, answer.wire
def test_post_call_masks_the_model_output_while_the_model_sees_the_raw_request(rig: Rig, model: str) -> None:
marker: Final = _marker()
local_part: Final = f"e2euser{marker}"
prompt: Final = _three_line_prompt(marker, local_part)
answer: Final = _ask(rig, model, CHAT, rig.guardrails.output_scope, prompt)
assert answer.sent_to_model == (prompt,), answer.sent_to_model
assert answer.text == f"{local_part}\n{MASKED_EMAIL}\n{MASKED_PHONE}", answer.wire
@pytest.mark.parametrize(
"surface", (CHAT, CHAT_STREAM, MESSAGES_STREAM), ids=("chat", "chat_stream", "messages_stream")
)
def test_default_scope_masks_a_card_number_the_model_generates(rig: Rig, model: str, surface: Surface) -> None:
answer: Final = _ask(rig, model, surface, rig.guardrails.default_scope, _card_prompt(_marker()))
assert answer.text == MASKED_CARD, answer.wire
assert "4111" not in answer.wire, answer.wire
@pytest.mark.parametrize(
"surface",
(CHAT, CHAT_STREAM, MESSAGES, MESSAGES_STREAM, RESPONSES),
ids=("chat", "chat_stream", "messages", "messages_stream", "responses"),
)
def test_spend_log_stores_the_masked_output_the_caller_received(rig: Rig, model: str, surface: Surface) -> None:
answer: Final = _ask(rig, model, surface, rig.guardrails.card_output_scope, _card_prompt(_marker()))
assert answer.text == MASKED_CARD, answer.wire
stored: Final = _spend_row(rig, model)["response"]
assert "\n".join(_content_strings(stored)) == MASKED_CARD, stored
assert "4111" not in json.dumps(stored), stored
def test_masking_run_is_recorded_on_the_spend_log_with_the_detected_entities(rig: Rig, model: str) -> None:
marker: Final = _marker()
email: Final = f"jane.doe.{marker}@example.com"
prompt: Final = _repeat_prompt(marker, email)
answer: Final = _ask(rig, model, CHAT, rig.guardrails.logged, prompt)
assert rig.guardrails.logged in answer.applied_guardrails, answer.applied_guardrails
metadata: Final = object_value(_spend_row(rig, model)["metadata"])
records: Final = tuple(
object_value(record)
for record in list_value(metadata["guardrail_information"])
if object_value(record)["guardrail_name"] == rig.guardrails.logged
)
assert len(records) == 1, metadata["guardrail_information"]
record: Final = records[0]
assert {
key: record[key]
for key in (
"guardrail_mode",
"guardrail_status",
"guardrail_provider",
"masked_entity_count",
)
} == {
"guardrail_mode": "pre_call",
"guardrail_status": "success",
"guardrail_provider": "presidio",
"masked_entity_count": {"EMAIL_ADDRESS": 1, "PHONE_NUMBER": 1},
}, record
detected: Final = sorted(
(_LoggedEntity.model_validate(entity) for entity in list_value(record["guardrail_response"])),
key=lambda entity: entity.start,
)
email_start: Final = prompt.index(email)
phone_start: Final = prompt.index(PHONE)
assert [(entity.entity_type, entity.start, entity.end, entity.score) for entity in detected] == [
("EMAIL_ADDRESS", email_start, email_start + len(email), 1.0),
("PHONE_NUMBER", phone_start, phone_start + len(PHONE), 0.4),
], record["guardrail_response"]

223
tests/integration/observability/test_weave_key_scoped_delivery.py Normal file
View file

@ -0,0 +1,223 @@
import base64
import json
import uuid
from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Final
import pytest
import yaml
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment
from integration._support.process import owned_proxy_process
from integration._support.wire import Reply, Request, Wire, wire_server
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest
from opentelemetry.proto.common.v1.common_pb2 import AnyValue
from opentelemetry.proto.trace.v1.trace_pb2 import Status
from pydantic import JsonValue
INPUT_RATE: Final = 0.001
OUTPUT_RATE: Final = 0.002
PROMPT_TOKENS: Final = 7
COMPLETION_TOKENS: Final = 2
KEY_WANDB_API_KEY: Final = "tenant-wandb-key"
KEY_WEAVE_PROJECT: Final = "tenant/key-project"
OPERATOR_WEAVE_PROJECT: Final = "operator/project"
FAILING_PREFIX: Final = "weave-fail-"
UPSTREAM_AUTH_ERROR: Final = "Incorrect API key provided for the scripted upstream"
LATE_DUPLICATE_WINDOW_SECONDS: Final = 3
@dataclass(frozen=True, slots=True)
class WeaveSpan:
status_code: int
attributes: Mapping[str, str]
project_id: str
authorization: str
@dataclass(frozen=True, slots=True)
class Rig:
proxy: Gateway
scenario: Scenario
model: str
weave: Wire
def _attribute_text(value: AnyValue) -> str:
match value.WhichOneof("value"):
case "string_value":
return value.string_value
case "int_value":
return str(value.int_value)
case "double_value":
return str(value.double_value)
case "bool_value":
return str(value.bool_value)
case _:
return ""
def _spans(request: Request) -> Iterator[WeaveSpan]:
if request.headers.get("content-type") != "application/x-protobuf":
return
batch: Final = ExportTraceServiceRequest.FromString(request.body)
for resource_spans in batch.resource_spans:
for scope_spans in resource_spans.scope_spans:
for span in scope_spans.spans:
yield WeaveSpan(
status_code=span.status.code,
attributes={attribute.key: _attribute_text(attribute.value) for attribute in span.attributes},
project_id=request.headers.get("project_id", ""),
authorization=request.headers.get("authorization", ""),
)
def _upstream_reply(request: Request) -> Reply:
prompt: Final = json.loads(request.body)["messages"][0]["content"]
if prompt.startswith(FAILING_PREFIX):
return Reply(
status=401,
body=json.dumps(
{
"error": {
"message": UPSTREAM_AUTH_ERROR,
"type": "invalid_request_error",
"param": None,
"code": "invalid_api_key",
}
}
).encode(),
)
return Reply(
body=json.dumps(
{
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1,
"model": "gpt-4o-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "OK"}, "finish_reason": "stop"}],
"usage": {
"prompt_tokens": PROMPT_TOKENS,
"completion_tokens": COMPLETION_TOKENS,
"total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS,
},
}
).encode()
)
def _collector(_request: Request) -> Reply:
return Reply(body=b"", content_type="application/x-protobuf")
def _weave_logging_for_the_key() -> list[JsonValue]:
return [
{
"callback_name": "weave_otel",
"callback_type": "success_and_failure",
"callback_vars": {"wandb_api_key": KEY_WANDB_API_KEY, "weave_project_id": KEY_WEAVE_PROJECT},
}
]
def _matching(weave: Wire, prompt: str) -> tuple[WeaveSpan, ...]:
return tuple(span for request in weave.drain() for span in _spans(request) if prompt in span.attributes.values())
def _delivered_once(weave: Wire, prompt: str) -> WeaveSpan:
first: Final = eventually(lambda: _matching(weave, prompt), bool, seconds=30)
late: Final = eventually(
lambda: _matching(weave, prompt),
bool,
seconds=LATE_DUPLICATE_WINDOW_SECONDS,
return_last_on_timeout=True,
)
assert len(first) + len(late) == 1, (first, late)
return first[0]
@contextmanager
def _started(directory: Path, callbacks: list[str]) -> Iterator[Rig]:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
path: Final = directory / "weave.yaml"
path.write_text(
yaml.safe_dump({**config, "litellm_settings": {**config["litellm_settings"], "callbacks": callbacks}})
)
with (
wire_server(_upstream_reply) as provider,
wire_server(_collector) as weave,
gateway_from_environment() as gateway,
owned_proxy_process(
gateway,
directory,
{
"LITELLM_OTEL_V2": "1",
"OTEL_BSP_SCHEDULE_DELAY": "100",
"OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT": "span_only",
"WANDB_API_KEY": "operator-wandb-key",
"WANDB_PROJECT_ID": OPERATOR_WEAVE_PROJECT,
"WANDB_HOST": weave.url,
},
config=path,
) as owned,
owned.gateway.scenario() as scenario,
):
model: Final = scenario.model(
api_base=provider.url + "/v1", input_cost_per_token=INPUT_RATE, output_cost_per_token=OUTPUT_RATE
)
yield Rig(owned.gateway, scenario, model, weave)
@pytest.fixture(scope="module", params=("weave_otel", "otel"))
def rig(request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
if request.param == "otel":
pytest.skip(
"BUG: a key-scoped weave_otel callback ships nothing when the OTEL v2 otel callback is the only one configured"
)
with _started(tmp_path_factory.mktemp("weave-key-scoped"), [request.param]) as started:
yield started
def _chat(rig: Rig, key: str, prompt: str) -> tuple[int, str, str]:
response: Final = rig.proxy.request(
"POST",
"/v1/chat/completions",
{"model": rig.model, "messages": [{"role": "user", "content": prompt}]},
key=key,
)
return response.status_code, response.headers.get("x-litellm-response-cost", ""), response.text
def test_key_scoped_weave_gets_one_span_with_the_request_cost_in_the_keys_project(rig: Rig) -> None:
prompt: Final = f"weave-ok-{uuid.uuid4().hex}"
key: Final = rig.scenario.key(models=[rig.model], metadata={"logging": _weave_logging_for_the_key()})
status, header_cost, text = _chat(rig, key, prompt)
assert status == 200, text
span: Final = _delivered_once(rig.weave, prompt)
assert span.project_id == KEY_WEAVE_PROJECT, span
assert span.authorization == "Basic " + base64.b64encode(f"api:{KEY_WANDB_API_KEY}".encode()).decode(), span
assert span.status_code != Status.STATUS_CODE_ERROR, span
expected_cost: Final = PROMPT_TOKENS * INPUT_RATE + COMPLETION_TOKENS * OUTPUT_RATE
assert float(header_cost) == pytest.approx(expected_cost), header_cost
assert float(span.attributes["litellm.cost.total"]) == pytest.approx(expected_cost), span
assert span.attributes["gen_ai.usage.input_tokens"] == str(PROMPT_TOKENS), span
assert span.attributes["gen_ai.usage.output_tokens"] == str(COMPLETION_TOKENS), span
def test_key_scoped_weave_gets_one_error_span_for_a_provider_failure_without_cost(rig: Rig) -> None:
prompt: Final = f"{FAILING_PREFIX}{uuid.uuid4().hex}"
key: Final = rig.scenario.key(models=[rig.model], metadata={"logging": _weave_logging_for_the_key()})
status, _, text = _chat(rig, key, prompt)
assert status == 401, text
span: Final = _delivered_once(rig.weave, prompt)
assert span.project_id == KEY_WEAVE_PROJECT, span
assert span.status_code == Status.STATUS_CODE_ERROR, span
assert span.attributes["error.type"] == "AuthenticationError", span
assert UPSTREAM_AUTH_ERROR in span.attributes["error.message"], span
assert float(span.attributes["litellm.cost.total"]) == 0, span

229
tests/integration/pricing/test_sail_service_tier_pricing.py Normal file
View file

@ -0,0 +1,229 @@
import json
import uuid
from dataclasses import dataclass
from typing import Final
import pytest
from pydantic import JsonValue
from tests.integration._support.client import Gateway, Scenario, eventually, object_value, string_value
from tests.integration._support.database import read_rows
from tests.integration._support.wire import Reply, Request, Wire, wire_server
BACKEND: Final = "zai-org/GLM-5.3"
PROMPT_TOKENS: Final = 1000
CACHED_TOKENS: Final = 400
COMPLETION_TOKENS: Final = 200
PRICES: Final[dict[str, JsonValue]] = {
"input_cost_per_token": 3e-06,
"output_cost_per_token": 9e-06,
"cache_read_input_token_cost": 1e-06,
"input_cost_per_token_balanced": 2e-06,
"output_cost_per_token_balanced": 6e-06,
"cache_read_input_token_cost_balanced": 7e-07,
"input_cost_per_token_flex": 1e-06,
"output_cost_per_token_flex": 3e-06,
"cache_read_input_token_cost_flex": 4e-07,
}
BASE_COST: Final = 600 * 3e-06 + 400 * 1e-06 + 200 * 9e-06
BALANCED_COST: Final = 600 * 2e-06 + 400 * 7e-07 + 200 * 6e-06
FLEX_COST: Final = 600 * 1e-06 + 400 * 4e-07 + 200 * 3e-06
def _chat_reply(_request: Request) -> Reply:
return Reply(
body=json.dumps(
{
"id": f"chatcmpl-{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1,
"model": BACKEND,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "OK"}, "finish_reason": "stop"}],
"usage": {
"prompt_tokens": PROMPT_TOKENS,
"completion_tokens": COMPLETION_TOKENS,
"total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS,
"prompt_tokens_details": {"cached_tokens": CACHED_TOKENS},
},
}
).encode()
)
def _responses_reply(_request: Request) -> Reply:
return Reply(
body=json.dumps(
{
"id": f"resp_{uuid.uuid4().hex}",
"object": "response",
"created_at": 1,
"status": "completed",
"model": BACKEND,
"output": [
{
"type": "message",
"id": f"msg_{uuid.uuid4().hex}",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": "OK", "annotations": []}],
}
],
"usage": {
"input_tokens": PROMPT_TOKENS,
"output_tokens": COMPLETION_TOKENS,
"total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS,
"input_tokens_details": {"cached_tokens": CACHED_TOKENS},
"output_tokens_details": {"reasoning_tokens": 0},
},
}
).encode()
)
def _messages_reply(_request: Request) -> Reply:
return Reply(
body=json.dumps(
{
"id": f"msg_{uuid.uuid4().hex}",
"type": "message",
"role": "assistant",
"model": BACKEND,
"content": [{"type": "text", "text": "OK"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": PROMPT_TOKENS, "output_tokens": COMPLETION_TOKENS},
}
).encode()
)
def _upstream_body(request: Request) -> dict[str, JsonValue]:
return object_value(json.loads(request.body))
def _sail_model(scenario: Scenario, wire: Wire) -> str:
return scenario.model(model=f"sail/{BACKEND}", api_base=wire.url, api_key="synthetic-sail-key", **PRICES)
def _assert_spend_row(gateway: Gateway, response_id: str, expected_cost: float) -> None:
rows: Final = eventually(
lambda: read_rows(
'SELECT spend, custom_llm_provider, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" '
"WHERE request_id = %s",
(response_id,),
),
lambda values: len(values) == 1,
seconds=70,
)
assert rows[0]["custom_llm_provider"] == "sail", rows
assert rows[0]["prompt_tokens"] == PROMPT_TOKENS, rows
assert rows[0]["completion_tokens"] == COMPLETION_TOKENS, rows
assert float(rows[0]["spend"]) == pytest.approx(expected_cost, rel=1e-6), rows
@dataclass(frozen=True, slots=True)
class Served:
upstream_body: dict[str, JsonValue]
header_cost: float
response_id: str
def _chat_through_sail(gateway: Gateway, extra: dict[str, JsonValue]) -> Served:
with wire_server(_chat_reply) as wire, gateway.scenario() as scenario:
model: Final = _sail_model(scenario, wire)
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": "Reply with one word."}], **extra},
)
assert response.status_code == 200, response.text
received: Final = wire.drain()
assert [(request.method, request.target) for request in received] == [("POST", "/chat/completions")]
return Served(
upstream_body=_upstream_body(received[0]),
header_cost=float(response.headers["x-litellm-response-cost"]),
response_id=string_value(object_value(response.json())["id"]),
)
@pytest.mark.parametrize(
("service_tier", "completion_window", "expected_cost"),
(
("balanced", "balanced", BALANCED_COST),
("flex", "flex", FLEX_COST),
("auto", None, BASE_COST),
),
ids=("balanced", "flex", "auto"),
)
def test_chat_service_tier_is_sent_as_completion_window_and_billed_at_that_window(
gateway: Gateway, service_tier: str, completion_window: str | None, expected_cost: float
) -> None:
served: Final = _chat_through_sail(gateway, {"service_tier": service_tier})
assert served.upstream_body == {
"model": BACKEND,
"messages": [{"role": "user", "content": "Reply with one word."}],
**({} if completion_window is None else {"metadata": {"completion_window": completion_window}}),
}, served.upstream_body
assert served.header_cost == pytest.approx(expected_cost, rel=1e-6)
_assert_spend_row(gateway, served.response_id, expected_cost)
@pytest.mark.parametrize("service_tier", ("bogus", 5), ids=("string", "integer"))
def test_chat_unknown_service_tier_under_drop_params_is_dropped_and_billed_at_the_base_rate(
gateway: Gateway, service_tier: str | int
) -> None:
served: Final = _chat_through_sail(gateway, {"service_tier": service_tier, "drop_params": True})
assert served.upstream_body == {
"model": BACKEND,
"messages": [{"role": "user", "content": "Reply with one word."}],
}, served.upstream_body
assert served.header_cost == pytest.approx(BASE_COST, rel=1e-6)
_assert_spend_row(gateway, served.response_id, BASE_COST)
def test_responses_caller_completion_window_reaches_sail_and_is_billed_at_that_window(gateway: Gateway) -> None:
with wire_server(_responses_reply) as wire, gateway.scenario() as scenario:
model: Final = _sail_model(scenario, wire)
response: Final = gateway.request(
"POST",
"/v1/responses",
{"model": model, "input": "Reply with one word.", "metadata": {"completion_window": "balanced"}},
)
assert response.status_code == 200, response.text
received: Final = wire.drain()
assert [(request.method, request.target) for request in received] == [("POST", "/responses")]
assert _upstream_body(received[0]) == {
"model": BACKEND,
"input": "Reply with one word.",
"metadata": {"completion_window": "balanced"},
}, received[0].body
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(BALANCED_COST, rel=1e-6), response.text
_assert_spend_row(gateway, string_value(object_value(response.json())["id"]), BALANCED_COST)
def test_messages_call_to_sail_returns_an_anthropic_message(gateway: Gateway) -> None:
with wire_server(_messages_reply) as wire, gateway.scenario() as scenario:
model: Final = _sail_model(scenario, wire)
response: Final = gateway.request(
"POST",
"/v1/messages",
{"model": model, "max_tokens": 64, "messages": [{"role": "user", "content": "Reply with one word."}]},
)
assert response.status_code == 200, response.text
received: Final = wire.drain()
assert [(request.method, request.target) for request in received] == [("POST", "/v1/messages")], received
assert _upstream_body(received[0]) == {
"model": BACKEND,
"max_tokens": 64,
"messages": [{"role": "user", "content": "Reply with one word."}],
"stream": False,
}, received[0].body
payload: Final = object_value(response.json())
assert payload["type"] == "message" and payload["role"] == "assistant", response.text
assert payload["content"] == [{"type": "text", "text": "OK"}], response.text
assert object_value(payload["usage"])["output_tokens"] == COMPLETION_TOKENS, response.text