mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
remove unused parameters and methods, simplify payload handling in tests
This commit is contained in:
rzeta-10
parent
05635459a6
commit
d8917243d3
2 changed files with 6 additions and 107 deletions
|
|
@ -58,8 +58,6 @@ class AktoGuardrail(CustomGuardrail):
|
|||
Args:
|
||||
akto_base_url: Akto API base URL. Falls back to AKTO_GUARDRAIL_API_BASE env var.
|
||||
akto_api_key: Akto API key. Falls back to AKTO_API_KEY env var.
|
||||
akto_account_id: Akto account ID. Falls back to AKTO_ACCOUNT_ID env var, then "1000000".
|
||||
akto_vxlan_id: Akto VXLAN ID. Falls back to AKTO_VXLAN_ID env var, then "0".
|
||||
unreachable_fallback: Behavior when Akto is unreachable — block or allow.
|
||||
guardrail_timeout: HTTP timeout in seconds for Akto API calls.
|
||||
"""
|
||||
|
|
@ -117,23 +115,6 @@ class AktoGuardrail(CustomGuardrail):
|
|||
route = metadata.get("user_api_key_request_route")
|
||||
return route if route else "/v1/chat/completions"
|
||||
|
||||
def prepare_headers(self) -> Dict[str, str]:
|
||||
"""Build HTTP headers for the Akto API call."""
|
||||
return {
|
||||
"content-type": "application/json",
|
||||
"Authorization": self.akto_api_key,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def build_query_params(*, guardrails: bool, ingest_data: bool) -> Dict[str, str]:
|
||||
"""Build query params that control Akto backend behavior (guardrail check and/or data ingestion)."""
|
||||
params: Dict[str, str] = {"akto_connector": AKTO_CONNECTOR_NAME}
|
||||
if guardrails:
|
||||
params["guardrails"] = "true"
|
||||
if ingest_data:
|
||||
params["ingest_data"] = "true"
|
||||
return params
|
||||
|
||||
@staticmethod
|
||||
def build_request_headers(request_data: dict) -> Dict[str, str]:
|
||||
"""Build the requestHeaders field from proxy request headers."""
|
||||
|
|
@ -225,23 +206,17 @@ class AktoGuardrail(CustomGuardrail):
|
|||
status_code: int = 200,
|
||||
include_response: bool = False,
|
||||
) -> Dict[str, Any]:
|
||||
"""Build the flat MIRRORING payload sent to Akto's HTTP proxy endpoint.
|
||||
|
||||
All body fields use double-encoding: json.dumps({"body": json.dumps(actual_body)})
|
||||
to match the canonical CLI hook format.
|
||||
"""
|
||||
"""Build the MIRRORING payload for Akto's HTTP proxy endpoint."""
|
||||
request_path = self.extract_request_path(request_data)
|
||||
request_headers = self.build_request_headers(request_data)
|
||||
request_body = self.build_request_body(inputs, request_data)
|
||||
tag = self.build_tag_metadata(request_data)
|
||||
|
||||
response_payload = json.dumps({}) # Empty body wrapper when no response yet
|
||||
response_payload = json.dumps({})
|
||||
response_headers: Dict[str, str] = {}
|
||||
if include_response:
|
||||
response_body = self.build_response_body(inputs, request_data)
|
||||
response_payload = json.dumps(
|
||||
{"body": json.dumps(response_body)}
|
||||
) # Double-encoded
|
||||
response_payload = json.dumps(response_body)
|
||||
response_headers = {"content-type": "application/json"}
|
||||
|
||||
# Extract client IP from proxy headers
|
||||
|
|
@ -264,9 +239,7 @@ class AktoGuardrail(CustomGuardrail):
|
|||
"requestHeaders": json.dumps(request_headers),
|
||||
"responseHeaders": json.dumps(response_headers),
|
||||
"method": "POST",
|
||||
"requestPayload": json.dumps(
|
||||
{"body": json.dumps(request_body)}
|
||||
), # Double-encoded
|
||||
"requestPayload": json.dumps(request_body),
|
||||
"responsePayload": response_payload,
|
||||
"ip": ip,
|
||||
"destIp": "127.0.0.1",
|
||||
|
|
@ -340,48 +313,6 @@ class AktoGuardrail(CustomGuardrail):
|
|||
str(guardrails_result.get("Reason", "")),
|
||||
)
|
||||
|
||||
def handle_unreachable(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
error: Exception,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
"""Handle Akto being unreachable based on fail_open/fail_closed config."""
|
||||
if self.unreachable_fallback == "fail_open":
|
||||
verbose_proxy_logger.critical(
|
||||
"Akto unreachable (fail-open): %s",
|
||||
str(error),
|
||||
exc_info=error,
|
||||
)
|
||||
return inputs
|
||||
|
||||
verbose_proxy_logger.error("Akto unreachable (fail-closed): %s", str(error))
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="Akto guardrail service unreachable",
|
||||
)
|
||||
|
||||
async def fire_and_forget_request(
|
||||
self,
|
||||
*,
|
||||
guardrails: bool,
|
||||
ingest_data: bool,
|
||||
payload: dict,
|
||||
) -> None:
|
||||
"""Send a request without awaiting it in the caller. Errors are logged, not raised."""
|
||||
try:
|
||||
response = await self.send_request(
|
||||
guardrails=guardrails,
|
||||
ingest_data=ingest_data,
|
||||
payload=payload,
|
||||
)
|
||||
if response.status_code != 200:
|
||||
verbose_proxy_logger.error(
|
||||
"Akto fire-and-forget returned HTTP %d",
|
||||
response.status_code,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error("Akto fire-and-forget error: %s", str(e))
|
||||
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -158,8 +158,7 @@ def test_build_akto_payload(akto_validate, sample_inputs, sample_request_data):
|
|||
req_headers = json.loads(payload["requestHeaders"])
|
||||
assert "content-type" in req_headers
|
||||
|
||||
req_wrapper = json.loads(payload["requestPayload"])
|
||||
req_body = json.loads(req_wrapper["body"])
|
||||
req_body = json.loads(payload["requestPayload"])
|
||||
assert req_body["model"] == "gpt-4"
|
||||
assert req_body["messages"][0]["content"] == "Hello, how are you?"
|
||||
|
||||
|
|
@ -173,8 +172,7 @@ def test_build_akto_payload(akto_validate, sample_inputs, sample_request_data):
|
|||
|
||||
def test_build_akto_payload_with_response(akto_validate, sample_inputs, sample_request_data):
|
||||
payload = akto_validate.build_akto_payload(sample_inputs, sample_request_data, include_response=True)
|
||||
resp_wrapper = json.loads(payload["responsePayload"])
|
||||
resp_body = json.loads(resp_wrapper["body"])
|
||||
resp_body = json.loads(payload["responsePayload"])
|
||||
assert "choices" in resp_body
|
||||
|
||||
|
||||
|
|
@ -192,20 +190,6 @@ def test_build_akto_payload_custom_ids(sample_request_data):
|
|||
assert payload["akto_vxlan_id"] == "7"
|
||||
|
||||
|
||||
def test_build_query_params():
|
||||
params = AktoGuardrail.build_query_params(guardrails=True, ingest_data=False)
|
||||
assert params == {"akto_connector": "litellm", "guardrails": "true"}
|
||||
|
||||
params = AktoGuardrail.build_query_params(guardrails=False, ingest_data=True)
|
||||
assert params == {"akto_connector": "litellm", "ingest_data": "true"}
|
||||
|
||||
params = AktoGuardrail.build_query_params(guardrails=True, ingest_data=True)
|
||||
assert params == {
|
||||
"akto_connector": "litellm",
|
||||
"guardrails": "true",
|
||||
"ingest_data": "true",
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Guardrail response handling
|
||||
|
|
@ -382,22 +366,6 @@ async def test_fail_closed_on_unreachable():
|
|||
assert exc_info.value.status_code == 503
|
||||
|
||||
|
||||
def test_fail_closed_generic_message():
|
||||
g = AktoGuardrail(
|
||||
akto_base_url="http://localhost:9090",
|
||||
akto_api_key="test-token",
|
||||
unreachable_fallback="fail_closed",
|
||||
guardrail_name="msg-test",
|
||||
event_hook="pre_call",
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
g.handle_unreachable(
|
||||
inputs=GenericGuardrailAPIInputs(texts=["test"], model="gpt-4"),
|
||||
error=Exception("http://internal-host:9090/secret-path"),
|
||||
)
|
||||
assert "internal-host" not in exc_info.value.detail
|
||||
assert exc_info.value.detail == "Akto guardrail service unreachable"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helper method tests
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue