mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
1dbde0845e
* test: move guardrails, image, audio, batches, search legacy tests to unit * test: restore the custom guardrail and remaining chunk 1 unit rows * test: move the ocr matrix to e2e * test: drop comments from moved tests and fix vertex ocr credentials * test: move litellm utils and unified google legacy tests to their tiers * test: move logging callback and pass-through unit legacy tests to their tiers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: move store_model_in_db legacy tests and start the admin ui migration Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: finish the proxy job rows and drop their CircleCI jobs Moves the remaining multi-instance, openai endpoint, otel, pass-through and top-level proxy tests into tests/integration with scripted upstreams, rewrites the weak replacements to assert the outbound wire and exact responses, and removes the build_and_test job group this PR owns from CircleCI * test: tighten the key generation and error log assertions, rename migrated unit files to their mapped names * test: shard the pagerduty unit tests and record the rust harness file on the coverage allowlist * test: drive the SSO sign-in tests through a scripted IdP and isolate the cost center stand-in * test: type the e2e beta header, prompt cache and secret manager values * test: assert the missing model error message on the OpenAI routes * test: drop the deleted otel base class import and restore mcp redaction and exact langfuse usage checks * test: make the langfuse fixture directory a package for the layout check * test: restore collected no-match cases, real langfuse and gcs payload checks, and gate the aws secret lane * test: keep the ocr e2e pdf fixture next to the test * test: put meta below covers on the streaming tool search e2e * test: run the OCR matrix against a scripted upstream and fix two runner-only e2e failures * test: count only OCR calls in the OCR matrix recorder * test: pin OCR parse bodies and billed cost, group new e2e cases into classes * ci: restore the python image anchor dropped in the main sync --------- Co-authored-by: yuneng <yuneng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
79 lines
3.1 KiB
Python
79 lines
3.1 KiB
Python
import json
|
|
from pathlib import Path
|
|
from typing import Final
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
import yaml
|
|
|
|
from tests.integration._support.client import Gateway
|
|
from tests.integration._support.process import owned_proxy
|
|
from tests.integration._support.wire import Reply, Request, wire_server
|
|
|
|
_REQUESTS: Final = 24
|
|
_COMPLETION: Final = {
|
|
"id": "chatcmpl-guarded",
|
|
"object": "chat.completion",
|
|
"created": 1,
|
|
"model": "gpt-4o-mini",
|
|
"choices": [{"index": 0, "message": {"role": "assistant", "content": "guarded"}, "finish_reason": "stop"}],
|
|
"usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4},
|
|
}
|
|
|
|
|
|
def _allow(request: Request) -> Reply:
|
|
return Reply(body=json.dumps({"action": "NONE"}).encode())
|
|
|
|
|
|
def _complete(request: Request) -> Reply:
|
|
return Reply(body=json.dumps(_COMPLETION).encode())
|
|
|
|
|
|
def _deployment(name: str, api_base: str) -> dict[str, object]:
|
|
return {
|
|
"guardrail_name": name,
|
|
"litellm_params": {
|
|
"guardrail": "generic_guardrail_api",
|
|
"mode": "pre_call",
|
|
"api_base": api_base,
|
|
"api_key": "synthetic-guardrail-key",
|
|
},
|
|
}
|
|
|
|
|
|
def test_requests_naming_a_shared_guardrail_are_spread_over_its_deployments(gateway: Gateway, tmp_path: Path) -> None:
|
|
pytest.skip(
|
|
"BUG: init_guardrails_v2 populates llm_router.guardrail_list from the proxy_server global, which is still "
|
|
"None during load_config, so every deployment sharing a guardrail_name checks every request"
|
|
)
|
|
name: Final = f"lb-guardrail-{uuid4().hex}"
|
|
prompts: Final = tuple(f"Hello request {index}" for index in range(_REQUESTS))
|
|
with (
|
|
wire_server(_allow) as first,
|
|
wire_server(_allow) as second,
|
|
wire_server(_complete) as upstream,
|
|
):
|
|
config: Final = {
|
|
**yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()),
|
|
"guardrails": [_deployment(name, first.url), _deployment(name, second.url)],
|
|
}
|
|
config_path: Final = tmp_path / "guardrail_load_balancing.yaml"
|
|
config_path.write_text(yaml.safe_dump(config))
|
|
with owned_proxy(gateway, tmp_path, {}, config=config_path) as proxy, proxy.scenario() as scenario:
|
|
model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=f"{upstream.url}/v1")
|
|
responses: Final = tuple(
|
|
proxy.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": prompt}], "guardrails": [name]},
|
|
)
|
|
for prompt in prompts
|
|
)
|
|
checked: Final = (first.drain(), second.drain())
|
|
|
|
assert [(response.status_code, response.headers["x-litellm-applied-guardrails"]) for response in responses] == [
|
|
(200, name)
|
|
] * _REQUESTS
|
|
seen: Final = tuple(frozenset(json.loads(call.body)["texts"][0] for call in calls) for calls in checked)
|
|
assert all(0 < len(prompts_seen) < _REQUESTS for prompts_seen in seen), [len(prompts_seen) for prompts_seen in seen]
|
|
assert seen[0] | seen[1] == frozenset(prompts)
|