mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
fix(guardrails): bound per_session dedup keys with a fixed-size hash
The dedup cache keyed sessions by the serialized caller, session id and side, so a caller sending large distinct litellm_session_id values grew it by their full size for an hour. Keys are now the SHA-256 of that tuple, 64 characters each, so the 100k-entry cap bounds the cache's memory
This commit is contained in:
parent
045b13fe21
commit
a37c292b5f
2 changed files with 22 additions and 1 deletions
|
|
@ -1,3 +1,4 @@
|
|||
import hashlib
|
||||
import json
|
||||
from typing import Final, Literal, NamedTuple
|
||||
|
||||
|
|
@ -49,6 +50,10 @@ def _jsonable(value: object) -> object:
|
|||
return to_jsonable_python(value, fallback=repr, bytes_mode="base64")
|
||||
|
||||
|
||||
def _session_key(caller: Caller, session_id: str, input_type: Literal["request", "response"]) -> str:
|
||||
return hashlib.sha256(json.dumps((*caller, session_id, input_type)).encode()).hexdigest()
|
||||
|
||||
|
||||
def _sent_value(sent: GenericGuardrailAPIInputs, key: str) -> object:
|
||||
return sent.get(key) if key in sent else ()
|
||||
|
||||
|
|
@ -87,7 +92,7 @@ class RecordScope:
|
|||
case "off":
|
||||
return False
|
||||
case "per_session":
|
||||
return session_id is None or self._claim_session(json.dumps((*caller, session_id, input_type)))
|
||||
return session_id is None or self._claim_session(_session_key(caller, session_id, input_type))
|
||||
return assert_never(self._scope)
|
||||
|
||||
def _claim_session(self, key: str) -> bool:
|
||||
|
|
|
|||
|
|
@ -96,3 +96,19 @@ def test_per_call_and_off_decide_without_claiming_a_session(scope: GuardrailInfo
|
|||
)
|
||||
|
||||
assert (decisions, session_still_unclaimed) == ([recorded, recorded], True)
|
||||
|
||||
|
||||
def test_a_huge_session_id_is_stored_under_a_fixed_size_key() -> None:
|
||||
sessions: Final = InMemoryCache()
|
||||
record_scope: Final = RecordScope("per_session", recorded_sessions=sessions)
|
||||
caller: Final = Caller(key_hash="hash-a", team_id=None, user_id=None)
|
||||
huge_session_id: Final = "s" * 1_000_000
|
||||
|
||||
decisions: Final = [
|
||||
record_scope.should_record_allow(session_id=huge_session_id, caller=caller, input_type="request")
|
||||
for _ in range(2)
|
||||
]
|
||||
|
||||
assert (decisions, [len(key) for key in sessions.cache_dict]) == ([True, False], [64]), (
|
||||
"a caller-chosen session id must not grow the in-memory dedup cache"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue