mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
fix(guardrails): store the masked output in spend logs when Presidio masks the response (#42441)
* fix(guardrails): store the masked output in spend logs when a post_call guardrail rewrites the response Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(guardrails): record served output without re-narrowing the logging object Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(guardrails): overlay the served output before message redaction so turn_off_message_logging still wins Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(logging): type the monkeypatch fixture in the redaction ordering regression Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): poll through raw card output until the guardrail reaches the serving worker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(guardrails): keep blanked, multi-choice and disconnected served output out of raw spend logs Served text keeps empty strings and tracks unavailable choices as None so a guardrail that blanks the output still overrides the raw provider text. Stream choices are sized from the highest choice index, served chunks are recorded before a client disconnect or stream failure propagates, and message-logging redaction drops the served text from callback kwargs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng <yucheng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
b682278aa9
commit
1a714548a4
14 changed files with 647 additions and 6 deletions
|
|
@ -107,6 +107,10 @@ from litellm.litellm_core_utils.redact_messages import (
|
|||
redact_streaming_responses_for_custom_logger,
|
||||
should_redact_message_logging,
|
||||
)
|
||||
from litellm.litellm_core_utils.served_output_texts import (
|
||||
SERVED_OUTPUT_TEXTS_KEY,
|
||||
overlay_served_output_texts,
|
||||
)
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
from litellm.llms.base_llm.search.transformation import SearchResponse
|
||||
from litellm.responses.utils import ResponseAPILoggingUtils
|
||||
|
|
@ -5843,15 +5847,12 @@ class StandardLoggingPayloadSetup:
|
|||
|
||||
modified_final_response_obj: Final = redact_message_input_output_from_logging(
|
||||
model_call_details=kwargs,
|
||||
result=final_response_obj,
|
||||
result=overlay_served_output_texts(final_response_obj, kwargs.get(SERVED_OUTPUT_TEXTS_KEY)),
|
||||
)
|
||||
|
||||
if modified_final_response_obj is not None and isinstance(modified_final_response_obj, BaseModel):
|
||||
final_response_obj = modified_final_response_obj.model_dump()
|
||||
else:
|
||||
final_response_obj = modified_final_response_obj
|
||||
|
||||
return final_response_obj
|
||||
return modified_final_response_obj.model_dump()
|
||||
return modified_final_response_obj
|
||||
|
||||
@staticmethod
|
||||
def get_additional_headers(
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from litellm.litellm_core_utils.classifier_logging import without_classifier_aud
|
|||
from litellm.litellm_core_utils.core_helpers import (
|
||||
get_metadata_variable_name_from_kwargs,
|
||||
)
|
||||
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
redact_vertex_ai_metadata_from_litellm_params,
|
||||
redact_vertex_ai_metadata_from_logged_object,
|
||||
|
|
@ -267,6 +268,7 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
|
|||
model_call_details["messages"] = [{"role": "user", "content": REDACTED_BY_LITELLM}]
|
||||
model_call_details["prompt"] = ""
|
||||
model_call_details["input"] = ""
|
||||
model_call_details.pop(SERVED_OUTPUT_TEXTS_KEY, None)
|
||||
standard_logging_object: Final = model_call_details.get("standard_logging_object")
|
||||
if isinstance(standard_logging_object, Mapping):
|
||||
model_call_details["standard_logging_object"] = _redact_standard_logging_object(standard_logging_object)
|
||||
|
|
|
|||
171
litellm/litellm_core_utils/served_output_texts.py
Normal file
171
litellm/litellm_core_utils/served_output_texts.py
Normal file
|
|
@ -0,0 +1,171 @@
|
|||
"""Assistant text the caller received, per choice, so the logging payload stores the response a
|
||||
post-call guardrail rewrote rather than the provider response the proxy assembled before it ran."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Final, Literal
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.types.utils import ModelResponse, ModelResponseStream
|
||||
|
||||
SERVED_OUTPUT_TEXTS_KEY: Final = "served_output_texts"
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, object])
|
||||
_JSON_LIST: Final = TypeAdapter(list[object])
|
||||
_TEXTS: Final = TypeAdapter(tuple[str | None, ...])
|
||||
|
||||
ServedTexts = tuple[str | None, ...]
|
||||
|
||||
|
||||
class _TextBlock(BaseModel):
|
||||
type: str
|
||||
text: str | None = None
|
||||
|
||||
|
||||
class _AnthropicMessage(BaseModel):
|
||||
type: Literal["message"]
|
||||
content: list[_TextBlock]
|
||||
|
||||
|
||||
class _ResponsesOutputItem(BaseModel):
|
||||
type: str
|
||||
content: list[_TextBlock] = []
|
||||
|
||||
|
||||
class _ResponsesResponse(BaseModel):
|
||||
object: Literal["response"]
|
||||
output: list[_ResponsesOutputItem]
|
||||
|
||||
|
||||
class _ChatChoices(BaseModel):
|
||||
choices: list[object]
|
||||
|
||||
|
||||
def _as_json_object(response: object) -> dict[str, object] | None:
|
||||
candidate: Final = response.model_dump() if isinstance(response, BaseModel) else response
|
||||
try:
|
||||
return _JSON_OBJECT.validate_python(candidate)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _joined_block_texts(blocks: Sequence[_TextBlock], *, text_type: str) -> str | None:
|
||||
texts: Final = tuple(block.text for block in blocks if block.type == text_type and block.text is not None)
|
||||
return "".join(texts) if texts else None
|
||||
|
||||
|
||||
def _chat_texts(response: ModelResponse) -> ServedTexts | None:
|
||||
texts: Final = tuple(
|
||||
choice.message.content if isinstance(choice.message.content, str) else None for choice in response.choices
|
||||
)
|
||||
return texts if any(text is not None for text in texts) else None
|
||||
|
||||
|
||||
def _anthropic_message_text(response: dict[str, object]) -> str | None:
|
||||
try:
|
||||
message: Final = _AnthropicMessage.model_validate(response)
|
||||
except ValidationError:
|
||||
return None
|
||||
return _joined_block_texts(message.content, text_type="text")
|
||||
|
||||
|
||||
def _responses_api_text(response: dict[str, object]) -> str | None:
|
||||
try:
|
||||
parsed: Final = _ResponsesResponse.model_validate(response)
|
||||
except ValidationError:
|
||||
return None
|
||||
texts: Final = tuple(
|
||||
text
|
||||
for item in parsed.output
|
||||
if item.type == "message" and (text := _joined_block_texts(item.content, text_type="output_text")) is not None
|
||||
)
|
||||
return "".join(texts) if texts else None
|
||||
|
||||
|
||||
def _chat_dict_texts(response: dict[str, object]) -> ServedTexts | None:
|
||||
try:
|
||||
_ChatChoices.model_validate(response)
|
||||
return _chat_texts(ModelResponse(**response))
|
||||
except (ValidationError, TypeError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def served_output_texts(response: object) -> ServedTexts | None:
|
||||
if isinstance(response, ModelResponse):
|
||||
return _chat_texts(response)
|
||||
mapping: Final = _as_json_object(response)
|
||||
if mapping is None:
|
||||
return None
|
||||
chat_texts: Final = _chat_dict_texts(mapping)
|
||||
if chat_texts is not None:
|
||||
return chat_texts
|
||||
anthropic_text: Final = _anthropic_message_text(mapping)
|
||||
text: Final = anthropic_text if anthropic_text is not None else _responses_api_text(mapping)
|
||||
return (text,) if text is not None else None
|
||||
|
||||
|
||||
def served_stream_output_texts(chunks: Sequence[object]) -> ServedTexts | None:
|
||||
if chunks and all(isinstance(chunk, ModelResponseStream) for chunk in chunks):
|
||||
return _chat_stream_texts(tuple(chunk for chunk in chunks if isinstance(chunk, ModelResponseStream)))
|
||||
from litellm.proxy.guardrails.anthropic_sse import assemble_anthropic_sse_stream, is_anthropic_sse_stream
|
||||
|
||||
if not is_anthropic_sse_stream(chunks):
|
||||
return None
|
||||
assembled: Final = assemble_anthropic_sse_stream(chunks)
|
||||
return _chat_texts(assembled) if assembled is not None else None
|
||||
|
||||
|
||||
def _chat_stream_choice_text(chunks: Sequence[ModelResponseStream], index: int) -> str | None:
|
||||
contents: Final = tuple(
|
||||
content
|
||||
for chunk in chunks
|
||||
for choice in chunk.choices
|
||||
if choice.index == index and isinstance(content := choice.delta.content, str)
|
||||
)
|
||||
return "".join(contents) if contents else None
|
||||
|
||||
|
||||
def _chat_stream_texts(chunks: Sequence[ModelResponseStream]) -> ServedTexts | None:
|
||||
choice_count: Final = max((choice.index + 1 for chunk in chunks for choice in chunk.choices), default=0)
|
||||
texts: Final = tuple(_chat_stream_choice_text(chunks, index) for index in range(choice_count))
|
||||
return texts if any(text is not None for text in texts) else None
|
||||
|
||||
|
||||
def record_served_output_texts(model_call_details: dict[str, object], texts: ServedTexts | None) -> None:
|
||||
if texts is None:
|
||||
return
|
||||
model_call_details[SERVED_OUTPUT_TEXTS_KEY] = texts # rebind-ok: model_call_details is the shared kwargs bag
|
||||
|
||||
|
||||
def overlay_served_output_texts(
|
||||
response_obj: dict[str, object] | str | list[object] | None, served_texts: object
|
||||
) -> dict[str, object] | str | list[object] | None:
|
||||
if not isinstance(response_obj, dict):
|
||||
return response_obj
|
||||
logged: Final = _as_json_object(response_obj)
|
||||
if logged is None:
|
||||
return response_obj
|
||||
try:
|
||||
texts: Final = _TEXTS.validate_python(served_texts)
|
||||
choices: Final = _JSON_LIST.validate_python(logged.get("choices"))
|
||||
except ValidationError:
|
||||
return response_obj
|
||||
return {
|
||||
**logged,
|
||||
"choices": [
|
||||
_choice_with_text(choice, texts[index]) if index < len(texts) else choice
|
||||
for index, choice in enumerate(choices)
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _choice_with_text(choice: object, text: str | None) -> object:
|
||||
choice_obj: Final = _as_json_object(choice)
|
||||
if choice_obj is None or text is None:
|
||||
return choice
|
||||
message: Final = _as_json_object(choice_obj.get("message"))
|
||||
if message is None or message.get("content") == text:
|
||||
return choice
|
||||
return {**choice_obj, "message": {**message, "content": text}}
|
||||
|
|
@ -68,6 +68,10 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import (
|
|||
get_response_headers,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.served_output_texts import (
|
||||
record_served_output_texts,
|
||||
served_output_texts,
|
||||
)
|
||||
from litellm.litellm_core_utils.streaming_handler import (
|
||||
backfill_missing_cache_usage_fields,
|
||||
)
|
||||
|
|
@ -2803,6 +2807,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
)
|
||||
record_served_output_texts(logging_obj.model_call_details, served_output_texts(response))
|
||||
except Exception:
|
||||
_exception_raised = True
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -144,6 +144,10 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.served_output_texts import (
|
||||
record_served_output_texts,
|
||||
served_stream_output_texts,
|
||||
)
|
||||
from litellm.litellm_core_utils.token_counter import offload_token_count
|
||||
from litellm.llms import load_guardrail_translation_mappings
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
|
|
@ -3809,12 +3813,16 @@ class ProxyLogging:
|
|||
translation=pipeline_translation,
|
||||
)
|
||||
|
||||
served_chunks: Final[list[object]] = [] # mutable-ok: accumulates while yielding to the client
|
||||
try:
|
||||
async for chunk in current_response:
|
||||
served_chunks.append(chunk)
|
||||
yield chunk
|
||||
except (GeneratorExit, asyncio.CancelledError):
|
||||
ProxyLogging._record_served_stream_output(request_data, served_chunks)
|
||||
raise
|
||||
except Exception as e:
|
||||
ProxyLogging._record_served_stream_output(request_data, served_chunks)
|
||||
if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, e):
|
||||
ProxyLogging._fire_deferred_stream_logging(request_data)
|
||||
raise
|
||||
|
|
@ -3823,6 +3831,7 @@ class ProxyLogging:
|
|||
# completed. unified_guardrail writes guardrail_information during
|
||||
# its end-of-stream block (inside current_response), so by the time
|
||||
# we reach this point the metadata is fully populated.
|
||||
ProxyLogging._record_served_stream_output(request_data, served_chunks)
|
||||
ProxyLogging._fire_deferred_stream_logging(request_data)
|
||||
|
||||
async def _pipeline_gated_stream(
|
||||
|
|
@ -3890,6 +3899,13 @@ class ProxyLogging:
|
|||
for buffered_item in buffered:
|
||||
yield buffered_item
|
||||
|
||||
@staticmethod
|
||||
def _record_served_stream_output(request_data: Mapping[str, object], served_chunks: Sequence[object]) -> None:
|
||||
logging_obj: Final = request_data.get("litellm_logging_obj")
|
||||
if not isinstance(logging_obj, Logging):
|
||||
return
|
||||
record_served_output_texts(logging_obj.model_call_details, served_stream_output_texts(served_chunks))
|
||||
|
||||
@staticmethod
|
||||
def _fire_deferred_stream_logging(request_data: dict) -> None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
- {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], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "AWS content guardrail blocks harmful input"}
|
||||
|
|
|
|||
|
|
@ -381,6 +381,26 @@ class GuardrailsClient:
|
|||
),
|
||||
)
|
||||
|
||||
def messages_raw(
|
||||
self,
|
||||
key: str,
|
||||
model: str,
|
||||
text: str,
|
||||
*,
|
||||
guardrails: list[str] | None = None,
|
||||
max_tokens: int = 64,
|
||||
) -> StreamingResponse:
|
||||
return self.proxy.transport.send(
|
||||
"/v1/messages",
|
||||
headers=self.proxy.transport.bearer(key),
|
||||
json=AnthropicMessagesBody(
|
||||
model=model,
|
||||
messages=[ChatMessage(role="user", content=text)],
|
||||
max_tokens=max_tokens,
|
||||
guardrails=guardrails,
|
||||
),
|
||||
)
|
||||
|
||||
def messages_stream_raw(
|
||||
self,
|
||||
key: str,
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ this suite deliberately requires the detected-entity details to remain visible.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
|
|
@ -481,6 +482,149 @@ class TestPresidioCreditCardOutputMasking:
|
|||
_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 json.dumps(row.response)
|
||||
|
||||
|
||||
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"])
|
||||
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"])
|
||||
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"])
|
||||
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"])
|
||||
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"])
|
||||
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])
|
||||
|
|
|
|||
|
|
@ -874,6 +874,8 @@ class SpendLogRow(BaseModel):
|
|||
request_tags: list[str] | None = None
|
||||
metadata: SpendLogMetadata | None = None
|
||||
proxy_server_request: JsonValue = None
|
||||
response: JsonValue = None
|
||||
litellm_call_id: str | None = None
|
||||
|
||||
|
||||
class SpendLogs(RootModel[list[SpendLogRow]]):
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import contextlib
|
||||
import datetime
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
|
@ -3013,6 +3014,54 @@ def test_get_final_response_obj_with_empty_response_obj_and_list_init():
|
|||
assert result[1].name == "Object2"
|
||||
|
||||
|
||||
def test_get_final_response_obj_stores_the_text_a_post_call_guardrail_served():
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
||||
|
||||
raw = {
|
||||
"id": "x",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"finish_reason": "stop",
|
||||
"message": {"role": "assistant", "content": "Card: 4111 1111 1111 1111"},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
logged = StandardLoggingPayloadSetup.get_final_response_obj(
|
||||
response_obj=raw, init_response_obj=raw, kwargs={SERVED_OUTPUT_TEXTS_KEY: ("Card: <CREDIT_CARD>",)}
|
||||
)
|
||||
untouched = StandardLoggingPayloadSetup.get_final_response_obj(response_obj=raw, init_response_obj=raw, kwargs={})
|
||||
|
||||
assert isinstance(logged, dict)
|
||||
assert logged["choices"][0]["message"]["content"] == "Card: <CREDIT_CARD>"
|
||||
assert logged["choices"][0]["finish_reason"] == "stop"
|
||||
assert untouched == raw
|
||||
|
||||
|
||||
def test_get_final_response_obj_redacts_the_served_text_when_message_logging_is_off(monkeypatch: pytest.MonkeyPatch):
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
||||
|
||||
monkeypatch.setattr(litellm, "turn_off_message_logging", True)
|
||||
raw = {
|
||||
"id": "x",
|
||||
"choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "Card: 4111"}}],
|
||||
}
|
||||
|
||||
logged = StandardLoggingPayloadSetup.get_final_response_obj(
|
||||
response_obj=raw,
|
||||
init_response_obj=raw,
|
||||
kwargs={SERVED_OUTPUT_TEXTS_KEY: ("Card: <CREDIT_CARD>",), "litellm_params": {}},
|
||||
)
|
||||
|
||||
assert isinstance(logged, dict)
|
||||
assert "<CREDIT_CARD>" not in json.dumps(logged), logged
|
||||
assert "4111" not in json.dumps(logged), logged
|
||||
|
||||
|
||||
def test_get_usage_as_dict():
|
||||
"""
|
||||
Test get_usage_as_dict returns usage as plain dict from response_obj or combined_usage_object.
|
||||
|
|
|
|||
|
|
@ -1032,3 +1032,11 @@ def test_a_callback_that_redacts_itself_keeps_its_messages_but_not_the_classifie
|
|||
assert "classifier_input" not in stored
|
||||
assert stored["messages"] == payload["messages"]
|
||||
assert stored["response"] == payload["response"]
|
||||
|
||||
|
||||
def test_perform_redaction_drops_the_served_output_texts_from_the_callback_kwargs() -> None:
|
||||
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
||||
|
||||
details: Final = {"litellm_params": {}, SERVED_OUTPUT_TEXTS_KEY: ("Card: <CREDIT_CARD>",)}
|
||||
perform_redaction(details, None)
|
||||
assert SERVED_OUTPUT_TEXTS_KEY not in details
|
||||
|
|
|
|||
|
|
@ -0,0 +1,114 @@
|
|||
from litellm.litellm_core_utils.served_output_texts import (
|
||||
SERVED_OUTPUT_TEXTS_KEY,
|
||||
overlay_served_output_texts,
|
||||
record_served_output_texts,
|
||||
served_output_texts,
|
||||
served_stream_output_texts,
|
||||
)
|
||||
from litellm.types.utils import Choices, Delta, Message, ModelResponse, ModelResponseStream, StreamingChoices
|
||||
|
||||
RAW = "Card: 4111 1111 1111 1111"
|
||||
MASKED = "Card: <CREDIT_CARD>"
|
||||
|
||||
|
||||
def _chat_response(*texts: str) -> ModelResponse:
|
||||
return ModelResponse(
|
||||
choices=[Choices(index=i, message=Message(content=text, role="assistant")) for i, text in enumerate(texts)]
|
||||
)
|
||||
|
||||
|
||||
def _chat_dict(*texts: str) -> dict[str, object]:
|
||||
return {
|
||||
"id": "x",
|
||||
"choices": [
|
||||
{"index": i, "finish_reason": "stop", "message": {"role": "assistant", "content": text}}
|
||||
for i, text in enumerate(texts)
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def _choice_texts(response: object) -> tuple[str | None, ...]:
|
||||
texts = served_output_texts(response)
|
||||
assert texts is not None, response
|
||||
return texts
|
||||
|
||||
|
||||
def _stream_chunk(text: str, index: int = 0) -> ModelResponseStream:
|
||||
return ModelResponseStream(choices=[StreamingChoices(index=index, delta=Delta(content=text))])
|
||||
|
||||
|
||||
def test_served_output_texts_reads_each_response_shape():
|
||||
assert served_output_texts(_chat_response(MASKED, "second")) == (MASKED, "second")
|
||||
assert served_output_texts(_chat_dict(MASKED)) == (MASKED,)
|
||||
assert served_output_texts(
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}],
|
||||
}
|
||||
) == ("ab",)
|
||||
responses_api: dict[str, object] = {
|
||||
"object": "response",
|
||||
"output": [
|
||||
{"type": "reasoning", "content": []},
|
||||
{"type": "message", "content": [{"type": "output_text", "text": MASKED}]},
|
||||
],
|
||||
}
|
||||
assert served_output_texts(responses_api) == (MASKED,)
|
||||
assert served_output_texts({"data": [{"embedding": [0.1]}]}) is None
|
||||
assert served_output_texts("plain") is None
|
||||
|
||||
|
||||
def test_served_stream_output_texts_joins_chat_chunks_and_reads_anthropic_sse():
|
||||
assert served_stream_output_texts([_stream_chunk("Card: "), _stream_chunk("<CREDIT_CARD>")]) == (MASKED,)
|
||||
sse = (
|
||||
'event: message_start\ndata: {"type":"message_start","message":{"id":"m","type":"message","role":"assistant",'
|
||||
'"content":[],"model":"x","usage":{"input_tokens":1,"output_tokens":1}}}\n\n',
|
||||
'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}\n\n',
|
||||
f'event: content_block_delta\ndata: {{"type":"content_block_delta","index":0,"delta":{{"type":"text_delta","text":"{MASKED}"}}}}\n\n',
|
||||
'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n',
|
||||
'event: message_stop\ndata: {"type":"message_stop"}\n\n',
|
||||
)
|
||||
assert served_stream_output_texts(tuple(chunk.encode() for chunk in sse)) == (MASKED,)
|
||||
assert served_stream_output_texts([]) is None
|
||||
assert served_stream_output_texts([b"not sse"]) is None
|
||||
|
||||
|
||||
def test_served_stream_output_texts_keeps_every_choice_index_when_chunks_carry_one_choice_each():
|
||||
chunks = [_stream_chunk("first ", 0), _stream_chunk(MASKED, 1), _stream_chunk("choice", 0)]
|
||||
assert served_stream_output_texts(chunks) == ("first choice", MASKED)
|
||||
|
||||
|
||||
def test_blanked_output_is_served_as_empty_text_and_overlaid():
|
||||
assert served_output_texts(_chat_response("")) == ("",)
|
||||
assert served_output_texts({"type": "message", "role": "assistant", "content": [{"type": "text", "text": ""}]}) == (
|
||||
"",
|
||||
)
|
||||
assert served_stream_output_texts([_stream_chunk("")]) == ("",)
|
||||
assert _choice_texts(overlay_served_output_texts(_chat_dict(RAW), ("",))) == ("",)
|
||||
|
||||
|
||||
def test_overlay_replaces_logged_choice_text_with_served_text():
|
||||
logged = _chat_dict(RAW, RAW)
|
||||
overlaid = overlay_served_output_texts(logged, (MASKED,))
|
||||
assert _choice_texts(overlaid) == (MASKED, RAW)
|
||||
assert isinstance(overlaid, dict)
|
||||
assert overlaid["id"] == "x"
|
||||
assert _choice_texts(logged) == (RAW, RAW)
|
||||
|
||||
|
||||
def test_overlay_leaves_unreadable_inputs_untouched():
|
||||
logged = _chat_dict(RAW)
|
||||
assert overlay_served_output_texts(logged, None) is logged
|
||||
assert overlay_served_output_texts(logged, "not a tuple") is logged
|
||||
assert overlay_served_output_texts(logged, (None,)) == logged
|
||||
assert overlay_served_output_texts("text", (MASKED,)) == "text"
|
||||
assert overlay_served_output_texts({"data": []}, (MASKED,)) == {"data": []}
|
||||
|
||||
|
||||
def test_record_served_output_texts_only_writes_readable_texts():
|
||||
details: dict[str, object] = {}
|
||||
record_served_output_texts(details, None)
|
||||
assert SERVED_OUTPUT_TEXTS_KEY not in details
|
||||
record_served_output_texts(details, (MASKED,))
|
||||
assert details[SERVED_OUTPUT_TEXTS_KEY] == (MASKED,)
|
||||
|
|
@ -9034,6 +9034,60 @@ class TestDetachedStreamFailureHook:
|
|||
assert [call["original_exception"] for call in recorder.calls] == [failure]
|
||||
|
||||
|
||||
class TestPostCallMaskedOutputReachesDeferredLogging:
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_streaming_records_the_masked_response_before_deferred_logging_fires(self, monkeypatch):
|
||||
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
||||
from litellm.types.utils import Choices, Message, ModelResponse
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_call_id = "lit-8325-call"
|
||||
logging_obj._defer_async_logging = False
|
||||
logging_obj._on_deferred_stream_complete = None
|
||||
logging_obj.cost_breakdown = None
|
||||
logging_obj.model_call_details = {}
|
||||
recorded_at_enqueue: dict[str, object] = {}
|
||||
logging_obj._enqueue_deferred_logging = lambda: recorded_at_enqueue.update(logging_obj.model_call_details)
|
||||
|
||||
processor = ProxyBaseLLMRequestProcessing(data={"model": "oa", "litellm_logging_obj": logging_obj})
|
||||
|
||||
def mask(data, user_api_key_dict, response):
|
||||
response.choices[0].message.content = "Card: <CREDIT_CARD>"
|
||||
return response
|
||||
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
|
||||
proxy_logging_obj.update_request_status = AsyncMock(return_value=None)
|
||||
proxy_logging_obj.post_call_success_hook = AsyncMock(side_effect=mask)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value=None)
|
||||
|
||||
async def fake_route_request(**kwargs):
|
||||
async def call():
|
||||
return ModelResponse(
|
||||
choices=[Choices(index=0, message=Message(content="Card: 4111 1111 1111 1111", role="assistant"))]
|
||||
)
|
||||
|
||||
return call()
|
||||
|
||||
monkeypatch.setattr(litellm.proxy.common_request_processing, "route_request", fake_route_request)
|
||||
|
||||
result = await processor.base_process_llm_request(
|
||||
request=Request(scope={"type": "http", "headers": []}),
|
||||
fastapi_response=Response(),
|
||||
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"),
|
||||
route_type="acompletion",
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
general_settings={},
|
||||
proxy_config=MagicMock(spec=ProxyConfig),
|
||||
select_data_generator=MagicMock(),
|
||||
is_streaming_request=False,
|
||||
skip_pre_call_logic=True,
|
||||
)
|
||||
|
||||
assert result.choices[0].message.content == "Card: <CREDIT_CARD>"
|
||||
assert recorded_at_enqueue[SERVED_OUTPUT_TEXTS_KEY] == ("Card: <CREDIT_CARD>",)
|
||||
|
||||
|
||||
class TestStreamingResponseHeadersFollowFallback:
|
||||
"""LIT-6767: the streaming branch has to publish the deployment that served the stream."""
|
||||
|
||||
|
|
|
|||
|
|
@ -347,6 +347,60 @@ async def test_post_call_stream_guardrail_keeps_own_iterator_on_chat_completions
|
|||
assert delivered_text != ""
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_stream_records_masked_text_for_deferred_logging(monkeypatch):
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
monkeypatch.setattr(litellm, "callbacks", [_content_filter_guardrail("MASK")])
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
||||
logging_obj = _streaming_logging_obj()
|
||||
|
||||
async def fake_stream():
|
||||
yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="the zebra runs"))])
|
||||
yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")])
|
||||
|
||||
delivered_text = ""
|
||||
async for chunk in proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
response=fake_stream(),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/chat/completions"),
|
||||
request_data={"model": "gpt-4o-mini", "metadata": {}, "litellm_logging_obj": logging_obj},
|
||||
):
|
||||
for choice in chunk.choices:
|
||||
delivered_text += choice.delta.content or ""
|
||||
|
||||
assert "zebra" not in delivered_text
|
||||
assert logging_obj.model_call_details[SERVED_OUTPUT_TEXTS_KEY] == (delivered_text,)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_stream_records_the_served_text_when_the_client_disconnects(monkeypatch):
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
monkeypatch.setattr(litellm, "callbacks", [_content_filter_guardrail("MASK")])
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
||||
logging_obj = _streaming_logging_obj()
|
||||
|
||||
async def fake_stream():
|
||||
yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="the zebra runs"))])
|
||||
yield ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content=" far"))])
|
||||
|
||||
stream = proxy_logging.async_post_call_streaming_iterator_hook(
|
||||
response=fake_stream(),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/chat/completions"),
|
||||
request_data={"model": "gpt-4o-mini", "metadata": {}, "litellm_logging_obj": logging_obj},
|
||||
)
|
||||
first = await stream.__anext__()
|
||||
await stream.aclose()
|
||||
|
||||
delivered_text = "".join(choice.delta.content or "" for choice in first.choices)
|
||||
assert "zebra" not in delivered_text
|
||||
assert logging_obj.model_call_details[SERVED_OUTPUT_TEXTS_KEY] == (delivered_text,)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unified_guardrail_iterator_accepts_explicit_guardrail():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue