mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
fix(memory): reuse shared credential redaction before persistence
This commit is contained in:
parent
f2026051a3
commit
6f024adfcc
3 changed files with 32 additions and 6 deletions
|
|
@ -3,21 +3,18 @@ from typing import Final
|
|||
|
||||
from rapidfuzz import fuzz, process
|
||||
|
||||
from litellm.litellm_core_utils.secret_redaction import redact_string
|
||||
from litellm.types.memory_v2 import MemoryEntry
|
||||
|
||||
_STOP_WORDS: Final = frozenset(
|
||||
"the and for how what why with this that does have from about our are was when should can you work team".split()
|
||||
)
|
||||
_TOKEN: Final = re.compile(r"[\w-]{2,}", re.UNICODE)
|
||||
_PRIVATE_KEY: Final = re.compile(r"-----BEGIN [A-Z ]*PRIVATE KEY-----.*?-----END [A-Z ]*PRIVATE KEY-----", re.DOTALL)
|
||||
_CREDENTIAL: Final = re.compile(r"\b(?:sk-|gh[pousr]_|github_pat_)[A-Za-z0-9_-]{12,}")
|
||||
_BEARER: Final = re.compile(r"(Bearer\s+)[A-Za-z0-9._~+/-]{12,}", re.IGNORECASE)
|
||||
|
||||
|
||||
def redact_memory(value: str) -> str:
|
||||
return _BEARER.sub(
|
||||
r"\1[REDACTED]", _CREDENTIAL.sub("[REDACTED TOKEN]", _PRIVATE_KEY.sub("[REDACTED PRIVATE KEY]", value))
|
||||
)
|
||||
return _CREDENTIAL.sub("[REDACTED TOKEN]", redact_string(value))
|
||||
|
||||
|
||||
def _similarity(term: str, text: str, words: tuple[str, ...]) -> float:
|
||||
|
|
|
|||
|
|
@ -55,4 +55,4 @@ def test_recognizable_credentials_are_redacted_without_removing_the_observation(
|
|||
assert "The integration failed" in redacted
|
||||
assert "abcdefghijklmnopqrst" not in redacted
|
||||
assert "private material" not in redacted
|
||||
assert "[REDACTED PRIVATE KEY]" in redacted
|
||||
assert "REDACTED" in redacted
|
||||
|
|
|
|||
|
|
@ -329,6 +329,35 @@ async def test_capture_records_authenticated_contributor_and_tenant(database: Ma
|
|||
assert data["organization_id"] == "org" and data["team_id"] == "team"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("credential", "secret"),
|
||||
(
|
||||
("AKIA" + "A" * 16, "A" * 16),
|
||||
("aws_secret_access_key=" + "B" * 40, "B" * 40),
|
||||
("AIza" + "C" * 35, "C" * 35),
|
||||
("postgres://test-user:memory-test-password@db.example.test/app", "memory-test-password"),
|
||||
("Authorization: Basic " + "D" * 24, "D" * 24),
|
||||
("Authorization: Bearer " + "E" * 24, "E" * 24),
|
||||
("ghp_" + "F" * 24, "F" * 24),
|
||||
("github_pat_" + "G" * 24, "G" * 24),
|
||||
),
|
||||
)
|
||||
async def test_capture_redacts_credentials_before_persisting_content_and_metadata(
|
||||
database: MagicMock, credential: str, secret: str
|
||||
) -> None:
|
||||
configure(database)
|
||||
text = "Deployment uses staging port 8123; credential: " + credential
|
||||
fields = ("title", "content", "evidence", "when_to_use", "scope", "source")
|
||||
await management.capture_entry(_CAPTURE.model_copy(update={field: text for field in fields}), auth())
|
||||
data = database.db.litellm_memorytable.create.call_args.kwargs["data"]
|
||||
stored = {"content": data["value"], **json.loads(data["metadata"])}
|
||||
for field in fields:
|
||||
assert secret not in stored[field]
|
||||
assert "REDACTED" in stored[field]
|
||||
assert "staging port 8123" in stored[field]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_search_finds_an_old_record_beyond_the_first_thousand(database: MagicMock) -> None:
|
||||
configure(database)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue