Skip to content

Commit

Permalink
guard against None in leakreplay attempt history management (#1081)
Browse files Browse the repository at this point in the history
resolves #879 

`NoneType` in attempt message history would cause a crash when
`leakreplay` rewrites that message history. Guard against `None` here.
It's unclear how a None would get in there in the first place, but the
original report hasn't had updates, so this may have been a transient
behaviour.

Thanks @bleszily


## Verification

- `python -m pytest
tests/probes/test_probes_leakreplay.py::test_leakreplay_handle_incomplete_attempt`
  • Loading branch information
leondz authored Jan 16, 2025
2 parents 0315908 + 70f8d9b commit 8f84963
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 3 deletions.
7 changes: 4 additions & 3 deletions garak/probes/leakreplay.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,9 +64,10 @@ def _attempt_prestore_hook(self, attempt: Attempt, seq: int) -> Attempt:

def _postprocess_hook(self, attempt: Attempt) -> Attempt:
for idx, thread in enumerate(attempt.messages):
attempt.messages[idx][-1]["content"] = re.sub(
"</?name>", "", thread[-1]["content"]
)
if thread[-1]["content"] is not None:
attempt.messages[idx][-1]["content"] = re.sub(
"</?name>", "", thread[-1]["content"]
)
return attempt


Expand Down
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ authors = [
{ name = "Shine-afk" },
{ name = "Rafael Sandroni" },
{ name = "Eric Hacker" },
{ name = "Blessed Uyo" },
]
license = { file = "LICENSE" }
description = "LLM vulnerability scanner"
Expand Down
8 changes: 8 additions & 0 deletions tests/probes/test_probes_leakreplay.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import garak._plugins
import garak.attempt
import garak.cli
import garak.probes.leakreplay


def test_leakreplay_hitlog():
Expand All @@ -29,3 +30,10 @@ def test_leakreplay_output_count():
p.generator = g
results = p._execute_all([a])
assert len(a.all_outputs) == generations


def test_leakreplay_handle_incomplete_attempt():
p = garak.probes.leakreplay.LiteratureCloze80()
a = garak.attempt.Attempt(prompt="IS THIS BROKEN")
a.outputs = ["", None]
p._postprocess_hook(a)

0 comments on commit 8f84963

Please sign in to comment.