scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

test(e2e): cover presidio post_call, tool_permission, and weave logging cells (#39279)

Browse source 
* test(e2e): cover presidio post_call, tool_permission, and weave logging cells

Five registry cells in Logging & Guardrails had no covering test. Each one now
has a live scenario read back from the real destination:

- guardrail.presidio.post_call.masks: an output-scoped Presidio guardrail
  anonymizes the PII the model repeats back. The prompt also asks for the
  address's local part, which Presidio does not mask, so one response proves the
  model saw the raw address (no pre-call masking) while the address itself comes
  back as <EMAIL_ADDRESS>
- guardrail.tool_permission.pre_call.blocks / .allows: an allow-list of one tool.
  A request declaring an unlisted tool is rejected 400 naming it; a request
  declaring the permitted tool is served and carries
  x-litellm-applied-guardrails, so the allow half cannot pass by the guardrail
  never running
- logging.niche_integrations.success.logs_spend / .failure.logs_spend: a
  key-scoped weave_otel callback delivers to the real Weave project, read back
  through Weave's query API. Success asserts exactly one call whose
  llm.response.cost equals the x-litellm-response-cost header; failure asserts
  one ERROR-status call naming the provider exception and carrying no cost

Logging & Guardrails coverage goes 24/59 to 29/59. No registry rows are added.

* test(e2e): make the tool-permission allow case deterministic and scope the Weave read-back

Review follow-ups on the coverage PR.

- the allow scenario forced the outcome to depend on whether the model felt like
  calling an optional tool, and checked for the tool name as a substring of the
  whole body, which a prose mention would satisfy. It now sends
  tool_choice="required" and asserts the parsed response carries exactly one tool
  call, for the permitted tool
- the Weave read-back queried the newest 200 calls of a shared project and
  filtered client-side, so busy traffic could push the target out of the window
  and read as a delivery failure. The query now scopes server-side to the
  litellm_request op and to calls started after the request, and pages through
  the window with offset
- the reader builds its results as tuples instead of accumulating into lists

Also unblocks the lint gate: `basedpyright tests/e2e` runs only on PRs that touch
tests/e2e, and it has been failing on staging for three FakeItem arguments in
test_junit_properties.py. The stand-in now goes through one typed adapter that
says why, so the gate is green without touching junit_properties.py itself.

* test(e2e): scope the presidio post_call guardrail to email and phone

Running the suite three times in a row caught a real flake: Presidio's broader
recognizers sometimes claim the email's local part as an NRP entity, so the
answer came back as `<NRP>\n<EMAIL_ADDRESS>\n<PHONE_NUMBER>` and the assertion
that the raw local part survives failed. That token is what tells output masking
apart from input masking, so it has to survive.

The post_call guardrail now registers pii_entities_config for EMAIL_ADDRESS and
PHONE_NUMBER only, which is also the narrower thing the scenario means. Verified
against the exact marker that failed, plus two others.

* test(e2e): mark weave logging cells stage red

* test(e2e): use per-test stage red skips for the weave logging cells
This commit is contained in:
yucheng-berri 2026-09-05 13:03:28 -07:00 • committed by GitHub
parent c091dd4608
commit 948e5755eb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 844 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
tests/e2e/CONTRIBUTING.md
View file

@ -50,7 +50,9 @@ 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
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`
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
### Record and replay

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

@ -18,6 +18,7 @@ from models import (
ChatBody,
ChatMessage,
ChatResponse,
ChatTool,
KeyGenerateBody,
LiteLLMParamsBody,
TeamDeleteBody,
@ -31,6 +32,8 @@ from proxy_client import ProxyClient
from pydantic import BaseModel
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"]
@ -81,6 +84,27 @@ class PresidioParamsBody(GuardrailParamsBase):
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`."""
id: str
tool_name: str
decision: Literal["allow", "deny"]
class ToolPermissionParamsBody(GuardrailParamsBase):
"""Tool-permission guardrail params. `default_action="deny"` makes the rules
an allow-list, and `on_disallowed_action="block"` turns a disallowed tool into
a 400 instead of rewriting the request; "rewrite" is a different product
promise and belongs to its own scenario."""
guardrail: Literal["tool_permission"] = "tool_permission"
rules: list[ToolPermissionRuleBody]
default_action: Literal["allow", "deny"] = "deny"
on_disallowed_action: Literal["block", "rewrite"] = "block"
GuardrailParamsBody = (
@ -89,6 +113,7 @@ GuardrailParamsBody = (
| OpenAIModerationParamsBody
| BlockCodeExecutionParamsBody
| PresidioParamsBody
| ToolPermissionParamsBody
)
@ -200,9 +225,7 @@ class GuardrailsClient:
self.proxy.transport.post(
"/guardrails",
headers=self.proxy.transport.master,
json=GuardrailCreateBody(
guardrail=GuardrailSpecBody(guardrail_name=name, litellm_params=params)
),
json=GuardrailCreateBody(guardrail=GuardrailSpecBody(guardrail_name=name, litellm_params=params)),
response_type=GuardrailCreateResponse,
)
).guardrail_id
@ -241,9 +264,7 @@ class GuardrailsClient:
)
def create_key_in_team(self, team_id: str) -> str:
return self.proxy.generate_key(
KeyGenerateBody(team_id=team_id, user_id="e2e-guardrails-user")
)
return self.proxy.generate_key(KeyGenerateBody(team_id=team_id, user_id="e2e-guardrails-user"))
def chat(
self,
@ -253,6 +274,7 @@ class GuardrailsClient:
*,
guardrails: list[str] | None = None,
max_tokens: int = 16,
tools: list[ChatTool] | None = None,
) -> Result[ChatResponse]:
"""Drive a chat call, optionally opting into named guardrails for this
request only (the per-request `guardrails` selector). With `guardrails`
@ -266,6 +288,35 @@ class GuardrailsClient:
messages=[ChatMessage(role="user", content=text)],
max_tokens=max_tokens,
guardrails=guardrails,
tools=tools,
),
)
def chat_raw(
self,
key: str,
model: str,
text: str,
*,
guardrails: list[str] | None = None,
max_tokens: int = 16,
tools: list[ChatTool] | None = None,
tool_choice: str | None = None,
) -> StreamingResponse:
"""Drive /chat/completions returning the raw HTTP outcome, for the
assertions a typed body cannot carry: the `x-litellm-applied-guardrails`
response header, which is how an ALLOW scenario proves the guardrail ran
rather than being absent."""
return self.proxy.transport.send(
"/chat/completions",
headers=self.proxy.transport.bearer(key),
json=ChatBody(
model=model,
messages=[ChatMessage(role="user", content=text)],
max_tokens=max_tokens,
guardrails=guardrails,
tools=tools,
tool_choice=tool_choice,
),
)
@ -323,9 +374,7 @@ class GuardrailsClient:
return self.proxy.transport.send(
"/v1/responses",
headers=self.proxy.transport.bearer(key),
json=_ResponsesGuardrailBody(
model=model, input=text, guardrails=guardrails
),
json=_ResponsesGuardrailBody(model=model, input=text, guardrails=guardrails),
)
def apply_guardrail(self, key: str, *, name: str, text: str) -> Result[ApplyGuardrailResponse]:
@ -349,9 +398,7 @@ class GuardrailsClient:
if isinstance(last, Success):
return
time.sleep(POLL_INTERVAL)
raise AssertionError(
f"team {team_id!r} was created but /team/info never returned it: {last}"
)
raise AssertionError(f"team {team_id!r} was created but /team/info never returned it: {last}")
def build_client(proxy: ProxyClient) -> GuardrailsClient:

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

@ -6,11 +6,19 @@ messages BEFORE the model runs, so the model only ever sees placeholders like
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 presidio_filter_scope="input" so only the
configured hook's callback exists (the default "both" adds a second post_call
output masker), and is deleted on teardown.
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.
"""
from __future__ import annotations
@ -18,13 +26,14 @@ from __future__ import annotations
import os
import time
from collections.abc import Callable
from typing import Literal
import pytest
from pydantic import BaseModel
from e2e_config import unique_marker
from e2e_http import Result, Success
from guardrails_client import GuardrailsClient, PresidioParamsBody
from guardrails_client import GuardrailMode, GuardrailsClient, PiiAction, PiiEntity, PresidioParamsBody
from lifecycle import ResourceManager
from models import AnthropicMessagesResponse, ChatResponse
@ -65,16 +74,20 @@ def _register_presidio(
resources: ResourceManager,
*,
name: str,
mode: GuardrailMode = "pre_call",
filter_scope: Literal["input", "output", "both"] = "input",
entities: dict[PiiEntity, PiiAction] | None = None,
) -> None:
analyzer, anonymizer = _presidio_bases()
guardrail_id = client.register(
name,
PresidioParamsBody(
mode="pre_call",
mode=mode,
default_on=False,
presidio_analyzer_api_base=analyzer,
presidio_anonymizer_api_base=anonymizer,
presidio_filter_scope="input",
presidio_filter_scope=filter_scope,
pii_entities_config=entities,
),
)
resources.defer(lambda: client.delete_guardrail(guardrail_id))
@ -182,3 +195,84 @@ class TestPresidioPreCallMasking:
_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"],
)
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)

167
tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py Normal file
View file

@ -0,0 +1,167 @@
"""Live e2e: the tool_permission guardrail gates which tools a request may declare.
The guardrail is registered `mode="pre_call"` with `default_action="deny"`, so its
rules are an allow-list applied to the tools the CALLER declares, before the model
runs. Two halves of one product promise:
- blocks: a request declaring a tool outside the allow-list is rejected with a 400
naming the denied tool, and never reaches the model
- allows: a request declaring only the permitted tool is served normally, comes
back with a real tool call for that tool, and carries an
`x-litellm-applied-guardrails` header naming the guardrail, which is what
separates "the guardrail ran and allowed it" from "the guardrail was never
attached". `tool_choice="required"` keeps the model from answering directly and
making the outcome depend on its mood
No vendor API is involved: `tool_permission` is a built-in guardrail, so the
verdict comes from the proxy itself.
"""
from __future__ import annotations
from typing import Final
import pytest
from e2e_config import unique_marker
from e2e_http import StreamingResponse, UnknownApiError
from guardrails_client import (
GuardrailsClient,
ToolPermissionParamsBody,
ToolPermissionRuleBody,
poll_until_blocked,
)
from lifecycle import ResourceManager
from models import ChatResponse, ChatTool, ChatToolFunction
pytestmark = pytest.mark.e2e
MODEL = "gemini-2.5-flash"
#: The one tool the guardrail permits, and one it does not. Both are declared by
#: the caller in the request body; the guardrail reads them there.
ALLOWED_TOOL: Final = ChatTool(
function=ChatToolFunction(
name="get_weather",
description="Get the current weather for a city",
parameters={
"type": "object",
"properties": {"city": {"type": "string"}},
"required": ["city"],
},
)
)
DENIED_TOOL: Final = ChatTool(
function=ChatToolFunction(
name="delete_customer_database",
description="Permanently delete the customer database",
parameters={"type": "object", "properties": {}},
)
)
TOOL_PROMPT: Final = "What is the weather in Paris right now?"
def _register_tool_permission(client: GuardrailsClient, resources: ResourceManager, *, name: str) -> None:
"""Allow-list exactly one tool: everything else falls to `default_action=deny`
and, with `on_disallowed_action=block`, is rejected outright."""
guardrail_id = client.register(
name,
ToolPermissionParamsBody(
mode="pre_call",
default_on=False,
default_action="deny",
on_disallowed_action="block",
rules=[
ToolPermissionRuleBody(
id="allow-get-weather",
tool_name=ALLOWED_TOOL.function.name,
decision="allow",
)
],
),
)
resources.defer(lambda: client.delete_guardrail(guardrail_id))
def _applied_guardrails(outcome: StreamingResponse) -> str:
return outcome.headers.get("x-litellm-applied-guardrails", "")
def _tool_call_names(response: ChatResponse) -> tuple[str, ...]:
return tuple(
call.function.name
for choice in response.choices
if choice.message
for call in choice.message.tool_calls or ()
if call.function.name
)
class TestToolPermissionPreCall:
@pytest.mark.covers("guardrail.tool_permission.pre_call.blocks", exercised_on=["chat_completions"])
def test_pre_call_blocks_tool_outside_the_allow_list(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
"""A request declaring a tool the guardrail does not permit must be
rejected with a 400 that names the denied tool. An unauthorized tool that
merely reaches the model is the whole failure mode this guardrail exists
to prevent, so a 200 here is a hard failure."""
name = f"e2e-toolperm-block-{unique_marker()}"
_register_tool_permission(client, resources, name=name)
result = poll_until_blocked(
lambda: client.chat(
scoped_key,
MODEL,
TOOL_PROMPT,
guardrails=[name],
max_tokens=128,
tools=[DENIED_TOOL],
)
)
match result:
case UnknownApiError(status_code=status, body=body):
assert status == 400, f"expected the guardrail block status 400, got {status}: {body[:400]}"
assert DENIED_TOOL.function.name in body, (
f"the block must name the denied tool so the caller can fix the request; got: {body[:400]}"
)
assert "guardrail" in body.lower(), (
f"the block body should identify itself as a guardrail verdict; got: {body[:400]}"
)
case _:
pytest.fail(f"tool_permission let a tool outside the allow-list through; got {result}")
@pytest.mark.covers("guardrail.tool_permission.pre_call.allows", exercised_on=["chat_completions"])
def test_pre_call_allows_permitted_tool(
self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str
) -> None:
"""The mirror half: a request declaring only the permitted tool is served
and the model calls it. Without the header check a guardrail that never
attached would pass this test for the wrong reason, so the 200 alone is
not the contract."""
name = f"e2e-toolperm-allow-{unique_marker()}"
_register_tool_permission(client, resources, name=name)
outcome = client.chat_raw(
scoped_key,
MODEL,
TOOL_PROMPT,
guardrails=[name],
max_tokens=128,
tools=[ALLOWED_TOOL],
tool_choice="required",
)
assert outcome.ok, f"the permitted tool must be served, got {outcome.status_code}: {outcome.body[:400]}"
applied = _applied_guardrails(outcome)
assert name in applied, (
"the allowed call must carry x-litellm-applied-guardrails naming the guardrail; "
f"without it the 200 only proves the guardrail never ran. Got {applied!r}"
)
called = _tool_call_names(ChatResponse.model_validate_json(outcome.body))
assert called == (ALLOWED_TOOL.function.name,), (
f"the served call must carry one tool call for the permitted tool, got {called!r}: {outcome.body[:400]}"
)

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

@ -189,6 +189,45 @@ 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")

192
tests/e2e/logging/test_weave_log_e2e.py Normal file
View file

@ -0,0 +1,192 @@
"""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
import pytest
from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker
from e2e_http import StreamingResponse
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
@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"])
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"])
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="anthropic/claude-haiku-4-5", 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

282
tests/e2e/logging/weave_reader.py Normal file
View file

@ -0,0 +1,282 @@
"""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
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
_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
@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())
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))
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

@ -35,6 +35,8 @@ class KeyLoggingCallbackVars(BaseModel):
langfuse_public_key: str | None = None
langfuse_secret_key: str | None = None
langfuse_host: str | None = None
wandb_api_key: str | None = None
weave_project_id: str | None = None
class KeyLoggingCallback(BaseModel):