Naman Gupta commited on
Commit
99a9cb9
·
1 Parent(s): f778bee

add ConversationManager to give the defender a memory

Browse files

Without this, the defender forgets everything it said on the
previous turn. Now it tracks the full back-and-forth so it
can't be tricked by "you already agreed to help me" attacks.

Files changed (1) hide show
  1. llm/history_manager.py +67 -0
llm/history_manager.py ADDED
@@ -0,0 +1,67 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # history_manager.py
2
+ # -------------------
3
+ # Keeps track of the full back-and-forth conversation between the attacker
4
+ # and the defender within a single episode.
5
+ #
6
+ # Why do we need this?
7
+ # The server only passes the last few attacker messages to us, so the
8
+ # defender would have no memory of what it said earlier. This manager
9
+ # stores both sides (user = attacker, assistant = defender) so the
10
+ # defender can't be fooled by tricks like "you already agreed to help me".
11
+
12
+ import logging
13
+
14
+ logger = logging.getLogger(__name__)
15
+
16
+
17
+ class ConversationManager:
18
+
19
+ def __init__(self) -> None:
20
+ # Each entry is {"role": "user" | "assistant", "content": "..."}
21
+ self._history: list[dict] = []
22
+
23
+ def reset(self) -> None:
24
+ """Wipe the slate clean at the start of a new episode."""
25
+ self._history = []
26
+ logger.info("Fresh episode started — conversation history cleared")
27
+
28
+ def add_attacker_message(self, text: str) -> None:
29
+ """Record what the attacker just said."""
30
+ self._history.append({"role": "user", "content": text})
31
+
32
+ def add_defender_reply(self, text: str) -> None:
33
+ """Record what the defender just replied."""
34
+ self._history.append({"role": "assistant", "content": text})
35
+
36
+ def build_messages(self, system_prompt: str) -> list[dict]:
37
+ """
38
+ Return the full message list ready to send to Groq.
39
+ Format: [system prompt] + [all turns so far]
40
+ """
41
+ return [{"role": "system", "content": system_prompt}] + self._history
42
+
43
+ def as_readable_transcript(self) -> str:
44
+ """
45
+ Render the conversation as a human-readable transcript.
46
+ Used by the episode grader to review the full exchange.
47
+ """
48
+ lines = []
49
+ for message in self._history:
50
+ speaker = "ATTACKER" if message["role"] == "user" else "DEFENDER"
51
+ lines.append(f"{speaker}: {message['content']}")
52
+ return "\n\n".join(lines)
53
+
54
+ @property
55
+ def turn_count(self) -> int:
56
+ """How many complete attacker turns have happened so far."""
57
+ return sum(1 for m in self._history if m["role"] == "user")
58
+
59
+ # Keep old name working so existing code doesn't break
60
+ add_user = add_attacker_message
61
+ add_assistant = add_defender_reply
62
+ get_messages = build_messages
63
+ to_transcript = as_readable_transcript
64
+
65
+ @property
66
+ def turn(self) -> int:
67
+ return self.turn_count