mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(guardrails): keep post_call guardrail info on streamed chat completions (#40806)
* fix(guardrails): sync logging_obj guardrail info on every record so post_call entries survive streamed chat completions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(guardrails): hoist regression test imports to module scope 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:
parent
e72d1cfa8d
commit
f84f986b4e
2 changed files with 61 additions and 1 deletions
|
|
@ -1217,6 +1217,7 @@ class CustomGuardrail(CustomLogger):
|
|||
|
||||
_, metadata_bucket = get_or_create_metadata_bucket(request_data)
|
||||
_append_guardrail_info(metadata_bucket)
|
||||
_sync_guardrail_info_to_logging_obj(request_data, request_data.get("litellm_logging_obj"))
|
||||
|
||||
_guardrail_self_recorded.set(True)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import asyncio
|
||||
import datetime as dt
|
||||
from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
|
|
@ -9,9 +10,16 @@ from litellm.integrations.custom_guardrail import (
|
|||
CustomGuardrail,
|
||||
log_guardrail_information,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
|
||||
from litellm.types.guardrails import GuardrailEventHooks, Mode
|
||||
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail
|
||||
from litellm.types.utils import (
|
||||
Choices,
|
||||
GenericGuardrailAPIInputs,
|
||||
GuardrailTracingDetail,
|
||||
Message,
|
||||
ModelResponse,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
|
@ -2345,6 +2353,57 @@ class TestUndecoratedApplyGuardrailIsLogged:
|
|||
|
||||
assert _Labelled.seen_label == "docs-style"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_recorded_outside_decorator_reaches_standard_logging_object(self):
|
||||
"""LIT-7608 regression: the auto-wrapped pre_call apply_guardrail copies the request bucket
|
||||
into logging_obj.litellm_params["metadata"]. A post_call entry recorded later without the
|
||||
decorator (the Bedrock streaming hook) must not be shadowed by that stale copy."""
|
||||
messages: Final = [{"role": "user", "content": "hello there"}]
|
||||
litellm_metadata: Final[dict] = {"user_api_key_user_id": "u1"}
|
||||
logging_obj: Final = Logging(
|
||||
model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
messages=messages,
|
||||
stream=True,
|
||||
call_type=CallTypes.acompletion.value,
|
||||
start_time=dt.datetime.now(),
|
||||
litellm_call_id="call-1",
|
||||
function_id="fn-1",
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
litellm_params={"litellm_metadata": litellm_metadata},
|
||||
optional_params={},
|
||||
model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
custom_llm_provider="bedrock",
|
||||
)
|
||||
request_data: Final = {
|
||||
"model": "bedrock-haiku",
|
||||
"messages": messages,
|
||||
"litellm_metadata": litellm_metadata,
|
||||
"litellm_logging_obj": logging_obj,
|
||||
}
|
||||
guardrail: Final = _UndecoratedGuardrail(guardrail_name="bedrock-pre", event_hook=GuardrailEventHooks.pre_call)
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
inputs=GenericGuardrailAPIInputs(texts=["hello there"]),
|
||||
request_data=request_data,
|
||||
input_type="request",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response={"action": "NONE"},
|
||||
request_data=request_data,
|
||||
guardrail_status="success",
|
||||
event_type=GuardrailEventHooks.post_call,
|
||||
)
|
||||
await logging_obj.async_success_handler(
|
||||
result=ModelResponse(choices=[Choices(message=Message(role="assistant", content="general kenobi"))]),
|
||||
start_time=dt.datetime.now(),
|
||||
end_time=dt.datetime.now(),
|
||||
)
|
||||
|
||||
entries: Final = logging_obj.model_call_details["standard_logging_object"]["guardrail_information"]
|
||||
assert [e["guardrail_mode"] for e in entries] == ["pre_call", "post_call"]
|
||||
|
||||
|
||||
class _ApplyOnlyObserver(CustomGuardrail):
|
||||
"""Overrides only apply_guardrail, like panw_prisma_airs; inherits async_logging_hook."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue